diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index b9c612621982..47536695b190 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 HfOverrides, ModelConfig from vllm.config.parallel import ParallelConfig @@ -118,6 +119,9 @@ class SpeculativeConfig: """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).""" + kv_cache_dtype: CacheDType | None = None + """KV cache dtype for the draft model. When `None`, the draft inherits the + target model's `--kv-cache-dtype`.""" 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.""" diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 742d62ac3698..07877b0812c7 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1754,6 +1754,10 @@ def create_speculative_config( if self.speculative_config is None: return None + self.speculative_config = { + k.replace("-", "_"): v for k, v in self.speculative_config.items() + } + # Note(Shangming): These parameters are not obtained from the cli arg # '--speculative-config' and must be passed in when creating the engine # config. diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 756c5f3b3717..f8b52d079a89 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1300,6 +1300,15 @@ def _create_draft_vllm_config(self) -> VllmConfig: ), ) + if spec_cfg.kv_cache_dtype is not None: + base = replace( + base, + cache_config=replace( + base.cache_config, + cache_dtype=spec_cfg.kv_cache_dtype, + ), + ) + return base def _get_model(self) -> nn.Module: diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 998a28b6aee0..478274a05671 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -27,6 +27,14 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dflash_head"): dflash_model = get_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index bbee614a1aeb..9ea2ff0f7363 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -27,6 +27,14 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dspark_head"): diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index bdd588e5786d..579652bc7d64 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA 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 speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + if speculative_config.kv_cache_dtype is not None: + vllm_config = replace( + vllm_config, + cache_config=replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ), + ) with set_model_tag("eagle_head"): eagle_model = get_model( vllm_config=vllm_config, model_config=draft_model_config