Fix Kimi EAGLE3 draft config under DCP

vllm/config/speculative.py

@@ -10,6 +10,7 @@
 from typing_extensions import Self
 
 from vllm.config import LoadConfig
+from vllm.config.cache import CacheDType
 from vllm.config.kernel import MoEBackend
 from vllm.config.model import ModelConfig
 from vllm.config.parallel import ParallelConfig
@@ -203,10 +204,14 @@ class SpeculativeConfig:
     inherits the target model's `--moe-backend` setting. Useful when the
     drafter and generator require different MoE kernels (e.g. quantized
     generator with unquantized drafter)."""
+    draft_kv_cache_dtype: CacheDType | None = None
+    """KV cache dtype to use for the draft model. When `None`, the draft model
+    inherits the target model's `--kv-cache-dtype` setting."""
+    draft_attention_backend: AttentionBackendEnum | Literal["auto"] | None = None
+    """Attention backend to use for the draft model. When `None`, the draft
+    model attention backend is independently auto-selected."""
     attention_backend: AttentionBackendEnum | None = None
-    """Attention backend to use for the draft model. When `None`, the backend is
-    automatically selected. Useful when the drafter requires a different attention
-    backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN)."""
+    """Alias for the draft attention backend used by upstream configs."""
     max_model_len: int | None = Field(default=None, ge=1)
     """The maximum model length of the draft model. Used when testing the
     ability to skip speculation for some sequences."""
@@ -1155,6 +1160,15 @@ def _maybe_apply_virtual_tp_to_draft(self) -> None:
             self.draft_parallel_config,
         )
 
+    @field_validator("draft_attention_backend", mode="before")
+    @classmethod
+    def _parse_draft_attention_backend(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            if value.lower() == "auto":
+                return "auto"
+            return AttentionBackendEnum[value.upper()]
+        return value
+
     @field_validator("attention_backend", mode="before")
     @classmethod
     def _parse_attention_backend(cls, value: Any) -> Any:
@@ -1185,6 +1199,9 @@ def _verify_args(self) -> Self:
                 f"than zero ({self.num_speculative_tokens})."
             )
 
+        if self.draft_attention_backend is None and self.attention_backend is not None:
+            self.draft_attention_backend = self.attention_backend
+
         if self.rejection_sample_method == "synthetic":
             # Consolidate to per-position rates
             self.synthetic_acceptance_rates = self._resolve_synthetic_acceptance_rates(

vllm/v1/spec_decode/llm_base_proposer.py

@@ -1314,14 +1314,28 @@ def _create_draft_vllm_config(self) -> VllmConfig:
                 ),
             )
 
+        if spec_cfg.draft_kv_cache_dtype is not None:
+            config = replace(
+                config,
+                cache_config=replace(
+                    config.cache_config,
+                    cache_dtype=spec_cfg.draft_kv_cache_dtype,
+                ),
+            )
+
         # Note (matt): Never inherit the attention backend from base, because there are
         # many opportunities for incompatibility, so we always independently autoselect
         # unless explicitly specified in the speculative config.
+        draft_backend = (
+            None
+            if spec_cfg.draft_attention_backend == "auto"
+            else spec_cfg.draft_attention_backend
+        )
         config = replace(
             config,
             attention_config=replace(
                 config.attention_config,
-                backend=spec_cfg.attention_backend,
+                backend=draft_backend,
             ),
         )
 

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -16,6 +16,7 @@
     init_attn_backend,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     AttentionStatePair,
     BatchExecutionDescriptor,
@@ -194,6 +195,20 @@ def _build_draft_attn_metadata(
             x[:num_reqs_padded] for x in self.block_tables.input_block_tables
         ]
         slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
+        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
+        dcp_local_seq_lens = None
+        if self.block_tables.cp_size > 1:
+            prepare_dcp_local_seq_lens(
+                self.input_buffers.dcp_local_seq_lens,
+                self.input_buffers.seq_lens,
+                num_reqs,
+                self.block_tables.cp_size,
+                self.block_tables.cp_rank,
+                self.block_tables.cp_interleave,
+            )
+            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[
+                :num_reqs_padded
+            ]
         with record_function_or_nullcontext(
             "vllm:v2/speculator/build_attn_metadata"
         ):
@@ -206,12 +221,13 @@ def _build_draft_attn_metadata(
                 ],
                 query_start_loc_cpu=query_start_loc_cpu,
                 max_query_len=num_query_per_req,
-                seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
+                seq_lens=seq_lens,
                 max_seq_len=self.draft_max_seq_len,
                 block_tables=block_tables,
                 slot_mappings=slot_mappings,
                 kv_cache_config=self.kv_cache_config,
                 causal=causal,
+                dcp_local_seq_lens=dcp_local_seq_lens,
             )
         return attn_metadata