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, 15:00 UTC
[Bugfix][MoE] Only unpad routed output before shared expert add or routed output transform (#40865)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Netanel Haber committed
Apr 25, 2026 at 23:50 UTC
12a3f6454b973d7cd8806d398ba287a7e1d22c63
1 file changed
+5
-1
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+5
-1
@@ -552,6 +552,8 @@ class MoERunner(MoERunnerInterface):
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
+ # so routed output can be trimmed before
556
+ # shared+routed add / latent up proj if needed.
557
routed_hidden_dim = hidden_states.shape[-1]
558
hidden_states, og_hidden_dim = self._maybe_pad_hidden_states(
559
shared_experts_input,
@@ -577,7 +579,9 @@ class MoERunner(MoERunnerInterface):
579
580
# Extract outputs from result
581
shared_output, fused_output = _unpack(result)
580
- if hidden_dim_was_padded:
582
+ if (
583
+ shared_output is not None or self.routed_output_transform is not None
584
+ ) and hidden_dim_was_padded:
585
fused_output = fused_output[..., :routed_hidden_dim]
586
587
# If combine kernel already reduced fused, reduce shared to match.