@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 726d89720c

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 07:30 UTC

[CI] Enable mypy import following for `vllm/spec_decode` (#33282)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

Lucas Kabela committed Jan 29, 2026 at 22:43 UTC 726d89720ca7ad4109c30fe7c5ce44456affc365

5 files changed +38 -20

tools/pre_commit/mypy.py

-1

@@ -32,7 +32,6 @@ SEPARATE_GROUPS = [

32 "vllm/model_executor",

33 # v1 related

34 "vllm/v1/kv_offload",

35 - "vllm/v1/spec_decode",

35 ]

37 # TODO(woosuk): Include the code from Megatron and HuggingFace.

vllm/v1/spec_decode/draft_model.py

+7 -6

@@ -5,7 +5,6 @@ from typing import Any

5 import torch

7 from vllm.config import VllmConfig, get_layers_from_vllm_config

8 -from vllm.config.speculative import SpeculativeConfig

8 from vllm.logger import init_logger

9 from vllm.model_executor.layers.attention import Attention

10 from vllm.model_executor.model_loader import get_model

@@ -56,7 +55,7 @@ class DraftModelProposer(SpecDecodeBaseProposer):

55 )

57 def _raise_if_padded_drafter_batch_disabled(self):

59 - if self.vllm_config.speculative_config.disable_padded_drafter_batch:

58 + if self.speculative_config.disable_padded_drafter_batch:

59 raise NotImplementedError(

60 "Speculative Decoding with draft models only supports "

61 "padded drafter batch. Please don't pass --disable-padded-drafter-batch"

@@ -64,7 +63,7 @@ class DraftModelProposer(SpecDecodeBaseProposer):

63 )

65 def _raise_if_vocab_size_mismatch(self):

67 - self.vllm_config.speculative_config.verify_equal_vocab_size_if_draft_model()

66 + self.speculative_config.verify_equal_vocab_size_if_draft_model()

68 def _raise_if_draft_tp_mismatch(self):

69 # Note(Tomas Ruiz) If we run the target model with TP > 1 and

@@ -73,7 +72,7 @@ class DraftModelProposer(SpecDecodeBaseProposer):

72 # (because TP=1), then the torch compile cache is overwritten and corrupted.

73 # We need a mechanism like this: https://github.com/vllm-project/vllm/pull/5414

74 # To prevent this error, we assert that both TP sizes must be the same.

76 - spec_cfg: SpeculativeConfig = self.vllm_config.speculative_config

75 + spec_cfg = self.speculative_config

76 tgt_tp = spec_cfg.target_parallel_config.tensor_parallel_size

77 draft_tp = spec_cfg.draft_parallel_config.tensor_parallel_size

78 if draft_tp != tgt_tp:

@@ -190,12 +189,14 @@ def create_vllm_config_for_draft_model(

189 The vllm_config is useful when loading the draft model with get_model().

190 """

191 old = target_model_vllm_config

193 - new_parallel_config = old.speculative_config.draft_parallel_config.replace(

192 + assert old.speculative_config is not None, "speculative_config is not set"

193 + old_spec_config = old.speculative_config

194 + new_parallel_config = old_spec_config.draft_parallel_config.replace(

195 rank=old.parallel_config.rank

196 )

197 new: VllmConfig = old.replace(

198 quant_config=None, # quant_config is recomputed in __init__()

198 - model_config=old.speculative_config.draft_model_config,

199 + model_config=old_spec_config.draft_model_config,

200 parallel_config=new_parallel_config,

201 )

202 return new

vllm/v1/spec_decode/eagle.py

+24 -9

@@ -3,6 +3,7 @@

3 import ast

4 from dataclasses import replace

5 from importlib.util import find_spec

6 +from typing import cast

8 import numpy as np

9 import torch

@@ -20,6 +21,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

21 from vllm.model_executor.model_loader import get_model

22 from vllm.model_executor.models import supports_multimodal

23 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache

24 +from vllm.model_executor.models.interfaces import SupportsMultiModal

25 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

26 from vllm.multimodal import MULTIMODAL_REGISTRY

27 from vllm.platforms import current_platform

@@ -62,8 +64,8 @@ class SpecDecodeBaseProposer:

64 runner=None,

65 ):

66 self.vllm_config = vllm_config

67 + assert vllm_config.speculative_config is not None

68 self.speculative_config = vllm_config.speculative_config

66 - assert self.speculative_config is not None

69 self.draft_model_config = self.speculative_config.draft_model_config

70 self.method = self.speculative_config.method

71 self.pass_hidden_states_to_model = pass_hidden_states_to_model

@@ -206,6 +208,7 @@ class SpecDecodeBaseProposer:

208

209 # Parse the speculative token tree.

210 spec_token_tree = self.speculative_config.speculative_token_tree

211 + assert spec_token_tree is not None

212 self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)

213 tree_depth = len(self.tree_choices[-1])

214 # Precompute per-level properties of the tree.

@@ -1077,9 +1080,12 @@ class SpecDecodeBaseProposer:

1080 return model.__class__.__name__

1081

1082 def load_model(self, target_model: nn.Module) -> None:

1080 - draft_model_config = self.vllm_config.speculative_config.draft_model_config

1083 + draft_model_config = self.speculative_config.draft_model_config

1084 target_attn_layer_names = set(

1082 - get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()

1085 + get_layers_from_vllm_config(

1086 + self.vllm_config,

1087 + AttentionLayerBase, # type: ignore[type-abstract]

1088 + ).keys()

1089 )

1090 # FIXME: support hybrid kv for draft model

1091 target_indexer_layer_names = set(

@@ -1096,7 +1102,10 @@ class SpecDecodeBaseProposer:

1102 )

1103

1104 draft_attn_layer_names = (

1099 - get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()

1105 + get_layers_from_vllm_config(

1106 + self.vllm_config,

1107 + AttentionLayerBase, # type: ignore[type-abstract]

1108 + ).keys()

1109 - target_attn_layer_names

1110 )

1111 indexer_layers = get_layers_from_vllm_config(

@@ -1136,6 +1145,7 @@ class SpecDecodeBaseProposer:

1145

1146 if supports_multimodal(target_model):

1147 # handle multimodality

1148 + assert hasattr(target_model, "config")

1149 if self.get_model_name(target_model) in [

1150 "Qwen2_5_VLForConditionalGeneration",

1151 "Qwen3VLForConditionalGeneration",

@@ -1152,16 +1162,21 @@ class SpecDecodeBaseProposer:

1162 self.model.config.image_token_index = (

1163 target_model.config.image_token_index

1164 )

1155 - target_language_model = target_model.get_language_model()

1165 + target_language_model = cast(

1166 + SupportsMultiModal, target_model

1167 + ).get_language_model()

1168 else:

1169 target_language_model = target_model

1170

1171 # share embed_tokens with the target model if needed

1172 if get_pp_group().world_size == 1:

1161 - if hasattr(target_language_model.model, "embed_tokens"):

1162 - target_embed_tokens = target_language_model.model.embed_tokens

1163 - elif hasattr(target_language_model.model, "embedding"):

1164 - target_embed_tokens = target_language_model.model.embedding

1173 + inner_model = getattr(target_language_model, "model", None)

1174 + if inner_model is None:

1175 + raise AttributeError("Target model does not have 'model' attribute")

1176 + if hasattr(inner_model, "embed_tokens"):

1177 + target_embed_tokens = inner_model.embed_tokens

1178 + elif hasattr(inner_model, "embedding"):

1179 + target_embed_tokens = inner_model.embedding

1180 else:

1181 raise AttributeError(

1182 "Target model does not have 'embed_tokens' or 'embedding' attribute"

vllm/v1/spec_decode/medusa.py

+6 -4

@@ -27,11 +27,13 @@ class MedusaProposer:

27 ):

28 # Save config parameters

29 self.vllm_config = vllm_config

30 + assert vllm_config.speculative_config is not None, (

31 + "Speculative config must be set"

32 + )

33 + self.spec_config = vllm_config.speculative_config

34 self.device = device

35 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

32 - self.hidden_size = (

33 - vllm_config.speculative_config.draft_model_config.get_hidden_size()

34 - )

36 + self.hidden_size = self.spec_config.draft_model_config.get_hidden_size()

37 self.dtype = vllm_config.model_config.dtype

39 def propose(

@@ -58,7 +60,7 @@ class MedusaProposer:

60 with set_model_tag("medusa_head"):

61 self.model = get_model(

62 vllm_config=self.vllm_config,

61 - model_config=self.vllm_config.speculative_config.draft_model_config,

63 + model_config=self.spec_config.draft_model_config,

64 )

65 assert not (

66 is_mixture_of_experts(self.model)

vllm/v1/spec_decode/suffix_decoding.py

@@ -15,6 +15,7 @@ class SuffixDecodingProposer:

16 def __init__(self, vllm_config: VllmConfig):

17 config = vllm_config.speculative_config

18 + assert config is not None, "Speculative config must be set"

19 self.num_speculative_tokens = config.num_speculative_tokens

20 self.max_tree_depth = config.suffix_decoding_max_tree_depth

21 self.max_spec_factor = config.suffix_decoding_max_spec_factor