From 3eec36066f484a0fed57b141217e73ca377bc4f2 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 23:29:56 -0500 Subject: [PATCH] perf(qwen): drop the redundant key sanitizing select Native-KV attention zeroes inactive key and value rows with a select over the whole cache on every step. The key select cannot change the result: the attention mask already replaces the score of every inactive position with the blocked value, and an elementwise select never propagates its unselected operand. Only the value select protects the context matmul from a zero probability times a non-finite row. Feed the unmasked keys to the score matmul and keep the value select. On an H200, Qwen3-0.6B fp16 at its default 40960-row capacity decodes 128 greedy tokens with identical token IDs in 6.89 instead of 7.73 ms per token, median of seven interleaved runs. Add a TensorRT test that fills unwritten cache rows with NaN and Inf and checks decode and prefill-chunk attention against a float64 reference. It fails when the value select is removed. Refs: #1409 Signed-off-by: RunguoLi --- families/qwen/native_kv_attention_builder.py | 12 +- .../qwen/tests/test_native_kv_attention.py | 143 ++++++++++++++++++ 2 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 families/qwen/tests/test_native_kv_attention.py diff --git a/families/qwen/native_kv_attention_builder.py b/families/qwen/native_kv_attention_builder.py index 071721d8a8..113b4acbb3 100644 --- a/families/qwen/native_kv_attention_builder.py +++ b/families/qwen/native_kv_attention_builder.py @@ -208,17 +208,15 @@ def add_explicit_masked_grouped_query_attention( np.array([0.0], dtype=np.float32), ) zero = _cast(network, zero, query.dtype) - safe_key = network.add_select( - active_grouped.get_output(0), - key_grouped.get_output(0), - zero, - ) + # Unwritten cache rows may hold non-finite values. Keys need no sanitizing: + # the attention mask replaces every inactive score with the blocked value. + # Values do, because a zero probability times a non-finite row is NaN. safe_value = network.add_select( active_grouped.get_output(0), value_grouped.get_output(0), zero, ) - if safe_key is None or safe_value is None: + if safe_value is None: raise RuntimeError("TensorRT failed to sanitize inactive native KV rows") if scale is None: @@ -241,7 +239,7 @@ def add_explicit_masked_grouped_query_attention( scores = network.add_matrix_multiply( scaled_query, trt.MatrixOperation.NONE, - safe_key.get_output(0), + key_grouped.get_output(0), trt.MatrixOperation.TRANSPOSE, ) if scores is None: diff --git a/families/qwen/tests/test_native_kv_attention.py b/families/qwen/tests/test_native_kv_attention.py new file mode 100644 index 0000000000..d6a1400d4f --- /dev/null +++ b/families/qwen/tests/test_native_kv_attention.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Masked native-KV attention ignores whatever unwritten cache rows contain.""" + +from __future__ import annotations + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt", reason="TensorRT is required for native KV attention") +cudart = pytest.importorskip("cuda.bindings.runtime", reason="cuda-python is required") + +from families.qwen.native_kv_attention_builder import ( # noqa: E402 + add_active_prefix_causal_masks, + add_explicit_masked_grouped_query_attention, +) + +pytestmark = [pytest.mark.gpu, pytest.mark.trt] + +HEADS, KV_HEADS, HEAD_DIM, CAPACITY = 4, 2, 8, 8 + + +def _check(result): + error = result[0] if isinstance(result, tuple) else result + if error != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"CUDA call failed: {error}") + return result[1] if isinstance(result, tuple) and len(result) > 1 else None + + +def _build(dtype: trt.DataType, query_length: int) -> bytes: + logger = trt.Logger(trt.Logger.ERROR) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + query = network.add_input("query", dtype, (1, HEADS, query_length, HEAD_DIM)) + key = network.add_input("key", dtype, (1, KV_HEADS, CAPACITY, HEAD_DIM)) + value = network.add_input("value", dtype, (1, KV_HEADS, CAPACITY, HEAD_DIM)) + rows = network.add_input("rows", trt.float32, (query_length, 1)) + write_index = network.add_input("write_index", trt.int32, (1,)) + length = network.add_input("length", trt.int32, (1,)) + masks = add_active_prefix_causal_masks(network, rows, write_index, length, CAPACITY) + context = add_explicit_masked_grouped_query_attention( + network, + query, + key, + value, + masks, + num_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + ) + context = network.add_cast(context, trt.float32).get_output(0) + context.name = "context" + network.mark_output(context) + config = builder.create_builder_config() + # Compare float32 against the float64 reference without TF32 matmul rounding. + config.clear_flag(trt.BuilderFlag.TF32) + plan = builder.build_serialized_network(network, config) + assert plan is not None + return bytes(plan) + + +def _run(plan: bytes, inputs: dict[str, np.ndarray], output_shape: tuple[int, ...]) -> np.ndarray: + runtime = trt.Runtime(trt.Logger(trt.Logger.ERROR)) + engine = runtime.deserialize_cuda_engine(plan) + context = engine.create_execution_context() + stream = _check(cudart.cudaStreamCreate()) + output = np.empty(output_shape, dtype=np.float32) + buffers = [] + try: + for name, array in [*inputs.items(), ("context", output)]: + pointer = _check(cudart.cudaMalloc(array.nbytes)) + buffers.append(pointer) + context.set_tensor_address(name, pointer) + if name != "context": + _check(cudart.cudaMemcpy( + pointer, array.ctypes.data, array.nbytes, + cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, + )) + assert context.execute_async_v3(stream) + _check(cudart.cudaMemcpyAsync( + output.ctypes.data, buffers[-1], output.nbytes, + cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream, + )) + _check(cudart.cudaStreamSynchronize(stream)) + finally: + for pointer in buffers: + _check(cudart.cudaFree(pointer)) + _check(cudart.cudaStreamDestroy(stream)) + return output + + +def _reference(query, key, value, write_index: int, length: int) -> np.ndarray: + groups = HEADS // KV_HEADS + query_length = query.shape[2] + result = np.zeros((1, HEADS, query_length, HEAD_DIM), dtype=np.float64) + for head in range(HEADS): + keys = key[0, head // groups].astype(np.float64) + values = value[0, head // groups].astype(np.float64) + for row in range(query_length): + visible = min(write_index + row + 1, length) + scores = keys[:visible] @ query[0, head, row].astype(np.float64) / np.sqrt(HEAD_DIM) + weights = np.exp(scores - scores.max()) + result[0, head, row] = (weights / weights.sum()) @ values[:visible] + return result + + +@pytest.mark.parametrize("dtype", [trt.float32, trt.float16]) +@pytest.mark.parametrize( + ("query_length", "write_index", "length"), + [(1, 4, 5), (3, 2, 5)], + ids=["decode", "prefill-chunk"], +) +def test_unwritten_cache_rows_cannot_reach_the_context(dtype, query_length, write_index, length): + numpy_dtype = np.float32 if dtype == trt.float32 else np.float16 + rng = np.random.default_rng(7) + query = rng.standard_normal((1, HEADS, query_length, HEAD_DIM)).astype(numpy_dtype) + key = rng.standard_normal((1, KV_HEADS, CAPACITY, HEAD_DIM)).astype(numpy_dtype) + value = rng.standard_normal((1, KV_HEADS, CAPACITY, HEAD_DIM)).astype(numpy_dtype) + # Rows at or beyond the active length were never written by this request. + key[:, :, length:] = np.nan + key[:, :, length:, 0] = np.inf + value[:, :, length:] = np.nan + value[:, :, length:, 1] = -np.inf + + actual = _run( + _build(dtype, query_length), + { + "query": query, + "key": key, + "value": value, + "rows": np.zeros((query_length, 1), dtype=np.float32), + "write_index": np.array([write_index], dtype=np.int32), + "length": np.array([length], dtype=np.int32), + }, + (1, HEADS, query_length, HEAD_DIM), + ) + + assert np.isfinite(actual).all() + tolerance = 1e-5 if dtype == trt.float32 else 5e-3 + np.testing.assert_allclose( + actual, _reference(query, key, value, write_index, length), rtol=0, atol=tolerance + )