[Bugfix][MLA] Fix BLHNC addressing for FlashInfer sparse MLA (#54465)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

tests/v1/attention/test_indexer_dcp_localize.py

@@ -642,7 +642,10 @@ def test_sparse_prefill_dcp_metadata_localizes_causal_bounds():
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
-def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():
+@pytest.mark.parametrize("block_stride_rows", [None, 12])
+def test_dcp_filter_compacts_valid_slots_for_sparse_kernel(
+    block_stride_rows: int | None,
+):
     block_size = 4
     num_topk = 128
     dcp_size = 2
@@ -658,6 +661,7 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():
         dcp_size=dcp_size,
         dcp_rank=0,
         BLOCK_SIZE=block_size,
+        BLOCK_STRIDE_ROWS=block_stride_rows,
         NUM_TOPK_TOKENS=num_topk,
         return_valid_counts=True,
     )
@@ -668,7 +672,11 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():
     assert (out[0, valid:] == -1).all()
     # In-kernel compaction packs valid slots to the front; prefix order is
     # unspecified, so compare as a set.
-    assert set(out[0, :valid].cpu().tolist()) == {40, 41, 42, 43}
+    block_stride_rows = block_stride_rows or block_size
+    expected_base = 10 * block_stride_rows
+    assert set(out[0, :valid].cpu().tolist()) == set(
+        range(expected_base, expected_base + block_size)
+    )
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

vllm/models/deepseek_v32/common/kernels.py

@@ -140,6 +140,7 @@ def _fused_norm_rope_kernel(
     mla_cache_ptr,
     mla_cache_block_stride,
     mla_cache_entry_stride,
+    MLA_CACHE_BLOCK_SIZE: tl.constexpr,
     MLA_CACHE_FP8: tl.constexpr,
     mla_cache_scale_ptr,
     # fp8_ds_mla cache views (block-scaled fp8 NoPE + unquantized bf16 RoPE).
@@ -242,7 +243,7 @@ def _fused_norm_rope_kernel(
                 return
 
             slot_idx = tl.load(slot_mapping_ptr + tok_idx)
-            mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
+            mla_block_size = MLA_CACHE_BLOCK_SIZE
             mla_block_idx = slot_idx // mla_block_size
             mla_block_off = slot_idx % mla_block_size
 
@@ -479,6 +480,7 @@ def fused_norm_rope(
     mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)
     mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device)
     if mla_kv_cache is not None:
+        mla_block_size = mla_kv_cache.shape[1]
         if mla_cache_ds_mla:
             # 656-byte custom layout addressed in bytes; mla_cache_ptr is the
             # 1-byte fp8 view, so block/entry strides are byte offsets and the
@@ -503,6 +505,7 @@ def fused_norm_rope(
         mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
         mla_block_stride = 0
         mla_entry_stride = 0
+        mla_block_size = 1
         mla_k_scale = _dummy((1,), torch.float32, device)
 
     if q_c_out is None:
@@ -570,6 +573,7 @@ def fused_norm_rope(
         mla_kv_cache,
         mla_block_stride,
         mla_entry_stride,
+        mla_block_size,
         mla_cache_fp8,
         mla_k_scale,
         mla_ds_scale_view,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -28,6 +28,7 @@ from vllm.v1.attention.backend import (
     MultipleOf,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    flat_kv_row_view,
     triton_convert_req_index_to_global_index,
     triton_filter_and_convert_dcp_index,
 )
@@ -377,6 +378,10 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_actual_toks]
 
+        _, block_stride_rows = flat_kv_row_view(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
+
         if self.dcp_world_size > 1:
             topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(
                 attn_metadata.req_id_per_token[:num_actual_toks],
@@ -386,6 +391,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
                 dcp_rank=self.dcp_rank,
                 cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size),
                 BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
                 return_valid_counts=True,
             )
@@ -395,6 +401,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
                 attn_metadata.block_table,
                 topk_indices,
                 BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
                 return_valid_counts=True,
             )

vllm/v1/attention/backends/mla/sparse_utils.py

@@ -307,6 +307,7 @@ def triton_filter_and_convert_dcp_index(
     dcp_rank: int,
     cp_kv_cache_interleave_size: int = 1,
     BLOCK_SIZE: int = 64,
+    BLOCK_STRIDE_ROWS: int | None = None,
     NUM_TOPK_TOKENS: int = 2048,
     BLOCK_N: int = 128,
     return_valid_counts: bool = False,
@@ -343,6 +344,7 @@ def triton_filter_and_convert_dcp_index(
             block_table,
             token_indices,
             BLOCK_SIZE=BLOCK_SIZE,
+            BLOCK_STRIDE_ROWS=BLOCK_STRIDE_ROWS,
             NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
             BLOCK_N=BLOCK_N,
             return_valid_counts=return_valid_counts,
@@ -390,7 +392,7 @@ def triton_filter_and_convert_dcp_index(
         None,
         max_num_blocks_per_req,
         BLOCK_SIZE,
-        BLOCK_SIZE,  # dense caches on the DCP path
+        BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
         block_n,
         False,  # HAS_PREFILL
         count_valid,