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 4, 16:15 UTC
[Bugfix][MoE] Unpad routed output before shared expert add [Fixes #35949] (#40794)
Signed-off-by: Netanel Haber <nhaber@nvidia.com>
Netanel Haber committed
Apr 24, 2026 at 14:53 UTC
e8eb0490ce098b1add05877363b185f3a7b570c5
1 file changed
+6
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+6
@@ -550,10 +550,14 @@ class MoERunner(MoERunnerInterface):
550
hidden_states
551
)
552
553
+ # Record before `_maybe_pad_hidden_states` pads activations to match
554
+ # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`
555
+ routed_hidden_dim = hidden_states.shape[-1]
556
hidden_states, og_hidden_dim = self._maybe_pad_hidden_states(
557
shared_experts_input,
558
hidden_states,
559
)
560
+ hidden_dim_was_padded = hidden_states.shape[-1] > routed_hidden_dim
561
562
result = self._forward_entry(
563
hidden_states,
@@ -573,6 +577,8 @@ class MoERunner(MoERunnerInterface):
577
578
# Extract outputs from result
579
shared_output, fused_output = _unpack(result)
580
+ if hidden_dim_was_padded:
581
+ fused_output = fused_output[..., :routed_hidden_dim]
582
583
# If combine kernel already reduced fused, reduce shared to match.
584
# See note above re: the two all-reduce points.