diff --git a/families/qwen/native_kv_attention_builder.py b/families/qwen/native_kv_attention_builder.py index 071721d8a..113b4acbb 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 000000000..d6a1400d4 --- /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 + )