执行摘要

• 一句话：为 MiniMax-H3 的 SubBlock 稀疏注意力添加 SM120 GPU 支持。
• 推荐动作：该 PR 值得精读。它是一个结构清晰、风险可控的硬件功能扩展，展示了良好的软件工程实践：严格的架构门控、全面的测试（调度、依赖、数值）、详尽的文档和性能基准。其设计模式（懒加载、适配器函数、平台解析器分支）是处理类似硬件适配任务的优秀参考。关注点在于对新 FlashInfer 内核的依赖管理和配置要求的明确传递。

功能与动机

SubBlock 稀疏注意力已支持 SM90（Hopper）和 SM100（数据中心 Blackwell），但缺少对计算能力 12.0 的 SM120 设备（如 RTX PRO 5000）的原生支持。为充分发挥新硬件性能并提供优化路径，需要添加 SM120 支持，利用其架构特定的 FlashInfer 内核实现更高的稀疏注意力加速。

实现拆解

变更遵循硬件扩展的标准模式，主要分为以下步骤：

1. 新增 SM120 内核加载器与适配器：在 router.py 中新增 load_bsa_attn_sm120_blk64_fwd 函数，通过懒加载方式引入 FlashInfer 的 bsa_attn_sm120_blk64_fwd 模块。在 subblock_sparse_attn.py 中新增 _sm120_sparse_attention 函数作为适配器，封装对加载器的调用，并添加一条一次性日志以确认 SM120 内核激活。
2. 扩展设备调度逻辑：在 subblock_sparse_attn.py 的设备调度函数 _get_subblock_sparse_attention_runner 中，增加对 (12, 0) 计算能力的判断，返回 _sm120_sparse_attention 适配器。同时更新相关文档字符串和错误提示信息，明确支持 9.0、10.0 和 12.0 三种架构。
3. 更新平台解析器与依赖验证：在 cuda.py 的 _SubBlockSparseAttentionBackendResolver 中，将 (12, 0) 加入 supported_capabilities 集合。其 resolve 方法被重构为根据计算能力（(9,0)，(10,0) 或其他）执行不同的依赖验证分支：对于 SM120，它调用 load_bsa_attn_sm120_blk64_fwd 来验证 FlashInfer 的 SM120 内核依赖是否就绪，并提供针对性的错误信息。
4. 更新公开接口与文档：在 subblock_sparse/__init__.py 中导出新的 load_bsa_attn_sm120_blk64_fwd 函数。更新了主模块和 README.md 文档，说明 SM120 的硬件要求、内核来源以及对 text_encoder 组件后端配置的特殊要求（在 SM120 上需使用 torch_sdpa）。
5. 补充全面测试覆盖：在 test_subblock_sparse_attention.py（CPU 测试）中，新增 test_dispatches_sm120、test_platform_resolver_loads_sm120_dependency 和 test_sm120_adapter_forwards_subblock_plan 等测试，覆盖调度、依赖加载和参数转发逻辑。在 test_subblock_sparse_attention.py（数值测试）中，新增 test_sm120_kernel_backed_mixed_query_mask 测试用例，并更新内核可用性检查函数 _subblock_kernel_available，以支持 SM120 跳过逻辑。现有的混合查询掩码测试也被扩展，覆盖了 SM120 路径。

关键文件：

• python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py（模块 稀疏注意力后端；类别 source；类型 core-logic；符号 _sm120_sparse_attention）: 核心变更文件，新增了 SM120 的稀疏注意力适配器函数 _sm120_sparse_attention 和设备调度逻辑，是功能实现的主体。
• python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py（模块 SubBlock路由器；类别 source；类型 entrypoint；符号 load_bsa_attn_sm120_blk64_fwd）: 新增 SM120 内核的懒加载函数 load_bsa_attn_sm120_blk64_fwd，是引入新硬件依赖的关键入口。
• python/sglang/multimodal_gen/runtime/platforms/cuda.py（模块 CUDA平台解析；类别 source；类型 dependency-wiring）: 平台解析器在此处扩展了 supported_capabilities 并重构了依赖验证逻辑，是将 SM120 纳入系统级后端选择流程的关键。
• test/registered/cpu/test_subblock_sparse_attention.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_dispatches_sm120, test_platform_resolver_loads_sm120_dependency, test_sm120_adapter_forwards_subblock_plan）: 核心测试文件，新增了三个关键的单元测试，覆盖了 SM120 的调度、平台解析依赖加载和适配器参数转发逻辑。

关键符号：_sm120_sparse_attention, _get_subblock_sparse_attention_runner, load_bsa_attn_sm120_blk64_fwd, _SubBlockSparseAttentionBackendResolver.resolve, test_dispatches_sm120, test_platform_resolver_loads_sm120_dependency, test_sm120_adapter_forwards_subblock_plan, test_sm120_kernel_backed_mixed_query_mask

关键源码片段

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py

核心变更文件，新增了 SM120 的稀疏注意力适配器函数 _sm120_sparse_attention 和设备调度逻辑，是功能实现的主体。

def _sm120_sparse_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q2k_block_index: torch.Tensor,
    topk: int,
    softmax_scale: float,
    block_counts: torch.Tensor | None = None,
) -> torch.Tensor:
    """Run a SubBlock routing plan through FlashInfer's SM120 kernel."""
    # 一次性日志，确认 SM120 内核被激活。
    logger.info_once("SubBlock sparse attention kernel active: FlashInfer SM120 blk64")
    # 调用新添加的加载器函数，获取 SM120 内核的入口点。
    out = load_bsa_attn_sm120_blk64_fwd()(
        q,
        k,
        v,
        q2k_block_index,
        topk,
        block_sizes=_cached_block_sizes(k.shape[1], k.device),
        q2k_block_nums=block_counts,
        softmax_scale=softmax_scale,
    )
    # FlashInfer 内核可能返回元组，这里统一处理。
    return out[0] if isinstance(out, tuple) else out
​
​
@functools.lru_cache(maxsize=None)
def _get_subblock_sparse_attention_runner(device: torch.device):
    """Resolve the architecture-specific kernel once per CUDA device."""
    capability = torch.cuda.get_device_capability(device)
    if capability == (9, 0):
        return _sm90_sparse_attention
    if capability == (10, 0):
        return _sm100_sparse_attention
    # 新增对 SM120 ( 计算能力 12.0) 的调度。
    if capability == (12, 0):
        return _sm120_sparse_attention
    # 错误信息已更新，明确支持三种架构。
    raise RuntimeError(
        "SubBlock sparse attention supports compute capability 9.0, 10.0, or 12.0; "
        f"this tensor is on a {capability[0]}.{capability[1]} device."
    )

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py

新增 SM120 内核的懒加载函数 load_bsa_attn_sm120_blk64_fwd，是引入新硬件依赖的关键入口。

@functools.lru_cache(maxsize=1)
def load_bsa_attn_sm120_blk64_fwd():
    """FlashInfer's CuTe-DSL 64-block sparse attention entry point for SM120."""
    try:
        # 尝试从 FlashInfer 中导入 SM120 专用的块稀疏注意力前向函数。
        from flashinfer.cute_dsl.sparse.bsa_attn_sm120 import (
            bsa_attn_sm120_blk64_fwd,
        )
    except Exception as exc:
        # 如果导入失败（如依赖未安装），抛出清晰的错误信息。
        raise ImportError(
            "SM120 SubBlock sparse attention requires FlashInfer's "
            "flashinfer.cute_dsl.sparse.bsa_attn_sm120 module"
        ) from exc
​
    return bsa_attn_sm120_blk64_fwd

评论区精华

审查讨论集中在变更的结构清晰度和风险可控性上。审阅者 BBuf 认可了变更的模块化设计，指出其清晰地遵循了现有的 SM90/SM100 模式，没有改变上层注意力契约。讨论中还引用了仓库历史中关于硬件能力严格门控、失败关闭依赖检查和真实硬件性能证据的模式，认为本 PR 符合这些良好实践。总体结论是变更结构良好，可以合并。

• SM120 硬件支持扩展的整体审查 (design): 变更结构良好，遵循了仓库现有的硬件扩展模式和安全检查实践（如失败关闭），可以合并。

风险与影响

• 风险：技术风险：
  1. 回归风险：变更局限于 SubBlock 稀疏注意力子系统，不修改其他模块，但新内核 (flashinfer.cute_dsl.sparse.bsa_attn_sm120) 的正确性依赖于外部依赖版本 (flashinfer-python==0.6.18rc2+g68693c2, nvidia-cutlass-dsl==4.7.0)。PR 中已通过位精确比较和 FP32 母版测试验证了内核正确性。
  2. 兼容性风险：变更严格限定了支持的计算能力（9.0、10.0、12.0），其他能力（如 10.3）会在启动时被正确拒绝，符合“失败关闭”原则，降低了意外运行在不支持硬件上的风险。
  3. 配置风险：文档明确指出在 SM120 设备上，文本编码器的注意力后端需配置为 torch_sdpa，而非 fa，这是一个用户需要遵循的配置要求，已通过文档和平台解析器（_FlashAttentionBackendResolver 会为 SM12.x 回退至 TORCH_SDPA）进行了引导。
  4. 测试风险：SM120 内核测试在非目标硬件上会被跳过，依赖 CI 环境（如 NVIDIA RTX PRO 5000）的可用性。PR 提供了详细的性能数据表明其已在目标硬件上验证。
• 影响：影响范围与程度：
  1. 用户影响：直接受益者是在 SM120 GPU（如 NVIDIA RTX PRO 5000）上运行 MiniMax-H3 等 DiT 模型进行视频生成的用户。该 PR 为其提供了显著的性能提升（延迟降低 18.23%-45.89%），是功能性的解锁。对 SM90/SM100 用户无负面影响。
  2. 系统影响：扩展了 SGLang 在 multimodal_gen 子系统中对 NVIDIA 新硬件架构的覆盖，遵循了既有的硬件适配模式，系统集成度高，未引入新的外部依赖（仅利用了已有的 FlashInfer 依赖）。
  3. 团队影响：为未来支持其他 Blackwell 架构变体或类似硬件扩展提供了清晰的代码模式和测试范例。
     - 风险标记：外部内核依赖, 硬件特定配置

关联脉络

• PR #37266 [diffusion] MiniMax-H3: tiered AdaLN plan cache (pinned-host tier + per-plan LRU): 同属 diffusion 模块下对 MiniMax-H3 模型的性能优化，#37266 优化了 AdaLN 计划缓存，而本 PR 优化了注意力计算内核，两者共同提升模型推理效率。
• PR #36735 [multimodal_gen] feat: support key masks on USPAttention's replicated-prefix path: 同属 multimodal_gen 模块的注意力层修复与增强，两者都涉及注意力机制的正确性或性能优化，属于同一功能线的不同方面。