minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
elvischenv
2026-08-25 22:09:09 +08:00
committed by GitHub
parent a1f9508dd4
commit 46d9427b91
2 changed files with 130 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8.py
+36 -29
View File
@@ -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,
)
test/registered/unit/layers/quantization/test_fp8_moe_weight_gating.py
+94
View File
@@ -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()