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

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"] = "flashinfer_mxfp4"
-            overrides["moe_runner_backend"] = "marlin"
             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

@@ -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":
+        elif self._fi_kernel in ("cutlass_sm90", "cutlass_sm120"):
-            # cutlass mixed-input GEMM contraction dim K must be % 128 == 0
+            # CUTLASS mixed-input GEMM dimensions must be % 128 == 0. The
-            # (interleave factor for MXFP4 group_size=32 is 4). The kernel
+            # kernels also expect ``fc1_expert_weights`` in halved
-            # also expects ``fc1_expert_weights`` in halved ``[up; gate]``
+            # ``[up; gate]`` layout, which means the padding boundary must
-            # layout, which means the padding boundary must fall on the
+            # fall on the gate / up split.
-            # 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
+            # rebuild a properly-padded buffer in the architecture-specific
-            # ``_process_weights_for_sm90_cutlass`` after the load completes.
+            # 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 (
+        elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
-            moe_runner_backend.is_flashinfer_mxfp4()
+            "cutlass_sm90",
-            and self._fi_kernel == "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

@@ -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"]))