CSA: Compressed Sparse Attention (toy implementation)¶
Paper Section 2.3.1 · Code: CSA in src/deepseek_v4_lab/attention/csa.py
Evidence tags: PAPER stated by DeepSeek · DERIVED mathematically implied · IMPLEMENTED reproduced in this project · MEASURED observed in our experiments · INFERRED our interpretation · UNKNOWN not established. Full definitions: evidence ledger
Why this exists
Problem: at 1M-token context, dense attention FLOPs grow quadratically and the uncompressed KV cache grows linearly into hundreds of GB. Both are unaffordable at once.
Response: compress KV m=4 tokens per entry via dual overlapped pooling, and let a lightning indexer select only the top-k relevant blocks per query.
New cost: pooling destroys information; top-k retrieval can miss the right block; and an untrained indexer selects noise, so sparsity itself must be trained.
Our experiment: Exp02 measures cache reduction (2.9% of a vanilla cache), needle-block recall vs k, and end-task accuracy under staged training; Exp04 places CSA on the accuracy-per-byte frontier.
Exchange rate: COMPRESSION buys memory and compute; spends information and retrieval accuracy.
What is implemented [IMPLEMENTED]¶
- Dual overlapped compression (Eqs 9–12): two KV series C^a/C^b with content-dependent pooling weights Z^a/Z^b and learnable positional biases; entry i pools its own block (a-series) plus the preceding block (b-series) under one joint softmax ⇒ net 1/m compression with smooth boundaries.
- Lightning indexer (Eqs 13–17): low-rank indexer queries, per-head ReLU-scored relevance against compressed indexer keys, learned head weights, top-k block selection restricted to causally visible blocks.
- Shared-KV MQA core attention (Eqs 18–19): selected compressed entries serve as both K and V; sliding-window branch (n_win uncompressed recent entries); RMSNorm on queries/entries; partial RoPE (last 8 dims toy / last 64 paper) including −i rotation of outputs; attention sink (Eq 27); grouped output projection.
Training curriculum [PAPER-motivated]¶
Dense-over-compressed warmup → sparse fine-tuning with a straight-through mixture so the indexer receives gradients, mirroring the paper's staged sparsity introduction (Section 4.2.2). Without the sparse phase the indexer is untrained and true top-k evaluation collapses.
Toy-scale result. Direction is informative; production magnitude is not.
The failure story this page tells [MEASURED]¶
dense attention
| quadratic FLOPs, linear KV bytes [DERIVED, Exp01]
v
compress KV entries (m=4 dual-overlap pooling)
| information loss: dense-over-compressed ceiling ~66% vs vanilla 82% [MEASURED]
v
add relevance indexer (top-k blocks per query)
| sparse retrieval misses: recall plateaus ~72%, never reaches 1.0 [MEASURED]
v
train the indexer with staged sparsity (dense -> STE mixture)
| new bottleneck: decoding pooled entries, not selecting them [MEASURED]
v
usable compressed attention: 2.9% of vanilla cache bytes at eval length [DERIVED]
Every arrow is priced somewhere on this page; nothing is asserted without a tag.
Measured results (Exp02, needle task, eval at 2× train length)¶
| k | accuracy | needle-block recall |
|---|---|---|
| 1 | 0.320 | 0.352 |
| 2 | 0.393 | 0.680 |
| 4 | 0.400 | 0.686 |
| 8 | 0.389 | 0.693 |
| 16 | 0.400 | 0.721 |
Recall saturates ≈ 72% while accuracy plateaus ≈ 40%: at toy scale the bottleneck shifts from selection to decoding the pooled entry. Dense-over-compressed ceiling: ≈ 66% ⇒ compression itself costs ~16 points vs vanilla 82%, selection costs the rest [MEASURED].
KV bytes at eval length: CSA m=4 = 2.9% of a GQA-like vanilla cache [DERIVED from analytic formulas, test-validated].
Deviations from paper (stated)¶
- Indexer key compression uses separate projections (paper says "same compression operation" without specifying weight sharing).
- RoPE positions for compressed entries use block-start indices (paper unspecified); STE training mixture skips entry-RoPE (position undefined).
- Sparse training uses a softmax-mixture relaxation instead of production sparse-attention kernels.
Limitations¶
- Random-noise sequences are the worst case for compression (every token carries maximal entropy); real text has exploitable redundancy.
- Toy dims (c=32, n_h=4, d_c=32) vs paper (c=512, n_h=64, d_c=1024).