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 12, 21:00 UTC
[Bugfix] Fix spec decode memory regression after #28549 (#28819)
Signed-off-by: zhewenli <zhewenli@meta.com>
Zhewen Li committed
Nov 20, 2025 at 03:05 UTC
93c8672ceb06f6e9c282a96fcd85a7ce41293693
4 files changed
+5
-19
vllm/model_executor/models/deepseek_eagle.py
-5
@@ -8,7 +8,6 @@ import torch.nn as nn
8
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import VllmConfig
11
-from vllm.distributed.parallel_state import get_pp_group
11
from vllm.model_executor.layers.fused_moe import FusedMoE
12
from vllm.model_executor.layers.layernorm import RMSNorm
13
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -172,10 +171,6 @@ class DeepseekV2Model(nn.Module):
171
)
172
break
173
else:
175
- # if PP disabled then draft will share embed with target
176
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
177
- continue
178
-
174
# Skip loading extra bias for GPTQ models.
175
if name.endswith(".bias") and name not in params_dict:
176
continue
vllm/model_executor/models/llama4_eagle.py
-7
@@ -23,7 +23,6 @@ import torch.nn as nn
23
24
from vllm.compilation.decorators import support_torch_compile
25
from vllm.config import VllmConfig
26
-from vllm.distributed.parallel_state import get_pp_group
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.layernorm import RMSNorm
28
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -127,17 +126,11 @@ class LlamaModel(nn.Module):
126
weight_loader(param, loaded_weight, shard_id)
127
break
128
else:
130
- # if PP disabled then draft will share embed with target
131
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
132
- continue
129
param = params_dict[name]
130
weight_loader = getattr(param, "weight_loader", default_weight_loader)
131
weight_loader(param, loaded_weight)
132
loaded_params.add(name)
133
for name in params_dict:
138
- # if PP disabled then draft will share embed with target
139
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
140
- continue
134
assert name in loaded_params, f"{name} is not loaded!"
135
return loaded_params
136
vllm/model_executor/models/llama_eagle.py
-5
@@ -9,7 +9,6 @@ from transformers import LlamaConfig
9
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
-from vllm.distributed.parallel_state import get_pp_group
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.linear import ReplicatedLinear
14
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -155,10 +154,6 @@ class LlamaModel(nn.Module):
154
weight_loader(param, loaded_weight, shard_id)
155
break
156
else:
158
- # if PP disabled then draft will share embed with target
159
- if get_pp_group().world_size == 1 and "embed_tokens." in name:
160
- continue
161
-
157
param = params_dict[name]
158
weight_loader = getattr(param, "weight_loader", default_weight_loader)
159
weight_loader(param, loaded_weight)
vllm/v1/spec_decode/eagle.py
+5
-2
@@ -1028,8 +1028,11 @@ class EagleProposer:
1028
elif (
1029
isinstance(target_embed_tokens.weight, torch.Tensor)
1030
and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)
1031
- and torch.equal(
1032
- target_embed_tokens.weight, self.model.model.embed_tokens.weight
1031
+ and torch.allclose(
1032
+ target_embed_tokens.weight.cpu(),
1033
+ self.model.model.embed_tokens.weight.cpu(),
1034
+ rtol=1e-5,
1035
+ atol=1e-7,
1036
)
1037
):
1038
share_embeddings = True