Kwai Summary Attention Technical Report
About
Long-context ability, has become one of the most important iteration direction of next-generation Large Language Models, 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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Code | HumanEval | HumanEval Accuracy31.71 | 118 | |
| Mathematics | GSM8K | GSM8K Score81.09 | 87 | |
| General Knowledge | MMLU-Pro | MMLU-Pro General Knowledge Score45.7 | 55 | |
| General Knowledge | CMMLU | Accuracy73.29 | 50 | |
| Long-context retrieval | RULER 16k | Score88.86 | 28 | |
| General Knowledge | MMLU | General Score70.73 | 25 | |
| Coding | MBPP | Score36.4 | 23 | |
| Math | CMath | Score84.58 | 22 | |
| Long-context retrieval | RULER 64K context | Accuracy76.09 | 19 | |
| General Knowledge | CEval | Score72.66 | 19 |