From f57f24ba688464ae54226535e8c1aa87b9195686 Mon Sep 17 00:00:00 2001 From: elvischenv <219235043+elvischenv@users.noreply.github.com> Date: Sun, 23 Aug 2026 23:40:47 +0800 Subject: [PATCH] Fix MXFP8 MoE weight sizing for non-gated models create_fp8_moe_weight_ hardcoded is_concat=True, so w13 was always sized as 2 * intermediate (gate + up). Non-gated MoEs fuse only up into w13, leaving the upper half as uninitialised torch.empty that no weight loader writes. NemotronH (relu2, checkpoint has up_proj/down_proj and no gate_proj) is the only MoE model in sglang with is_gated=False. It scored 0.066 gsm8k under --quantization mxfp8 versus 0.943 in bf16; with this fix it scores 0.945. Derive is_concat from the layer's gating instead. Add a CPU unit test covering both the sizing primitive and the call site. It fails on the unfixed tree and passes with the fix. --- python/sglang/srt/layers/quantization/fp8.py | 2 +- .../test_fp8_moe_weight_gating.py | 79 +++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 3dee0a97dc7b..55fd078984f5 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1148,7 +1148,7 @@ def create_fp8_moe_weight_( w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes( intermediate_size_per_partition, is_aiter_moe=_use_aiter, - is_concat=True, + is_concat=layer.moe_runner_config.is_gated, is_packed=False, ) diff --git a/test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py b/test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py new file mode 100644 index 000000000000..6d16a5dda8af --- /dev/null +++ b/test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py @@ -0,0 +1,79 @@ +"""w13 must be sized by the layer's gating, not assumed to be gate+up fused. + +A non-gated MoE (e.g. NemotronH: relu2, checkpoint carries up_proj/down_proj and +no gate_proj) has a single projection fused into w13. Sizing it as 2*intermediate +leaves the upper half as uninitialised ``torch.empty`` that no weight loader ever +writes, which silently corrupts quantized MoE weights. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import unittest +from unittest.mock import MagicMock, patch + +from sglang.srt.layers.moe.utils import get_moe_weight_sizes +from sglang.test.test_utils import CustomTestCase + +INTERMEDIATE = 640 + + +class TestGetMoeWeightSizes(CustomTestCase): + """The sizing primitive itself.""" + + def test_gated_fuses_gate_and_up(self): + w13_up_dim, w2_down_dim, _ = get_moe_weight_sizes( + INTERMEDIATE, is_concat=True, is_packed=False, is_aiter_moe=False + ) + self.assertEqual(w13_up_dim, 2 * INTERMEDIATE) + self.assertEqual(w2_down_dim, INTERMEDIATE) + + def test_non_gated_holds_up_only(self): + w13_up_dim, w2_down_dim, _ = get_moe_weight_sizes( + INTERMEDIATE, is_concat=False, is_packed=False, is_aiter_moe=False + ) + self.assertEqual(w13_up_dim, INTERMEDIATE) + self.assertEqual(w2_down_dim, INTERMEDIATE) + + +class TestFp8MoEWeightGating(CustomTestCase): + """Fp8MoEMethod must forward the layer's gating into the sizing call.""" + + def _is_concat_used_for(self, is_gated: bool) -> bool: + from sglang.srt.layers.quantization import fp8 as fp8_quant + + layer = MagicMock() + layer.moe_runner_config.is_gated = is_gated + + # Stop right after the sizing decision; we only assert on its argument. + with patch.object( + fp8_quant, "get_moe_weight_sizes", return_value=(0, 0, False) + ) as sizes, patch.object(fp8_quant, "get_parallel") as parallel: + parallel.return_value.tp_size = 1 + with self.assertRaises(Exception): + fp8_quant.Fp8MoEMethod.create_fp8_moe_weight_( + layer=layer, + num_experts=8, + hidden_size=128, + intermediate_size_per_partition=INTERMEDIATE, + block_quant=True, + quant_config=MagicMock(weight_block_size=[1, 32]), + use_mxfp8=True, + is_checkpoint_fp8_serialized=False, + is_fp4_expert=False, + params_dtype=None, + ) + self.assertTrue(sizes.called, "get_moe_weight_sizes was never reached") + return sizes.call_args.kwargs["is_concat"] + + def test_gated_layer_requests_concat(self): + self.assertTrue(self._is_concat_used_for(is_gated=True)) + + def test_non_gated_layer_does_not_request_concat(self): + # Regression: this was hardcoded True, over-allocating w13 for NemotronH. + self.assertFalse(self._is_concat_used_for(is_gated=False)) + + +if __name__ == "__main__": + unittest.main()