执行摘要

• 一句话：MiniMax-H3 SubBlock 稀疏注意力从 SM100 扩展到 SM90，Hopper 端到端提速最高 2×
• 推荐动作：值得精读。这是“同一路由方案、双 kernel 后端”的教科书式实现：重点看 (1) _get_subblock_sparse_attention_runner 的按设备一次性缓存分发，避免热路径开销；(2) _sm90_sparse_attention 中显式排序与 None full-block 的推理链条——为何排序是正确性的前提、为何 None 比空张量更优；(3) block_sparse_utils.py 的编译期守卫，区分“运行时值”与“编译期常量”是 CuTe DSL 调试的高频难点；(4) _tile_size_fwd_sm90 的 64×64 特例边界。Review 讨论本身也具有很高的学习价值，特别是对 kernel 内部消费顺序契约、compile key 隐式依赖的识别。

功能与动机

PR body 明确说明这是 #34148 的 follow-up：该 PR 为 MiniMax-H3 在 SM100 上引入了 SubBlock 稀疏注意力，本 PR 的目标是“extends the backend to SM90 GPUs using SGLang's CuTe-DSL block-sparse FlashAttention kernel without changing the existing SM100 path”。SM100 硬件成本高、可获得性受限，H100/H200 才是视频生成推理的主流部署机型，因此在 SM90 上复用仓库已有的 CuTe-DSL 稀疏 kernel 承载同一套 64×64 SubBlock 路由方案，可以在不改动 SM100 行为的前提下把稀疏加速的收益覆盖面显著扩大。

实现拆解

1. 统一分发入口（subblock_sparse_attn.py）：新增 _sm90_sparse_attention / _sm100_sparse_attention，分别封装 SM90 CuTe-DSL 稀疏 kernel 与 SM100 FlashInfer blk64 kernel；_get_subblock_sparse_attention_runner 用 functools.lru_cache(maxsize=None) 按 device 缓存架构判定（9.0 → SM90、10.0 → SM100、其余抛 RuntimeError）；_run_subblock_sparse_attention 成为 SubBlockSparseAttentionImpl._sparse_attention 的唯一出口，替换原先硬编码 load_bsa_attn_blk64_fwd 的 SM100 逻辑。
2. SM90 adapter 的排序语义（_sm90_sparse_attention）：显式对 q2k_block_index 做 sort(dim=-1) 升序排序。原因是 router 的公共契约允许 block id 任意顺序，而 SM90 稀疏消费者从最高槽位向下消费，且 mask_seqlen 只作用于第一个被消费的 block；升序排序保证最大 block id（可能的 ragged tail）永远落在被掩码的最高槽位。BlockSparseTensorsTorch 的 full_block_cnt/full_block_idx 传 None，因为 SubBlock 方案没有 always-dense block，mask-only 形式让 full-block 分支在编译期整段消除。
3. SM90 kernel 配套（interface.py / block_sparse_utils.py）：_tile_size_fwd_sm90 新增 sparse_block_size_kv 参数，并在 head_dim == 128 且 Q/K sparse block 均为 64 时返回 FwdConfig(64, 64, True, True)（num_wg_mma == 1，本仓库首次编译运行的配置）；head_dim=96 等其他情况保持原 128/192 tile 选择。block_sparse_utils.py 的 producer/consumer 增加 const_expr(blocksparse_tensors.full_block_cnt is None) 守卫，因为 mask_empty 是运行时值、CuTe 仍会编译两条动态分支，不特判会在 full-list 一侧对 None 下标而编译失败。
4. 平台 resolver（platforms/cuda.py）：_SubBlockSparseAttentionBackendResolver 从 required_capability = (10, 0) 改为 supported_capabilities = {(9, 0), (10, 0)}，并按架构分别校验依赖：SM90 检查 CuTe-DSL/Quack 导入，SM100 检查 FlashInfer blk64；10.3/12.x 仍保持 fail-closed。capability 先转 tuple 再查集合，避免对 None 取 major 的隐患。
5. 测试与文档：test/registered/cpu/test_subblock_sparse_attention.py（新增）验证分发缓存只解析一次、SM100 分发、10.3 拒绝；test_subblock_sparse_sm90.py（新增）验证 64×64 tile 选择仅在 head_dim=128 生效；test_subblock_sparse_attention.py 把 requires_sm100 泛化为 requires_subblock_kernel，SM90/SM100 数值测试共用，并新增 adapter 排序断言测试与 test_unsorted_ragged_tail_oversubscribes_sms（516 个 query-head tile、跨多波次、故意乱序）以及随机路由负对照；原仅在 SM100 CI 上跑的回归测试现在会落到 1-gpu-h100。README 与模块 docstring 同步为双路径描述。

关键文件：

• python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 SubBlockSparseAttentionBackend, _load_sm90_block_sparse_attention, _sm90_sparse_attention, _sm100_sparse_attention）: PR 核心：新增 SM90/SM100 双路径分发与 SM90 adapter。_get_subblock_sparse_attention_runner 按计算能力 9.0/10.0 每设备缓存一次选择 kernel；_sm90_sparse_attention 完成 block id 显式排序与 mask-only BlockSparseTensorsTorch 构造，是 SM90 正确性的关键逻辑所在。
• python/sglang/kernels/ops/attention/flash_attn/cute/interface.py（模块 SM90 内核；类别 infra；类型 core-logic；符号 _tile_size_fwd_sm90, _flash_attn_fwd）: SM90 tile 选择的关键改动：_tile_size_fwd_sm90 新增 sparse_block_size_kv 参数，并在 head_dim == 128 的 64×64 sparse block 场景返回 FwdConfig(64, 64, True, True)，这是该配置首次被编译运行；同时修复 _flash_attn_fwd 中 sparse_kv 变量遮蔽 compile key 的风险。
• python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py（模块 SM90 内核；类别 infra；类型 core-logic；符号 produce_block_sparse_loads, consume_block_sparse_loads）: mask-only 稀疏模式能否编译运行的关键：producer/consumer 需要编译期守卫，避免 full-block 列表为 None 时 CuTe 对 None 下标导致编译失败；守卫键统一为 full_block_cnt 以保证与上游一致性。
• python/sglang/multimodal_gen/runtime/platforms/cuda.py（模块 设备解析；类别 source；类型 configuration；符号 _SubBlockSparseAttentionBackendResolver）: 启动期 fail-closed 策略的调整点：resolver 从仅接受 10.0 改为接受 9.0/10.0，并按架构分别校验 CuTe-DSL 或 FlashInfer 依赖，10.3/12.x 仍拒绝。这是用户能否在 H100/H200 上启动 subblock_sparse_attn 的入口。
• python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py（模块 单元测试；类别 test；类型 test-coverage；符号 _subblock_kernel_available, test_sm90_adapter_sorts_indices_and_uses_64x64_blocks, test_unsorted_ragged_tail_oversubscribes_sms, _FakeBlockSparseTensors）: 测试主战场：requires_sm100 泛化为 requires_subblock_kernel 让原 SM100 专用数值测试（含 ragged tail 复现、随机路由负对照）落到 H100 CI；新增 adapter 排序断言测试与乱序超订多波次测试，把排序失效从静默错误变成可观测失败。
• test/registered/cpu/test_subblock_sparse_attention.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestSubBlockSparseAttentionDispatch, test_dispatch_is_resolved_once_per_device, test_dispatches_sm100, test_rejects_unsupported_compute_capability）: 新增的 CPU 注册测试，验证 _get_subblock_sparse_attention_runner 的缓存语义（每设备只解析一次）、SM100 分发与 10.3 拒绝路径，把分发正确性变成无 GPU 依赖的确定性回归。
• test/registered/kernels/ops/attention/test_subblock_sparse_sm90.py（模块 内核测试；类别 test；类型 test-coverage；符号 TestSubBlockSparseSM90, test_64x64_routing_mask_uses_matching_compute_tile, test_64x64_special_case_is_limited_to_head_dim_128）: 新增的 SM90 kernel 回归测试，固定 64×64 tile 选择只在 head_dim=128 生效、head_dim=96 仍保持 128×128 的既有结论，防止未来 tile 选择逻辑回归。
• python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md（模块 文档；类别 docs；类型 documentation）: 文档同步：从“10.0 only”改为“9.0 or 10.0”，明确 H100/H200 用 CuTe-DSL、B200/GB200 用 FlashInfer，并保留 10.3/12.x fail-closed 的说明。

关键符号：_run_subblock_sparse_attention, _get_subblock_sparse_attention_runner, _sm90_sparse_attention, _sm100_sparse_attention, _load_sm90_block_sparse_attention, _tile_size_fwd_sm90, produce_block_sparse_loads, consume_block_sparse_loads, _SubBlockSparseAttentionBackendResolver.resolve

关键源码片段

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

PR 核心：新增 SM90/SM100 双路径分发与 SM90 adapter。_get_subblock_sparse_attention_runner 按计算能力 9.0/10.0 每设备缓存一次选择 kernel；_sm90_sparse_attention 完成 block id 显式排序与 mask-only BlockSparseTensorsTorch 构造，是 SM90 正确性的关键逻辑所在。

# 按 CUDA 设备做一次架构判定并缓存结果，避免每次前向都走 capability 查询。
@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 # Hopper 走 SGLang CuTe-DSL kernel
    if capability == (10, 0):
        return _sm100_sparse_attention # Blackwell 走 FlashInfer blk64 kernel
    raise RuntimeError(
        "SubBlock sparse attention supports compute capability 9.0 or 10.0; "
        f"this tensor is on a {capability[0]}.{capability[1]} device."
    )
​
​
def _run_subblock_sparse_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    q2k_block_index: torch.Tensor, topk: int, softmax_scale: float,
) -> torch.Tensor:
    """把同一份 64x64 路由方案分发给 Hopper 或 Blackwell 的实现。"""
    runner = _get_subblock_sparse_attention_runner(q.device)
    return runner(q, k, v, q2k_block_index, topk, softmax_scale)
​
​
def _sm90_sparse_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    q2k_block_index: torch.Tensor, topk: int, softmax_scale: float,
) -> torch.Tensor:
    """用 SGLang 已有的 SM90 CuTe 稀疏 kernel 执行 SubBlock 路由方案。"""
    BlockSparseTensorsTorch, flash_attn_func = _load_sm90_block_sparse_attention()
​
    # router 契约允许 block id 任意顺序，而 SM90 稀疏消费者从最高槽位向下消费，
    # 且 mask_seqlen 只作用于第一个被消费的 block。显式升序排序保证可能的 ragged
    # tail（最大 block id）永远落在被掩码的最高槽位，不依赖 fused top-k kernel
    # 当前恰好升序输出的内部行为。
    ordered_index = q2k_block_index.sort(dim=-1).values
    block_counts = torch.full(
        ordered_index.shape[:-1], topk, dtype=torch.int32, device=ordered_index.device,
    )
    sparse_tensors = BlockSparseTensorsTorch(
        mask_block_cnt=block_counts,
        mask_block_idx=ordered_index,
        # SubBlock 方案没有 always-dense block，full 侧两个张量传 None。
        # block_sparse_broadcast_pattern 会把 None 原样保留在 broadcast 模式里，
        # 而该模式参与 compile key，因此 mask-only 与 mask+full 调用不会复用
        # 同一个已编译 kernel——这就是传 None 安全的依据。
        full_block_cnt=None,
        full_block_idx=None,
        block_size=(SUBBLOCK_SPARSE_BLOCK_SIZE, SUBBLOCK_SPARSE_BLOCK_SIZE),
    )
    out, _ = flash_attn_func(
        q, k, v,
        softmax_scale=softmax_scale,
        causal=False,
        num_splits=1,
        block_sparse_tensors=sparse_tensors,
    )
    return out

python/sglang/kernels/ops/attention/flash_attn/cute/interface.py

SM90 tile 选择的关键改动：_tile_size_fwd_sm90 新增 sparse_block_size_kv 参数，并在 head_dim == 128 的 64×64 sparse block 场景返回 FwdConfig(64, 64, True, True)，这是该配置首次被编译运行；同时修复 _flash_attn_fwd 中 sparse_kv 变量遮蔽 compile key 的风险。

def _tile_size_fwd_sm90(
    head_dim, head_dim_v, is_causal, is_local,
    sparse_block_size_q=None, sparse_block_size_kv=None,
):
    """SM90 forward 的 tile 选择；出现 sparse block 时计算 tile 必须贴合 mask 两个轴。"""
    # SubBlock 场景特例：每个 64 行 query block 拥有自己独立路由的 64 行 KV
    # block 列表，因此不能像常规 sparse mask 那样粗化成 128x128 tile——
    # 那样会把某一行 query 的 mask 错误套到相邻行上。注意 FwdConfig(64, 64, True, True)
    # 对应 num_wg_mma == 1，是本仓库首次被编译运行的配置，head_dim == 128 是
    # 额外硬约束，避免让从未使用 tile_m=64 的 192/256 head_dim 静默继承该配置。
    if (
        head_dim == 128
        and sparse_block_size_q == 64
        and sparse_block_size_kv == 64
    ):
        return FwdConfig(64, 64, True, True)
    # 其余 sparse mask 只要求 tile_m 整除 sparse_block_size_q；head_dim <= 96
    # 时优先用 192，不兼容再回退 128（沿用 C++ tile_size.h 的既有策略）。
    if head_dim <= 64:
        ...

python/sglang/kernels/ops/attention/flash_attn/cute/block_sparse_utils.py

mask-only 稀疏模式能否编译运行的关键：producer/consumer 需要编译期守卫，避免 full-block 列表为 None 时 CuTe 对 None 下标导致编译失败；守卫键统一为 full_block_cnt 以保证与上游一致性。

# produce_block_sparse_loads 中的 mask-only 特判分支。normalize 保证 full 的
# count 与 index 要么同时存在、要么同时为 None；mask_empty 是运行时值，CuTe
# 仍会编译两条动态分支，因此若不在此处特判，full-list 一侧会对 None 做下标
# 运算而直接编译失败——这里不是性能优化，而是可编译性的前提。
if const_expr(blocksparse_tensors.full_block_cnt is None):
    kv_producer_state = load_block_list(
        curr_mask_block_idx,
        mask_begin,
        mask_end,
        first_block_preloaded=False,
        kv_producer_state=kv_producer_state,
        load_K=load_K,
        load_V=load_V,
        pipeline_k=pipeline_k,
        pipeline_v=pipeline_v,
        intra_wg_overlap=intra_wg_overlap,
    )
    if const_expr(intra_wg_overlap) and not mask_empty:
        kv_producer_state = finish_overlap_v_load(
            curr_mask_block_idx,
            mask_begin,
            mask_end,
            load_V,
            pipeline_v,
            kv_producer_state,
        )
    return kv_producer_state
​
# consumer 侧的守卫统一用 full_block_cnt 做编译期存在性信号，与
# get_curr_blocksparse_tensors 保持一致；normalize_block_sparse_tensors 已保证
# full 的 count/index 成对出现，所以两个字段不可能在 fixed-length 或 varlen
# 布局中出现分歧。
if (
    const_expr(blocksparse_tensors.full_block_cnt is not None)
    and split_full_block_cnt > 0
):
    full_n_block = curr_full_block_idx[full_end - 1]

评论区精华

核心 review 由 triple-mu 主导，整体设计获得确认：“On SM90 q_stage = 1，加上 block_sparsity.py 的两处检查，(64, 64) 是 64 行路由粒度下唯一能通过的 tile 组合；升序排序确实必要，因为 block_sparse_utils.py 从最高槽位向下消费，mask_seqlen 只作用于第一个被消费的 block。” 主要交锋点：

• sort 注释与实现不符：triple-mu 指出 _topk_kernel 实际已升序输出，sort 是防御性的而非纠正性的，建议保留 sort 但改写注释，避免依赖 Triton kernel 未文档化的内部行为；作者采纳并改写。
• full-block 传 None 而非空张量：triple-mu 论证 block_sparse_utils.py:73-79 显式支持 None，可让整个 full-block 分支编译期消除并省掉两次张量分配；作者采纳，并因此暴露并修复了 producer/consumer 对 None 的编译期访问。
• 变量遮蔽风险：_flash_attn_fwd 中新增 sparse_kv 局部变量与原布尔 sparse_kv（参与 compile key）同名，triple-mu 警告“900 行函数里一个名字两个含义”会静默污染缓存键；作者改名 sparse_block_size_kv。
• 64x64 特例作用域：triple-mu 指出特例分支在 head_dim 分支之前，会匹配到从未使用 tile_m=64 的 192/256 head_dim；作者加上 head_dim == 128 限制并补回归测试。
• 守卫键一致性：consumer 守卫原来挂在 curr_full_block_idx is not None，与上游 get_curr_blocksparse_tensors 的 full_block_cnt is not None 只在“同时置空”时成立；作者改为统一用 full_block_cnt，并确认 normalize_block_sparse_tensors 已保证 count/index 成对出现。

• 测试覆盖缺口：triple-mu 指出 test_varlen_routes_each_document 两个文档单波次即完成，排序失效会静默通过；作者用新的乱序超订测试补上了该场景。

  • SM90 adapter 排序的必要性与注释准确性 (correctness): 保留显式 sort，注释改为“router 公共契约允许任意顺序，SM90 consumer 要求最大 block id 落在最高槽位，排序不依赖 top-k kernel 的内部输出顺序”。
  • full-block 传 None 优于空张量及编译期裁剪 (design): full_block_cnt/full_block_idx 均传 None，并在 CuTe producer/consumer 增加 const_expr(... is None) 特判；后续评论澄清该守卫是编译前提而非性能优化。
  • sparse_kv 局部变量遮蔽 compile key 布尔值 (style): 作者重命名为 sparse_block_size_kv，规避歧义。
  • 64x64 tile 特例的作用域限制 (design): 作者加 head_dim == 128 硬约束，并新增 head_dim=96 回归测试确认仍走 128×128。
  • consumer 编译期守卫键统一为 full_block_cnt (correctness): 作者统一改为 blocksparse_tensors.full_block_cnt，并说明 normalize_block_sparse_tensors 已保证 count/index 成对出现。
  • varlen 测试无法暴露排序失效 (testing): 作者新增 test_unsorted_ragged_tail_oversubscribes_sms：516 个 query-head tile、跨多波次、故意把 ragged block 放在最低槽位，用零 Q/K + 单位 V 使输出幅度偏差无法被 assert_close 忽略。
  • 模块 docstring 与 README 遗漏 FlashInfer-only 描述 (documentation): 作者在 follow-up commit 8133e9f7 中更新 docstring 并补上 None 安全的依据注释（broadcast pattern 保留 None 且参与 compile key）。

风险与影响

• 风险：
  • 首次编译运行的新 kernel 配置：reviewer 明确警告 FwdConfig(64, 64, True, True) 对应 num_wg_mma == 1，此前 _tile_size_fwd_sm90 只产出 128/192 且无 caller 传 tile_mn，该配置在本仓库从未被编译运行过。作者通过让原有 test_full_budget_reproduces_dense、test_ragged_tail_reproduces_dense（8192+37，正是排序保护的目标路径）等数值测试落到 H100 CI 上缓解，但 64×64 tile 在真实负载下的 kernel 稳定性仍需关注。
  • mask-only 与 compile key 的隐式耦合：mask-only 调用不与 mask+full 复用编译 kernel，依赖 block_sparse_broadcast_pattern 把 None 保留进 compile key 的隐含行为，而 compile_key 本身没有 full-block 位；未来若有人“顺手”把 None 改成空张量，会导致错误 kernel 被静默复用。PR 内已用注释记录该依赖，但仍是结构性脆弱点。
  • SM90 消费顺序契约：排序正确性建立在“从高槽位向下消费 + 首个 block 被掩码”的 kernel 内部约定上，上游 kernel 若改消费方向，排序保护即失效；外部文档并未固化这一契约。
  • CI 未真正运行：作者与 reviewer 都明确指出 approval 时缺 run-ci 标签，84 个测试 job 全部跳过；BBuf 最后执行了 /rerun-failed-ci，但合并前完整 CI 结果未被确认。
  • 900 行大函数集中改动：interface.py 的 _flash_attn_fwd 是超长函数，本次新增 sparse_block_size_kv 读取与 tile 选择调用，变量遮蔽问题虽已修复，但后续维护仍需谨慎。
• 影响：
  • 用户侧：H100/H200 上部署 MiniMax-H3 的用户可直接使用 --attention-backend subblock_sparse_attn，端到端提速 1.26×-2.01×（长视频收益更大：15 s Ref2AV 达 2.01×），SSIM 0.82-0.98 的可接受精度损失换得显著吞吐提升。
  • 系统侧：sglang/kernels 的 SM90 CuTe-DSL block-sparse 路径新增 mask-only 模式与 64×64 tile 配置，所有使用该 kernel 的调用都会受影响；但 reviewer 确认“(64, 64) mask 在 SM90 上此前会在 normalize_block_sparse_config 直接报错，没有既有 caller 受影响”，因此向后兼容性有保障。
  • 团队侧：SM100 路径逐字保留，双架构共享同一套 SubBlock router 与调度配置，后续架构扩展（如 SM100a 变体）只需在 _get_subblock_sparse_attention_runner 增加分支；此前被跳过的 SM90 数值测试现在进入 H100 CI，提升了回归覆盖质量。
  • 风险标记：首次编译 FwdConfig(64,64) 配置, mask-only compile key 隐式依赖, CI 未在合并前完整运行, SM90 消费顺序契约未固化, 900 行大函数集中改动

关联脉络

• PR #34148 引入 SM100 SubBlock 稀疏注意力（PR body 直接引用）: 本 PR 明确声明是其 follow-up：把同一 SubBlock 路由方案从 SM100 扩展到 SM90，SM100 路径逐字保留。
• PR #34581 [Diffusion] Optimizing MiniMax-H3 for consumer-level GPUs: INT8 Linear + pluggable DiT attention backends: 同一 MiniMax-H3 diffusion 优化线：引入可插拔 DiT attention backends 体系，本 PR 的 subblock_sparse_attn 正是该体系下新增的 SM90 支持。
• PR #35339 [diffusion] Per-request lossy accelerations: Cache-DiT, CFG gating, attention backend override: 后续演进：把有损加速改为按请求开关并支持 attention backend 覆盖，与 SubBlock 这类有损注意力共同构成 MiniMax-H3 感知质量/吞吐权衡的完整能力栈。