From c5e5c33fcd510da53270f0845bc7d3ce1cb3ef64 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:06:28 +0300 Subject: [PATCH] [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> --- tests/quantization/test_blackwell_moe.py | 11 ++-------- .../layers/fused_moe/runner/moe_runner.py | 20 ++++++++++--------- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index 652748c668f..da70491bbc2 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -283,7 +283,7 @@ def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch): ) -def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch): +def test_nemotron_fp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", hf_overrides=HF_OVERRIDE_TEXT, @@ -291,14 +291,7 @@ def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch) ) -@pytest.mark.skip( - reason=( - "FP4 MoE backend FLASHINFER_TRTLLM does not support the " - "deployment configuration since kernel does not support " - "hidden_dim % 512 != 0." - ) -) -def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch): +def test_nemotron_fp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): can_initialize( "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", hf_overrides=HF_OVERRIDE_TEXT, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index cf8bd73fbf6..b638db13fd2 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -499,11 +499,13 @@ class MoERunner(MoERunnerInterface): # pre_xform: applied to fused_output BEFORE routed_output_transform # post_xform: applied to the final result AFTER all-reduce # - # Latent MoE with shared experts (NemotronH): - # - pre_xform strips padding from the latent dim so - # routed_output_transform receives the correct input size - # - post_xform truncates to shared_experts_hidden_dim (full hidden) - # after shared + routed outputs are combined and all-reduced + # MoE with routed output transform or shared experts: + # - pre_xform applies if the transform needs unpadded routed output + # or shared+routed add needs matching hidden dims. For Nemotron-3 + # Nano, TRTLLM NVFP4 pads routed MoE hidden dim 2688->2816, while + # shared output stays 2688. + # - post_xform uses shared_experts_hidden_dim when transform and shared + # experts make the final output full hidden dim. # # Standard MoE / MoE without transforms (GPT-OSS, Mixtral): # - pre_xform is None (no early truncation) @@ -511,12 +513,12 @@ class MoERunner(MoERunnerInterface): if transformed_hidden_dim == hidden_states.shape[-1]: transformed_hidden_dim = None - if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: + pre_xform_trunc_size = None + if self.routed_output_transform is not None or shared_experts_hidden_dim > 0: pre_xform_trunc_size = transformed_hidden_dim + post_xform_trunc_size = transformed_hidden_dim + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: post_xform_trunc_size = shared_experts_hidden_dim - else: - pre_xform_trunc_size = None - post_xform_trunc_size = transformed_hidden_dim return hidden_states, pre_xform_trunc_size, post_xform_trunc_size