From 46d9427b913b20ce0d72b95a9209717448b37368 Mon Sep 17 00:00:00 2001 From: elvischenv <219235043+elvischenv@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:09:09 +0800 Subject: [PATCH] Fix MXFP8 MoE weight sizing for non-gated models (#36097) --- python/sglang/srt/layers/quantization/fp8.py | 65 +++++++------ .../test_fp8_moe_weight_gating.py | 94 +++++++++++++++++++ 2 files changed, 130 insertions(+), 29 deletions(-) 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 9b528516e..6d18ba2c7 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1144,11 +1144,12 @@ class Fp8MoEMethod(FusedMoEMethodBase): params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn tp_size = get_parallel().tp_size + w13_num_shards = 2 if layer.moe_runner_config.is_gated else 1 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, ) @@ -1182,7 +1183,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_weight = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, + w13_num_shards * intermediate_size_per_partition, hidden_size // 2, dtype=torch.int8, ), @@ -1202,7 +1203,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_weight = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, + w13_num_shards * intermediate_size_per_partition, hidden_size // 8, dtype=params_dtype, ), @@ -1249,13 +1250,12 @@ class Fp8MoEMethod(FusedMoEMethodBase): # BIAS (optional, e.g. GPT-OSS) if with_bias: - w13_up_dim = ( - 2 * intermediate_size_per_partition - if layer.moe_runner_config.is_gated - else intermediate_size_per_partition - ) w13_weight_bias = torch.nn.Parameter( - torch.empty(num_experts, w13_up_dim, dtype=torch.float32), + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + dtype=torch.float32, + ), requires_grad=False, ) layer.register_parameter("w13_weight_bias", w13_weight_bias) @@ -1276,7 +1276,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_weight_scale = torch.nn.Parameter( torch.ones( num_experts, - 2 * intermediate_size_per_partition, + w13_num_shards * intermediate_size_per_partition, hidden_size // fp4_block_k, dtype=fp4_scale_dtype, ), @@ -1299,7 +1299,8 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_weight_scale = torch.nn.Parameter( scale_init( num_experts, - 2 * ((intermediate_size_per_partition + block_n - 1) // block_n), + w13_num_shards + * ((intermediate_size_per_partition + block_n - 1) // block_n), (hidden_size + block_k - 1) // block_k, dtype=scale_dtype, ), @@ -1323,10 +1324,11 @@ class Fp8MoEMethod(FusedMoEMethodBase): assert quant_config.activation_scheme == "dynamic" else: - # Allocate 2 scales for w1 and w3 respectively. - # They will be combined to a single scale after weight loading. + # One scale per w13 shard; a gated layer combines its two into a + # single scale after weight loading. w13_weight_scale = torch.nn.Parameter( - torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False + torch.ones(num_experts, w13_num_shards, dtype=torch.float32), + requires_grad=False, ) w2_weight_scale = torch.nn.Parameter( torch.ones(num_experts, dtype=torch.float32), requires_grad=False @@ -1339,7 +1341,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_weight_scale1 = torch.nn.Parameter( torch.ones( num_experts, - 2 * intermediate_size_per_partition, + w13_num_shards * intermediate_size_per_partition, dtype=torch.float32, ), requires_grad=False, @@ -2106,20 +2108,23 @@ class Fp8MoEMethod(FusedMoEMethodBase): # Fp8 moe kernel needs single weight scale for w13 per expert. # We take the max then dequant and requant each expert. assert layer.w13_weight_scale is not None + w13_num_shards = 2 if layer.moe_runner_config.is_gated else 1 shard_size = layer.intermediate_size_per_partition max_w13_scales = layer.w13_weight_scale.max(dim=1).values - for expert_id in range(layer.num_local_experts): - start = 0 - for shard_id in range(2): - dq_weight = per_tensor_dequantize( - layer.w13_weight[expert_id][start : start + shard_size, :], - layer.w13_weight_scale[expert_id][shard_id], - ) - ( - layer.w13_weight[expert_id][start : start + shard_size, :], - _, - ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) - start += shard_size + # A single shard already carries one scale per expert; nothing to fuse. + if w13_num_shards > 1: + for expert_id in range(layer.num_local_experts): + start = 0 + for shard_id in range(w13_num_shards): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id]) + start += shard_size layer.w13_weight_scale = torch.nn.Parameter( max_w13_scales, requires_grad=False @@ -2263,12 +2268,13 @@ class Fp8MoEMethod(FusedMoEMethodBase): # We won't do requant each expert's fp8 weight (not direct available), # instead we adjust half of INT4 w13_weight_scale1 numbers assert layer.w13_weight_scale is not None + w13_num_shards = 2 if layer.moe_runner_config.is_gated else 1 shard_size = layer.intermediate_size_per_partition max_w13_scales = layer.w13_weight_scale.max(dim=1).values for expert_id in range(layer.num_local_experts): start = 0 max_w13_scale_fp8 = max_w13_scales[expert_id] - for shard_id in range(2): + for shard_id in range(w13_num_shards): if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8: int4_rescale = ( layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8 @@ -2624,13 +2630,14 @@ class Fp8MoEMethod(FusedMoEMethodBase): num_experts = layer.w13_weight.shape[0] hidden_size = layer.w2_weight.shape[1] intermediate_size_per_partition = layer.intermediate_size_per_partition + w13_num_shards = 2 if layer.moe_runner_config.is_gated else 1 self.ab_strides1 = torch.full( (num_experts,), hidden_size, device=device, dtype=torch.int64 ) self.c_strides1 = torch.full( (num_experts,), - 2 * intermediate_size_per_partition, + w13_num_shards * intermediate_size_per_partition, device=device, dtype=torch.int64, ) 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 000000000..e68d1bd51 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py @@ -0,0 +1,94 @@ +"""w13 buffers 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) fuses a single projection into w13. Sizing w13 as 2*intermediate +leaves the upper half as uninitialised ``torch.empty`` that no weight loader ever +writes, which silently corrupts quantized MoE weights. The weight scales must +follow the same shard count, or the two describe different tensors. +""" + +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 + +import torch + +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.test.test_utils import CustomTestCase + +NUM_EXPERTS = 4 +HIDDEN = 256 +INTERMEDIATE = 640 +BLOCK_N = 128 +BLOCK_K = 128 + + +class _RecordingLayer: + """Collects the parameters ``create_fp8_moe_weight_`` registers.""" + + def __init__(self, is_gated: bool): + self.moe_runner_config = MoeRunnerConfig(is_gated=is_gated) + self.params = {} + + def register_parameter(self, name, param): + self.params[name] = param + + +def _create_weights(is_gated: bool, block_quant: bool): + from sglang.srt.layers.quantization import fp8 as fp8_quant + + layer = _RecordingLayer(is_gated) + quant_config = MagicMock( + weight_block_size=[BLOCK_N, BLOCK_K], + activation_scheme="dynamic", + is_checkpoint_fp8_serialized=False, + ) + + with patch.object(fp8_quant, "get_parallel") as parallel: + parallel.return_value.tp_size = 1 + fp8_quant.Fp8MoEMethod.create_fp8_moe_weight_( + layer=layer, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN, + intermediate_size_per_partition=INTERMEDIATE, + block_quant=block_quant, + quant_config=quant_config, + use_mxfp8=False, + is_checkpoint_fp8_serialized=False, + is_fp4_expert=False, + params_dtype=torch.bfloat16, + ) + return layer.params + + +class TestFp8MoEWeightGating(CustomTestCase): + def test_gated_fuses_gate_and_up(self): + params = _create_weights(is_gated=True, block_quant=True) + self.assertEqual(params["w13_weight"].shape[1], 2 * INTERMEDIATE) + self.assertEqual( + params["w13_weight_scale_inv"].shape[1], 2 * (INTERMEDIATE // BLOCK_N) + ) + + def test_non_gated_w13_holds_up_only(self): + # Regression: w13 was always sized 2*intermediate, so the upper half + # stayed uninitialised for NemotronH. + params = _create_weights(is_gated=False, block_quant=True) + self.assertEqual(params["w13_weight"].shape[1], INTERMEDIATE) + + def test_non_gated_block_scale_matches_weight(self): + params = _create_weights(is_gated=False, block_quant=True) + weight_rows = params["w13_weight"].shape[1] + scale_rows = params["w13_weight_scale_inv"].shape[1] + self.assertEqual(scale_rows * BLOCK_N, weight_rows) + + def test_non_gated_per_tensor_scale_is_single(self): + # One shard means one scale per expert; nothing to fuse afterwards. + params = _create_weights(is_gated=False, block_quant=False) + self.assertEqual(params["w13_weight_scale"].shape, (NUM_EXPERTS, 1)) + + +if __name__ == "__main__": + unittest.main()