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