[None][fix] Follow up PrimTS refresh backport review tests and docs

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,

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.

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

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()

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:

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,