[Bugfix][Spec Decode] Only create draft_id_to_target_id when draft vocab differs (#53458)
Signed-off-by: drslark <slarksblood@qq.com> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: GLM-5.3 <glm@z.ai> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
drslark committed
Sep 15, 2026 at 22:59 UTC
9446ea168067c141992f86c3a796e8551f37139e
3 files changed
+24
-12
vllm/model_executor/models/deepseek_eagle3.py
+8
-4
@@ -305,10 +305,14 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
305
self.logits_processor = LogitsProcessor(
306
self.config.draft_vocab_size, scale=logit_scale
307
)
308
- self.draft_id_to_target_id = nn.Parameter(
309
- torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
310
- requires_grad=False,
311
- )
308
+ target_vocab_size = vllm_config.model_config.get_vocab_size()
309
+ if self.config.draft_vocab_size != target_vocab_size:
310
+ self.draft_id_to_target_id = nn.Parameter(
311
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
312
+ requires_grad=False,
313
+ )
314
+ else:
315
+ self.draft_id_to_target_id = None
316
317
def embed_input_ids(
318
self,
vllm/model_executor/models/llama_eagle3.py
+8
-4
@@ -335,10 +335,14 @@ class Eagle3LlamaForCausalLM(_Eagle3LlamaForCausalLMBase):
335
self.logits_processor = LogitsProcessor(
336
self.config.draft_vocab_size, scale=logit_scale
337
)
338
- self.draft_id_to_target_id = nn.Parameter(
339
- torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
340
- requires_grad=False,
341
- )
338
+ target_vocab_size = vllm_config.model_config.get_vocab_size()
339
+ if self.config.draft_vocab_size != target_vocab_size:
340
+ self.draft_id_to_target_id = nn.Parameter(
341
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
342
+ requires_grad=False,
343
+ )
344
+ else:
345
+ self.draft_id_to_target_id = None
346
347
self.use_parallel_drafting = speculative_config.parallel_drafting
348
vllm/model_executor/models/qwen3_eagle3.py
+8
-4
@@ -309,10 +309,14 @@ class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
309
self.logits_processor = LogitsProcessor(
310
self.config.draft_vocab_size, scale=logit_scale
311
)
312
- self.draft_id_to_target_id = nn.Parameter(
313
- torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
314
- requires_grad=False,
315
- )
312
+ target_vocab_size = vllm_config.model_config.get_vocab_size()
313
+ if self.config.draft_vocab_size != target_vocab_size:
314
+ self.draft_id_to_target_id = nn.Parameter(
315
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
316
+ requires_grad=False,
317
+ )
318
+ else:
319
+ self.draft_id_to_target_id = None
320
321
self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting
322