Summary-token architectures such as Kwai Summary Attention (KSA) compress distant context into learnable Summary token per chunk, and hybrid variants interleave these layers with full attention to retain exact access to the original text. In such hybrid KSA-Full architecture, the interleaving inserted summary tokens are served as a streaming flow role across different layers. Although effective, however, it also introduce some bottlenecks: (1) the different layers summary compression chunk must be same, therefore hard to reach a flexible KV Cache compression, e.g., different layer set different summary compression chunk window. (2) in the full-attention layers, the summary token introduced additional computation cost, but such layer ignore them and still pay quadratic prefilling, therefore unfriendly to the computation side. Motivated by them, we developed the KSA V2, which not only achieves a more flexible KV cache memory compression but also reduces the computation cost in Full attention layer. Specifically, we found that the summary tokens can serve two roles: the KV memory and sparse indexer.
Figures & tables
Figure 1: Left: long-context and general capabilities of 1.9B C-KSA+S-KSA and the reference models reported in KSA. Right: per token inference FLOPs and cache size ( 1GiB=230 bytes).
Figure 2: Architecture overview: from Hybrid KSA to KSA V2.
Figure 3: Compressed KSA with ks=4 and km=2 . Complete remote Summary groups are compressed into Meta KV, while boundary Summary KV and local text KV remain directly accessible.
Figure 4: Sparse KSA with ks=4 and b=2 .
Figure 5: C-KSA cache layout and decoding updates with ks=2 , W=4 , and km=4 . Gray ( T ), orange ( S ), and purple ( M ) denote text, Summary, and Meta KV. Blue marks indicate reads or writes.
Mechanism
Effective Context
KV Cache Size ( N→∞ )
Compression Rate
MHA
N (exactly)
2⋅N⋅H⋅dh
—
GQA
N (exactly)
2⋅N⋅HKV⋅dh
HHKV≈6.25%
MLA
N (exactly)
N⋅(dc+dr)
2Hdhdc+dr≈1.76%
GDN
N (ambiguous)
H⋅dh2
2Ndh≈0%
SWA
w (exactly)
2⋅w⋅HKV⋅dh
Nw⋅HHKV≈0%
KSA
N (summarized)
2⋅ksN⋅H⋅dh
ks1≈12.5%
Table 1: KV-cache cost comparison for a sequence of N text tokens. Example settings use ks=8 , km=128 , H=128 , dh=128 , HKV=8 , dc=512 , and dr=64 . Compression rates are relative to MHA.
Figure 6: KV cache size comparison across different mechanisms as sequence length grows. The inset uses a logarithmic vertical axis to separate the low-cache curves and labels C-KSA values at 128K.
Figure 7: Kernel design for C-KSA and S-KSA: (a) direct KV addressing, (b) fused block scoring and selection under a K -token budget, and (c) shared tiled prefill with KV reuse.
Configuration
1.9B
4B
Shared backbone
Number of layers
24
36
Hidden size
2048
2560
Intermediate size
6144
9728
Attention heads (Q/KV)
16/16
32/8
Head dimension
128
128
Table 2: Configurations of the selected 1.9B and 4B models. Module settings apply where the corresponding modules are used.
Figure 8: Sequential adaptation and parameter selection for KSA V2 at 1.9B.
Configuration
C-KSA
S-KSA
Module-only
Full-parameter
Module-only
Trainable parameters
Aggregation modules
All
Aggregation and Indexer
Training tokens (B), 1.9B
6.7
3.4
18.5
Training tokens (B), 4B
6.7
15.9
4.2
Target learning rate
10−4
10−4
10−3
LR warmup (updates)
200
400
200
Table 3: Adaptation hyperparameters and token budgets for the selected models.
Table 4: Long-context retrieval and general capabilities at 1.9B (%; pass@1 for code). Bold and underlined values mark the best and second-best scores, respectively, here and in subsequent tables.
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.
Large language models (LLMs) face severe memory bottlenecks in long-context inference due to the linearly growing size of key-value (KV) caches. Existing KV cache compression techniques typically rely on simple heuristics, overlooking the distinct functional roles of different attention heads. We present SGD-KV (Summarization-Guided KV Cache Compression), a head-aware framework that leverages a novel chunk-summarization diagnostic task to systematically identify and prioritize attention heads specialized in hierarchical information aggregation. Experiments on Qwen2.5-7B-1M and Qwen3-32B across diverse long-context benchmarks demonstrate that SGD-KV achieves state-of-the-art performance with contexts up to 1M tokens, while reducing KV cache memory usage by up to 75%. Our findings show that strategically allocating the KV cache budget based on the summarization score distribution of attention heads yields a superior efficiency-accuracy trade-off for long-context inference.
Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches, yet existing compression methods often apply uniform budgets across layers or tokens and degrade retrieval when lexical cues and semantic states require different preservation. We introduce DepthWeave-KV, a token-adaptive cache compression method that factorizes key and value states across neighboring transformer layers using shared low-rank channel bases while retaining lightweight token-specific residuals where attention behavior is sensitive. DepthWeave-KV combines cross-depth residual factorization with a token-conditional depth router that allocates higher reconstruction rank to instruction-bearing and retrieval-critical tokens, and uses calibration-free online error tracking from attention-output probes to adapt compression during generation without retraining the base model. A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic. Across LongBench, Needle-in-a-Haystack, L-Eval, and long-form QA and summarization benchmarks, DepthWeave-KV achieves near-full-cache task quality with substantially lower memory use, improving average score and retrieval accuracy over prior compressed caches while reaching 8.3x KV memory reduction and 72.8 tokens per second at 64K context.
Anna Cordoba, Adam Puente Tercero, Nerea Angulo Hijo +4