obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-07 14:28:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix][MoE] Restore routed output unpadding before shared expert add (#45707)

Browse Source 
Signed-off-by: Netanel Haber <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Netanel Haber
2026-06-16 16:06:28 +03:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent a8c86eeb16
commit c5e5c33fcd
2 changed files with 13 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/quantization/test_blackwell_moe.py
+2 -9
View File
@@ -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,
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+11 -9
View File
@@ -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