From 4c3f01e9a94e5369c4c2535dd19441f1935d697e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 16 Jun 2026 10:04:08 +0000 Subject: [PATCH 1/2] spec-decode: restore draft Kimi config overrides --- vllm/config/speculative.py | 23 ++++++++++++++++++++--- vllm/v1/spec_decode/llm_base_proposer.py | 16 +++++++++++++++- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index abdf8ca9ac08..d565bf612d36 100644 --- a/vllm/config/speculative.py +++ b/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( diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 3cdf31e429a6..66b1ab30249f 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/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, ), ) From 58f130b5aba090dd6deebcff9ba21dd6f0b5d98e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 16 Jun 2026 10:21:27 +0000 Subject: [PATCH 2/2] spec-decode: pass DCP seq lens to draft MLA metadata --- vllm/v1/worker/gpu/spec_decode/speculator.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index bd41e77c9a58..942130929133 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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