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, 03:45 UTC
[Bugfix] Restore moe_forward output shape invariant on TRTLLM MXFP4 path (#41646)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Stefano Castagnetta committed
May 8, 2026 at 00:26 UTC
54f548e9e58087f0155e4e164e416ad7efdfde6d
1 file changed
+43
-6
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+43
-6
@@ -93,6 +93,7 @@ def _moe_forward(
93
shared_experts_input: torch.Tensor | None,
94
input_ids: torch.Tensor | None,
95
layer_name: _layer_name_type,
96
+ hidden_dim_unpadded: int,
97
) -> torch.Tensor:
98
layer = get_layer_from_name(_resolve_layer_name(layer_name))
99
return layer.runner._forward_impl(
@@ -110,7 +111,14 @@ def _moe_forward_fake(
111
shared_experts_input: torch.Tensor | None,
112
input_ids: torch.Tensor | None,
113
layer_name: _layer_name_type,
114
+ hidden_dim_unpadded: int,
115
) -> torch.Tensor:
116
+ # `hidden_dim_unpadded > 0` only on the TRT-LLM MXFP4 path, where the
117
+ # real kernel writes narrower than `hidden_states.shape[-1]`. Plumbed
118
+ # as an op arg (not peeked from the layer registry) to keep the fake
119
+ # a pure shape function of its inputs and preserve subgraph dedup.
120
+ if hidden_dim_unpadded > 0:
121
+ return hidden_states.new_empty((*hidden_states.shape[:-1], hidden_dim_unpadded))
122
return torch.empty_like(hidden_states)
123
124
@@ -120,6 +128,7 @@ def _moe_forward_shared(
128
shared_experts_input: torch.Tensor | None,
129
input_ids: torch.Tensor | None,
130
layer_name: _layer_name_type,
131
+ hidden_dim_unpadded: int,
132
) -> tuple[torch.Tensor, torch.Tensor]:
133
layer = get_layer_from_name(_resolve_layer_name(layer_name))
134
return layer.runner._forward_impl(
@@ -137,13 +146,17 @@ def _moe_forward_shared_fake(
146
shared_experts_input: torch.Tensor | None,
147
input_ids: torch.Tensor | None,
148
layer_name: _layer_name_type,
149
+ hidden_dim_unpadded: int,
150
) -> tuple[torch.Tensor, torch.Tensor]:
141
- # Output shapes:
142
- # - fused_out: same as hidden_states (routed experts use transformed size)
143
- # - shared_out: same as shared_experts_input if provided, else same as
144
- # hidden_states
145
- # (For latent MoE: shared experts use original hidden_size, not latent size)
146
- fused_out = torch.empty_like(hidden_states)
151
+ # `fused_out`: see `_moe_forward_fake` for hidden_dim_unpadded semantics.
152
+ # `shared_out`: matches `shared_experts_input` if provided (latent MoE),
153
+ # else `hidden_states`.
154
+ if hidden_dim_unpadded > 0:
155
+ fused_out = hidden_states.new_empty(
156
+ (*hidden_states.shape[:-1], hidden_dim_unpadded)
157
+ )
158
+ else:
159
+ fused_out = torch.empty_like(hidden_states)
160
if shared_experts_input is not None:
161
shared_out = torch.empty_like(shared_experts_input)
162
else:
@@ -389,6 +402,29 @@ class MoERunner(MoERunnerInterface):
402
return "from_forward_context"
403
return self.layer_name
404
405
+ def _trtllm_mxfp4_unpadded_dim(self) -> int:
406
+ """Return ``hidden_dim_unpadded`` when the active backend is TRT-LLM
407
+ MXFP4 (whose kernel writes narrower than the padded
408
+ ``hidden_states.shape[-1]``), else 0. Other MXFP4 backends (notably
409
+ Cutlass MXFP4 MXFP8) write the full padded width, so
410
+ ``moe_config.hidden_dim_unpadded`` alone is insufficient: it encodes
411
+ the model's logical hidden, not whether the kernel narrows. Computed
412
+ caller-side and passed as an op arg; doing the isinstance check
413
+ inside the fake would specialize per ``layer_name`` and break
414
+ subgraph dedup for identical-architecture models (e.g. Phi-MoE).
415
+ """
416
+ from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
417
+ TrtLlmMxfp4ExpertsBase,
418
+ )
419
+
420
+ moe_kernel = getattr(self._quant_method, "moe_kernel", None)
421
+ fused_experts = getattr(
422
+ getattr(moe_kernel, "impl", None), "fused_experts", None
423
+ )
424
+ if isinstance(fused_experts, TrtLlmMxfp4ExpertsBase):
425
+ return self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim
426
+ return 0
427
+
428
def _maybe_pad_hidden_states(
429
self,
430
shared_experts_input: torch.Tensor | None,
@@ -585,6 +621,7 @@ class MoERunner(MoERunnerInterface):
621
shared_experts_input,
622
input_ids,
623
self._encode_layer_name(),
624
+ self._trtllm_mxfp4_unpadded_dim(),
625
)
626
627
#