From 5b0228a95410b0dd8ded00b61ddb0344f4bad002 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 15 Oct 2026 05:24:59 +0000 Subject: [PATCH] [None][fix] Address PrimTS refresh backport review tests and docs Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../sparse-attention-development-guide.md | 4 ++ docs/source/features/sparse-attention.md | 5 ++ .../attention/ATTENTION_DEVELOPER_GUIDE.md | 5 ++ .../sparse/test_prims_ts_block_sparse.py | 62 +++++++++++++++---- .../test_prims_ts_attention_backend.py | 62 +++++++++++++++++++ .../_torch/attention/test_prims_ts_fmha.py | 23 ++++--- 6 files changed, 142 insertions(+), 19 deletions(-) diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index 4b939446199d..44c5803b0400 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -330,6 +330,10 @@ three prediction methods. A `VanillaAttention` implementation instead overrides (BSR `block_indptr`/`block_indices` or a packed `exact_block_bits` bitmask), optional K/V summaries for proxy routes, and optional `kv_valid_bits` masking ragged KV tails. +- **Summary contract**: both `k_summary` and `v_summary` are arithmetic means + over the tokens present in each KV block, excluding padding in a partial + final block. The kernel applies the represented token count to each proxy; + do not pass V sums or pre-weight the summaries. - **Default**: hands through `SparseBackendForwardArgs.block_sparse_inputs`, so modules that predict before the core forward do not override it. Override it to predict inside the backend from the flattened Q/K/V, diff --git a/docs/source/features/sparse-attention.md b/docs/source/features/sparse-attention.md index acd9609c491a..d31a005bf703 100644 --- a/docs/source/features/sparse-attention.md +++ b/docs/source/features/sparse-attention.md @@ -163,6 +163,11 @@ routes is never served by a dense kernel. Both the contiguous prefill path, used by diffusion models that keep no KV cache, and the paged generation path are provided by the vendored PrimTS kernels. +Both `k_summary` and `v_summary` are arithmetic means over the tokens present +in each KV block. A partial final block excludes padding from its mean. The +kernel weights each proxy by its represented token count; do not pass V sums +or pre-weight the summaries. + This is an attention capability, not a standalone public `SparseAttentionConfig` algorithm. A user-facing algorithm must also provide the selector, metadata, and backend integration. diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index a44632837d7c..f82b0faf1801 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -224,6 +224,11 @@ routes, and optional token-validity bits mask ragged KV tails. Plans contain only static format, proxy, geometry, and capacity choices; every run receives the live routes, summaries, validity bits, page tables, and sequence lengths. +Both `k_summary` and `v_summary` are arithmetic means over the tokens present +in each KV block, excluding padding in a partial final block. The kernel +weights each proxy by its represented token count; do not pass V sums or +pre-weight the summaries. + `PrimsTSBlockSparseFmha` owns its wrapper-plan cache by default. Integrations whose attention layers execute serially may explicitly bind a model-scoped cache to reuse graph-stable route workspaces across compatible layers. The diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 124fa905b67f..b594c70f9f80 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -38,8 +38,6 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType -pytestmark = pytest.mark.cpu_only - _REQUIRES_PRIMTS_GPU = pytest.mark.skipif( getSMVersion() not in (100, 103), reason="PrimTS block-sparse attention requires SM100 or SM103", @@ -102,23 +100,47 @@ def _proxy_reference( block_size = 64 exact_tokens = torch.arange( exact_block * block_size, - (exact_block + 1) * block_size, + min((exact_block + 1) * block_size, k.shape[1]), device=q.device, ) proxy_blocks = [block for block in range(k_summary.shape[1]) if block != exact_block] + proxy_token_counts = torch.tensor( + [min(block_size, k.shape[1] - block * block_size) for block in proxy_blocks], + device=q.device, + dtype=torch.float32, + ) q_rows = q[0, :, 0].float() exact_logits = q_rows @ k[0, exact_tokens, 0].float().T proxy_logits = q_rows @ k_summary[0, proxy_blocks, 0].float().T logits = torch.cat((exact_logits, proxy_logits), dim=1) / math.sqrt(q.shape[-1]) weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) - exact_weights, proxy_weights = weights.split((block_size, len(proxy_blocks)), dim=1) + exact_weights, proxy_weights = weights.split((exact_tokens.numel(), len(proxy_blocks)), dim=1) + proxy_weights = proxy_weights * proxy_token_counts numerator = exact_weights @ v[0, exact_tokens, 0].float() - numerator += (proxy_weights * block_size) @ v_summary[0, proxy_blocks, 0].float() + numerator += proxy_weights @ v_summary[0, proxy_blocks, 0].float() denominator = exact_weights.sum(dim=1, keepdim=True) - denominator += proxy_weights.sum(dim=1, keepdim=True) * block_size + denominator += proxy_weights.sum(dim=1, keepdim=True) return (numerator / denominator).to(q.dtype)[None, :, None] +@pytest.mark.cpu_only +@pytest.mark.parametrize("kv_seq_len", [160, 192]) +@pytest.mark.parametrize("exact_block", [0, 2]) +def test_proxy_reference_weights_present_tokens(kv_seq_len: int, exact_block: int) -> None: + q = torch.zeros((1, 64, 1, 128)) + k = torch.zeros((1, kv_seq_len, 1, 128)) + v = torch.zeros_like(k) + v[:, 128:] = 1 + k_summary = torch.stack([block.mean(dim=1) for block in k.split(64, dim=1)], dim=1) + v_summary = torch.stack([block.mean(dim=1) for block in v.split(64, dim=1)], dim=1) + + actual = _proxy_reference(q, k, v, k_summary, v_summary, exact_block) + + # Uniform logits give every present token equal weight, including the tail. + expected = v.mean(dim=1, keepdim=True).expand_as(q) + torch.testing.assert_close(actual, expected) + + class _Attention: def __init__(self) -> None: self.sparse_params = None @@ -207,6 +229,7 @@ def _paged_case(): return attention, fmha, q, metadata, args +@pytest.mark.cpu_only def test_block_sparse_route_mode_is_derived_from_payload() -> None: bsr = _bsr_inputs() exact = _bitmask_inputs(proxy=False) @@ -229,6 +252,7 @@ def test_block_sparse_route_mode_is_derived_from_payload() -> None: ({"k_summary": torch.empty(0)}, "k_summary and v_summary"), ], ) +@pytest.mark.cpu_only def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) -> None: kwargs = { "q_block_size": 64, @@ -243,6 +267,7 @@ def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) BlockSparseForwardInputs(**kwargs) +@pytest.mark.cpu_only def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( monkeypatch, ) -> None: @@ -260,6 +285,7 @@ def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( assert reason == "paged block-sparse attention only supports BSR exact routes" +@pytest.mark.cpu_only def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) @@ -280,6 +306,7 @@ def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) @pytest.mark.parametrize("paged", [False, True]) +@pytest.mark.cpu_only def test_block_sparse_support_rejects_invalid_static_kernel_profile( monkeypatch, paged, @@ -319,6 +346,7 @@ def test_block_sparse_support_rejects_invalid_static_kernel_profile( assert reason == "block-sparse requires head_dim=128" +@pytest.mark.cpu_only def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: _attention, fmha, q, k, v, _metadata, args = _contiguous_case() wrapper = Mock() @@ -362,6 +390,7 @@ def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypat assert call.kwargs["v_summary"] is inputs.v_summary +@pytest.mark.cpu_only def test_block_sparse_plan_key_includes_attention_head_topology() -> None: inputs = _bitmask_inputs(proxy=True) q = torch.empty((128, 256), dtype=torch.bfloat16) @@ -385,6 +414,7 @@ def _key(fmha): assert _key(first) != _key(second) +@pytest.mark.cpu_only def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: first = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) second = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) @@ -404,6 +434,7 @@ def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: } +@pytest.mark.cpu_only def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() @@ -480,6 +511,7 @@ def snapshot(*_args, **kwargs): assert snapshots[1][3] is _get_block_sparse_inputs(args).block_indptr +@pytest.mark.cpu_only def test_paged_block_tables_remain_live_across_graph_replay(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() @@ -527,6 +559,7 @@ def snapshot(*_args, **kwargs): torch.testing.assert_close(seen[1][1], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) +@pytest.mark.cpu_only def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: _attention, fmha, _q, _metadata, _args = _paged_case() query_device = torch.device("cuda:1") @@ -560,6 +593,7 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: assert workspace.numel() == 0 +@pytest.mark.cpu_only def test_prepare_workspace_skips_generation_layout_for_contiguous_requests(monkeypatch) -> None: _attention, fmha, q, _k, _v, metadata, args = _contiguous_case() layout = Mock() @@ -618,16 +652,19 @@ def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: @_REQUIRES_PRIMTS_GPU +@pytest.mark.parametrize("kv_seq_len", [160, 192]) @torch.no_grad() -def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: +def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries(kv_seq_len: int) -> None: torch.manual_seed(20260901) q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.bfloat16) - k = torch.randn((1, 192, 1, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn((1, kv_seq_len, 1, 128), device="cuda", dtype=torch.bfloat16) v = torch.randn_like(k) - k_blocks = k.float().view(1, 3, 64, 1, 128) - v_blocks = v.float().view(1, 3, 64, 1, 128) - initial_k_summary = k_blocks.mean(dim=2).to(k.dtype) - initial_v_summary = v_blocks.mean(dim=2).to(v.dtype) + initial_k_summary = torch.stack( + [block.float().mean(dim=1) for block in k.split(64, dim=1)], dim=1 + ).to(k.dtype) + initial_v_summary = torch.stack( + [block.float().mean(dim=1) for block in v.split(64, dim=1)], dim=1 + ).to(v.dtype) live_k_summary = initial_k_summary.clone() live_v_summary = initial_v_summary.clone() live_exact_bits = torch.tensor([[[[1]]]], device="cuda", dtype=torch.uint32) @@ -885,6 +922,7 @@ def test_real_gpu_paged_routes_match_reference_across_topologies( torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) +@pytest.mark.cpu_only def test_paged_support_rejects_pages_smaller_than_a_route_fragment() -> None: page_size = 32 _attention, fmha, q, metadata, args = _paged_case() diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index 2adfac26c11c..b3737e937fcd 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -563,6 +563,68 @@ def check_output() -> None: check_output() +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +@pytest.mark.parametrize("log2_jump", [7.0, 9.0], ids=["freeze-row-max", "rescale-row-max"]) +@torch.no_grad() +def test_prims_ts_context_lazy_rescale_with_late_key_block( + dtype: torch.dtype, log2_jump: float +) -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import BatchPrefillPagedTSWrapper + + device = torch.device("cuda") + q_len, kv_len, page_size, head_dim = 64, 512, 64, 128 + sm_scale = head_dim**-0.5 + query = torch.zeros((q_len, 1, head_dim), device=device, dtype=dtype) + query[..., 0] = 1 + k_cache = torch.zeros((kv_len // page_size, 1, page_size, head_dim), device=device, dtype=dtype) + v_cache = torch.ones_like(k_cache) + # Only the last page raises the row maximum, after earlier KV tiles have + # contributed to the accumulator. The skip threshold is 8 in log2 logits. + k_cache[-1, :, :, 0] = log2_jump / (sm_scale * math.log2(math.e)) + v_cache[-1] = -1 + qo_indptr = torch.tensor([0, q_len], device=device, dtype=torch.int32) + block_tables = torch.arange(kv_len // page_size, device=device, dtype=torch.int32)[None] + seq_lens = torch.tensor([kv_len], device=device, dtype=torch.int32) + output = torch.empty_like(query) + + logical_k = k_cache[:, 0].flatten(0, 1).float() + logical_v = v_cache[:, 0].flatten(0, 1).float() + scores = (query[:, 0].float() @ logical_k.T) * sm_scale + actual_log2_jump = (scores[0, -1] - scores[0, 0]).item() * math.log2(math.e) + assert (actual_log2_jump > 8.0) == (log2_jump > 8.0) + expected = (torch.softmax(scores, dim=-1) @ logical_v).to(dtype)[:, None] + + wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND") + wrapper.plan( + device=device, + batch_size=1, + max_seq_len_q=q_len, + max_kv_len=kv_len, + num_qo_heads=1, + num_kv_heads=1, + head_dim=head_dim, + q_dtype=dtype, + k_dtype=dtype, + v_dtype=dtype, + out_dtype=dtype, + page_size=page_size, + mask_type="dense", + sm_scale=sm_scale, + ) + wrapper.run(query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + wrapper.run( + query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output, validate=False + ) + output.zero_() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + + def test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 8902890021ce..8b705bf20564 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -1156,7 +1156,7 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( False, ) ) - wrapper = Mock() + wrapper = create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) wrapper._plan_state = SimpleNamespace( policy=( ("use_split_kv", use_split_kv), @@ -1324,7 +1324,10 @@ def test_decode_layer_adapters_bind_the_same_shared_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_decode_module, @@ -1410,7 +1413,10 @@ def test_decode_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_decode_module, @@ -1510,7 +1516,7 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( dtype=torch.int32, ) build_metadata = Mock(return_value=(kv_cache, block_tables, None)) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_ts_module.thop, @@ -1704,7 +1710,10 @@ def test_mla_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention(head_dim=576, is_mla=True, num_heads=4) fmha = PrimsTSFmha(attn) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_mla_module, @@ -1770,7 +1779,7 @@ def test_mla_wrapper_capture_uses_cached_plan_and_rejects_plan_miss( "is_current_stream_capturing", lambda: capturing, ) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_mla_module, @@ -1816,7 +1825,7 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( dtype=torch.int32, ) build_metadata = Mock(return_value=(kv_cache, block_tables, None)) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_ts_module.thop,