CommunityKV: Efficient Long-Context Decoding via Graph Partitioning
Organizations: Amazon AGI
Abstract
Scaling Transformers to long contexts is constrained by the quadratic cost of self-attention and the linear growth of key-value cache memory transfer. Sparse attention mitigates this by retrieving only relevant tokens, but current approaches either require large-scale training or, within the training-free regime, rely on semantically coarse heuristics or expensive clustering that is difficult to update efficiently during decoding. We introduce CommunityKV, a framework that formulates sparse attention as a community detection problem. CommunityKV constructs a token graph from the scores already computed during standard prefill, and partitions the graph into communities to enable retrieval of semantically coherent token groups. A local update rule assigns newly generated tokens to communities in constant time, enabling sparse retrieval throughout streaming decoding without global re-partitioning. We evaluate CommunityKV on Qwen3 and Llama-3.1 models across three long-context benchmarks. With one graph per query head, CommunityKV delivers up to the end-to-end generation throughput of dense attention, while query-group graph aggregation yields up to with comparable accuracy.
Figures & tables
| Method | Initial Clustering | Per-step Update | Per-Step Retrieval |
|---|---|---|---|
| Quest | |||
| Multipole | |||
| CommunityKV |
| Qwen3-8B (Dense: 31.4) | ||||
|---|---|---|---|---|
| Method | 4096 | 8192 | 16384 | 32768 |
| GraphKV | 3.4 | 10.5 | 18.5 | 22.5 |
| CommunityKV | 31.1 | 31.2 | 31.4 | 31.4 |
| Model | Method | 128 | 256 | 512 | 1024 | 2048 | 4096 |
|---|---|---|---|---|---|---|---|
| Qwen3-4B (Dense: 27.4) | Quest | 0.2 | 0.8 | 3.0 | 6.6 | 12.5 | 14.1 |
| Multipole | 25.8 | 26.2 | 25.6 | 25.8 | 25.6 | 25.8 | |
| CommunityKV | 25.4 | 26.1 | 25.8 | 26.6 | 26.2 | 27.4 | |
| Qwen3-8B (Dense: 31.4) | Quest | 5.6 | 9.1 | 15.1 | 17.5 | 22.5 | 24.7 |
| Multipole | 29.4 | 30.0 | 28.4 | 29.8 | 28.2 | 28.8 | |
| CommunityKV | 28.5 | 28.8 | 29.5 | 30.1 | 29.9 | 31.1 |
| Qwen3-8B | Llama-3.1-8B-Instruct | |||||||||||
| Method | QA1 | QA2 | QA3 | QA4 | QA5 | Avg. | QA1 | QA2 | QA3 | QA4 | QA5 | Avg. |
| Dense | 63 | 32 | 30 | 55 | 63 | 48.6 | 62 | 26 | 23 | 43 | 58 | 42.4 |
| CommunityKV | 48 | 31 | 30 | 53 | 65 | 45.4 | 63 | 24 | 21 | 44 | 59 | 42.2 |
| TokenSelect | 58 | 28 | 28 | 48 | 56 | 43.6 | 57 | 21 | 16 | 44 | 52 | 38.0 |
| FreeKV | 52 | 28 | 23 | 48 | 58 | 41.8 | 58 | 23 | 22 | 47 | 57 | 41.4 |
| SparQ | 49 | 32 | 17 | 51 | 57 | 41.2 | 59 | 15 | 15 | 44 | 54 | 37.4 |
| Model | Context | Graph Agg | Tok/s | Speedup | Aux | Aux/KV | Peak | Peak |
|---|---|---|---|---|---|---|---|---|
| Qwen3-4B | 64k | Dense | 88.95 | 1.00 | 32.80 | |||
| Per-query head | 99.73 | 1.12 | 9.07 | 67.2% | 39.69 | 21.0% | ||
| Query group | 133.86 | 1.50 | 2.03 | 15.0% | 33.83 | 3.1% | ||
| 128k | Dense | 73.13 | 1.00 | 47.19 | ||||
| Per-query head | 90.77 | 1.24 | 12.36 | 54.9% | 60.65 | 28.5% | ||
| Query group | 125.37 | 1.71 | 2.69 | 12.0% | 51.22 | 8.5% |
| Graph Sparsity ( ) | |||||
|---|---|---|---|---|---|
| 30.3 | 31.1 | 32.0 | 28.7 | 27.8 | |
| Graph Weight ( ) | |||||
| 28.8 | 29.0 | 32.0 | 29.0 | 28.5 | |
| Centroid Budget | 0 | 4 | 16 | 64 | 256 |
| 32.0 | 27.5 | 29.7 | 28.4 | 28.6 | |
| Graph Aggregation | No Agg. | Query Group | Layer-Wise | – | – |
| Model | Method | 128 | 256 | 512 | 1024 | 2048 | 4096 |
|---|---|---|---|---|---|---|---|
| Qwen3-4B | Connected components | 1.1 | 2.7 | 3.5 | 7.3 | 12.1 | 20.3 |
| Top- neighborhoods | 23.5 | 23.3 | 23.3 | 21.9 | 19.0 | 18.2 | |
| CommunityKV | 25.4 | 26.1 | 25.8 | 26.6 | 26.2 | 27.4 | |
| Qwen3-8B | Connected components | 3.2 | 5.5 | 6.7 | 11.6 | 15.8 | 21.8 |
| Top- neighborhoods | 27.6 | 28.4 | 27.6 | 26.4 | 28.8 | 26.2 | |
| CommunityKV | 28.5 | 28.8 | 29.5 | 30.1 | 29.9 | 31.1 |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| NLL | KL | Top-1 agreement | ||||||
|---|---|---|---|---|---|---|---|---|
| Domain | Horizon | Dense | Incr. | Refresh | Incr. | Refresh | Incr. | Refresh |
| Python | 2k | 0.7045 | 0.7534 | 0.7375 | 0.0600 | 0.0311 | 94.34% | 95.63% |
| 4k | 0.4952 | 0.5478 | 0.5288 | 0.0694 | 0.0369 | 94.53% | 95.59% | |
| 8k | 0.7264 | 0.8470 | 0.7729 | 0.1372 | 0.0418 | 92.30% | 95.59% | |
| 16k | 0.7042 | 0.8297 | 0.7414 | 0.2510 | 0.0702 | 87.50% | 94.30% | |
| PG-19 | 2k | 2.5129 | 2.6743 | 2.6158 | 0.1525 | 0.1144 | 82.81% | 86.48% |
| KL | Top-1 agreement | ||||
|---|---|---|---|---|---|
| Domain | Horizon | Incr. | Refresh | Incr. | Refresh |
| PG-19 | 2k | 0.0519 | 0.0521 | 93.96% | 93.91% |
| 4k | 0.0083 | 0.0196 | 99.34% | 97.87% | |
| 8k | 0.0045 | 0.0097 | 99.75% | 99.13% | |
| Python | 2k | 0.0290 | 0.0150 | 98.13% | 99.01% |
| 4k | 0.0146 | 0.0105 | 99.21% | 99.11% | |
| Prefill – Total Wall Time (s) | Decoding – Avg Per-Step Wall Time (ms) | ||||||
| 4B | 8B | 14B | 4B | 8B | 14B | ||
| Dense | 16.5 | 18.2 | 26.3 | Dense | 63 | 64 | 71 |
| CommunityKV | 17.8 | 19.6 | 28.2 | CommunityKV | 37 | 34 | 51 |
| Top- Sel. | 0.94 | 1.09 | 1.49 | Retrieval | 5.3 | 4.7 | 7.2 |
| Leiden | 0.37 | 0.45 | 0.42 | Attention | 6.1 | 5.5 | 8.2 |
| Overhead | +7.8% | +7.7% | +7.2% | Speedup | 1.69 | 1.86 | 1.39 |
| Method | 128 | 256 | 512 | 1024 | 2048 | 4096 |
|---|---|---|---|---|---|---|
| Qwen3-4B (Dense: 27.4) | ||||||
| Quest | 0.2 | 0.8 | 3.0 | 6.6 | 12.5 | 14.1 |
| Multipole | 25.8 | 26.2 | 25.6 | 25.8 | 25.6 | 25.8 |
| TokenSelect | 26.1 | 26.1 | 25.9 | 25.7 | 25.9 | 25.3 |
| FreeKV | 22.6 | 24.2 | 25.7 | 26.7 | 26.1 | 26.7 |
| SparQ | 25.6 | 25.0 | 24.8 | 26.4 | 26.6 | 27.2 |
| Qwen3-8B | Llama-3.1-8B-Instruct | |||||||
|---|---|---|---|---|---|---|---|---|
| Method | MK-2 | MQ | VT | Avg. | MK-2 | MQ | VT | Avg. |
| CommunityKV | 78.0 | 92.5 | 22.0 | 64.2 | 95.0 | 97.8 | 96.8 | 96.5 |
| TokenSelect | 81.0 | 91.8 | 17.4 | 63.4 | 92.0 | 90.0 | 93.2 | 91.7 |
| FreeKV | 64.0 | 93.5 | 46.4 | 68.0 | 91.0 | 97.0 | 98.4 | 95.5 |
| SparQ | 69.0 | 88.5 | 26.8 | 61.4 | 47.0 | 85.5 | 92.0 | 74.8 |
| Model | Dense | StreamingLLM | H2O | SnapKV | CommunityKV |
|---|---|---|---|---|---|
| Qwen3-8B | 76 | 17 | 17 | 52 | 53 |
| Llama-3.1-8B-Instruct | 86 | 20 | 21 | 76 | 73 |