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

2026-08-10

AI Summary

This work addresses the high communication overhead in tensor-parallel training caused by Query-Key normalization, which traditionally requires transmitting full hidden vectors. The authors propose SwiftQK, the first communication-efficient QK-Norm mechanism that exchanges only scalar normalization statistics instead of entire vectors. SwiftQK integrates a deadlock-free persistent GPU kernel that overlaps point-to-point reductions with element-wise computations and fuses the RMSNorm operator to minimize latency. Experimental results demonstrate substantial improvements: compared to an All-Gather baseline, SwiftQK reduces QK-Norm communication latency by 81.4%–93.9%, and achieves an average 29.5% reduction in time per output token (TPOT) during end-to-end inference, outperforming existing scalar aggregation approaches by 14.3%.

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.

Problem

• Query-Key Normalization
• Tensor Parallelism
• cross-GPU communication
• Large Language Models
• RMSNorm

Innovation

• Tensor Parallelism
• Query-Key Normalization
• Communication-Efficient
• RMSNorm
• Persistent Kernel