diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 48f3ffd5ac..7f7c20d53e 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/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") diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 3131517248..8f3afc28d0 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 512627559a..3798000d7d 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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, ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 8ffc31ca3b..00679d5e26 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/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,