执行摘要

• 一句话：为AMD MI300/MI355 GPU优化稀疏注意力内核，提升性能。
• 推荐动作：建议精读此PR以学习AMD GPU内核优化策略，特别是条件内存分配和网格划分设计。关注安全漏洞修复的实现细节，以及性能基准测试方法。

功能与动机

根据PR body，动机是“Enable the faster/new tilelang kernel on MI300”和“Improve longer-context kernel performance on MI355”，旨在通过优化内核提升AMD GPU上的推理速度，特别是针对DeepSeek-V3.2等模型的长上下文场景。

实现拆解

实现主要集中在修改sparse_mla_fwd_decode_partial函数：1）添加inner_iter参数，将网格计算改为N_GROUPS = topk // (block_I * inner_iter)，以处理多个KV tile；2）根据inner_iter值条件分配Q缓冲区（inner_iter == 1时在共享内存，inner_iter > 1时在fragment），以重用Q提升占用率；3）添加安全索引逻辑处理负索引，避免越界读取；4）保留v1内核用于文档目的，作为设计权衡。

关键文件：

• python/sglang/srt/layers/attention/nsa/tilelang_kernel.py（模块 attention/nsa）: 这是唯一修改的文件，包含稀疏注意力内核的核心优化，直接影响AMD GPU性能。

关键符号：sparse_mla_fwd_decode_partial

评论区精华

Review讨论核心包括：1）gemini-code-assist[bot]指出安全漏洞，索引未验证可能导致GPU越界读取，作者回应添加安全索引逻辑；2）gemini-code-assist[bot]识别冗余计算循环，作者移除以优化性能；3）HaiShaw建议回滚v1内核删除以保留文档用途，作者执行。决策包括安全修复和性能优化，所有疑虑已解决。

• 安全漏洞：索引未验证导致越界读取 (security): 作者添加T.if_then_else逻辑处理负索引，漏洞已修复。
• 冗余计算循环优化 (performance): 作者移除冗余循环，代码简化并优化。
• 保留v1内核用于文档目的 (design): v1内核被保留，设计决策基于维护和文档考虑。

风险与影响

• 风险：技术风险：1）安全漏洞修复可能引入错误，如果安全索引逻辑不正确；2）内核参数变更（如inner_iter）可能影响现有配置的兼容性，但添加了断言验证；3）性能调优针对特定硬件（MI300/MI355），在其他环境可能未充分测试。风险局部于单个文件，回归测试通过基准和准确性验证。
• 影响：影响范围：用户端，AMD GPU用户将体验到推理速度提升（MI300高达2倍加速）；系统端，优化资源利用率，尤其长上下文场景；团队端，展示硬件特定内核优化模式，促进类似工作。影响程度中等，针对性强但限于稀疏注意力模块。
• 风险标记：核心路径变更, 安全漏洞修复, 性能调优

关联脉络

• PR #21188 [AMD] Add fused GemmaRMSNorm forward_hip to use aiter/vllm kernels for qwen3.5: 同属AMD GPU性能优化，涉及内核代码和JIT-kernel标签，展示硬件特定调优趋势。
• PR #20438 [Perf] Overlap NSA-CP key all-gather with query computation for DeepSeek-V3.2: 涉及NSA（稀疏注意力）性能优化，关联相同模块，揭示注意力内核演进方向。