Kwai Summary Attention V2 Technical Report

What it claims

KSA V2 is a retrofit for pretrained Hybrid KSA models, which interleave 3 summary-attention layers with 1 full-attention layer. It reuses the existing per-chunk Summary tokens in two ways:

• C-KSA: a second-level compression of distant KV.
• S-KSA: a block router that replaces dense full attention.

On a 1.9B MHA backbone at 128K context, the paper reports these results:

RULER retrieval is roughly unchanged. Neither new layer adds tokens or changes the Summary layout, which is why it can be retrofitted.

How it works

flowchart LR
  S[Summary token per k_s text tokens] --> C[C-KSA: pool k_m Summaries into Meta KV]
  S --> R[S-KSA: Summary-guided block keys]
  R --> I[Lightning Indexer top-K blocks]
  I --> A[Attend sink + selected + local original KV]
  C --> L[Attend Meta + boundary Summaries + local text]

C-KSA

• Each layer pools every k_m consecutive Summary KV entries into one "Meta" entry.
• The pooling uses per-feature softmax weights: a learned projection A plus a within-group position bias. This follows DeepSeek-V4.
• RoPE is applied to the pooled key at the position of the group's last Summary.
• Meta entries are layer-local and never fed forward. Later layers still see full-granularity Summaries, so k_m can differ per layer.
• Distant-history compression becomes k_s·k_m, for example 1024× with k_s=8 and k_m=128.
• log k_m is added to each Meta logit. This keeps attention from shifting toward the local window just because the softmax has fewer remote terms.
• A contiguous KV slice plus a Summary ring buffer support incremental decoding.

S-KSA

• Each chunk is split into blocks of b tokens. The chunk's Summary acts as the query in auxiliary attention heads, which pool each block's tokens into one routing key.
• A Lightning Indexer (from DSA) scores the blocks and selects up to K/b remote blocks. Sink blocks are kept within that budget.
• Main attention reads the original-text KV of the selected blocks plus a local window, in a single softmax.
• Compared with token-level indexing, routing needs b× fewer candidates, but prefill routing cost is still O(N²/b).
• Unlike CSA, main attention never reads compressed text.
• The motivation is an oracle test on a frozen Hybrid KSA at 128K. Selecting about 4% of text positions kept about 82% of decode attention mass and near-dense retrieval.

Evidence

All numbers below are from the 1.9B models in Figure 1.

RULER mean

General benchmarks, KSA V2 vs Hybrid KSA

All of these are above Full attention, for example 48.3 on GSM8K.

Ablations

• C-KSA alone matches its parent.
• Sparse adaptation loses less retrieval accuracy on the compressed backbone than on the original.

4B GQA backbone: retrieval stays close to the parent model. No specific numbers appear in the fetched text.

How to use it

• The paper describes a retrofit recipe: start from a pretrained Hybrid KSA, keep the 3:1 layout as C-KSA ×3 plus S-KSA ×1, and mostly train only the new compression and sparse modules.
• The authors say they provide GPU kernels.
• The fetched text gives no repository, weights, licence or training budget.
• For 4B/GQA, apply RMSNorm to the indexer queries and keys before RoPE, which the authors needed for convergence.

Limits

• Aggregation-heavy tasks, such as frequent-word extraction, degrade under block selection (Table 7).
• After the change, most of the remaining cache is original-text KV in the S-KSA layers.
• It only applies to existing Hybrid KSA models. It is tested at 1.9B and 4B only, and only up to 128K.
• The fetched page is cut off partway through the KV-cache section, so the detailed results tables and the 4B numbers were not visible.