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

arXiv: https://arxiv.org/abs/2608.09160 · HTML: https://mlsanity.com/paper/2608.09160/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 12 references that are in the map.
- First posted on arXiv 10 Aug 2026.

## Builds on (most cited first, 9 of 9)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [2 OLMo 2 Furious](https://mlsanity.com/paper/2501.00656/) (2024)
- [OLMoE: Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2409.02060/) (2024)
- [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024)
- [SimpleGPT: Improving GPT via A Simple Normalization Strategy](https://mlsanity.com/paper/2602.01212/) (2026)

## Abstract (arXiv)

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.