From eb0160e8f6b5ec36d46dd87dbb1d1afe96d749cc Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Thu, 6 Aug 2026 20:13:21 -0400 Subject: [PATCH] fix(ds4): declare compressed MLA decode row capacity Co-authored-by: OpenAI Codex --- tests/kernels/test_compressor_kv_cache.py | 72 ++++++ .../test_b12x_compressed_mla_workspace.py | 230 +++++++++++++++++- vllm/models/deepseek_v4/nvidia/b12x.py | 48 ++++ .../backends/mla/compressor_utils.py | 3 + 4 files changed, 347 insertions(+), 6 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index b8f1cb8bfa73..636074553714 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -21,6 +21,7 @@ from vllm.models.deepseek_v4.common.ops import ( dequantize_and_gather_k_cache, quantize_and_insert_k_cache, + save_partial_states, ) from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, @@ -61,6 +62,77 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): # ── Test A: DeepseekV4 Attention path ────────────────────────────────────────────── +@pytest.mark.parametrize("draft_tokens", [5, 7]) +def test_save_partial_states_writes_every_verifier_row(draft_tokens: int) -> None: + """The compressor state write covers all target rows for each request.""" + device = "cuda" + verifier_width = draft_tokens + 1 + num_reqs = 2 + num_actual = num_reqs * verifier_width + head_size = 16 + block_size = 16 + compress_ratio = 4 + + kv = ( + torch.arange( + (num_actual + 1) * head_size, + dtype=torch.float32, + device=device, + ) + .view(num_actual + 1, head_size) + .to(torch.bfloat16) + ) + score = (kv + 1000).to(torch.bfloat16) + ape = ( + torch.arange( + compress_ratio * head_size, + dtype=torch.float32, + device=device, + ) + .view(compress_ratio, head_size) + .to(torch.bfloat16) + ) + positions = torch.arange(num_actual + 1, dtype=torch.int64, device=device) + slot_mapping = torch.cat( + ( + torch.arange(verifier_width, dtype=torch.int64, device=device), + torch.arange( + block_size, + block_size + verifier_width, + dtype=torch.int64, + device=device, + ), + torch.tensor([-1], dtype=torch.int64, device=device), + ) + ) + state_cache = torch.full( + (2, block_size, 2 * head_size), + fill_value=-1, + dtype=torch.bfloat16, + device=device, + ) + + save_partial_states( + kv=kv, + score=score, + ape=ape, + positions=positions, + state_cache=state_cache, + slot_mapping=slot_mapping, + block_size=block_size, + state_width=head_size, + compress_ratio=compress_ratio, + ) + torch.cuda.synchronize() + + actual_slots = slot_mapping[:num_actual] + actual = state_cache.view(-1, 2 * head_size)[actual_slots] + torch.testing.assert_close(actual[:, :head_size], kv[:num_actual]) + expected_score = score[:num_actual] + ape[positions[:num_actual] % compress_ratio] + torch.testing.assert_close(actual[:, head_size:], expected_score) + assert torch.all(state_cache.view(-1, 2 * head_size)[verifier_width] == -1) + + @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size: int): diff --git a/tests/models/deepseek_v4/test_b12x_compressed_mla_workspace.py b/tests/models/deepseek_v4/test_b12x_compressed_mla_workspace.py index bb228c9396fa..7cd83e5dde64 100644 --- a/tests/models/deepseek_v4/test_b12x_compressed_mla_workspace.py +++ b/tests/models/deepseek_v4/test_b12x_compressed_mla_workspace.py @@ -9,8 +9,10 @@ import pytest import torch +from vllm.models.deepseek_v4 import compressor as compressor_mod from vllm.models.deepseek_v4.nvidia import b12x as b12x_mod from vllm.utils.math_utils import round_up +from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla.compressor_utils import ( get_c128a_topk_width, get_compressed_mla_max_q_chunks, @@ -52,9 +54,16 @@ def plan(caps_value): shapes_and_dtypes=lambda: (((1,), torch.uint8),), ) - def split_chunks_for_contract(*, rows: int, width: int, max_chunks: int) -> int: + def split_chunks_for_contract( + *, + rows: int, + width: int, + max_chunks: int, + decode_row_capacity: int | None = None, + ) -> int: del width - return min(max_chunks, 8 if rows <= 256 else 1) + decode_split_max_rows = max(256, int(decode_row_capacity or 0)) + return min(max_chunks, 8 if rows <= decode_split_max_rows else 1) compressed_mla.Caps = caps # type: ignore[attr-defined] compressed_mla.plan = plan # type: ignore[attr-defined] @@ -76,6 +85,8 @@ def _make_layer( compress_ratio: int, dspark: bool, dcp_world_size: int = 1, + max_num_batched_tokens: int = _MAX_ROWS, + max_num_seqs: int = 64, ): layer = object.__new__(b12x_mod.DeepseekV4B12xMLAAttention) torch.nn.Module.__init__(layer) @@ -88,7 +99,7 @@ def _make_layer( layer.indexer = None layer.max_model_len = 524288 layer.window_size = _WINDOW_SIZE - layer.max_num_batched_tokens = _MAX_ROWS + layer.max_num_batched_tokens = max_num_batched_tokens layer.swa_cache_layer = SimpleNamespace(block_size=_PAGE_SIZE) speculative_config = ( SimpleNamespace( @@ -100,6 +111,10 @@ def _make_layer( ) layer.vllm_config = SimpleNamespace( speculative_config=speculative_config, + scheduler_config=SimpleNamespace( + max_num_seqs=max_num_seqs, + max_num_batched_tokens=max_num_batched_tokens, + ), parallel_config=SimpleNamespace( decode_context_parallel_size=dcp_world_size, ), @@ -136,13 +151,16 @@ def test_reserve_uses_full_runtime_width( caps = caps_calls[-1] split_cap = get_compressed_mla_split_cap(expected_width) + decode_row_capacity = 64 * (1 + 5) if dspark else None assert caps.max_width == expected_width assert caps.max_q_rows == _MAX_ROWS + assert caps.decode_row_capacity == decode_row_capacity assert caps.max_q_chunks == get_compressed_mla_max_q_chunks( _MAX_ROWS, expected_width, split_cap, split_chunks, + decode_row_capacity=decode_row_capacity, ) assert workspace.specs is not None @@ -170,6 +188,105 @@ def test_workspace_width_helpers_match_reporter_geometry() -> None: assert get_dspark_swa_index_width(512, 5) == 1024 +@pytest.mark.parametrize( + ("max_num_seqs", "num_speculative_tokens", "expected"), + [ + pytest.param(24, 5, 144, id="mns24-k5"), + pytest.param(64, 5, 384, id="mns64-k5"), + pytest.param(128, 5, 768, id="mns128-k5"), + pytest.param(64, 7, 512, id="mns64-k7"), + ], +) +def test_dspark_decode_row_capacity_comes_from_scheduler_contract( + max_num_seqs: int, + num_speculative_tokens: int, + expected: int, +) -> None: + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + use_dspark=lambda: True, + num_speculative_tokens=num_speculative_tokens, + ), + scheduler_config=SimpleNamespace( + max_num_seqs=max_num_seqs, + max_num_batched_tokens=8192, + ), + ) + assert b12x_mod._get_dspark_decode_row_capacity(vllm_config) == expected + + +def test_non_dspark_has_no_declared_decode_row_capacity() -> None: + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace(use_dspark=lambda: False), + scheduler_config=SimpleNamespace( + max_num_seqs=128, + max_num_batched_tokens=8192, + ), + ) + assert b12x_mod._get_dspark_decode_row_capacity(vllm_config) is None + + +def test_dspark_decode_row_capacity_fails_closed() -> None: + b12x_mod._validate_compressed_mla_decode_row_capacity( + rows=384, + mode="decode", + decode_row_capacity=384, + ) + b12x_mod._validate_compressed_mla_decode_row_capacity( + rows=8192, + mode="extend", + decode_row_capacity=384, + ) + with pytest.raises(ValueError, match="rows 385 exceed.*capacity 384"): + b12x_mod._validate_compressed_mla_decode_row_capacity( + rows=385, + mode="decode", + decode_row_capacity=384, + ) + + +@pytest.mark.parametrize("draft_tokens", [5, 7]) +def test_compressor_metadata_keeps_every_verifier_row( + monkeypatch, + draft_tokens: int, +) -> None: + verifier_width = draft_tokens + 1 + num_reqs = 2 + num_tokens = num_reqs * verifier_width + builder = object.__new__(compressor_mod.CompressorMetadataBuilder) + builder.block_size = 16 + builder.token_to_req_indices = torch.empty(num_tokens, dtype=torch.int32) + common = CommonAttentionMetadata( + query_start_loc=torch.tensor( + [0, verifier_width, num_tokens], dtype=torch.int32 + ), + query_start_loc_cpu=torch.tensor( + [0, verifier_width, num_tokens], dtype=torch.int32 + ), + seq_lens=torch.tensor([100, 200], dtype=torch.int32), + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=verifier_width, + max_seq_len=200, + block_table_tensor=torch.zeros((num_reqs, 1), dtype=torch.int32), + slot_mapping=torch.arange(num_tokens, dtype=torch.int64), + causal=False, + ) + monkeypatch.setattr(compressor_mod, "_prefer_two_stage_compressor", lambda: False) + + metadata = builder.build(0, common) + + assert metadata.slot_mapping.shape == (num_tokens,) + assert metadata.token_to_req_indices is not None + torch.testing.assert_close( + metadata.token_to_req_indices, + torch.repeat_interleave( + torch.arange(num_reqs, dtype=torch.int32), + verifier_width, + ), + ) + + def test_q_chunk_envelope_is_cached() -> None: call_count = 0 @@ -193,6 +310,57 @@ def split_chunks_for_contract(**kwargs) -> int: assert call_count == _MAX_ROWS +def test_production_q_chunk_envelope_only_reserves_declared_capacity() -> None: + def split_chunks_for_contract( + *, + rows: int, + width: int, + max_chunks: int, + decode_row_capacity: int | None = None, + ) -> int: + decode_split_max_rows = max(256, int(decode_row_capacity or 0)) + chunk_size = 64 if rows <= decode_split_max_rows else 1024 + return min(max_chunks, math.ceil(width / chunk_size)) + + max_rows = 8192 + width = 4608 + split_cap = get_compressed_mla_split_cap(width) + get_compressed_mla_max_q_chunks.cache_clear() + + legacy = get_compressed_mla_max_q_chunks( + max_rows, + width, + split_cap, + split_chunks_for_contract, + ) + mns24 = get_compressed_mla_max_q_chunks( + max_rows, + width, + split_cap, + split_chunks_for_contract, + decode_row_capacity=144, + ) + mns64 = get_compressed_mla_max_q_chunks( + max_rows, + width, + split_cap, + split_chunks_for_contract, + decode_row_capacity=384, + ) + mns128 = get_compressed_mla_max_q_chunks( + max_rows, + width, + split_cap, + split_chunks_for_contract, + decode_row_capacity=768, + ) + + assert legacy == 40960 + assert mns24 == legacy + assert mns64 == legacy + assert mns128 == 55296 + + def _aligned_nbytes( specs: tuple[tuple[tuple[int, ...], torch.dtype], ...], ) -> int: @@ -233,14 +401,14 @@ def _runtime_plan_nbytes(compressed_mla, *, rows: int, width: int, heads: int) - ), [ pytest.param(1, False, 1, (128,), 64.502929688, id="swa-causal"), - pytest.param(1, True, 1, (128, 512), 64.502929688, id="swa-dspark"), + pytest.param(1, True, 1, (128, 512), 96.627929688, id="swa-dspark"), pytest.param(4, False, 1, (2176,), 273.315429688, id="c4-causal"), pytest.param( 4, True, 1, (2176, 2560), - 321.502929688, + 482.127929688, id="c4-dspark", ), pytest.param(128, False, 1, (4224,), 530.315429688, id="c128-causal"), @@ -249,7 +417,7 @@ def _runtime_plan_nbytes(compressed_mla, *, rows: int, width: int, heads: int) - True, 1, (4224, 4608), - 578.502929688, + 867.627929688, id="c128-dspark", ), pytest.param(128, False, 2, (4224,), 1060.627929688, id="c128-dcp2"), @@ -291,3 +459,53 @@ def test_real_sparkinfer_reserve_dominates_runtime_envelope( for rows in range(1, _MAX_ROWS + 1) ) assert reserve_bytes >= runtime_bytes + + +def test_mns64_contract_does_not_grow_production_mnb8192_scratch( + monkeypatch, +) -> None: + compressed_mla = pytest.importorskip("sparkinfer.attention.compressed_mla") + workspace = _RecordingWorkspaceManager() + monkeypatch.setattr(b12x_mod, "current_workspace_manager", lambda: workspace) + max_rows = 8192 + layer = _make_layer( + compress_ratio=128, + dspark=True, + max_num_batched_tokens=max_rows, + max_num_seqs=64, + ) + + layer._reserve_dummy_compressed_mla_scratch( + torch.empty((1, _LOCAL_HEADS, 512), dtype=torch.bfloat16) + ) + + assert workspace.specs is not None + candidate_bytes = _aligned_nbytes(workspace.specs) + width = 4608 + split_cap = get_compressed_mla_split_cap(width) + max_q_chunks = get_compressed_mla_max_q_chunks( + max_rows, + width, + split_cap, + compressed_mla.split_chunks_for_contract, + ) + splits = compressed_mla.split_chunks_for_contract( + rows=max_rows, + width=width, + max_chunks=split_cap, + ) + legacy_plan = compressed_mla.plan( + compressed_mla.Caps( + device="cpu", + num_q_heads=_LOCAL_HEADS, + max_q_rows=max_rows, + max_width=width, + head_dim=512, + v_head_dim=512, + page_size=_PAGE_SIZE, + max_chunks_per_row=splits, + max_q_chunks=max_q_chunks, + ) + ) + + assert candidate_bytes == _aligned_nbytes(legacy_plan.shapes_and_dtypes()) diff --git a/vllm/models/deepseek_v4/nvidia/b12x.py b/vllm/models/deepseek_v4/nvidia/b12x.py index 4d244c087ae7..1afba68d0f18 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x.py +++ b/vllm/models/deepseek_v4/nvidia/b12x.py @@ -20,6 +20,7 @@ import torch +from vllm.config import VllmConfig from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.common.ops import ( compute_dcp_global_topk_indices_and_lens, @@ -53,6 +54,36 @@ _VALIDATE_DCP_INDICES_ENV = "VLLM_DSV4_DCP_VALIDATE_INDICES" +def _get_dspark_decode_row_capacity(vllm_config: VllmConfig) -> int | None: + """Return the largest target-verifier row count the scheduler can emit.""" + speculative_config = vllm_config.speculative_config + if speculative_config is None or not speculative_config.use_dspark(): + return None + num_speculative_tokens = int(speculative_config.num_speculative_tokens or 0) + if num_speculative_tokens <= 0: + return None + scheduler_config = vllm_config.scheduler_config + return min( + int(scheduler_config.max_num_batched_tokens), + int(scheduler_config.max_num_seqs) * (1 + num_speculative_tokens), + ) + + +def _validate_compressed_mla_decode_row_capacity( + *, + rows: int, + mode: Literal["decode", "extend"], + decode_row_capacity: int | None, +) -> None: + if mode != "decode" or decode_row_capacity is None: + return + if int(rows) > int(decode_row_capacity): + raise ValueError( + f"compressed MLA decode rows {rows} exceed the declared " + f"capacity {decode_row_capacity}" + ) + + def _cdiv(x: int, y: int) -> int: return (int(x) + int(y) - 1) // int(y) @@ -286,6 +317,7 @@ def _run_compressed_mla( indexed_lens: torch.Tensor | None, indexed_page_size: int | None, mode: Literal["decode", "extend"] = "decode", + decode_row_capacity: int | None = None, return_lse: bool = False, lse_scale: Literal["natural", "base2"] = "natural", ) -> torch.Tensor | None: @@ -309,6 +341,11 @@ def _run_compressed_mla( ) rows, heads = int(q.shape[0]), int(q.shape[1]) + _validate_compressed_mla_decode_row_capacity( + rows=rows, + mode=mode, + decode_row_capacity=decode_row_capacity, + ) q = q.contiguous() swa_indices = swa_indices.contiguous() swa_lens = swa_lens.contiguous() @@ -341,6 +378,7 @@ def _run_compressed_mla( rows=max(1, rows), width=width, max_chunks=decode_split_cap, + decode_row_capacity=decode_row_capacity, ) plan = plan_compressed_mla_scratch( @@ -353,6 +391,7 @@ def _run_compressed_mla( v_head_dim=_DSV4_V_HEAD_DIM, page_size=int(swa_page_size), max_chunks_per_row=num_splits_cap, + decode_row_capacity=decode_row_capacity, ) ) scratch = current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) @@ -411,6 +450,7 @@ def _run_dcp_compressed_mla( indexed_lens: torch.Tensor | None, indexed_page_size: int | None, mode: Literal["decode", "extend"] = "decode", + decode_row_capacity: int | None = None, ) -> None: from vllm.distributed.parallel_state import get_dcp_group @@ -443,6 +483,7 @@ def _run_dcp_compressed_mla( indexed_lens=indexed_lens, indexed_page_size=indexed_page_size, mode=mode, + decode_row_capacity=decode_row_capacity, return_lse=True, lse_scale="natural", ) @@ -547,6 +588,7 @@ def _reserve_dummy_compressed_mla_scratch(self, q: torch.Tensor) -> None: swa_width = int(self.window_size) speculative_config = self.vllm_config.speculative_config + decode_row_capacity = _get_dspark_decode_row_capacity(self.vllm_config) if speculative_config is not None and speculative_config.use_dspark(): swa_width = max( swa_width, @@ -563,12 +605,14 @@ def _reserve_dummy_compressed_mla_scratch(self, q: torch.Tensor) -> None: rows=rows, width=width, max_chunks=decode_split_cap, + decode_row_capacity=decode_row_capacity, ) max_q_chunks = get_compressed_mla_max_q_chunks( rows, width, decode_split_cap, compressed_mla_split_chunks_for_contract, + decode_row_capacity=decode_row_capacity, ) dcp_world_size = max( int(self.vllm_config.parallel_config.decode_context_parallel_size), @@ -587,6 +631,7 @@ def _reserve_dummy_compressed_mla_scratch(self, q: torch.Tensor) -> None: page_size=int(self.swa_cache_layer.block_size), max_chunks_per_row=num_splits_cap, max_q_chunks=max_q_chunks, + decode_row_capacity=decode_row_capacity, ) ) current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) @@ -693,6 +738,7 @@ def _forward_b12x_decode( num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + decode_row_capacity = _get_dspark_decode_row_capacity(vllm_config) dcp_rank = 0 if dcp_world_size > 1: from vllm.distributed.parallel_state import get_dcp_group @@ -770,6 +816,7 @@ def _forward_b12x_decode( indexed_lens=topk_lens, indexed_page_size=indexed_page_size, mode="decode", + decode_row_capacity=decode_row_capacity, ) else: _run_compressed_mla( @@ -786,6 +833,7 @@ def _forward_b12x_decode( indexed_lens=topk_lens, indexed_page_size=indexed_page_size, mode="decode", + decode_row_capacity=decode_row_capacity, ) def _forward_b12x_prefill( diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 7f140a8dfe23..cfa5f62d0748 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -62,6 +62,7 @@ def get_compressed_mla_max_q_chunks( width: int, max_chunks: int, split_chunks_for_contract: Callable[..., int], + decode_row_capacity: int | None = None, ) -> int: """Return the q-chunk cap over every reachable row count. @@ -70,6 +71,7 @@ def get_compressed_mla_max_q_chunks( width: Combined SWA and indexed width. max_chunks: Maximum chunks allowed for each row. split_chunks_for_contract: Kernel contract split-count function. + decode_row_capacity: Integration-declared decode/verifier row capacity. Returns: The largest reachable product of rows and chunks per row. @@ -83,6 +85,7 @@ def get_compressed_mla_max_q_chunks( rows=rows, width=width, max_chunks=max_chunks, + decode_row_capacity=decode_row_capacity, ) for rows in range(1, max_rows + 1) )