A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][MoE] Restore routed output unpadding before shared expert add (#45707)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Netanel Haber committed
Jun 16, 2026 at 16:06 UTC
c5e5c33fcd510da53270f0845bc7d3ce1cb3ef64
2 files changed
+13
-18
tests/quantization/test_blackwell_moe.py
+2
-9
@@ -283,7 +283,7 @@ def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):
283
)
284
285
286
-def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):
286
+def test_nemotron_fp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
287
can_initialize(
288
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
289
hf_overrides=HF_OVERRIDE_TEXT,
@@ -291,14 +291,7 @@ def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch)
291
)
292
293
294
-@pytest.mark.skip(
295
- reason=(
296
- "FP4 MoE backend FLASHINFER_TRTLLM does not support the "
297
- "deployment configuration since kernel does not support "
298
- "hidden_dim % 512 != 0."
299
- )
300
-)
301
-def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):
294
+def test_nemotron_fp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
295
can_initialize(
296
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
297
hf_overrides=HF_OVERRIDE_TEXT,
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+11
-9
@@ -499,11 +499,13 @@ class MoERunner(MoERunnerInterface):
499
# pre_xform: applied to fused_output BEFORE routed_output_transform
500
# post_xform: applied to the final result AFTER all-reduce
501
#
502
- # Latent MoE with shared experts (NemotronH):
503
- # - pre_xform strips padding from the latent dim so
504
- # routed_output_transform receives the correct input size
505
- # - post_xform truncates to shared_experts_hidden_dim (full hidden)
506
- # after shared + routed outputs are combined and all-reduced
502
+ # MoE with routed output transform or shared experts:
503
+ # - pre_xform applies if the transform needs unpadded routed output
504
+ # or shared+routed add needs matching hidden dims. For Nemotron-3
505
+ # Nano, TRTLLM NVFP4 pads routed MoE hidden dim 2688->2816, while
506
+ # shared output stays 2688.
507
+ # - post_xform uses shared_experts_hidden_dim when transform and shared
508
+ # experts make the final output full hidden dim.
509
#
510
# Standard MoE / MoE without transforms (GPT-OSS, Mixtral):
511
# - pre_xform is None (no early truncation)
@@ -511,12 +513,12 @@ class MoERunner(MoERunnerInterface):
513
if transformed_hidden_dim == hidden_states.shape[-1]:
514
transformed_hidden_dim = None
515
514
- if self.routed_output_transform is not None and shared_experts_hidden_dim > 0:
516
+ pre_xform_trunc_size = None
517
+ if self.routed_output_transform is not None or shared_experts_hidden_dim > 0:
518
pre_xform_trunc_size = transformed_hidden_dim
519
+ post_xform_trunc_size = transformed_hidden_dim
520
+ if self.routed_output_transform is not None and shared_experts_hidden_dim > 0:
521
post_xform_trunc_size = shared_experts_hidden_dim
517
- else:
518
- pre_xform_trunc_size = None
519
- post_xform_trunc_size = transformed_hidden_dim
522
523
return hidden_states, pre_xform_trunc_size, post_xform_trunc_size
524