@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5707355209

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

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

Lucas Wilkinson committed Aug 30, 2026 at 21:13 UTC 570735520942b6428ae98f1eb442c6d88adae55e

4 files changed +25 -4

tests/v1/attention/test_indexer_dcp_localize.py

+10 -2

@@ -642,7 +642,10 @@ def test_sparse_prefill_dcp_metadata_localizes_causal_bounds():

642

643

644 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

645 -def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():

645 +@pytest.mark.parametrize("block_stride_rows", [None, 12])

646 +def test_dcp_filter_compacts_valid_slots_for_sparse_kernel(

647 + block_stride_rows: int | None,

648 +):

649 block_size = 4

650 num_topk = 128

651 dcp_size = 2

@@ -658,6 +661,7 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():

661 dcp_size=dcp_size,

662 dcp_rank=0,

663 BLOCK_SIZE=block_size,

664 + BLOCK_STRIDE_ROWS=block_stride_rows,

665 NUM_TOPK_TOKENS=num_topk,

666 return_valid_counts=True,

667 )

@@ -668,7 +672,11 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel():

672 assert (out[0, valid:] == -1).all()

673 # In-kernel compaction packs valid slots to the front; prefix order is

674 # unspecified, so compare as a set.

671 - assert set(out[0, :valid].cpu().tolist()) == {40, 41, 42, 43}

675 + block_stride_rows = block_stride_rows or block_size

676 + expected_base = 10 * block_stride_rows

677 + assert set(out[0, :valid].cpu().tolist()) == set(

678 + range(expected_base, expected_base + block_size)

679 + )

680

681

682 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")

vllm/models/deepseek_v32/common/kernels.py

+5 -1

@@ -140,6 +140,7 @@ def _fused_norm_rope_kernel(

140 mla_cache_ptr,

141 mla_cache_block_stride,

142 mla_cache_entry_stride,

143 + MLA_CACHE_BLOCK_SIZE: tl.constexpr,

144 MLA_CACHE_FP8: tl.constexpr,

145 mla_cache_scale_ptr,

146 # fp8_ds_mla cache views (block-scaled fp8 NoPE + unquantized bf16 RoPE).

@@ -242,7 +243,7 @@ def _fused_norm_rope_kernel(

243 return

244

245 slot_idx = tl.load(slot_mapping_ptr + tok_idx)

245 - mla_block_size = mla_cache_block_stride // mla_cache_entry_stride

246 + mla_block_size = MLA_CACHE_BLOCK_SIZE

247 mla_block_idx = slot_idx // mla_block_size

248 mla_block_off = slot_idx % mla_block_size

249

@@ -479,6 +480,7 @@ def fused_norm_rope(

480 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)

481 mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device)

482 if mla_kv_cache is not None:

483 + mla_block_size = mla_kv_cache.shape[1]

484 if mla_cache_ds_mla:

485 # 656-byte custom layout addressed in bytes; mla_cache_ptr is the

486 # 1-byte fp8 view, so block/entry strides are byte offsets and the

@@ -503,6 +505,7 @@ def fused_norm_rope(

505 mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)

506 mla_block_stride = 0

507 mla_entry_stride = 0

508 + mla_block_size = 1

509 mla_k_scale = _dummy((1,), torch.float32, device)

510

511 if q_c_out is None:

@@ -570,6 +573,7 @@ def fused_norm_rope(

573 mla_kv_cache,

574 mla_block_stride,

575 mla_entry_stride,

576 + mla_block_size,

577 mla_cache_fp8,

578 mla_k_scale,

579 mla_ds_scale_view,

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

@@ -28,6 +28,7 @@ from vllm.v1.attention.backend import (

28 MultipleOf,

29 )

30 from vllm.v1.attention.backends.mla.sparse_utils import (

31 + flat_kv_row_view,

32 triton_convert_req_index_to_global_index,

33 triton_filter_and_convert_dcp_index,

34 )

@@ -377,6 +378,10 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):

378 assert self.topk_indices_buffer is not None

379 topk_indices = self.topk_indices_buffer[:num_actual_toks]

380

381 + _, block_stride_rows = flat_kv_row_view(

382 + kv_c_and_k_pe_cache, attn_metadata.block_size

383 + )

384 +

385 if self.dcp_world_size > 1:

386 topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(

387 attn_metadata.req_id_per_token[:num_actual_toks],

@@ -386,6 +391,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):

391 dcp_rank=self.dcp_rank,

392 cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size),

393 BLOCK_SIZE=attn_metadata.block_size,

394 + BLOCK_STRIDE_ROWS=block_stride_rows,

395 NUM_TOPK_TOKENS=topk_indices.shape[1],

396 return_valid_counts=True,

397 )

@@ -395,6 +401,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):

401 attn_metadata.block_table,

402 topk_indices,

403 BLOCK_SIZE=attn_metadata.block_size,

404 + BLOCK_STRIDE_ROWS=block_stride_rows,

405 NUM_TOPK_TOKENS=topk_indices.shape[1],

406 return_valid_counts=True,

407 )

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

+3 -1

@@ -307,6 +307,7 @@ def triton_filter_and_convert_dcp_index(

307 dcp_rank: int,

308 cp_kv_cache_interleave_size: int = 1,

309 BLOCK_SIZE: int = 64,

310 + BLOCK_STRIDE_ROWS: int | None = None,

311 NUM_TOPK_TOKENS: int = 2048,

312 BLOCK_N: int = 128,

313 return_valid_counts: bool = False,

@@ -343,6 +344,7 @@ def triton_filter_and_convert_dcp_index(

344 block_table,

345 token_indices,

346 BLOCK_SIZE=BLOCK_SIZE,

347 + BLOCK_STRIDE_ROWS=BLOCK_STRIDE_ROWS,

348 NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,

349 BLOCK_N=BLOCK_N,

350 return_valid_counts=return_valid_counts,

@@ -390,7 +392,7 @@ def triton_filter_and_convert_dcp_index(

392 None,

393 max_num_blocks_per_req,

394 BLOCK_SIZE,

393 - BLOCK_SIZE, # dense caches on the DCP path

395 + BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,

396 block_n,

397 False, # HAS_PREFILL

398 count_valid,