# PR #32668 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Enable GPT-OSS FlashInfer MXFP4 on SM120
- 合并时间：2026-07-30 08:04
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/32668

---

# 执行摘要

- 一句话：在 SM120 上启用 GPT-OSS FlashInfer MXFP4 MoE
- 推荐动作：值得精读。本 PR 展示了如何为特定 GPU 架构添加新的 MoE 内核后端，包括初始化检测、对齐约束处理、权重后处理以及自动切换。对于涉及多架构内核支持的开发者有很好的参考价值。

# 功能与动机

在 SM120（Blackwell 架构）上，FlashInfer CUTLASS MXFP4 内核比现有 Marlin 内核提供显著性能优势（最高 30% 吞吐提升）。PR body 中的基准测试显示从低并发到高并发全面超越 Marlin，因此希望为 GPT-OSS 模型默认启用该内核。

# 实现拆解

1. **在 `mxfp4.py` 中注册 SM120 内核路径**：在 `Mxfp4MoEMethod.__init__()` 中新增 `is_sm120_supported()` 分支，设置 `_fi_kernel = 'cutlass_sm120'`；在 `create_weights()` 中将 SM120 与 SM90 的 padding 策略合并（要求 dimension % 128 == 0），因为这些 CUTLASS 内核有相同的对齐约束。

2. **在 `mxfp4.py` 中新增权重后处理函数**：在 `process_weights_after_loading()` 中为 `cutlass_sm120` 分支调用新增的 `_process_weights_for_sm120_cutlass()`，该函数包含 `_stack_up_gate_w13()`、`_pad_w2_3d()` 和 `_apply_sm120_cutlass()` 等辅助函数，负责将模型权重重新排列为 FlashInfer CUTLASS 内核所需的布局（halved `[up; gate]` 布局），并执行 padding。

3. **在 `overrides.py` 中修改 SM120 自动选择**：当检测到 SM120 且模型使用 MXFP4 量化格式时，将 `moe_runner_backend` 从 `'marlin'` 改为 `'flashinfer_mxfp4'`，使得 SM120 用户无需手动指定即可启用新内核。

4. **新增单元测试 `test_mxfp4_sm120_cutlass.py`**：添加 `test_gpt_oss_sm120_padding_layout_and_kernel` 函数，在 SM120 上构建模拟层并调用 Mxfp4MoEMethod 的权重处理，然后通过 FlashInfer `cutlass_fused_moe` 执行前向，验证结果与显式调用 FlashInfer 直接路径一致，确保 padding 和内核行为正确。

关键文件：
- `python/sglang/srt/layers/quantization/mxfp4.py`（模块 量化层；类别 source；类型 core-logic；符号 _process_weights_for_sm120_cutlass, _stack_up_gate_w13, _pad_w2_3d, _apply_sm120_cutlass）: 核心变更，新增 SM120 内核路径和权重处理函数，是功能实现的主文件。
- `test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py`（模块 集成测试；类别 test；类型 test-coverage；符号 test_gpt_oss_sm120_padding_layout_and_kernel）: 新增的集成测试，验证 padding 布局和内核在 SM120 上的正确性。
- `python/sglang/srt/arg_groups/overrides.py`（模块 参数覆盖；类别 source；类型 core-logic）: 修改 SM120 自动选择后端策略，使新内核成为默认。

关键符号：_process_weights_for_sm120_cutlass, _stack_up_gate_w13, _pad_w2_3d, _apply_sm120_cutlass

## 关键源码片段

### `python/sglang/srt/layers/quantization/mxfp4.py`

核心变更，新增 SM120 内核路径和权重处理函数，是功能实现的主文件。

```python
# Inside Mxfp4MoEMethod.__init__():
self._fi_kernel: Optional[str] = None
if self.use_flashinfer:
    if is_sm100_supported():
        self._fi_kernel = "trtllm_sm100"
    elif is_sm120_supported():
        # SM120 -> use FlashInfer CUTLASS MXFP8 x MXFP4 MoE kernel
        self._fi_kernel = "cutlass_sm120"
    elif is_sm90_supported():
        if not _FI_HAS_SM90_CUTLASS_MXFP4:
            raise RuntimeError(...)
        self._fi_kernel = "cutlass_sm90"
    else:
        raise NotImplementedError(
            "moe_runner_backend=flashinfer_mxfp4 requires SM90, SM100, or SM120."
        )

# Later in process_weights_after_loading():
if self._fi_kernel == "cutlass_sm120":
    self._process_weights_for_sm120_cutlass(layer)
    return

```

# 评论区精华

在 PR 审核中，b8zhong 询问 FlashInfer 是否已支持 SM90，如果支持是否可以完全删除 triton MoE 内核。mmangkad 回复 SM90 已支持，但由于 minimax m3（MXFP8）和 ROCm 平台仍依赖 triton 内核的某些部分，目前无法完全删除。这一讨论凸显了内核维护的复杂性：多架构、多精度格式共存时，不能简单用单一后端替代所有场景。

- FlashInfer 对 SM90 的支持与 triton 内核删除可行性 (design): mmangkad 回复 SM90 已支持，但由于 minimax m3（MXFP8）和 ROCm 平台仍依赖 triton 内核的某些部分，目前无法完全删除。

# 风险与影响

- 风险：
 1. **回归风险**：`overrides.py` 中的自动选择修改可能影响其他架构或非 GPT-OSS 模型的默认后端选择，需确保条件判断互斥且完备。
 2. **正确性风险**：新增的 `_process_weights_for_sm120_cutlass` 函数涉及复杂的权重重排和 padding，如果输入维度不满足内核约束（如 %128 != 0），可能导致静默错误或数值异常。单元测试覆盖了特定尺寸（hidden=160, intermediate=160），但真实模型维度可能触发未测试的边界条件。
 3. **依赖风险**：依赖 FlashInfer 版本，需确保 `cutlass_fused_moe` 支持 MXFP8×MXFP4 的 SM120 入口，否则抛出 NotImplementedError。
 4. **性能风险**：N/A（基准测试已显示积极收益）。
 - 影响：对用户：SM120 用户使用 GPT-OSS 模型时将自动获得 FlashInfer MXFP4 内核，无需任何配置更改，性能提升显著。对其他架构无影响。对系统：代码量增加约 330 行，主要集中在权重处理函数，未引入新的外部依赖。对团队：维护成本略有增加，但内核选择逻辑更清晰，且与 SM90 复用 padding 策略降低了长期维护负担。
 - 风险标记：量化层核心路径变更 , SM120 依赖硬件可用性 , 权重处理逻辑新增 , 依赖 FlashInfer 版本

# 关联脉络

- PR #32818 Route asymmetric-KV models to fa4 on SM100 and pin MiMoV2 FP8 MoE to flashinfer_trtllm: 同样修改了 overrides.py 中的架构特定后端路由逻辑，本 PR 新增 SM120 路由，形成对 SM100/SM120 的统一覆盖。