SwiftQK：用于查询键规范化的快速且通信高效的张量并行性

查询键标准化（QK-Norm）提高了现代 大语言模型（LLM）的训练稳定性和质量。然而，在张量并行（TP）下，分层 QK-Norm 引入了额外的跨 GPU 通信，因为归一化因子取决于完整的隐藏向量。我们提出了 SwiftQK，一个多 GPU RMSNorm 内核，它仅交换标量归一化统计数据，并在死锁安全的持久内核中将剩余的点对点缩减与独立的元素计算重叠。对最近 LLM 的评估表明，相对于使用全向量 All-Gather 的标准 TP QK-Norm，SwiftQK 将 QK-Norm 延迟降低了 81.4--93.9%。在端到端服务中，SwiftQK 与基于 All-Gather 的基线相比，平均将 TPOT 减少了 29.5%，与优化的标量聚合实现相比，平均减少了 14.3%。