Fix MXFP8 MoE weight sizing for non-gated models (#36097)

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.
+            # One scale per w13 shard; a gated layer combines its two into a
-            # They will be combined to a single scale after weight loading.
+            # 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):
+            # A single shard already carries one scale per expert; nothing to fuse.
-                start = 0
+            if w13_num_shards > 1:
-                for shard_id in range(2):
+                for expert_id in range(layer.num_local_experts):
-                    dq_weight = per_tensor_dequantize(
+                    start = 0
-                        layer.w13_weight[expert_id][start : start + shard_size, :],
+                    for shard_id in range(w13_num_shards):
-                        layer.w13_weight_scale[expert_id][shard_id],
+                        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])
+                            layer.w13_weight[expert_id][start : start + shard_size, :],
-                    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,
         )

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()