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

[quantization] NVFP4 MoE: split fused w13 gate/up global scales (#27588)

Browse Source 
Co-authored-by: prajjwal1 <prajjwal1@protonmail.com>
This commit is contained in:
Prajj
2026-06-14 21:18:36 -07:00
committed by GitHub
co-authored by prajjwal1
parent ce9fad7196
commit 441b75ee69
3 changed files with 389 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+25 -10
View File
@@ -476,6 +476,26 @@ def _align_fp4_moe_weights(
return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
def _compute_g1_scale_c(
w2_input_scale_quant: torch.Tensor,
g1_alphas: torch.Tensor,
g1_alphas_up: torch.Tensor,
is_gated: bool,
) -> torch.Tensor:
"""TRT-LLM GEMM1-output scale for the up (w3) half.
TRT-LLM dequantizes the two halves of the fused GEMM1 separately: g1_alphas
covers the gate half, this scalar the up half (hence g1_alphas_up). The
1/a2_scale factor (w2_input_scale_quant) requantizes GEMM2's input. A shared
scale passes g1_alphas as g1_alphas_up and recovers the single-scale value;
non-gated (Relu2) has no gate half, so it is just 1/a2_scale per expert.
"""
if is_gated:
return (w2_input_scale_quant * g1_alphas_up).to(torch.float32)
num_experts = g1_alphas.shape[0]
return w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous()
def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
"""Prepare FP4 MoE weights/scales for FlashInfer TRT-LLM kernels.
@@ -535,18 +555,13 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
layer, "w2_weight_scale", gemm2_scales_fp4_shuffled.contiguous()
)
# Compute additional scaling factor needed for TRT-LLM.
# For gated (SwiGLU): g1_scale_c = g1_alphas * a2_gscale
# For non-gated (Relu2): g1_scale_c = a2_gscale (no gate dequant contribution)
# Extra GEMM1-output scalar that TRT-LLM needs (up-half dequant).
w2_input_scale_quant = cast(torch.Tensor, layer.w2_input_scale_quant)
g1_alphas = cast(torch.Tensor, layer.g1_alphas)
if layer.moe_runner_config.is_gated:
g1_scale_c = (w2_input_scale_quant * g1_alphas).to(torch.float32)
else:
num_experts = g1_alphas.shape[0]
g1_scale_c = (
w2_input_scale_quant.to(torch.float32).expand(num_experts).contiguous()
)
g1_alphas_up = cast(torch.Tensor, getattr(layer, "g1_alphas_up", g1_alphas))
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, g1_alphas, g1_alphas_up, layer.moe_runner_config.is_gated
)
copy_or_rebind_param(layer, "g1_scale_c", g1_scale_c)
# Update intermediate_size_per_partition to reflect any padding applied
python/sglang/srt/layers/quantization/modelopt_quant.py
+56 -24
View File
@@ -1679,6 +1679,32 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
return out.view(*output_shape)
def _compute_gemm1_alphas(
w13_weight_scale_2: torch.Tensor,
w13_input_scale: torch.Tensor,
is_gated: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
"""GEMM1 weight x input alphas for the gate (w1) and up (w3) halves of w13.
w13 fuses the gate and up projections, which may carry separate NVFP4 weight
scales stored as [num_experts, 2] (col 0 = gate, col 1 = up). A 1-D (or
[num_experts, 1]) scale, and any non-gated layer, shares one scale across
both halves; the col-1 read is guarded so those cases stay in bounds.
Returns (g1_alphas, g1_alphas_up), equal for a shared scale. Single-alpha
backends use g1_alphas; the TRT-LLM path also uses g1_alphas_up.
"""
if is_gated and w13_weight_scale_2.dim() == 2 and w13_weight_scale_2.shape[1] >= 2:
gate_scale = w13_weight_scale_2[:, 0]
up_scale = w13_weight_scale_2[:, 1]
else:
gate_scale = w13_weight_scale_2.reshape(w13_weight_scale_2.shape[0])
up_scale = gate_scale
g1_alphas = (w13_input_scale * gate_scale).to(torch.float32)
g1_alphas_up = (w13_input_scale * up_scale).to(torch.float32)
return g1_alphas, g1_alphas_up
class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
"""
MoE Method for FP4 Quantization with Blockscales and PerTensorScales
@@ -1919,29 +1945,32 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
Only supports pre-quantized checkpoints with FP8 weights and scales.
"""
# GEMM 1 scale processing
if layer.moe_runner_config.is_gated:
if layer.w13_weight_scale_2.dim() == 1:
# Some checkpoints store a shared scale for w1/w3.
w13_weight_scale_2 = layer.w13_weight_scale_2
else:
if layer.w13_weight_scale_2.shape[1] >= 2 and not torch.allclose(
layer.w13_weight_scale_2[:, 0],
layer.w13_weight_scale_2[:, 1],
):
logger.warning_once(
"w1_weight_scale_2 must match w3_weight_scale_2. "
"Accuracy may be affected."
)
w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0]
else:
w13_weight_scale_2 = layer.w13_weight_scale_2[:]
# GEMM1 scale processing is deferred until the input scale is known;
# see _compute_gemm1_alphas, which splits w13's gate/up weight scales.
moe_runner_backend = getattr(
self, "_moe_runner_backend", get_moe_runner_backend()
)
if moe_runner_backend.is_marlin():
# Marlin supports only a single shared w1/w3 weight scale, so collapse
# the gate/up columns to the gate scale here. Other backends keep the
# raw scale and split the halves later (see _compute_gemm1_alphas).
if layer.moe_runner_config.is_gated:
if layer.w13_weight_scale_2.dim() == 1:
# Some checkpoints store a shared scale for w1/w3.
w13_weight_scale_2 = layer.w13_weight_scale_2
else:
if layer.w13_weight_scale_2.shape[1] >= 2 and not torch.allclose(
layer.w13_weight_scale_2[:, 0],
layer.w13_weight_scale_2[:, 1],
):
logger.warning_once(
"w1_weight_scale_2 must match w3_weight_scale_2. "
"Accuracy may be affected."
)
w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0]
else:
w13_weight_scale_2 = layer.w13_weight_scale_2[:]
copy_or_rebind_param(
layer,
"w13_weight_scale_2",
@@ -1988,12 +2017,15 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w13_input_scale = torch.ones_like(w13_input_scale, dtype=torch.float32)
w2_input_scale = torch.ones_like(w2_input_scale, dtype=torch.float32)
# Create shared parameters
copy_or_rebind_param(
layer,
"g1_alphas",
(w13_input_scale * w13_weight_scale_2).to(torch.float32),
# Create shared parameters. g1_alphas / g1_alphas_up are the gate (w1)
# and up (w3) GEMM1 scales (equal for shared-scale checkpoints).
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
layer.w13_weight_scale_2,
w13_input_scale,
layer.moe_runner_config.is_gated,
)
copy_or_rebind_param(layer, "g1_alphas", g1_alphas)
copy_or_rebind_param(layer, "g1_alphas_up", g1_alphas_up)
copy_or_rebind_param(
layer,
"g2_alphas",