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 2, 08:30 UTC
[Bug] Fix gemma4 MTP IMA issue when TP>1, `CUDA error: an illegal memory access was encountered` (#43909)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
May 30, 2026 at 10:34 UTC
e1105064b282bb807ba9c309741b40a3b64e2261
1 file changed
+9
-1
vllm/model_executor/models/gemma4_mtp.py
+9
-1
@@ -501,6 +501,7 @@ class Gemma4MTP(nn.Module):
501
config = vllm_config.speculative_config.draft_model_config.hf_config
502
text_config = _get_text_config(config)
503
self.config = config
504
+ self._stable_full_lm_head_weight: torch.Tensor | None = None
505
506
self.model = Gemma4MultiTokenPredictor(
507
vllm_config=vllm_config,
@@ -567,6 +568,8 @@ class Gemma4MTP(nn.Module):
568
)
569
570
def _get_full_lm_head_weight(self) -> torch.Tensor:
571
+ if self._stable_full_lm_head_weight is not None:
572
+ return self._stable_full_lm_head_weight
573
lm_head_weight = self.lm_head.weight
574
tp_size = get_tensor_model_parallel_world_size()
575
if tp_size > 1:
@@ -574,7 +577,11 @@ class Gemma4MTP(nn.Module):
577
lm_head_weight,
578
dim=0,
579
)
577
- return lm_head_weight[: self.masked_embedding.vocab_size]
580
+ lm_head_weight = lm_head_weight[: self.masked_embedding.vocab_size]
581
+ if tp_size > 1:
582
+ lm_head_weight = lm_head_weight.contiguous()
583
+ self._stable_full_lm_head_weight = lm_head_weight
584
+ return lm_head_weight
585
586
def compute_logits(
587
self,
@@ -599,5 +606,6 @@ class Gemma4MTP(nn.Module):
606
)
607
608
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
609
+ self._stable_full_lm_head_weight = None
610
loader = AutoWeightsLoader(self)
611
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)