执行摘要
本PR针对GLM-5模型在AMD平台（特别是TP8配置）下的Aiter稀疏注意力进行性能优化，通过填充头数以适配内核要求，并移除硬编码参数让内核内部优化。影响范围限于使用Aiter后端的AMD平台，能提升GLM-5推理速度，但需等待外部Aiter更新后方可合并。
功能与动机
优化GLM-5模型性能，解决在AMD平台TP8配置下（每设备头数=8）运行Aiter稀疏注意力时，因内核要求头数为16的倍数而导致的性能问题。PR body中明确目标为“Optimize GLM-5 model performance”，并通过填充和参数简化来适配内核约束。
实现拆解
主要修改集中在两个文件：
-
python/sglang/srt/layers/attention/nsa_backend.py：
- 在
__init__中添加need_pad_heads和head_repeat_factor逻辑，判断是否需要填充。
- 在
_forward_aiter和_forward_aiter_extend中实现张量填充与还原：
python
if self.need_pad_heads:
q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim).repeat_interleave(self.head_repeat_factor, dim=1)
o_kernel = q.new_empty((q.shape[0], layer.tp_q_head_num * self.head_repeat_factor, layer.v_head_dim))
调用mla_decode_fwd后，通过o = o_kernel[:, :: self.head_repeat_factor, :]还原输出。
-
python/sglang/srt/layers/attention/nsa/nsa_indexer.py：
- 删除
deepgemm.fp8_paged_mqa_logits调用中的硬编码参数ChunkK=128, TotalCuCount=256, WavePerEU=5，让内核内部配置优化参数。
评论区精华
review讨论较少，但有两个关键点：
gemini-code-assist[bot]："The padding and unpadding of query and output tensors are handled correctly within the _forward_aiter and _forward_aiter_extend methods. The changes are well-explained by comments and appear to be correct and efficient for the stated motivation."
HaiShaw："Merge later, see https://github.com/sgl-project/sglang/issues/21302" 并说明等待ROCm/aiter#2213合并。
风险与影响
- 风险：填充操作引入额外张量开销，但为满足内核必要；移除硬编码参数依赖内核内部优化，需确保Aiter更新后配置最优。
- 影响：仅影响使用Aiter稀疏注意力的AMD平台，提升GLM-5 TP8配置性能，对其他模型或后端无影响。
关联脉络
- 与PR #21986（AMD平台简化API）相似，都涉及移除硬编码/白名单以提升维护性。
- 与PR #22361（Whisper批量编码）和#22567（Tokenizer消除O(n²)复制）同属性能优化类别，反映团队持续关注关键路径性能调优。
- 依赖外部Aiter更新（ROCm/aiter#2213），需跟踪issue #21302以确定合并时机。
参与讨论