[Suggestion for #57169] Map kernel blocks per attention group instead of re-paging in builders

tests/kernels/attention/test_deepgemm_attention.py

@@ -321,14 +321,13 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
     not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
 )
 def test_deepgemm_paged_mqa_packed_manager_block_stride():
-    from vllm.models.glm5next.common.sparse_indexer import _kpool_flat_page_view
+    from vllm.v1.worker.utils import map_kv_cache_to_kernel_blocks
 
     torch.manual_seed(0)
     num_blocks, block_size, packed_rows = 4, 256, 320
     num_heads, head_dim, row_bytes = 64, 128, 132
-    page_size = 64
-    pages_per_block = block_size // page_size
-    dense_stride = pages_per_block
+    page_size, pages_per_block = 64, block_size // 64
+    dense_stride, packed_stride = pages_per_block, packed_rows // page_size
     compact_pages = kv_cache_cast_to_fp8(
         torch.randn(
             num_blocks, block_size, 1, head_dim, device="cuda", dtype=torch.bfloat16
@@ -343,8 +342,9 @@ def test_deepgemm_paged_mqa_packed_manager_block_stride():
     )
     packed_manager = packed_backing[:, :block_size]
     packed_manager.copy_(compact_pages.squeeze(2).view(num_blocks, block_size, -1))
-    packed_stride = packed_manager.stride(0) // (page_size * row_bytes)
-    packed_pages = _kpool_flat_page_view(packed_manager).unsqueeze(2)
+    packed_pages = map_kv_cache_to_kernel_blocks(
+        packed_manager, pages_per_block
+    ).unsqueeze(2)
     offsets = torch.arange(pages_per_block, device="cuda", dtype=torch.int32)
     tables = [
         (stride + offsets).unsqueeze(0) for stride in (dense_stride, packed_stride)

tests/kernels/attention/test_rocm_aiter_mla_sink.py

@@ -455,8 +455,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):
     set_random_seed(412)
     num_tokens, num_heads, block_size = 2, 8, 16
     q = torch.randn(num_tokens, num_heads, Q_HEAD_DIM, device="cuda").to(dtype)
-    backing = torch.randn(2, 2, block_size, Q_HEAD_DIM, device="cuda").to(dtype)
-    kv = backing[:, 1]
+    kv = torch.randn(2, 2, block_size, Q_HEAD_DIM, device="cuda").to(dtype)[:, 1]
     selected = torch.tensor([[1, 7, 18], [0, 3, 20]], dtype=torch.int32, device="cuda")
     sinks = torch.linspace(-3.0, 5.0, num_heads, device="cuda")
     impl = object.__new__(ROCMAiterMLASparseImpl)

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

@@ -29,15 +29,7 @@
 TOPK = 128  # triton convert requires width % 128 == 0
 
 
-def ref_convert(
-    req_id,
-    block_table,
-    token_indices,
-    BLOCK_SIZE=64,
-    BLOCK_STRIDE_ROWS=None,
-    **_,
-):
-    block_stride_rows = BLOCK_STRIDE_ROWS or BLOCK_SIZE
+def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **kw):
     out = torch.full_like(token_indices, -1)
     counts = torch.zeros(token_indices.shape[0], dtype=torch.int32)
     for t in range(token_indices.shape[0]):
@@ -49,7 +41,7 @@ def ref_convert(
             blk = int(block_table[int(req_id[t]), pos // BLOCK_SIZE])
             if blk < 0:
                 continue
-            vals.append(blk * block_stride_rows + pos % BLOCK_SIZE)
+            vals.append(blk * kw["BLOCK_STRIDE_ROWS"] + pos % BLOCK_SIZE)
         out[t, : len(vals)] = torch.tensor(vals, dtype=out.dtype)
         counts[t] = len(vals)
     return out, counts
@@ -115,10 +107,7 @@ def make_batch(rows, topk_rows, own_blocks):
     for t, row in enumerate(topk_rows):
         topk[t, : len(row)] = torch.tensor(row, dtype=torch.int32)
     return SimpleNamespace(
-        req_id_per_token=req_id,
-        block_table=block_table,
-        block_size=BLOCK_SIZE,
-        topk_indices=topk,
+        req_id_per_token=req_id, block_table=block_table, block_size=BLOCK_SIZE
     )
 
 
@@ -141,18 +130,11 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     state.kv_indptr[1:5] = torch.tensor([8, 13, 16, 17], dtype=torch.int32)
     state.kv_indptr[5:] = 17
     meta = make_batch(rows, topk_rows, [3])
-    impl.topk_indices_buffer.copy_(meta.topk_indices)
     meta.state = state
     q_nope = torch.randn(rows, impl.num_heads, HEAD)
     q_rope = torch.randn(rows, impl.num_heads, qk_rope)
-    backing = torch.zeros(
-        8,
-        2,
-        BLOCK_SIZE,
-        impl.head_size,
-        dtype=torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16,
-    )
-    cache = backing[:, 0]
+    dtype = torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16
+    cache = torch.zeros(8, 2, BLOCK_SIZE, impl.head_size, dtype=dtype)[:, 0]
 
     out, lse = impl.forward_mqa(
         (q_nope, q_rope), cache, meta, SimpleNamespace(_k_scale_float=0.5)
@@ -308,21 +290,6 @@ def test_pack_indices_replays_with_updated_offsets(width):
         assert (state.kv_indices[expected.numel() :] == -99).all()
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
-def test_pack_indices_padding_uses_valid_slot():
-    """Planned rows may read padding; it must not point at unrelated cache rows."""
-    state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
-    state.kv_indptr = torch.tensor([0, 6, 10], dtype=torch.int32, device="cuda")
-    state.kv_indices = torch.empty(10, dtype=torch.int32, device="cuda")
-    slots = torch.full((2, TOPK), -1, dtype=torch.int32, device="cuda")
-    slots[0, :2] = torch.tensor([7, 3], device="cuda")
-    slots[1, 0] = 13
-
-    state.pack_indices(slots)
-
-    assert state.kv_indices.tolist() == [7, 3, 7, 7, 7, 7, 13, 13, 13, 13]
-
-
 @pytest.mark.parametrize("adaptive", [False, True])
 def test_kv_lens_host_formula(adaptive):
     """Per-row host lengths: context == position + 1; capped at

tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py

@@ -20,6 +20,10 @@
 import torch
 
 from vllm.platforms.interface import DeviceCapability
+from vllm.v1.attention.backend import MultipleOf
+from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
+    FlashAttnMLASparseBackend,
+)
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     QUANTIZED_DS_MLA_CACHE_FORMATS,
     FlashMLASparseBackend,
@@ -96,6 +100,12 @@ def test_supported_head_sizes_include_512():
     assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
 
 
+def test_flash_attn_sparse_accepts_manager_blocks():
+    (supported,) = FlashAttnMLASparseBackend.get_supported_kernel_block_sizes()
+    assert isinstance(supported, MultipleOf)
+    assert supported.base == 64
+
+
 def test_flashmla_bf16_nope_accepts_packed_manager_blocks():
     spec = MLAAttentionSpec(
         block_size=1152,
@@ -105,8 +115,6 @@ def test_flashmla_bf16_nope_accepts_packed_manager_blocks():
         cache_dtype_str="bfloat16",
         block_stride_alignment=1024,
     )
-    assert FlashMLASparseBackend.get_supported_kernel_block_sizes() == [64]
-    assert FlashMLASparseBackend.get_strided_block_page_rows(spec) == 1
     num_blocks = 2
     layers = ["layer.0", "layer.1"]
     page_size = spec.page_size_bytes
@@ -147,10 +155,8 @@ def test_flashmla_quantized_cache_keeps_fixed_kernel_pages():
         state_content_bytes=656,
     )
 
-    assert FlashMLASparseBackend.get_supported_kernel_block_sizes() == [64]
     assert FlashMLASparseBackend.get_supported_kernel_block_sizes(spec) == [64]
-    assert FlashMLASparseBackend.get_strided_block_page_rows(spec) is None
-    assert select_common_block_size(1152, [FlashMLASparseBackend]) == 64
+    assert select_common_block_size(1152, [FlashMLASparseBackend], [spec]) == 64
 
 
 def test_quantized_ds_mla_formats_are_the_envelope_set():