执行摘要

• 一句话：修复 SM120 设备上 MoE tile scatter 回退
• 推荐动作：该 PR 解决了一个明确的兼容性问题，但模块级别的 CUDA 上下文初始化存在潜在风险。建议后续提交一个惰性初始化的优化 PR 来解决该问题。整体上值得合入。

功能与动机

用户在 RTX Pro 6000 (SM120) 上运行 lmsys/gpt-oss-20b-bf16 模型时，使用 --moe-runner-backend triton_kernel 参数，遇到 tile scatter not supported for sm120 错误。此 PR 旨在通过回退到非持久化实现来解决该兼容性问题。

实现拆解

1. 新增导入与条件回退在 python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py 中，新增对 triton_kernels.matmul_ogs_details.opt_flags 中 update_opt_flags_constraints 的导入，以及从 sglang.srt.utils.common 导入 is_sm120_supported。
2. 模块级别检测：在模块加载时，通过 if is_sm120_supported() 判断当前设备的计算能力是否为 SM120。若为真，则调用 update_opt_flags_constraints({"is_persistent": False})，强制后续 Triton 内核使用非持久化的 gather/scatter 实现，以此绕过 SM120 对 tile scatter 的原生不支持。

关键文件：

• python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py（模块 MoE内核；类别 source；类型 dependency-wiring；符号 update_opt_flags_constraints, is_sm120_supported）: 唯一的变更文件，包含了修复 SM120 tile scatter 回退的全部逻辑。

关键符号：未识别

关键源码片段

python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py

唯一的变更文件，包含了修复 SM120 tile scatter 回退的全部逻辑。

# python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py
​
from triton_kernels.matmul_ogs_details.opt_flags import update_opt_flags_constraints
from sglang.srt.utils.common import is_sm120_supported
​
# 在模块加载时检查是否为 SM120 设备（如 RTX Pro 6000）
# 如果是，则强制关闭持久化 scatter，使用非持久化回退
if is_sm120_supported():
    # 注意：此处在模块级别调用会提前初始化 CUDA 上下文
    # 在多 GPU 环境中可能引发 VRAM 泄漏（详见 review 讨论）
    update_opt_flags_constraints({"is_persistent": False})

评论区精华

review 中 gemini-code-assist[bot] 提出了一个关键问题：在模块级别（import 时）调用 is_sm120_supported() 会提前初始化 CUDA 上下文（默认在 GPU 0 上），导致多 GPU 环境中 VRAM 泄漏和运行时错误。建议采用惰性初始化来避免。但该 PR 最终由 ch-wan 审核并合并，未采纳此建议。

• 模块级别 CUDA 上下文初始化 (performance): 未采纳该建议，PR 被合并。

风险与影响

• 风险：当前实现在模块加载时调用 is_sm120_supported()，会触发 torch.cuda.get_device_capability()，从而初始化 CUDA 上下文。在分布式多 GPU 环境中（如张量并行），所有工作进程在调用 torch.cuda.set_device() 之前就会初始化 GPU 0 的 CUDA 上下文，可能导致 VRAM 浪费或 OOM 错误。
• 影响：对用户的影响：SM120 GPU（如 RTX Pro 6000）用户现在可以正常使用 triton_kernel 后端运行 MoE 模型。影响范围限于该文件和相关 MoE 内核路径，改动量小（+6行），风险较低。
• 风险标记：模块级别 CUDA 上下文初始化, 多 GPU 环境潜在 VRAM 泄漏

关联脉络

• 暂无明显关联 PR