viable/strict/1788331551: [ROCm][CK] Remove host-side FMHA sequence padding workaround (#195498)
PyTorch 在 ROCm 上移除了 FMHA 的 host-side sequence padding workaround，将原始逻辑张量直接传给 CK，并增加了对未对齐序列长度、dense bias、dropout replay 和 grouped-query attention 的算子覆盖。
发展脉络
- 首次出现viable/strict/1788331551: [ROCm][CK] Remove host-side FMHA sequence padding workaround (#195498)PyTorch Core
- 行业反馈trunk/3ee461fd8f0bc4b7b10c578a1677ee70c587fab0: [ROCm][CK] Remove host-side FMHA sequence padding workaround (#195498)PyTorch Core
- 当前判断PyTorch 对 ROCm 的持续优化表明 AMD GPU 在 AI 训练和推理中的地位日益重要。移除 workaround 并依赖 CK 的 tail-safe 实例，反映了开源生态在硬件适配上的协作加深。未来可观察 AMD 在软件栈上的投入是否加速其市场份额增长。Agent Pulse · 分析
PyTorch 在 ROCm 上移除了 FMHA 的 host-side sequence padding workaround。此前，gfx950 上未填充的 CK FMHA TRLOAD 实例在 tile 未对齐的 Q/K 序列长度上出错，因此 #187152 添加了 host-side pad/copy 的临时方案。现在，ROCm/rocm-libraries#11519 为生成的 dispatch 分支添加了保护，使得未对齐的请求会选择 CK 现有的 tail-safe kPadSeqLenQ/K 实例。因此，PyTorch CK FMHA wrapper 移除了外部的 Q/K/V 序列填充、O/LSE 临时分配和输出拷贝，直接将原始逻辑张量传给 CK。同时增加了对未对齐序列长度、dense bias、dropout replay 和 grouped-query attention 的算子覆盖。测试计划包括在 MI350X/gfx950 上运行相关测试。
该改动表明 CK 的 FMHA 实现已能通过生成的分发器处理未对齐的序列长度，无需在 host 端进行填充。这减少了内存分配和拷贝开销，可能提升推理性能。未来可关注 CK 是否进一步支持更多未对齐场景，以及 PyTorch 是否在其他后端采用类似策略。
PyTorch 对 ROCm 的持续优化表明 AMD GPU 在 AI 训练和推理中的地位日益重要。移除 workaround 并依赖 CK 的 tail-safe 实例，反映了开源生态在硬件适配上的协作加深。未来可观察 AMD 在软件栈上的投入是否加速其市场份额增长。
该改动降低了 ROCm 上 FMHA 的内存和计算开销，可能提升 AMD GPU 上 PyTorch 模型的推理效率，增强 AMD 在 AI 推理市场的竞争力。对于使用 AMD GPU 的云服务商和企业，这可能带来成本效益。
后续可关注该改动是否合入稳定版本，以及是否在更多 ROCm 架构上启用。同时，CK 库的更新可能带来更多性能优化，值得跟踪。