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",
test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py
+308
View File
@@ -0,0 +1,308 @@
"""CPU unit tests for NVFP4 fused-MoE GEMM1 global-scale handling.
w13 fuses the gate (w1) and up (w3) projections, which can carry separate NVFP4
weight scales. These tests pin the two helpers that handle that: the gate/up
split in _compute_gemm1_alphas and the TRT-LLM up-half output scalar in
_compute_g1_scale_c. Both are plain tensor ops, so they run on CPU.
The full TRT-LLM path only runs on Blackwell, and most checkpoints ship
near-equal gate/up scales, so this regression wouldn't surface in an accuracy
eval -- hence the direct contract tests. The wiring around the helpers (Marlin's
single-scale collapse, the w1/w3 mismatch warning, and registering g1_alphas /
g1_alphas_up in process_weights_after_loading) needs the GPU kernels and is
covered on-device.
Shapes and scale magnitudes follow real NVFP4 MoE checkpoints rather than toy
sizes, so a failure looks like one a real checkpoint would hit.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
import torch
# Import modelopt_quant before flashinfer_trtllm. Importing flashinfer_trtllm
# first hits a pre-existing circular import through the compressed_tensors
# schemes package; the quantization-package-first order masks it. The isort
# guards keep that order.
# isort: off
from sglang.srt.layers.quantization.modelopt_quant import _compute_gemm1_alphas
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import _compute_g1_scale_c
# isort: on
from sglang.test.test_utils import CustomTestCase
# Representative NVFP4 MoE shapes (name, num_experts, hidden, intermediate,
# is_gated). Only num_experts and is_gated matter to the scale helpers; hidden
# and intermediate just feed the reconstruct GEMM.
REAL_MOE_CONFIGS = [
("deepseek-r1", 256, 7168, 2048, True),
("llama4-maverick", 128, 5120, 8192, True),
("qwen3-235b", 128, 4096, 1536, True),
("mixtral-8x7b", 8, 4096, 14336, True),
("relu2-nongated", 128, 4096, 4096, False),
]
GATED_CONFIGS = [c for c in REAL_MOE_CONFIGS if c[-1]]
NONGATED_CONFIG = next(c for c in REAL_MOE_CONFIGS if not c[-1])
def _global_scales(num_experts: int, num_cols: int, seed: int) -> torch.Tensor:
"""Small positive fp32 weight scales, distinct per expert and column.
Real NVFP4 global scales are amax/(448*6)-style values, ~O(1e-2). Keeping
them distinct makes a gate/up swap visible. A local generator avoids
touching the global RNG.
"""
g = torch.Generator().manual_seed(seed)
scales = torch.rand((num_experts, num_cols), generator=g) * 0.19 + 0.01
return scales if num_cols > 1 else scales.squeeze(1)
class TestGemm1Alphas(CustomTestCase):
"""w13 gate/up global-scale split feeding g1_alphas / g1_alphas_up."""
def test_distinct_gate_up_columns(self):
# col 0 is gate, col 1 is up; each alpha must read its own column.
w13_input_scale = torch.tensor(0.05)
for i, (name, num_experts, _, _, _) in enumerate(GATED_CONFIGS):
with self.subTest(config=name):
w13_weight_scale_2 = _global_scales(num_experts, 2, seed=i)
gate_col = w13_weight_scale_2[:, 0]
up_col = w13_weight_scale_2[:, 1]
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
w13_weight_scale_2, w13_input_scale, is_gated=True
)
self.assertEqual(g1_alphas.shape, (num_experts,))
self.assertEqual(g1_alphas_up.shape, (num_experts,))
torch.testing.assert_close(g1_alphas, w13_input_scale * gate_col)
torch.testing.assert_close(g1_alphas_up, w13_input_scale * up_col)
self.assertFalse(torch.allclose(g1_alphas, g1_alphas_up))
def test_columns_are_decoupled(self):
# Bump one column and the other alpha must not move. Catches a column
# swap or coupling that a formula-mirroring check would wave through.
w13_input_scale = torch.tensor(0.05)
for i, (name, num_experts, _, _, _) in enumerate(GATED_CONFIGS):
with self.subTest(config=name):
base = _global_scales(num_experts, 2, seed=i)
gate_col, up_col = base[:, 0], base[:, 1]
bumped_up = torch.stack([gate_col, up_col * 10], dim=1)
bumped_gate = torch.stack([gate_col * 10, up_col], dim=1)
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
base, w13_input_scale, is_gated=True
)
g1_alphas_bu, g1_alphas_up_bu = _compute_gemm1_alphas(
bumped_up, w13_input_scale, is_gated=True
)
g1_alphas_bg, g1_alphas_up_bg = _compute_gemm1_alphas(
bumped_gate, w13_input_scale, is_gated=True
)
torch.testing.assert_close(g1_alphas, g1_alphas_bu)
self.assertFalse(torch.allclose(g1_alphas_up, g1_alphas_up_bu))
torch.testing.assert_close(g1_alphas_up, g1_alphas_up_bg)
self.assertFalse(torch.allclose(g1_alphas, g1_alphas_bg))
def test_per_expert_input_scale(self):
# The non-TRT-LLM backends pass a per-expert w13_input_scale (rather than
# the scalar the TRT-LLM path uses); it should scale each column
# element-wise.
num_experts = GATED_CONFIGS[0][1]
w13_weight_scale_2 = _global_scales(num_experts, 2, seed=0)
w13_input_scale = _global_scales(num_experts, 1, seed=5) # per-expert [E]
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
w13_weight_scale_2, w13_input_scale, is_gated=True
)
torch.testing.assert_close(
g1_alphas, w13_input_scale * w13_weight_scale_2[:, 0]
)
torch.testing.assert_close(
g1_alphas_up, w13_input_scale * w13_weight_scale_2[:, 1]
)
def test_alphas_reconstruct_each_gemm1_half(self):
# The alphas dequantize the GEMM1 accumulator back to full precision, one
# per half. Fake a GEMM in the quantized domain (block scales omitted --
# only the global-scale split is under test) and check each half
# reconstructs with its own alpha; the gate alpha would mis-scale the up
# half. Real expert counts, but k/n are capped (the reconstruction math
# doesn't depend on them) and tokens kept small to fit CPU CI. atol
# covers fp32 accumulation, not the split -- a wrong alpha is off by an
# O(1) factor, far above it.
m, k_cap, n_cap = 8, 128, 128
a_scale = torch.tensor(0.05)
for i, (name, num_experts, hidden, inter, _) in enumerate(GATED_CONFIGS):
with self.subTest(config=name):
k, n = min(hidden, k_cap), min(inter, n_cap)
gen = torch.Generator().manual_seed(100 + i)
a = torch.randn(num_experts, m, k, generator=gen)
w_gate = torch.randn(num_experts, n, k, generator=gen)
w_up = torch.randn(num_experts, n, k, generator=gen)
w13_weight_scale_2 = _global_scales(num_experts, 2, seed=i)
gate_wscale = w13_weight_scale_2[:, 0]
up_wscale = w13_weight_scale_2[:, 1]
a_q = a / a_scale
acc_gate = torch.einsum(
"emk,enk->emn", a_q, w_gate / gate_wscale[:, None, None]
)
acc_up = torch.einsum(
"emk,enk->emn", a_q, w_up / up_wscale[:, None, None]
)
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
w13_weight_scale_2, a_scale, is_gated=True
)
ref_gate = torch.einsum("emk,enk->emn", a, w_gate)
ref_up = torch.einsum("emk,enk->emn", a, w_up)
torch.testing.assert_close(
acc_gate * g1_alphas[:, None, None], ref_gate, rtol=1e-3, atol=1e-3
)
torch.testing.assert_close(
acc_up * g1_alphas_up[:, None, None], ref_up, rtol=1e-3, atol=1e-3
)
def test_shared_scale_forms_are_equivalent(self):
# A shared w1/w3 scale can arrive 1-D, [E, 1], or [E, 2] with equal
# columns. All three should give gate == up, and the single-column form
# must not reach for the missing column 1.
w13_input_scale = torch.tensor(0.05)
for i, (name, num_experts, _, _, _) in enumerate(GATED_CONFIGS):
shared = _global_scales(num_experts, 1, seed=i)
expected = w13_input_scale * shared
forms = {
"1d": shared,
"single_column": shared.unsqueeze(1),
"equal_columns": torch.stack([shared, shared], dim=1),
}
for form_name, w13_weight_scale_2 in forms.items():
with self.subTest(config=name, form=form_name):
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
w13_weight_scale_2, w13_input_scale, is_gated=True
)
torch.testing.assert_close(g1_alphas, expected)
torch.testing.assert_close(g1_alphas_up, expected)
def test_non_gated_shares_scale(self):
# Non-gated layers carry one 1-D scale and never split.
_, num_experts, _, _, _ = NONGATED_CONFIG
scale = _global_scales(num_experts, 1, seed=0)
w13_input_scale = torch.tensor(0.05)
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
scale, w13_input_scale, is_gated=False
)
self.assertEqual(g1_alphas.shape, (num_experts,))
torch.testing.assert_close(g1_alphas, w13_input_scale * scale)
torch.testing.assert_close(g1_alphas, g1_alphas_up)
def test_alphas_are_float32(self):
# Checkpoints store the scale as fp32 and kernels read fp32 alphas; a
# lower-precision input is still upcast.
num_experts = GATED_CONFIGS[0][1]
fp32_scale = _global_scales(num_experts, 2, seed=0)
for dtype in (torch.float32, torch.float16):
with self.subTest(dtype=dtype):
w13_weight_scale_2 = fp32_scale.to(dtype)
w13_input_scale = torch.tensor(0.05, dtype=dtype)
g1_alphas, g1_alphas_up = _compute_gemm1_alphas(
w13_weight_scale_2, w13_input_scale, is_gated=True
)
self.assertEqual(g1_alphas.dtype, torch.float32)
self.assertEqual(g1_alphas_up.dtype, torch.float32)
class TestG1ScaleC(CustomTestCase):
"""TRT-LLM up-half output scalar derived from g1_alphas_up."""
def test_gated_tracks_up_and_ignores_gate(self):
# g1_scale_c scales the up half, so it tracks g1_alphas_up and ignores the
# gate alpha (the old bug used the gate alpha here). The real path passes a
# scalar w2_input_scale_quant = 1 / w2_input_scale.max(), which broadcasts
# to one value per expert.
w2_input_scale_quant = torch.tensor(20.0) # ~ 1 / 0.05, scalar
for i, (name, num_experts, _, _, _) in enumerate(GATED_CONFIGS):
with self.subTest(config=name):
gate = _global_scales(num_experts, 1, seed=i)
up = _global_scales(num_experts, 1, seed=100 + i)
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, gate, up, is_gated=True
)
self.assertEqual(g1_scale_c.shape, (num_experts,))
torch.testing.assert_close(g1_scale_c, w2_input_scale_quant * up)
# Moving only the gate alpha must not move the result.
g1_scale_c_other_gate = _compute_g1_scale_c(
w2_input_scale_quant, gate * 5, up, is_gated=True
)
torch.testing.assert_close(g1_scale_c, g1_scale_c_other_gate)
def test_gated_accepts_per_expert_a2_scale(self):
# A per-expert (not scalar) w2_input_scale_quant should still broadcast.
num_experts = GATED_CONFIGS[0][1]
w2_input_scale_quant = _global_scales(num_experts, 1, seed=7)
gate = _global_scales(num_experts, 1, seed=8)
up = _global_scales(num_experts, 1, seed=9)
g1_scale_c = _compute_g1_scale_c(w2_input_scale_quant, gate, up, is_gated=True)
self.assertEqual(g1_scale_c.shape, (num_experts,))
torch.testing.assert_close(g1_scale_c, w2_input_scale_quant * up)
def test_gated_reduces_to_single_scale_when_up_equals_gate(self):
# Shared-scale checkpoints pass g1_alphas as g1_alphas_up, recovering the
# old single-scale value.
num_experts = GATED_CONFIGS[0][1]
w2_input_scale_quant = torch.tensor(20.0)
g1_alphas = _global_scales(num_experts, 1, seed=3)
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, g1_alphas, g1_alphas, is_gated=True
)
torch.testing.assert_close(g1_scale_c, w2_input_scale_quant * g1_alphas)
def test_non_gated_is_reciprocal_a2_per_expert(self):
# Relu2: no gate dequant, just 1/a2_scale as one contiguous value per
# expert. The alphas are ignored apart from the count.
_, num_experts, _, _, _ = NONGATED_CONFIG
w2_input_scale_quant = torch.tensor(20.0)
g1_alphas = torch.zeros(num_experts)
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, g1_alphas, g1_alphas, is_gated=False
)
self.assertEqual(g1_scale_c.shape, (num_experts,))
self.assertTrue(g1_scale_c.is_contiguous())
torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
def test_scale_c_is_float32(self):
# Lower-precision inputs are upcast to fp32 for the kernel.
num_experts = GATED_CONFIGS[0][1]
w2_input_scale_quant = torch.ones(num_experts, dtype=torch.float16)
g1_alphas_up = torch.ones(num_experts, dtype=torch.float16)
g1_scale_c = _compute_g1_scale_c(
w2_input_scale_quant, g1_alphas_up, g1_alphas_up, is_gated=True
)
self.assertEqual(g1_scale_c.dtype, torch.float32)
if __name__ == "__main__":
unittest.main()