SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization

📄 arXiv: 2608.09160v1 📥 PDF

作者: Gyudong Kim, Wonjun Han, Young Geun Kim

分类: cs.LG, cs.DC

发布日期: 2026-08-10

💡 一句话要点

提出SwiftQK以解决QK-Norm在多GPU下的通信效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 张量并行 QK归一化 多GPU计算 大型语言模型 通信效率 训练优化 机器学习

📋 核心要点

  1. 现有的QK-Norm在张量并行环境下引入了额外的跨GPU通信，影响了训练效率。
  2. SwiftQK通过仅交换标量归一化统计量，重叠计算与通信，显著提高了效率。
  3. 实验结果显示，SwiftQK在QK-Norm延迟上减少了81.4%至93.9%，在端到端服务中TPOT平均减少29.5%。

📝 摘要（中文）

Query-Key Normalization (QK-Norm) 提升了现代大型语言模型（LLMs）的训练稳定性和质量。然而，在张量并行（TP）下，逐层的QK-Norm引入了额外的跨GPU通信，因为归一化因子依赖于完整的隐藏向量。本文提出了SwiftQK，这是一种多GPU RMSNorm内核，仅交换标量归一化统计量，并在死锁安全的持久内核中重叠剩余的点对点归约与独立的逐元素计算。对近期LLMs的评估表明，SwiftQK相较于标准TP QK-Norm使用全向量All-Gather，减少了81.4%至93.9%的QK-Norm延迟。在端到端服务中，SwiftQK平均减少了TPOT 29.5%，相较于基于All-Gather的基线，减少了14.3%，相较于优化的标量聚合实现。

🔬 方法详解

问题定义：本文解决的是在多GPU环境下，QK-Norm因依赖完整隐藏向量而导致的额外跨GPU通信问题，这影响了训练效率和速度。

核心思路：SwiftQK的核心思想是通过仅交换标量归一化统计量，避免全向量的通信开销，并在计算过程中重叠通信与计算，以提高整体效率。

技术框架：SwiftQK的整体架构包括一个多GPU RMSNorm内核，采用持久化的计算模式，确保在进行点对点归约时不会发生死锁，同时实现独立的逐元素计算。

关键创新：SwiftQK的主要创新在于其通信策略的优化，显著减少了QK-Norm的延迟，尤其是在多GPU设置下，突破了传统方法的瓶颈。

关键设计：在设计上，SwiftQK采用了优化的标量聚合方法，确保在计算和通信之间的高效重叠，具体的参数设置和网络结构细节在实验中进行了验证。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示，SwiftQK在QK-Norm延迟上减少了81.4%至93.9%，在端到端服务中TPOT平均减少29.5%。与基于All-Gather的基线相比，SwiftQK的性能提升显著，展示了其在多GPU环境下的高效性。

🎯 应用场景

SwiftQK的研究成果在大型语言模型的训练和推理中具有广泛的应用潜力，尤其是在需要高效并行计算的场景中，如自然语言处理、机器翻译和对话系统等。其优化的通信策略将有助于提升模型训练的速度和稳定性，推动相关领域的技术进步。

📄 摘要（原文）

Query-Key Normalization (QK-Norm) improves the training stability and quality of modern Large Language Models (LLMs). However, under Tensor Parallelism (TP), layerwise QK-Norm introduces additional cross-GPU communication because the normalization factor depends on the full hidden vector. We present SwiftQK, a multi-GPU RMSNorm kernel that exchanges only scalar normalization statistics and overlaps the remaining Peer-to-Peer reduction with independent element-wise computation in a deadlock-safe persistent kernel. Evaluations on recent LLMs show that SwiftQK reduces QK-Norm latency by 81.4--93.9% relative to the standard TP QK-Norm using full-vector All-Gather. In end-to-end serving, SwiftQK reduces TPOT on average by 29.5% over the All-Gather-based baseline and by 14.3% over an optimized scalar-aggregation implementation.