执行摘要

• 一句话：修复 GB200 上 Triton attention 寄存器耗尽问题
• 推荐动作：建议合并以修复 GB200 上的崩溃问题。注意后续可能需要对 sm_120a 和其他 Blackwell 变体进行调整，并考虑是否优化 Lq <= 128 的配置。开发者可关注 flashinfer PR#2959 的进展。

功能与动机

Gemma4 NVFP4 检查点在 GB200 上运行失败，错误为 'Register allocation failed with register count of 255'。根本原因是 _get_block_sizes_for_extend_attention 没有处理 sm_100a，误用了 Hopper 的块大小，加上 fp8 量化导致寄存器压力超出限制。

实现拆解

1. 在 extend_attention.py 的 _get_block_sizes_for_extend_attention 函数中，在 Hopper 分支（CUDA_CAPABILITY[0] >= 9）之前插入针对 Blackwell 数据中心架构（sm_100a）的分支。
2. 对于 Lq > 256 的情况，设置 BLOCK_M=16, BLOCK_N=64（原 Hopper 配置为 BLOCK_M=32, BLOCK_N=64），减少寄存器占用。
3. 对于 Lq <= 256 的情况，保持 BLOCK_M=64, BLOCK_N=64，与 Hopper 相同。
4. 此分支仅在 _is_cuda 且 CUDA_CAPABILITY[0] == 10 时触发，不影响其他架构。

关键文件：

• python/sglang/srt/layers/attention/triton_ops/extend_attention.py（模块 注意力层；类别 source；类型 core-logic；符号 _get_block_sizes_for_extend_attention）: 唯一修改的文件，为 Blackwell sm_100a 添加专用分支以避免寄存器耗尽。

关键符号：_get_block_sizes_for_extend_attention

关键源码片段

python/sglang/srt/layers/attention/triton_ops/extend_attention.py

唯一修改的文件，为 Blackwell sm_100a 添加专用分支以避免寄存器耗尽。

# python/sglang/srt/layers/attention/triton_ops/extend_attention.py
​
def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
    """根据查询长度和值长度选择块大小。
​
    新增 sm_100a (Blackwell) 分支，使用较小 tile 避免寄存器溢出。
    """
    # ... 前面的代码 ...
    elif _is_cuda and CUDA_CAPABILITY[0] == 10:
        # Blackwell data-center architecture (GB200, B200, sm_100a)
        # sm_100a 的寄存器约束与 Hopper 不同，Hopper 的块大小
        # 在大 head dim (Lq=512) 时会导致 PTX 寄存器耗尽 (>255 regs)。
        if Lq <= 256:
            BLOCK_M, BLOCK_N = (64, 64)
        else:
            BLOCK_M, BLOCK_N = (16, 64) # 使用更小的 BLOCK_M 减少寄存器压力
    elif _is_cuda and CUDA_CAPABILITY[0] >= 9:
        # Hopper architecture (H100, etc.)
        if Lq <= 256:
            BLOCK_M, BLOCK_N = (64, 64)
        else:
            BLOCK_M, BLOCK_N = (32, 64)
    # ... 后续代码 ...

评论区精华

Review 中 alexnails 询问是否也应处理 Lq <= 128 的情况（建议 BLOCK_M=128），但未得到回复。Issue 评论中 jeremylea 询问是否处理 sm_120a（RTX 6000），但该 PR 仅针对 sm_100a。baoskee 报告在旧容器中仍存在问题，nvpohanh 建议使用最新 dev-cu13 容器。

• Lq <= 128 的块大小配置 (design): 未处理，PR 保持当前实现。
• sm_120a 兼容性 (question): 未涉及，仅针对 sm_100a。
• 容器兼容性 (question): 鼓励使用新容器，PR 本身无法解决容器旧问题。

风险与影响

• 风险：该修改仅影响 CUDA capability 10（Blackwell 数据中心架构）的块大小选择，其他架构行为不变。风险极低，但未处理 Lq <= 128 的优化可能潜在性能损失。另需注意，该修复假定 _is_cuda 和 CUDA_CAPABILITY 全局变量可用，之前的代码已有类似假设，因此无新增风险。
• 影响：直接影响在 GB200/B200 上使用 Triton attention 后端的用户，尤其是运行 Gemma4 NVFP4 或明确设置 kv_cache_dtype=fp8_e4m3 的模型。修复后避免崩溃，但可能因 tile 变小导致预填充阶段吞吐降低。对其他硬件无影响。
• 风险标记：仅修改配置逻辑, 新增架构分支, 未处理 Lq≤128 优化

关联脉络

• PR #21952 Previous NVFP4 attempt (依赖的基础): 本 PR 基于 #21952 修改，依赖其引入的 NVFP4 支持。