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

Enable GPT-OSS FlashInfer MXFP4 on SM120 (#32668)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-07-30 00:04:23 +00:00
committed by GitHub
parent e5c46ff07d
commit a55e1764a2
3 changed files with 329 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+3 -3
View File
@@ -631,10 +631,10 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict:
"Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."
)
elif is_sm120_supported() and is_mxfp4_quant_format:
# trtllm-gen only supports SM100
overrides["moe_runner_backend"] = "marlin"
overrides["moe_runner_backend"] = "flashinfer_mxfp4"
logger.warning(
"Detected SM120 and MXFP4 quantization format for GPT-OSS model, enabling Marlin MOE kernel."
"Detected SM120 and MXFP4 quantization format for GPT-OSS model, "
"enabling FlashInfer CUTLASS MXFP4 MOE kernel."
)
elif (is_hip() and envs.SGLANG_USE_AITER.get()) and is_mxfp4_quant_format:
overrides["moe_runner_backend"] = "auto"
python/sglang/srt/layers/quantization/mxfp4.py
+139 -13
View File
@@ -335,15 +335,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.flashinfer_mxfp4_moe_precision = (
get_server_args().flashinfer_mxfp4_moe_precision
)
# When `flashinfer_mxfp4` is enabled, dispatch to one of two FlashInfer
# 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)
self._fi_kernel: Optional[str] = None
if self.use_flashinfer:
if is_sm100_supported():
self._fi_kernel = "trtllm_sm100"
elif is_sm120_supported():
self._fi_kernel = "cutlass_sm120"
elif is_sm90_supported():
if not _FI_HAS_SM90_CUTLASS_MXFP4:
raise RuntimeError(
@@ -355,7 +358,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self._fi_kernel = "cutlass_sm90"
else:
raise NotImplementedError(
"moe_runner_backend=flashinfer_mxfp4 requires SM90 or SM100."
"moe_runner_backend=flashinfer_mxfp4 requires SM90, SM100, "
"or SM120."
)
def create_weights(
@@ -398,12 +402,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
intermediate_size_per_partition_after_pad = round_up(
intermediate_size_per_partition, triton_kernels_padding_alignment
)
elif self._fi_kernel == "cutlass_sm90":
# cutlass mixed-input GEMM contraction dim K must be % 128 == 0
# (interleave factor for MXFP4 group_size=32 is 4). The kernel
# also expects ``fc1_expert_weights`` in halved ``[up; gate]``
# layout, which means the padding boundary must fall on the
# gate / up split.
elif self._fi_kernel in ("cutlass_sm90", "cutlass_sm120"):
# CUTLASS mixed-input GEMM dimensions must be % 128 == 0. The
# kernels also expect ``fc1_expert_weights`` in halved
# ``[up; gate]`` layout, which means the padding boundary must
# fall on the gate / up split.
#
# The mxfp4 weight loader (FusedMoE.weight_loader fast path) does
# a NAIVE copy of HF's ``[2*intermediate_size, hidden_packed]``
@@ -411,8 +414,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
# buffer here would push the gate/up boundary, so HF's "up"
# rows would land in the buffer's "gate" half and vice versa.
# Marlin sidesteps this by not padding; we do the same and
# rebuild a properly-padded buffer in
# ``_process_weights_for_sm90_cutlass`` after the load completes.
# rebuild a properly-padded buffer in the architecture-specific
# 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)
# create_weights below uses the *unpadded* sizes so the loader's
@@ -532,6 +535,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
if self._fi_kernel == "cutlass_sm90":
self._process_weights_for_sm90_cutlass(layer)
return
if self._fi_kernel == "cutlass_sm120":
self._process_weights_for_sm120_cutlass(layer)
return
if self.use_flashinfer:
# TODO: these values are hardcoded for now, we need to get them from the model
layer.gemm1_alpha = Parameter(
@@ -1007,6 +1013,99 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
torch.cuda.empty_cache()
def _process_weights_for_sm120_cutlass(self, layer):
"""Prepare GPT-OSS MXFP4 experts for FlashInfer CUTLASS on SM120.
GPT-OSS stores gate/up rows pair-wise as
``[gate_0, up_0, gate_1, up_1, ...]``. FlashInfer's fused MoE consumes
two contiguous halves in ``[up; gate]`` order. Build that layout after
checkpoint loading so padding cannot move the split, pad both GEMMs to
CUTLASS's 128-element alignment, and swizzle the native E8M0 scales for
the SM120 MXFP8-by-MXFP4 kernels. Packed FP4 weight bytes themselves do
not need an SM120 permutation.
"""
from flashinfer import block_scale_interleave
sf_block_size = 32
N_un = layer.w13_weight.shape[1] // 2
K_un = layer.w13_weight.shape[2] * 2
N_pad = self._padded_intermediate
K_pad = self._padded_hidden
E = layer.num_local_experts
device = layer.w13_weight.device
def _stack_up_gate_w13(unpadded, last_pad, last_un):
gate_rows = unpadded[:, 0::2, :]
up_rows = unpadded[:, 1::2, :]
out = torch.zeros(
E, 2 * N_pad, last_pad, dtype=unpadded.dtype, device=device
)
out[:, :N_un, :last_un] = up_rows
out[:, N_pad : N_pad + N_un, :last_un] = gate_rows
return out
w13_padded = _stack_up_gate_w13(layer.w13_weight.data, K_pad // 2, K_un // 2)
w13_scale_padded = _stack_up_gate_w13(
layer.w13_weight_scale.data,
K_pad // sf_block_size,
K_un // sf_block_size,
)
bias_dtype = layer.w13_weight_bias.dtype
w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device)
w13_bias_padded[:, :N_un] = layer.w13_weight_bias.data[:, 1::2]
w13_bias_padded[:, N_pad : N_pad + N_un] = layer.w13_weight_bias.data[:, 0::2]
def _pad_w2_3d(unpadded, last_pad, last_un):
out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
out[:, :K_un, :last_un] = unpadded[:, :K_un, :]
return out
w2_padded = _pad_w2_3d(layer.w2_weight.data, N_pad // 2, N_un // 2)
w2_scale_padded = _pad_w2_3d(
layer.w2_weight_scale.data,
N_pad // sf_block_size,
N_un // sf_block_size,
)
w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data
w13_scale_interleaved = block_scale_interleave(w13_scale_padded)
w2_scale_interleaved = block_scale_interleave(w2_scale_padded)
layer.w13_weight = Parameter(w13_padded, requires_grad=False)
layer.w2_weight = Parameter(w2_padded, requires_grad=False)
layer.w13_weight_scale = Parameter(
w13_scale_interleaved.reshape_as(w13_scale_padded),
requires_grad=False,
)
layer.w2_weight_scale = Parameter(
w2_scale_interleaved.reshape_as(w2_scale_padded),
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)
layer.swiglu_alpha = Parameter(
torch.full((E,), 1.702, dtype=torch.float32, device=device),
requires_grad=False,
)
layer.swiglu_beta = Parameter(
torch.ones(E, dtype=torch.float32, device=device),
requires_grad=False,
)
layer.swiglu_limit = Parameter(
torch.full((E,), 7.0, dtype=torch.float32, device=device),
requires_grad=False,
)
# The MXFP4 ABI uses a neutral global weight scale for each GEMM.
layer.mxfp4_weight_global_scale = Parameter(
torch.ones(E, dtype=torch.float32, device=device),
requires_grad=False,
)
layer._mxfp4_backend = "flashinfer_cutlass_sm120"
torch.cuda.empty_cache()
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
):
@@ -1032,9 +1131,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
or moe_runner_backend.is_marlin()
):
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
elif (
moe_runner_backend.is_flashinfer_mxfp4()
and self._fi_kernel == "cutlass_sm90"
elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
"cutlass_sm90",
"cutlass_sm120",
):
# Register the fused func at runner construction so the FusedOpPool
# lookup at `MoeRunner.__init__` finds it.
@@ -1073,6 +1172,31 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)
return self.runner.run(dispatch_output, quant_info)
def _apply_sm120_cutlass(self, layer, dispatch_output):
"""SM120 GPT-OSS MXFP8 x MXFP4 MoE via FlashInfer CUTLASS."""
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
FlashInferCutlassMxfp4MoeQuantInfo,
)
quant_info = FlashInferCutlassMxfp4MoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_scale=layer.w2_weight_scale,
mxfp4_weight_global_scale=layer.mxfp4_weight_global_scale,
w13_bias=layer.w13_weight_bias,
w2_bias=layer.w2_weight_bias,
swiglu_alpha=layer.swiglu_alpha,
swiglu_beta=layer.swiglu_beta,
swiglu_limit=layer.swiglu_limit,
moe_tp_size=layer.moe_tp_size,
moe_tp_rank=layer.moe_tp_rank,
moe_ep_size=layer.moe_ep_size,
moe_ep_rank=layer.moe_ep_rank,
padded_hidden=self._padded_hidden,
)
return self.runner.run(dispatch_output, quant_info)
def apply(
self,
layer: torch.nn.Module,
@@ -1148,6 +1272,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
if self._fi_kernel == "cutlass_sm90":
return self._apply_sm90_cutlass(layer, dispatch_output)
if self._fi_kernel == "cutlass_sm120":
return self._apply_sm120_cutlass(layer, dispatch_output)
if self.use_flashinfer:
# When bf16 mode is enabled, we don't need to quantize the input,
# TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
+187 -1
View File
@@ -13,7 +13,7 @@ import torch
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small")
def _random_weights(num_experts: int, hidden: int, intermediate: int):
@@ -247,5 +247,191 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
assert torch.equal(actual, expected)
def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if torch.cuda.get_device_capability() != (12, 0):
pytest.skip("SM120 required")
pytest.importorskip("flashinfer.fused_moe")
from flashinfer import block_scale_interleave, mxfp8_quantize
from flashinfer.fused_moe import cutlass_fused_moe
from flashinfer.fused_moe.core import ActivationType
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as runner_module
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
from sglang.srt.layers.moe.topk import StandardTopKOutput
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
monkeypatch.setattr(
runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
)
monkeypatch.setattr(runner_module, "is_allocation_symmetric", lambda: False)
monkeypatch.setattr(runner_module, "get_tp_group", lambda: None)
num_experts, hidden, intermediate = 4, 160, 160
padded_hidden = padded_intermediate = 256
w13, w2, w13_scale, w2_scale = _random_weights(num_experts, hidden, intermediate)
generator = torch.Generator(device="cuda").manual_seed(2)
w13_bias = torch.randn(
num_experts,
2 * intermediate,
dtype=torch.bfloat16,
device="cuda",
generator=generator,
)
w2_bias = torch.randn(
num_experts,
hidden,
dtype=torch.bfloat16,
device="cuda",
generator=generator,
)
layer = SimpleNamespace(
w13_weight=torch.nn.Parameter(
w13.view(torch.uint8).clone(), requires_grad=False
),
w2_weight=torch.nn.Parameter(w2.view(torch.uint8).clone(), requires_grad=False),
w13_weight_scale=torch.nn.Parameter(
w13_scale.view(torch.uint8).clone(), requires_grad=False
),
w2_weight_scale=torch.nn.Parameter(
w2_scale.view(torch.uint8).clone(), requires_grad=False
),
w13_weight_bias=torch.nn.Parameter(w13_bias.clone(), requires_grad=False),
w2_weight_bias=torch.nn.Parameter(w2_bias.clone(), requires_grad=False),
num_local_experts=num_experts,
moe_tp_size=1,
moe_tp_rank=0,
moe_ep_size=1,
moe_ep_rank=0,
)
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
method._fi_kernel = "cutlass_sm120"
method.num_experts = num_experts
method.hidden_size = hidden
method.intermediate_size_per_partition = intermediate
method._padded_hidden = padded_hidden
method._padded_intermediate = padded_intermediate
config = MoeRunnerConfig(
num_experts=num_experts,
num_local_experts=num_experts,
hidden_size=hidden,
intermediate_size_per_partition=intermediate,
top_k=4,
activation="silu",
is_gated=True,
gemm1_alpha=1.702,
gemm1_clamp_limit=7.0,
)
method.moe_runner_config = config
method.runner = MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, config)
method._process_weights_for_sm120_cutlass(layer)
expected_w13 = torch.zeros(
num_experts,
2 * padded_intermediate,
padded_hidden // 2,
dtype=torch.uint8,
device="cuda",
)
expected_w13[:, :intermediate, : hidden // 2] = w13[:, 1::2]
expected_w13[
:, padded_intermediate : padded_intermediate + intermediate, : hidden // 2
] = w13[:, 0::2]
expected_w13_scale = torch.zeros(
num_experts,
2 * padded_intermediate,
padded_hidden // 32,
dtype=torch.uint8,
device="cuda",
)
expected_w13_scale[:, :intermediate, : hidden // 32] = w13_scale.view(torch.uint8)[
:, 1::2
]
expected_w13_scale[
:,
padded_intermediate : padded_intermediate + intermediate,
: hidden // 32,
] = w13_scale.view(torch.uint8)[:, 0::2]
expected_w13_scale = block_scale_interleave(expected_w13_scale).reshape_as(
expected_w13_scale
)
assert torch.equal(layer.w13_weight, expected_w13)
assert torch.equal(layer.w13_weight_scale, expected_w13_scale)
assert torch.equal(layer.w13_weight_bias[:, :intermediate], w13_bias[:, 1::2])
assert torch.equal(
layer.w13_weight_bias[
:, padded_intermediate : padded_intermediate + intermediate
],
w13_bias[:, 0::2],
)
assert torch.all(layer.swiglu_alpha == 1.702)
assert torch.all(layer.swiglu_beta == 1.0)
assert torch.all(layer.swiglu_limit == 7.0)
assert layer._mxfp4_backend == "flashinfer_cutlass_sm120"
x = torch.randn(
8,
hidden,
dtype=torch.bfloat16,
device="cuda",
generator=generator,
)
logits = torch.randn(
8,
num_experts,
dtype=torch.float32,
device="cuda",
generator=generator,
)
topk_weights, topk_ids = torch.topk(torch.softmax(logits, dim=-1), 4, dim=-1)
topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
dispatch_output = StandardDispatchOutput(
x,
None,
StandardTopKOutput(topk_weights, topk_ids.to(torch.int32), logits),
)
actual = method._apply_sm120_cutlass(layer, dispatch_output).hidden_states
x_padded = torch.nn.functional.pad(x, (0, padded_hidden - hidden))
x_quant, x_scale = mxfp8_quantize(
x_padded, is_sf_swizzled_layout=True, alignment=32
)
expected = torch.empty(
x.shape[0], padded_hidden, dtype=torch.bfloat16, device="cuda"
)
cutlass_fused_moe(
input=x_quant,
token_selected_experts=topk_ids.to(torch.int32),
token_final_scales=topk_weights,
fc1_expert_weights=layer.w13_weight.view(torch.int64),
fc2_expert_weights=layer.w2_weight.view(torch.int64),
output_dtype=torch.bfloat16,
quant_scales=[
layer.w13_weight_scale.view(torch.int32),
layer.mxfp4_weight_global_scale,
layer.w2_weight_scale.view(torch.int32),
layer.mxfp4_weight_global_scale,
],
input_sf=x_scale,
fc1_expert_biases=layer.w13_weight_bias,
fc2_expert_biases=layer.w2_weight_bias,
swiglu_alpha=layer.swiglu_alpha,
swiglu_beta=layer.swiglu_beta,
swiglu_limit=layer.swiglu_limit,
use_w4_group_scaling=False,
use_mxfp8_act_scaling=True,
activation_type=ActivationType.Swiglu,
tune_max_num_tokens=8,
output=expected,
)
assert torch.equal(actual, expected[:, :hidden].contiguous())
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))