执行摘要

• 一句话：修复Gemma 4 NVFP4模型在GB200上Triton attention kernel因PTX寄存器耗尽导致的崩溃问题。
• 推荐动作：建议工程师精读此PR以了解Triton kernel硬件适配模式，关注块大小调优对寄存器压力的影响。设计决策中值得注意：为不同CUDA能力添加专用分支以避免寄存器耗尽，但可考虑扩展更细粒度优化以适应不同场景。

功能与动机

根据PR body，Gemma 4 NVFP4 checkpoints在GB200上使用Triton attention后端时崩溃，错误为'PTXAS error: Register allocation failed with register count of 255'。根本原因是_get_block_sizes_for_extend_attention函数缺少针对CUDA_CAPABILITY[0] == 10（Blackwell架构）的分支，导致使用Hopper的块大小配置时寄存器压力过大，特别是在KV缓存为fp8时加剧。

实现拆解

在python/sglang/srt/layers/attention/triton_ops/extend_attention.py的_get_block_sizes_for_extend_attention函数中，添加一个新的条件分支：当CUDA_CAPABILITY[0] == 10（Blackwell架构）时，根据查询长度Lq设置块大小：若Lq <= 256，BLOCK_M, BLOCK_N = (64, 64)；否则为(16, 64)。这替代了原先的Hopper分支，以适配sm_100a的寄存器约束，避免PTX寄存器耗尽。

关键文件：

• python/sglang/srt/layers/attention/triton_ops/extend_attention.py（模块 attention/triton_ops）: 核心修复文件，修改了Triton attention kernel的块大小选择逻辑以适配Blackwell架构，避免PTX寄存器耗尽，直接影响Gemma 4 NVFP4模型在GB200上的运行。

关键符号：_get_block_sizes_for_extend_attention

评论区精华

reviewer alexnails在代码第77行评论：'can you also include Lq <= 128 case? (e.g I believe 128x64, but it could be 128x128 if I am missing something from Blackwell tuning guide)'。此建议旨在进一步优化块大小选择以提升性能，但最终代码未采纳该修改，可能因修复优先级或测试覆盖不足，无其他深入讨论。

• 为Blackwell架构添加更优块大小分支 (design): 建议未采纳，代码保持原修改，可能因时间紧迫或已有配置足够。

风险与影响

• 风险：技术风险包括：1) 回归风险：新分支可能影响其他Blackwell模型或配置的性能，需测试覆盖；2) 兼容性风险：仅针对sm_100a，可能未覆盖其他Blackwell变体如sm_120a（如评论提及）；3) 性能风险：缺少更细粒度分支可能导致某些场景性能次优。
• 影响：影响范围：直接受益于Gemma 4 NVFP4模型在GB200上的用户，确保模型可运行，提升硬件兼容性。系统层面，修复了Triton attention后端在Blackwell架构上的一个崩溃bug，增强对新兴硬件的适配能力。团队需注意此硬件特定调优可能需后续迭代优化。
• 风险标记：硬件特定调优, 缺少细粒度优化, 潜在回归风险

关联脉络

• PR #21952 未知: PR body中提及基于此PR，但历史分析中未提供详细信息，可能为相关前置修复或依赖。
• PR #22323 [Lora] Lora quat info re-factor and support deepseekv3 mla lora: 共享quant标签，涉及量化相关优化，反映团队对量化模型的持续关注。