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 8, 00:00 UTC
[ROCm][Bugfix]: Only save unpadded sizes for shared_experts in MoERunner to fix rmsnorm pad fusion (#34636)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Feb 20, 2026 at 21:56 UTC
ded333fb9b903e9de9f1cc5d82d2b5c5ab726750
1 file changed
+6
-3
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+6
-3
@@ -384,8 +384,11 @@ class DefaultMoERunner(MoERunner):
384
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
385
# For latent MoE: save ORIGINAL hidden_states before transform
386
# (shared_experts need original dimension, routed experts use transformed)
387
- original_hidden_states = hidden_states
388
- original_hidden_dim = hidden_states.shape[-1]
387
+ if self.shared_experts is not None:
388
+ original_hidden_states = hidden_states
389
+ original_hidden_dim = hidden_states.shape[-1]
390
+ else:
391
+ original_hidden_states = None
392
393
# Apply transform for routed experts (e.g., latent projection for latent MoE)
394
hidden_states = self.apply_routed_input_transform(hidden_states)
@@ -407,7 +410,7 @@ class DefaultMoERunner(MoERunner):
410
self._encode_layer_name(),
411
)
412
410
- if isinstance(fused_output, tuple):
413
+ if self.shared_experts is not None:
414
orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim]
415
else:
416
orig_hidden_dims = [transformed_hidden_dim]