diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index aceffbece71f..ccc8cd1b43d7 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -4,17 +4,20 @@ import torch -def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): +def test_deepseek_v4_c128a_metadata_preserves_capacity_stride(): from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata device = torch.device("cuda") capacity_width = 256 active_width = 128 - global_decode_buffer = torch.empty( - (2, capacity_width), dtype=torch.int32, device=device + untouched = -77 + global_decode_buffer = torch.full( + (2, capacity_width), untouched, dtype=torch.int32, device=device ) decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device) - prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device) + prefill_buffer = torch.full( + (2, capacity_width), untouched, dtype=torch.int32, device=device + ) global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( positions=torch.tensor([255, 511], dtype=torch.int64, device=device), @@ -31,15 +34,17 @@ def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): max_compressed_tokens=active_width, ) - assert global_decode.shape == (1, active_width) - assert prefill_local.shape == (1, active_width) - assert global_decode.stride() == (active_width, 1) - assert prefill_local.stride() == (active_width, 1) + assert global_decode.shape == (1, capacity_width) + assert prefill_local.shape == (1, capacity_width) + assert global_decode.stride() == (capacity_width, 1) + assert prefill_local.stride() == (capacity_width, 1) assert global_decode[0, :2].cpu().tolist() == [768, 769] assert decode_lens.cpu().tolist() == [2] assert prefill_local[0, :4].cpu().tolist() == list(range(4)) - assert torch.all(global_decode[0, 2:] == -1) - assert torch.all(prefill_local[0, 4:] == -1) + assert torch.all(global_decode[0, 2:active_width] == -1) + assert torch.all(prefill_local[0, 4:active_width] == -1) + assert torch.all(global_decode[0, active_width:] == untouched) + assert torch.all(prefill_local[0, active_width:] == untouched) def test_sparse_flashmla_metadata_smoke(): diff --git a/tests/v1/attention/test_deepseek_v4_sparse_mla_metadata.py b/tests/v1/attention/test_deepseek_v4_sparse_mla_metadata.py index 5b7b012c70b0..6589fd0c5a64 100644 --- a/tests/v1/attention/test_deepseek_v4_sparse_mla_metadata.py +++ b/tests/v1/attention/test_deepseek_v4_sparse_mla_metadata.py @@ -14,6 +14,63 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec +def test_c128a_builder_uses_capture_time_row_capacity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import vllm.models.deepseek_v4.sparse_mla as sparse_mla + + capacity_width = 256 + builder = object.__new__(DeepseekV4FlashMLAMetadataBuilder) + builder.deepseek_v4_decode_threshold = 1 + builder.compress_ratio = 128 + builder.c128a_max_compressed = capacity_width + builder.kv_cache_spec = SimpleNamespace(block_size=capacity_width * 128) + builder.c128a_global_decode_buffer = torch.empty(2, capacity_width) + builder.c128a_decode_lens_buffer = torch.empty(2, dtype=torch.int32) + builder.c128a_prefill_buffer = torch.empty(2, capacity_width) + + monkeypatch.setattr( + sparse_mla, + "split_decodes_and_prefills", + lambda *_args, **_kwargs: (2, 0, 2, 0), + ) + observed: dict[str, int] = {} + + def record_width(*_args, max_compressed_tokens: int, **_kwargs): + observed["width"] = max_compressed_tokens + return ( + torch.empty(2, capacity_width, dtype=torch.int32), + torch.empty(2, dtype=torch.int32), + torch.empty(0, capacity_width, dtype=torch.int32), + ) + + monkeypatch.setattr(sparse_mla, "build_c128a_topk_metadata", record_width) + common_metadata = SimpleNamespace( + # This active context would select a 128-wide packed row in the + # incompatible implementation. + max_seq_len=512, + positions=torch.arange(2), + block_table_tensor=torch.empty(2, 1, dtype=torch.int32), + slot_mapping=torch.arange(2), + ) + + metadata = builder._build_c128a_metadata( + common_metadata, + req_id_per_token=torch.zeros(2, dtype=torch.int32), + actual_num_query_tokens=2, + dcp_world_size=1, + dcp_rank=0, + cp_kv_cache_interleave_size=1, + ) + + assert observed["width"] == capacity_width + assert metadata["c128a_global_decode_topk_indices"].shape == ( + 2, + 1, + capacity_width, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize( ("dcp_rank", "valid_offset", "expected_slot"), diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 38043dbba616..0b4cdd22f4a0 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -311,13 +311,6 @@ def _build_c128a_metadata( assert cm.positions is not None, ( "positions is required for C128A metadata build" ) - active_topk_width = min( - max( - triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), - _C128A_TOPK_ALIGNMENT, - ), - self.c128a_max_compressed, - ) block_size = self.kv_cache_spec.block_size // self.compress_ratio global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( cm.positions[:num_total], @@ -331,7 +324,7 @@ def _build_c128a_metadata( self.c128a_global_decode_buffer, self.c128a_decode_lens_buffer, self.c128a_prefill_buffer, - max_compressed_tokens=active_topk_width, + max_compressed_tokens=self.c128a_max_compressed, dcp_world_size=dcp_world_size, dcp_rank=dcp_rank, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, @@ -370,30 +363,26 @@ def build_c128a_topk_metadata( Decode tokens: position → block_table lookup → global slot ids + topk_lens. Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. - Writes into packed views of pre-allocated buffers for CUDA graph stability. + Writes into capacity-strided views of pre-allocated buffers. FULL CUDA + graphs capture consumers with this row stride, so runtime metadata must + preserve it independently of the active sequence length. """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens - # view(-1) as 1-d array and then expanded to - # [num_decode_tokens, max_compressed_tokens] - global_decode = global_decode_buffer.view(-1)[ - : num_decode_tokens * max_compressed_tokens - ].view(num_decode_tokens, max_compressed_tokens) + global_decode = global_decode_buffer[:num_decode_tokens] decode_lens = decode_lens_buffer[:num_decode_tokens] - prefill_local = prefill_buffer.view(-1)[ - : num_prefill_tokens * max_compressed_tokens - ].view(num_prefill_tokens, max_compressed_tokens) + prefill_local = prefill_buffer[:num_prefill_tokens] if num_tokens == 0: return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( global_decode_buffer, - max_compressed_tokens, + global_decode_buffer.stride(0), decode_lens_buffer, prefill_buffer, - max_compressed_tokens, + prefill_buffer.stride(0), positions, compress_ratio, max_compressed_tokens,