[Spec decode] Support Kimi-K3 DCP with DSpark

tests/transformers_utils/test_dspark_mla_config.py

@@ -140,26 +140,3 @@ def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):
     assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]
     assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"
     assert speculative_config.draft_model_config.use_mla
-
-
-def test_dspark_mla_rejects_decode_context_parallelism(tmp_path):
-    target_path = tmp_path / "target"
-    draft_path = tmp_path / "draft"
-    _write_target_config(target_path)
-    _write_dspark_config(draft_path)
-    target_config = ModelConfig(
-        model=str(target_path), tokenizer_mode="skip", max_model_len=32768
-    )
-
-    with pytest.raises(ValueError, match="does not currently support decode context"):
-        SpeculativeConfig(
-            model=str(draft_path),
-            method="dspark",
-            num_speculative_tokens=8,
-            target_model_config=target_config,
-            target_parallel_config=ParallelConfig(
-                tensor_parallel_size=2,
-                decode_context_parallel_size=2,
-                distributed_executor_backend="external_launcher",
-            ),
-        )

tests/v1/attention/test_flashinfer_mla_dcp.py

@@ -45,6 +45,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
     impl.bmm1_scale = 1.0
     impl.bmm2_scale = 1.0
     impl.need_to_return_lse_for_decode = True
+    impl.dcp_world_size = 2
     impl.num_heads = 6
     impl.qk_nope_head_dim = 128
     impl.kv_lora_rank = 512

tests/v1/attention/test_mla_backends.py

@@ -837,9 +837,86 @@ class _AttnMeta:
 def test_tokenspeed_mla_noncausal_capability():
     builder = tokenspeed_mla_module.TokenspeedMLAMetadataBuilder
     assert builder.supports_non_causal_multi_token_decode
+    assert builder.supports_non_causal_multi_token_dcp
     assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()
 
 
+def test_flashinfer_mla_dcp_multi_token_decode_uses_per_query_bounds(monkeypatch):
+    flashinfer_mla_module = pytest.importorskip(
+        "vllm.v1.attention.backends.mla.flashinfer_mla"
+    )
+
+    decode_call = None
+
+    def fake_decode(**kwargs):
+        nonlocal decode_call
+        decode_call = kwargs
+        query = kwargs["query"]
+        output = torch.empty(*query.shape[:-1], 512, dtype=torch.bfloat16)
+        lse = torch.empty(query.shape[0], query.shape[-2], dtype=torch.float32)
+        return output, lse
+
+    monkeypatch.setattr(
+        flashinfer_mla_module,
+        "trtllm_batch_decode_with_kv_cache_mla",
+        fake_decode,
+    )
+    monkeypatch.setattr(
+        flashinfer_mla_module,
+        "_get_workspace_buffer",
+        lambda return_lse: torch.empty(1, dtype=torch.int8),
+    )
+
+    impl = object.__new__(flashinfer_mla_module.FlashInferMLAImpl)
+    impl.dcp_world_size = 2
+    impl.dcp_rank = 1
+    impl.cp_kv_cache_interleave_size = 1
+    impl.need_to_return_lse_for_decode = True
+    impl.kv_lora_rank = 512
+    impl.qk_nope_head_dim = 128
+    impl.qk_rope_head_dim = 64
+    impl.bmm1_scale = 1.0
+    impl.bmm2_scale = 1.0
+
+    block_table = torch.tensor([[1], [2]], dtype=torch.int32)
+    metadata = SimpleNamespace(
+        num_decodes=2,
+        num_decode_tokens=6,
+        max_seq_len=7,
+        causal=True,
+        decode=SimpleNamespace(
+            block_table=block_table,
+            seq_lens=torch.tensor([5, 6], dtype=torch.int32),
+            dcp_tot_seq_lens=torch.tensor([10, 13], dtype=torch.int32),
+            flattened_block_table=None,
+            flattened_seq_lens=None,
+            query_len=0,
+        ),
+    )
+    query = torch.empty(6, 2, 576, dtype=torch.bfloat16)
+    kv_cache = torch.empty(3, 16, 576, dtype=torch.bfloat16)
+
+    output, lse = impl.forward_mqa(
+        query,
+        kv_cache,
+        metadata,
+        SimpleNamespace(),
+    )
+
+    assert output.shape == (6, 2, 512)
+    assert lse is not None
+    assert lse.shape == (6, 2)
+    assert decode_call is not None
+    assert decode_call["query"].shape == (6, 1, 2, 576)
+    torch.testing.assert_close(
+        decode_call["seq_lens"],
+        torch.tensor([4, 4, 5, 5, 6, 6], dtype=torch.int32),
+    )
+    torch.testing.assert_close(
+        decode_call["block_tables"], block_table.repeat_interleave(3, dim=0)
+    )
+
+
 @pytest.mark.parametrize(
     ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"),
     [

tests/v1/spec_decode/test_dflash_prepare_inputs.py

@@ -17,7 +17,14 @@
 )
 
 
-def _run_prepare(*, target_positions: list[int], block_table_values: list[int]):
+def _run_prepare(
+    *,
+    target_positions: list[int],
+    block_table_values: list[int],
+    cp_rank: int = 0,
+    cp_size: int = 1,
+    cp_interleave: int = 1,
+):
     device = torch.device("cuda")
     max_num_reqs = 4
     max_num_tokens = 16
@@ -86,6 +93,9 @@ def _run_prepare(*, target_positions: list[int], block_table_values: list[int]):
         input_seeds,
         block_table,
         4,
+        cp_rank,
+        cp_size,
+        cp_interleave,
         123,
         num_speculative_steps,
         num_speculative_steps,
@@ -131,6 +141,20 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
     assert out.seeds[2].item() == 17
 
 
+def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
+    out = _run_prepare(
+        target_positions=[10, 11, 12, 13],
+        block_table_values=[0, 7, 8, 9],
+        cp_rank=1,
+        cp_size=2,
+        cp_interleave=2,
+    )
+
+    assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
+    assert out.context_slot_mapping[:4].tolist() == [28, 29, PAD_SLOT_ID, PAD_SLOT_ID]
+    assert out.query_slot_mapping[:3].tolist() == [PAD_SLOT_ID, PAD_SLOT_ID, 30]
+
+
 def test_prepare_dflash_inputs_never_writes_the_null_block():
     # The valid context uses logical block 0 and the replacement query uses
     # logical block 1. Both map to the null block and must remain unwritable.

vllm/config/speculative.py

@@ -1046,16 +1046,6 @@ def __post_init__(self):
                 if self.method in ("dflash", "dspark"):
                     self.parallel_drafting = True
 
-                if (
-                    self.method == "dspark"
-                    and "K3DSparkModel" in self.draft_model_config.architectures
-                    and self.target_parallel_config.decode_context_parallel_size > 1
-                ):
-                    raise ValueError(
-                        "MLA DSpark does not currently support decode context "
-                        "parallelism; set decode_context_parallel_size=1."
-                    )
-
                 if self.num_speculative_tokens is not None and hasattr(
                     self.draft_model_config.hf_config, "num_lookahead_tokens"
                 ):

vllm/model_executor/layers/attention/mla_attention.py

@@ -1984,13 +1984,41 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
     # Whether this builder can flatten a non-causal query block into decode rows.
     supports_non_causal_multi_token_decode: ClassVar[bool] = False
 
+    # Whether can support non-causal multi-token decode with DCP KV cache.
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = False
+
     # The threshold for reordering the batch into decode and prefill requests.
     # If > 1, the batch will be reordered such that requests with
     # query length <= threshold are classified as decode requests.
     # Use `query_len_support` (above) to set this automatically
     # when speculative decoding is enabled.
     reorder_batch_threshold: int = 1
 
+    def _validate_dspark_dcp_support(self, supports_dcp_with_varlen: bool) -> None:
+        speculative_config = getattr(self.vllm_config, "speculative_config", None)
+        parallel_config = self.vllm_config.parallel_config
+        if (
+            speculative_config is None
+            or getattr(speculative_config, "method", None) != "dspark"
+            or parallel_config.decode_context_parallel_size <= 1
+        ):
+            return
+
+        if self.non_causal_multi_token_decode:
+            supported = self.supports_non_causal_multi_token_dcp
+            query_mode = "non-causal draft"
+        else:
+            supported = supports_dcp_with_varlen
+            query_mode = "causal multi-token"
+
+        if not supported:
+            raise ValueError(
+                f"{type(self).__name__} does not support {query_mode} MLA "
+                "attention for DSpark with decode context parallelism. Select "
+                "a backend with explicit DSpark DCP support or set "
+                "decode_context_parallel_size=1."
+            )
+
     @staticmethod
     def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:
         scheduler_config = vllm_config.scheduler_config
@@ -2084,6 +2112,7 @@ def __init__(
         self.non_causal_multi_token_decode = getattr(
             kv_cache_spec, "non_causal_multi_token_decode", False
         )
+        self._validate_dspark_dcp_support(supports_dcp_with_varlen)
 
         # A draft cache group can have a different head count from the target.
         self.num_heads = get_num_attention_heads_from_layers(

vllm/models/kimi_k3/nvidia/mla.py

@@ -341,11 +341,6 @@ def __init__(
             "parallelism."
         )
         self.dcp_world_size = parallel_config.decode_context_parallel_size
-        assert self.dcp_world_size <= 1 or self.rotary_emb is None, (
-            "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode "
-            "context parallelism because gathered queries require gathered "
-            "positions."
-        )
         self.dcp_manager: MLADCPManager | None = None
         if self.dcp_world_size > 1:
             query_dtype = (

vllm/v1/attention/backend.py

@@ -330,6 +330,11 @@ def supports_pcp(cls) -> bool:
         except NotImplementedError:
             return False
 
+    @classmethod
+    def supports_non_causal_dcp(cls) -> bool:
+        builder_cls = cls.get_builder_cls()
+        return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False))
+
     @classmethod
     def supports_attn_type(cls, attn_type: str) -> bool:
         """Check if backend supports a given attention type.
@@ -378,6 +383,7 @@ def validate_configuration(
         use_kv_connector: bool = False,
         use_pcp: bool = False,
         use_adaptive_verification: bool = False,
+        use_dcp: bool = False,
     ) -> list[str]:
         invalid_reasons = []
         if not cls.supports_head_size(head_size):
@@ -414,6 +420,8 @@ def validate_configuration(
             invalid_reasons.append("sliding window not supported")
         if use_non_causal and not cls.supports_non_causal():
             invalid_reasons.append("non-causal attention not supported")
+        if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():
+            invalid_reasons.append("non-causal MLA attention with DCP not supported")
         if use_batch_invariant and not cls.supports_batch_invariance():
             invalid_reasons.append("batch invariance not supported")
         if use_kv_connector and not cls.supports_kv_connector():