biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix] Fix spec decode memory regression after #28549 (#28819)

Browse Source 
Signed-off-by: zhewenli <zhewenli@meta.com>
This commit is contained in:
Zhewen Li
2025-11-20 03:05:50 -08:00
committed by GitHub
parent 371b1d4c61
commit 93c8672ceb
4 changed files with 5 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/models/llama4_eagle.py
View File

@@ -23,7 +23,6 @@ import torch.nn as nn
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -127,17 +126,11 @@ class LlamaModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
# if PP disabled then draft will share embed with target
if get_pp_group().world_size == 1 and "embed_tokens." in name:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
for name in params_dict:
# if PP disabled then draft will share embed with target
if get_pp_group().world_size == 1 and "embed_tokens." in name:
continue
assert name in loaded_params, f"{name} is not loaded!"
return loaded_params