SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization
2026-08-10 • Machine Learning
Machine LearningDistributed, Parallel, and Cluster Computing
AI summaryⓘ
The authors improved a step called Query-Key Normalization that helps large language models learn better. Normally, when using multiple GPUs, this step needs a lot of communication between GPUs, which slows things down. They created SwiftQK, which reduces the communication by only sharing small numbers instead of large vectors and overlaps tasks to avoid delays. Their tests show SwiftQK makes this process much faster, speeding up the overall model running time compared to older methods.
Query-Key NormalizationLarge Language ModelsTensor ParallelismGPU CommunicationRMSNormAll-GatherPeer-to-Peer ReductionLatencyMulti-GPU Kernels
Authors
Gyudong Kim, Wonjun Han, Young Geun Kim
Abstract
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.