Kwai Summary Attention Technical Report

Abstract: Long-context ability, has become one of the most important iteration direction of next-generation LLMs, particularly in semantic understanding/reasoning, code agentic intelligence and recommendation system. However, the standard softmax attention exhibits quadratic time complexity with respect to sequence length. As the sequence length increases, this incurs substantial overhead in long-context settings, leading the training and inference costs of extremely long sequences deteriorate rapidly. Existing solutions mitigate this issue through two technique routings: i) Reducing the KV cache per layer, such as from the head-level compression GQA, and the embedding dimension-level compression MLA, but the KV cache remains linearly dependent on the sequence length at a 1:1 ratio. ii) Interleaving with KV Cache friendly architecture, such as local attention SWA, linear kernel GDN, but often involve trade-offs among KV Cache and long-context modeling effectiveness. Besides the two technique routings, we argue that there exists an intermediate path not well explored: {Maintaining a linear relationship between the KV cache and sequence length, but performing semantic-level compression through a specific ratio $k$}. This $O(n/k)$ path does not pursue a ``minimum KV cache'', but rather trades acceptable memory costs for complete, referential, and interpretable retention of long distant dependency. Motivated by this, we propose Kwai Summary Attention (KSA), a novel attention mechanism that reduces sequence modeling cost by compressing historical contexts into learnable summary tokens.

• The paper introduces a hybrid attention mechanism that uses learnable summary tokens to semantically compress long contexts, preserving interpretability and high retrieval fidelity.
• It demonstrates superior long-context retrieval and convergence efficiency on benchmarks like RULER, MMLU, and GSM8K compared to full attention models.
• The approach reduces KV cache size by up to 2.5× and can be combined with other methods, paving the way for scalable, efficient LLM deployment.

Kwai Summary Attention: An Efficient and Expressive Mechanism for Long-Context LLMs

Introduction

Kwai Summary Attention (KSA), introduced in the "Kwai Summary Attention Technical Report" (2604.24432), proposes a structured attention mechanism designed for efficient and effective long-context modeling in LLMs. The work addresses the core issue of quadratic inference and training cost in Transformer-based models, specifically the linear growth of key-value (KV) cache and quadratic attention computation with respect to sequence length. Existing approaches either compress KV caches (e.g., GQA, MLA) or interleave full, local, and linear attention mechanisms but suffer from either maintainable but suboptimal scaling or from lossy context representation. KSA strategically compresses the sequence at a semantic level while maintaining interpretability and high retrieval fidelity.

Methodology

KSA introduces learnable summary tokens into the input sequence at fixed intervals, dividing it into non-overlapping chunks. Each summary token is responsible for semantically distilling and representing the content of its chunk. The resulting architecture enables an efficient combination: local token interactions are modeled through sliding chunk attention (SCA), while long-range dependencies are captured via the sequence of summary tokens.

Figure 2: The input sequence is partitioned into chunks with learnable summary tokens injected at boundaries, facilitating chunk-wise semantic compression.

The model's attention visibility is designed to guarantee clean and non-overlapping access to context. Summary tokens attend only to the tokens in their own chunk, while text tokens attend to neighboring tokens in their local window and to summary tokens of previous chunks, excluding partial or ambiguous chunk coverage. This strict routing prevents information loss due to boundary crossing, a limitation in classical sliding window attention.

Figure 1: Sliding Window Attention may lose information at chunk boundaries, while Sliding Chunk Attention in KSA cleanly isolates chunk and summary information.

The KSA framework is implemented with a block-sparse attention kernel in training/prefill and a highly efficient, contiguous summary KV cache for decoding. The KV cache is structured so each step can load the relevant contiguous block of memory without any dynamic concatenation, maximizing throughput.

Figure 3: The organized contiguous buffer layout of KSA permits efficient, non-gather read operations during decoding.

KSA can be combined orthogonally with other compression methods, such as GQA and MLA, for compound improvements in KV cache scalability.

Figure 4: KSA provides a substantial reduction in KV cache size across growing sequence lengths compared to other mechanisms.

Experimental Results

The report presents comprehensive evaluations in both from-scratch and continual pre-training (CPT) regimes. Models are compared against strong baselines, including Hybrid-SWA, Hybrid-SCA, and Hybrid-GDN as well as Full attention. Key benchmarks include long-context retrieval (RULER up to 128K), general language understanding (MMLU, CMMLU, C-Eval, MMLU-Pro), mathematical reasoning (GSM8K, CMath, MATH), and code synthesis (MBPP, HumanEval).

KSA, particularly in its hybrid architecture (interleaving KSA and Full layers, typically at a 3:1 ratio), demonstrates:

• Superior or state-of-the-art long-context retrieval, outperforming or closely matching Full attention at 128K context length on RULER (e.g., 71.67 vs. 65.86 for KSA vs. Full in CPT).
• Competitive or better general task performance, exhibiting minimal or positive gaps relative to Full attention in general knowledge, math, and coding.
• Robust performance at extreme context lengths with only minor degradation, enabling reliable retrieval and reasoning over 128K-long inputs (Needle-in-a-Haystack tasks).
• Best convergence efficiency and stability during training.

  Figure 5: Hybrid-KSA demonstrates leading performance among efficient attention mechanisms on long-context and general benchmarks.

  

Figure 7: Hybrid-KSA achieves the fastest and most stable convergence among model variants in from-scratch training.

Inference Efficiency

KSA achieves significant KV cache reduction (up to $2.5\times$ over Full attention) and maintains competitive or higher decoding throughput compared to other state-of-the-art efficient mechanisms.

Figure 9: Hybrid-KSA delivers substantial KV cache reductions while maintaining or improving decode throughput.

Ablation and Analysis

Ablation experiments dissect the effects of summary chunk number, chunk size, and hybrid ratio between KSA and Full attention:

• Chunk size: Smaller sizes boost long-context retrieval, larger sizes favor general benchmarks.
• Hybrid ratio: Increasing summary layers improves context retrieval but may slightly reduce general performance; an intermediate value (e.g., 3:1) achieves the best balance.

Visualizations of per-layer attention reveal that Hybrid-KSA develops "comb-like" patterns, with shallow layers efficiently attending summary tokens and Full-attention layers acting as cross-chunk integrators, enhancing retrieval robustness.

Figure 6: Per-layer patterns show Hybrid-KSA fosters cross-chunk retrieval via structured summary attention and integration by Full layers.

Practical and Theoretical Implications

KSA offers a new axis of design: sequence-level semantic compression with interpretable retention, filling the gap between minimal cache approaches and unscalable full attention. Its orthogonality to other compression methods means it can be directly composed with existing state-of-the-art mechanisms. The block-sparse kernel and efficient buffer layout further address real-world inference constraints, making KSA practical for deployment in agentic, dialogue, and recommendation scenarios where long context, low latency, and high throughput are critical.

The approach empirically narrows or eliminates the expressivity-efficiency trade-off typically observed in efficient attention research and provides a scalable pathway to extend LLM memory windows without bandwidth-prohibitive costs.

Future Directions

Promising avenues for extension include learned sparse selection of summary tokens (efficient, query-aware retrieval at extreme lengths), post-training adaptation (task/prompt-optimized summaries), theoretical scaling analyses for compression ratios, and the integration of KSA-style memory into foundation models for recommendation and agentic behavior.

Conclusion

KSA demonstrates that structured, chunk-level semantic compression via summary tokens enables sublinear scaling of the KV cache while preserving or exceeding the long-context modeling of Full attention. The hybrid KSA architecture, block-sparse kernel, and efficient cache innovations together yield a practical and theoretically grounded long-context mechanism suitable for next-generation LLM applications requiring expansive context comprehension and efficient deployment.