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

[MoE] Add FlashInfer SM90 MXFP4 W4A8 CUTLASS MoE (#34967)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-08-31 20:04:41 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent 22337e9c56
commit 5b04408784
13 changed files with 662 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+5
View File
@@ -401,6 +401,11 @@ class FusedMoE(torch.nn.Module):
self.quant_config = quant_config
self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4()
# TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic
# Keep the pre-round value: the mxfp4 SM90 CUTLASS post-load processor
# needs to know which trailing K columns are padding, because the
# loader never writes them and their scale entries keep the buffer's
# _UE8M0_ONE fill.
self.hidden_size_unpadded = hidden_size
if (
self.quant_config is not None
and self.quant_config.get_name() == "mxfp4"
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+42 -2
View File
@@ -65,8 +65,9 @@ class FlashInferCutlassMoeQuantInfo(MoeQuantInfo):
class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
"""Quantization payload for CUTLASS MXFP4 MoE.
SM90 consumes W4A16-interleaved weights and scales. SM120 consumes packed
MXFP4 weights and block-interleaved scales with MXFP8 activations.
SM90 consumes either W4A16-interleaved weights/scales or the Humming-style
W4A8 layouts. SM120 consumes packed MXFP4 weights and block-interleaved
scales with MXFP8 activations.
"""
# SM90 weights are interleaved; SM120 weights remain checkpoint-packed.
@@ -80,6 +81,13 @@ class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
# A non-None global scale selects the SM120 MXFP8 activation path.
mxfp4_weight_global_scale: Optional[torch.Tensor] = None
# A complete non-None triplet selects the SM90 Humming W4A8 path. The
# residuals are FP32 [num_local_experts] and already include the fixed 2^6
# compensation required by FlashInfer's epilogue.
w13_humming_residual_scale: Optional[torch.Tensor] = None
w2_humming_residual_scale: Optional[torch.Tensor] = None
humming_fc2_act_scale: Optional[torch.Tensor] = None
# Per-expert bias. GPT-OSS has both; DSv4 leaves both None.
w13_bias: Optional[torch.Tensor] = None # bf16 [E, 2*N]
w2_bias: Optional[torch.Tensor] = None # bf16 [E, K]
@@ -340,6 +348,22 @@ def fused_experts_none_to_flashinfer_mxfp4(
weight_global_scale = quant_info.mxfp4_weight_global_scale
use_mxfp8_act_scaling = weight_global_scale is not None
w13_humming_residual_scale = quant_info.w13_humming_residual_scale
w2_humming_residual_scale = quant_info.w2_humming_residual_scale
humming_fc2_act_scale = quant_info.humming_fc2_act_scale
humming_scales = (
w13_humming_residual_scale,
w2_humming_residual_scale,
humming_fc2_act_scale,
)
use_wfp4afp8_humming = any(scale is not None for scale in humming_scales)
if use_wfp4afp8_humming and not all(scale is not None for scale in humming_scales):
raise ValueError(
"SM90 Humming MXFP4 MoE requires both expert residual scales "
"and the FC2 activation scale."
)
if use_wfp4afp8_humming and use_mxfp8_act_scaling:
raise ValueError("SM90 Humming and SM120 MXFP8 scaling are mutually exclusive.")
input_sf = None
fc1_expert_weights = quant_info.w13_weight
fc2_expert_weights = quant_info.w2_weight
@@ -359,6 +383,17 @@ def fused_experts_none_to_flashinfer_mxfp4(
quant_info.w2_weight_scale.view(torch.int32),
weight_global_scale,
]
elif use_wfp4afp8_humming:
assert w13_humming_residual_scale is not None
assert w2_humming_residual_scale is not None
assert humming_fc2_act_scale is not None
quant_scales = [
quant_info.w13_weight_scale.view(torch.int32),
w13_humming_residual_scale,
humming_fc2_act_scale,
quant_info.w2_weight_scale.view(torch.int32),
w2_humming_residual_scale,
]
else:
quant_scales = [
quant_info.w13_weight_scale.view(torch.int32),
@@ -367,6 +402,10 @@ def fused_experts_none_to_flashinfer_mxfp4(
out_hidden = padded_hidden if do_pad else origin_hidden
output_dtype = torch.bfloat16
# FlashInfer 0.6.17 intentionally reverted the Humming API. Do not pass the
# new keyword at all on the existing W4A16/MXFP8 paths, so those paths keep
# working with SGLang's currently pinned release.
humming_kwargs = {"use_wfp4afp8_humming": True} if use_wfp4afp8_humming else {}
with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
@@ -398,6 +437,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
tune_max_num_tokens=next_power_of_2(x.shape[0]),
output=out,
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
**humming_kwargs,
)
if do_pad:
python/sglang/srt/layers/quantization/mxfp4.py
+109 -29
View File
@@ -405,14 +405,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.flashinfer_mxfp4_moe_precision = (
get_exec().moe.flashinfer_mxfp4_moe_precision
)
self._use_sm90_humming = False
# When `flashinfer_mxfp4` is enabled, dispatch to one of three FlashInfer
# entry points depending on the GPU:
# - SM100 (Blackwell) -> trtllm_fp4_block_scale_moe (existing)
# - SM120 (Blackwell) -> cutlass_fused_moe(MXFP8 x MXFP4)
# - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True)
# (FlashInfer PR #3084, post-0.6.10)
# - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True),
# W4A16 by default (PR #3084) or opt-in
# Humming W4A8 (PR #3738/#4431)
self._fi_kernel: Optional[str] = None
if self.use_flashinfer:
# precision=fp8 is an SM90 knob (Humming W4A8). The Blackwell
# paths already run MXFP8 activations, so the flag is inert there
# rather than an error -- one config can move across hardware.
if get_platform().is_sm100:
self._fi_kernel = "trtllm_sm100"
elif get_platform().is_sm120:
@@ -425,6 +430,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
"from FlashInfer PR #3084 (>= 0.6.11). Upgrade flashinfer-python "
"or pick a different backend (e.g. marlin / triton_kernel)."
)
self._use_sm90_humming = self.flashinfer_mxfp4_moe_precision == "fp8"
self._fi_kernel = "cutlass_sm90"
else:
raise NotImplementedError(
@@ -497,6 +503,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
# CUTLASS post-load processor after the load completes.
self._padded_intermediate = round_up(intermediate_size_per_partition, 128)
self._padded_hidden = round_up(hidden_size, 128)
# `hidden_size` here may ALREADY be FusedMoE's rounded value (GPT-OSS
# 2880 -> 3072). Remember the checkpoint's K so the post-load
# processor can exclude the never-written tail from Humming's
# per-expert scale range.
self._unpadded_hidden = getattr(layer, "hidden_size_unpadded", hidden_size)
# create_weights below uses the *unpadded* sizes so the loader's
# naive-copy fast path is correct.
intermediate_size_per_partition_after_pad = intermediate_size_per_partition
@@ -1113,7 +1124,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
_interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True)
def _stack_up_gate_w13(unpadded_w13, last_pad, last_un):
# FusedMoE may have rounded hidden up (GPT-OSS 2880 -> 3072) BEFORE
# create_weights, so K_un above is that rounded value, not the
# checkpoint's K. The loader never writes the trailing columns, so they
# keep the scale buffer's _UE8M0_ONE (2^0) fill -- far above a real
# per-expert max. Humming derives its residual from each expert's
# min/max E8M0 exponent, so letting those columns through would shift
# the residual and perturb the REAL weights. Copy only the checkpoint's
# columns and let the preserve_expert_range fill cover the rest; the
# packed weights there are zero, so the scale is numerically inert.
K_real = min(getattr(self, "_unpadded_hidden", None) or K_un, K_un)
# ceil: a partial trailing group is still real and must be kept.
w13_scale_real = -(-K_real // sf_block_size)
def _stack_up_gate_w13(
unpadded_w13, last_pad, last_un, preserve_expert_range=False, last_real=None
):
# unpadded_w13: [E, 2*N_un, last_un]
# Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order.
if _interleaved:
@@ -1126,10 +1152,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
out = torch.zeros(
E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device
)
if preserve_expert_range:
# Humming derives one residual from each expert's min/max E8M0
# exponents. Fill padding with an existing expert value so
# padding cannot change that range.
out.copy_(unpadded_w13[:, :1, :1])
# When protecting the expert range, stop at the checkpoint's K so
# the stale tail is covered by the fill instead of copied through.
copy_un = (
min(last_real, last_un)
if (preserve_expert_range and last_real is not None)
else last_un
)
# First half: up (with row + col padding zeros).
out[:, :N_un, :last_un] = up_rows
out[:, :N_un, :copy_un] = up_rows[:, :, :copy_un]
# Second half: gate.
out[:, N_pad : N_pad + N_un, :last_un] = gate_rows
out[:, N_pad : N_pad + N_un, :copy_un] = gate_rows[:, :, :copy_un]
return out
w13_padded = _stack_up_gate_w13(
@@ -1139,6 +1177,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
layer.w13_weight_scale.data,
K_pad // sf_block_size,
K_un // sf_block_size,
preserve_expert_range=self._use_sm90_humming,
last_real=w13_scale_real,
)
# Bias: same de-interleave on dim=-1.
if _interleaved:
@@ -1151,9 +1191,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
w13_bias_padded[:, :N_un] = w13_bias_up
w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate
def _pad_w2_3d(unpadded, last_pad, last_un):
def _pad_w2_3d(
unpadded, last_pad, last_un, preserve_expert_range=False, k_real=None
):
out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
out[:, :K_un, :last_un] = unpadded[:, :K_un, :]
if preserve_expert_range:
out.copy_(unpadded[:, :1, :1])
# Same stale-tail exclusion as _stack_up_gate_w13, on w2's K rows.
k_copy = (
min(k_real, K_un)
if (preserve_expert_range and k_real is not None)
else K_un
)
out[:, :k_copy, :last_un] = unpadded[:, :k_copy, :]
return out
# ---- w2 (no halving, just pad to [E, K_pad, N_pad/2]) ----------------
@@ -1164,6 +1214,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
layer.w2_weight_scale.data,
N_pad // sf_block_size,
N_un // sf_block_size,
preserve_expert_range=self._use_sm90_humming,
k_real=K_real,
)
w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data
@@ -1187,26 +1239,47 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
# ---- FlashInfer SM90 byte / scale interleave -----------------------
# The padded buffers above are contiguous by construction (allocated
# via torch.zeros + slice assignment), so we feed them straight in.
layer.w13_weight = Parameter(
interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4"),
requires_grad=False,
)
layer.w2_weight = Parameter(
interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4"),
requires_grad=False,
)
layer.w13_weight_scale = Parameter(
interleave_moe_scales_for_sm90_mixed_gemm(
if self._use_sm90_humming:
from flashinfer.fused_moe import (
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
)
w13_il, w13_scale_il, w13_residual = (
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
w13_padded, w13_scale_padded
)
)
w2_il, w2_scale_il, w2_residual = (
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
w2_padded, w2_scale_padded
)
)
# Humming keeps the FP4->FP8 exponent-bias compensation in the
# epilogue. FlashInfer #4431 consumes these in local expert order.
layer.w13_humming_residual_scale = Parameter(
(w13_residual * 64.0).contiguous(), requires_grad=False
)
layer.w2_humming_residual_scale = Parameter(
(w2_residual * 64.0).contiguous(), requires_grad=False
)
layer.humming_fc2_act_scale = Parameter(
torch.ones((), dtype=torch.float32, device=device),
requires_grad=False,
)
else:
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4")
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4")
w13_scale_il = interleave_moe_scales_for_sm90_mixed_gemm(
w13_scale_padded, group_size=sf_block_size
),
requires_grad=False,
)
layer.w2_weight_scale = Parameter(
interleave_moe_scales_for_sm90_mixed_gemm(
)
w2_scale_il = interleave_moe_scales_for_sm90_mixed_gemm(
w2_scale_padded, group_size=sf_block_size
),
requires_grad=False,
)
)
layer.w13_weight = Parameter(w13_il, requires_grad=False)
layer.w2_weight = Parameter(w2_il, requires_grad=False)
layer.w13_weight_scale = Parameter(w13_scale_il, requires_grad=False)
layer.w2_weight_scale = Parameter(w2_scale_il, requires_grad=False)
layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
@@ -1349,10 +1422,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
pass
def _apply_sm90_cutlass(self, layer, dispatch_output):
"""SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input
path (PR #3084). Routed through the unified ``MoeRunner`` -- this
helper only builds the quant_info; the actual kernel call lives in
:mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`."""
"""SM90 MXFP4 x BF16/FP8 MoE via FlashInfer's mixed-input kernels.
Routed through the unified ``MoeRunner``; this helper only builds the
quant_info. The actual kernel call lives in
:mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`.
"""
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
FlashInferCutlassMxfp4MoeQuantInfo,
)
@@ -1362,6 +1437,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
w13_humming_residual_scale=getattr(
layer, "w13_humming_residual_scale", None
),
w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None),
humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None),
w13_bias=layer.w13_weight_bias,
w2_bias=layer.w2_weight_bias,
swiglu_alpha=layer.swiglu_alpha,
python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
+65 -20
View File
@@ -14,7 +14,7 @@ import torch
from torch.nn import Module
from torch.nn.parameter import Parameter
from sglang.srt.runtime_context import get_platform
from sglang.srt.runtime_context import get_exec, get_platform
from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0
# Suppress TRT-LLM CUTLASS trace logs without overriding user configuration.
@@ -30,7 +30,7 @@ _GROUP_SIZE = 32
class Mxfp4FlashinferCutlassMoEMethod:
"""FlashInfer MXFP4 MoE: W4A16 on SM90 and W4A8 on SM120."""
"""FlashInfer MXFP4 MoE: W4A16/W4A8 on SM90 and W4A8 on SM120."""
fuse_routed_scaling_factor_in_topk = True
@@ -38,6 +38,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
if not is_flashinfer_available():
raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.")
self._use_mxfp8_act_scaling = get_platform().is_sm120
precision = get_exec().moe.flashinfer_mxfp4_moe_precision
# precision=fp8 is an SM90 knob (Humming W4A8); on SM120 the MXFP8
# activation path already computes in FP8, so the flag is simply inert
# there rather than an error -- one config can move across hardware.
self._use_sm90_humming = not self._use_mxfp8_act_scaling and precision == "fp8"
self._fp8 = fp8_method
self.prefix = prefix
self._swiglu_limit_tensor: torch.Tensor | None = None
@@ -124,9 +129,12 @@ class Mxfp4FlashinferCutlassMoEMethod:
return
arch = "SM120" if self._use_mxfp8_act_scaling else "SM90"
precision = (
"W4A8" if self._use_sm90_humming or self._use_mxfp8_act_scaling else "W4A16"
)
log_info_on_rank0(
logger,
f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS "
f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS {precision} "
f"(layer: {self.prefix})...",
)
@@ -152,23 +160,51 @@ class Mxfp4FlashinferCutlassMoEMethod:
for scale_u8 in (w13_scale_u8, w2_scale_u8):
scale_u8.copy_(block_scale_interleave(scale_u8).reshape_as(scale_u8))
else:
from flashinfer.fused_moe import (
interleave_moe_scales_for_sm90_mixed_gemm,
interleave_moe_weights_for_sm90_mixed_gemm,
)
if self._use_sm90_humming:
from flashinfer.fused_moe import (
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
)
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
)
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
)
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
w13_scale_u8, group_size=_GROUP_SIZE
)
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
w2_scale_u8, group_size=_GROUP_SIZE
)
w13_il, w13_s_il, w13_residual = (
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
layer.w13_weight.data.view(torch.uint8).contiguous(),
w13_scale_u8,
)
)
w2_il, w2_s_il, w2_residual = (
preprocess_moe_weights_for_sm90_mixed_gemm_humming(
layer.w2_weight.data.view(torch.uint8).contiguous(),
w2_scale_u8,
)
)
layer.w13_humming_residual_scale = Parameter(
(w13_residual * 64.0).contiguous(), requires_grad=False
)
layer.w2_humming_residual_scale = Parameter(
(w2_residual * 64.0).contiguous(), requires_grad=False
)
layer.humming_fc2_act_scale = Parameter(
torch.ones((), dtype=torch.float32, device=w13_scale_u8.device),
requires_grad=False,
)
else:
from flashinfer.fused_moe import (
interleave_moe_scales_for_sm90_mixed_gemm,
interleave_moe_weights_for_sm90_mixed_gemm,
)
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
)
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
)
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
w13_scale_u8, group_size=_GROUP_SIZE
)
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
w2_scale_u8, group_size=_GROUP_SIZE
)
layer.w13_weight = Parameter(w13_il, requires_grad=False)
layer.w2_weight = Parameter(w2_il, requires_grad=False)
layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False)
@@ -177,7 +213,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
layer._dsv4_mxfp4_backend = (
"flashinfer_cutlass_sm120"
if self._use_mxfp8_act_scaling
else "flashinfer_cutlass_sm90"
else (
"flashinfer_cutlass_sm90_fp8"
if self._use_sm90_humming
else "flashinfer_cutlass_sm90"
)
)
# SM90 creates full-size interleaved copies; release old layouts per layer.
if not self._use_mxfp8_act_scaling:
@@ -198,6 +238,11 @@ class Mxfp4FlashinferCutlassMoEMethod:
w13_weight_scale=layer.w13_weight_scale_inv,
w2_weight_scale=layer.w2_weight_scale_inv,
mxfp4_weight_global_scale=self._mxfp4_weight_global_scale_tensor,
w13_humming_residual_scale=getattr(
layer, "w13_humming_residual_scale", None
),
w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None),
humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None),
w13_bias=None,
w2_bias=None,
swiglu_alpha=None,
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+3
View File
@@ -54,6 +54,9 @@ class Mxfp4FlashinferTrtllmMoEMethod:
def __init__(self, fp8_method, prefix: str):
self._fp8 = fp8_method
self.prefix = prefix
# precision=fp8 is an SM90 knob (Humming W4A8); this SM100 trtllm path
# already runs MXFP8 activations, so the flag is inert here rather than
# an error -- one config can move across hardware.
self.flashinfer_mxfp4_moe_precision = (
get_exec().moe.flashinfer_mxfp4_moe_precision
)
python/sglang/srt/server_args.py
+4 -2
View File
@@ -2422,8 +2422,10 @@ class ServerArgs:
NS("exec.moe"),
] = "auto"
flashinfer_mxfp4_moe_precision: A[
Literal["default", "bf16"],
"Choose the computation precision of flashinfer mxfp4 moe",
Literal["default", "bf16", "fp8"],
"Choose the computation precision of flashinfer mxfp4 moe. "
"On SM90, `fp8` selects the Humming-style MXFP4-weight x FP8-activation "
"path introduced by FlashInfer #3738 and requires FlashInfer >= 0.6.18.",
NS("exec.moe"),
] = "default"
deepep_mode: A[