From a55e1764a2e12a757b8a5f98054d41c82ccb4b3a Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:04:23 +0800 Subject: [PATCH] Enable GPT-OSS FlashInfer MXFP4 on SM120 (#32668) --- python/sglang/srt/arg_groups/overrides.py | 6 +- .../sglang/srt/layers/quantization/mxfp4.py | 152 ++++++++++++-- .../quantization/test_mxfp4_sm120_cutlass.py | 188 +++++++++++++++++- 3 files changed, 329 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index b9dc7da89..08a236345 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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" diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index ba0f998e6..332ad7eda 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -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, diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py index 38079a140..435c6d7fe 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py @@ -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"]))