[Spec Decode] Add kv_cache_dtype to speculative_config to control separately from target (#48787)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Michael Goin committed
Jul 16, 2026 at 01:49 UTC
8bfd6839016f9670508cf07c16c68891f83ea117
6 files changed
+42
-1
vllm/config/speculative.py
+4
@@ -10,6 +10,7 @@ from pydantic import Field, SkipValidation, field_validator, model_validator
10
from typing_extensions import Self
11
12
from vllm.config import LoadConfig
13
+from vllm.config.cache import CacheDType
14
from vllm.config.kernel import MoEBackend
15
from vllm.config.model import HfOverrides, ModelConfig
16
from vllm.config.parallel import ParallelConfig
@@ -118,6 +119,9 @@ class SpeculativeConfig:
119
"""Attention backend to use for the draft model. When `None`, the backend is
120
automatically selected. Useful when the drafter requires a different attention
121
backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN)."""
122
+ kv_cache_dtype: CacheDType | None = None
123
+ """KV cache dtype for the draft model. When `None`, the draft inherits the
124
+ target model's `--kv-cache-dtype`."""
125
max_model_len: int | None = Field(default=None, ge=1)
126
"""The maximum model length of the draft model. Used when testing the
127
ability to skip speculation for some sequences."""
vllm/engine/arg_utils.py
+4
@@ -1754,6 +1754,10 @@ class EngineArgs:
1754
if self.speculative_config is None:
1755
return None
1756
1757
+ self.speculative_config = {
1758
+ k.replace("-", "_"): v for k, v in self.speculative_config.items()
1759
+ }
1760
+
1761
# Note(Shangming): These parameters are not obtained from the cli arg
1762
# '--speculative-config' and must be passed in when creating the engine
1763
# config.
vllm/v1/spec_decode/llm_base_proposer.py
+9
@@ -1300,6 +1300,15 @@ class SpecDecodeBaseProposer:
1300
),
1301
)
1302
1303
+ if spec_cfg.kv_cache_dtype is not None:
1304
+ base = replace(
1305
+ base,
1306
+ cache_config=replace(
1307
+ base.cache_config,
1308
+ cache_dtype=spec_cfg.kv_cache_dtype,
1309
+ ),
1310
+ )
1311
+
1312
return base
1313
1314
def _get_model(self) -> nn.Module:
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+8
@@ -27,6 +27,14 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
27
use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config),
28
backend=speculative_config.attention_backend,
29
),
30
+ cache_config=(
31
+ replace(
32
+ vllm_config.cache_config,
33
+ cache_dtype=speculative_config.kv_cache_dtype,
34
+ )
35
+ if speculative_config.kv_cache_dtype is not None
36
+ else vllm_config.cache_config
37
+ ),
38
)
39
with set_model_tag("dflash_head"):
40
dflash_model = get_model(
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+8
@@ -27,6 +27,14 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
27
use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config),
28
backend=speculative_config.attention_backend,
29
),
30
+ cache_config=(
31
+ replace(
32
+ vllm_config.cache_config,
33
+ cache_dtype=speculative_config.kv_cache_dtype,
34
+ )
35
+ if speculative_config.kv_cache_dtype is not None
36
+ else vllm_config.cache_config
37
+ ),
38
)
39
40
with set_model_tag("dspark_head"):
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+9
-1
@@ -3,7 +3,7 @@
3
import torch
4
import torch.nn as nn
5
6
-from vllm.config import VllmConfig
6
+from vllm.config import VllmConfig, replace
7
from vllm.distributed.parallel_state import get_pp_group
8
from vllm.lora.layers.base import BaseLayerWithLoRA
9
from vllm.model_executor.model_loader import get_model
@@ -39,6 +39,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
39
speculative_config = vllm_config.speculative_config
40
assert speculative_config is not None
41
draft_model_config = speculative_config.draft_model_config
42
+ if speculative_config.kv_cache_dtype is not None:
43
+ vllm_config = replace(
44
+ vllm_config,
45
+ cache_config=replace(
46
+ vllm_config.cache_config,
47
+ cache_dtype=speculative_config.kv_cache_dtype,
48
+ ),
49
+ )
50
with set_model_tag("eagle_head"):
51
eagle_model = get_model(
52
vllm_config=vllm_config, model_config=draft_model_config