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

[Model] Supplement to PR 24862: Pass param prefix to LLMHead (#25805)

Browse Source 
Signed-off-by: whx-sjtu <2952154980@qq.com>
This commit is contained in:
whx
2025-10-03 21:34:53 +08:00
committed by GitHub
parent 5f42fc53b6
commit cbf9221992
8 changed files with 35 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/model_executor/models/deepseek_mtp.py
View File

@@ -28,13 +28,15 @@ class SharedHead(nn.Module):
def __init__(
self,
config: PretrainedConfig,
prefix: str,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.head = ParallelLMHead(config.vocab_size,
config.hidden_size,
quant_config=quant_config)
quant_config=quant_config,
prefix=maybe_prefix(prefix, "head"))
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
return self.norm(hidden_states)
@@ -64,7 +66,9 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
device="cuda")
else:
topk_indices_buffer = None
self.shared_head = SharedHead(config=config, quant_config=quant_config)
self.shared_head = SharedHead(config=config,
prefix=prefix,
quant_config=quant_config)
self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix,
topk_indices_buffer)