CommunityKV: Efficient Long-Context Decoding via Graph Partitioning

Joe McKenna; Anastasios Alexandridis; Nathan Susanj; Jing Liu

cs.LGSep 30, 2026

CommunityKV: Efficient Long-Context Decoding via Graph Partitioning

Authors: Joe McKenna, Anastasios Alexandridis, Nathan Susanj, Jing Liu

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 $QK^T$ 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 $1.25\times$ the end-to-end generation throughput of dense attention, while query-group graph aggregation yields up to $1.71\times$ with comparable accuracy.

Figures & tables

Figure 1: Overview of the CommunityKV framework. (1) Graph Construction. During prefill, we induce a sparse graph from the attention matrix by connecting each query to its top- $\kappa$ keys (gray cells, $\kappa=3$ shown). Sink tokens (hatched) are excluded from the top- $\kappa$ pool. (2) Initial Partitioning. We partition the graph into semantic communities (colored regions) using the Leiden algorithm and compute a centroid (star) per community. (3) Incremental Update. At each decoding step (token 16 shown), we score the query against community centroids to retrieve a candidate subset, run exact attention on this subset together with sink tokens, and use the resulting top- $\kappa$ keys to update the new token’s edges and assign it to a community in constant time. For clarity, only direct attention $w^{(1)}$ relations are depicted; this corresponds to $\lambda=1$ in the combined graph.

Method	Initial Clustering	Per-step Update	Per-Step Retrieval
Quest	$\mathcal{O}(nd)$	$\mathcal{O}(d)$	$\mathcal{O}(nd)$
Multipole	$\mathcal{O}(n^{2}d)$	$\mathcal{O}(nd)$	$\mathcal{O}(nd)$
CommunityKV	$\mathcal{O}(n\log n+nd)$	$\mathcal{O}(d)$	$\mathcal{O}(\sqrt{n}\,d)$

Table 1: Computational complexity comparison.

n

is sequence length,

d

is hidden dimension. Note that our method achieves sub-linear retrieval

\mathcal{O}(\sqrt{n}d)

without expensive clustering or update costs.

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

Table 2: Accuracy on LongBench v2 across token budgets (4k–32k) for GraphKV and CommunityKV. Dynamic community retrieval matches dense accuracy at

8\times

lower budget than static eviction.

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

Table 3: Accuracy on LongBench v2 across token budgets (128–4096) at cluster size 16 for Qwen3 models (4B, 8B, 14B). CommunityKV accuracy improves with budget and reaches near dense parity at 4096 tokens across all model scales.

	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

Table 4: Accuracy on BABILong QA1–QA5 across model families and sparse-attention baselines. Results use 64k contexts and a matched 4,096-token budget. We compare CommunityKV with direct retrieval methods (TokenSelect, FreeKV, and SparQ) and eviction/compression methods (SnapKV, H2O, and StreamingLLM).

Model	Context	Graph Agg	Tok/s	Speedup	Aux	Aux/KV	Peak	$\Delta$ Peak
Qwen3-4B	64k	Dense	88.95	1.00 $\times$			32.80
		Per-query head	99.73	1.12 $\times$	9.07	67.2%	39.69	21.0%
		Query group	133.86	1.50 $\times$	2.03	15.0%	33.83	3.1%
	128k	Dense	73.13	1.00 $\times$			47.19
		Per-query head	90.77	1.24 $\times$	12.36	54.9%	60.65	28.5%
		Query group	125.37	1.71 $\times$	2.69	12.0%	51.22	8.5%

Table 5: Generation of 32,768 tokens on an H200 with batch size 1. Timing spans prefill through final decode and includes all model and CommunityKV operations. Auxiliary storage and peak memory are reported in GiB. Auxiliary storage comprises graphs and centroids, with percentages relative to KV-cache size; peak-memory percentages are increases over dense attention.

Graph Sparsity ( $\kappa$ )	$\kappa=2$	$\kappa=4$	$\kappa=8$	$\kappa=16$	$\kappa=32$
	30.3	31.1	32.0	28.7	27.8
Graph Weight ( $\lambda$ )	$\lambda=0.0$	$\lambda=0.25$	$\lambda=0.50$	$\lambda=0.75$	$\lambda=1.0$
	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	–	–

Table 6: Ablation studies on LongBench v2 (Qwen3-14B, token budget=4096). The default configuration (

\kappa{=}8

\lambda{=}0.5

, per-head graphs, sinks excluded) is marked with gray shading and achieves the best accuracy in every category.

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- $k$ 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- $k$ 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

Table 7: CommunityKV outperforms connected-components and top-

k

-neighborhood controls at every budget and model scale, highlighting the importance of Leiden partitioning.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Empirical analysis of graph partition properties on LongBench v2 (Qwen3-14B). Left: Average community size as a function of the resolution parameter $\gamma$ . Right: Number of communities vs. context length, confirming $\mathcal{O}(\sqrt{n})$ sub-linear scaling.

Figure 3: Partition stability over a 4096-token generation horizon (Qwen3-14B, LongBench v2). Left: Graph modularity $Q$ remains stable and increases slightly as new tokens reinforce community structure. Right: Attention mass recall $\tilde{\mu}_{\kappa}$ at varying retrieval budgets. At the default 4096-token budget, recall stays above 85% throughout generation.

		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%

Appendix

Table 8: Fixed-continuation stability over generation length. Lower NLL and KL are better; higher top-1 agreement is better.

		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%

Appendix

Table 9: Sampled-generation stability over generation length. Values compare dense attention with each method on the same generated history.

Figure 4: Pipeline schedule for graph partitioning during prefill. Leiden partitioning for each layer is dispatched asynchronously after its attention scores are computed, overlapping with subsequent layer execution. The red segment indicates the residual Leiden latency that remains exposed on the critical path when the final layer’s partitioning has not completed by the time it is needed.

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- $\kappa$ 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 $\times$	1.86 $\times$	1.39 $\times$

Appendix

Table 10: Latency comparison between FlashAttention-2 (Dense) and CommunityKV using Hugging Face’s DynamicCache on an NVIDIA H100 80GB at 131k context. Left: total prefill wall time with breakdown into fused top-

\kappa

selection and exposed Leiden partitioning latency. Right: average per-step decoding wall time (128 generated tokens) with breakdown into token retrieval and sparse attention. Only attention-related kernels are reported; the remaining per-step time is MLP and layer normalization.

Qwen3-4B (Dense: 27.4)
Method	128	256	512	1024	2048	4096
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

Appendix

Table 11: Accuracy on LongBench v2 across matched token budgets for direct retrieval, eviction/compression, and CommunityKV methods.

	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

Appendix

Table 12: RULER accuracy for direct retrieval methods. MK-2, MQ, and VT denote Multikey-2, Multiquery, and variable tracking, respectively. Results use a matched 4,096-token retrieval budget. The average covers the three reported tasks only. Qwen3 tasks use 64k contexts; for Llama-3.1, the NIAH tasks use 64k contexts and variable tracking uses 32k.

Model	Dense	StreamingLLM	H2O	SnapKV	CommunityKV
Qwen3-8B	76	17	17	52	53
Llama-3.1-8B-Instruct	86	20	21	76	73

Appendix

Table 13: RULER accuracy across the complete 13-task suite. Results use a matched 4,096-token active-token budget. Dense attention is included as a reference; reported values are unweighted averages over all 13 tasks.

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunDynamic Sparse Attention Efficient Long-Context Inference
KAIST
ResidualKV: Residual-Based KV Cache Compression for Efficient Long-Context Inference
Feb 8, 2026Jitai Hao, Qiang Huang, Yaowei Wang +2Efficient Long-Context Inference Key-Value Cache Compression
Harbin Institute of Technology (Shenzhen), China
SeKV: Resolution-Adaptive KV Cache with Hierarchical Semantic Memory for Long-Context LLM Inference
Jun 30, 2026Amirhossein Abaskohi, Giuseppe Carenini, Peter West +1Key-Value Cache Compression Key-Value Cache
University of British Columbia · Microsoft Research

Jun 22, 2026, cs.LGJ/K move · Enter open · S save

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Huzama Ahmad, Se-Young Yun

KAIST

Dynamic Sparse Attention Efficient Long-Context Inference Prefill Plug-And-Play

Open paper AskPDF arXiv

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.

Abstract

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

ResidualKV: Residual-Based KV Cache Compression for Efficient Long-Context Inference

SeKV: Resolution-Adaptive KV Cache with Hierarchical Semantic Memory for Long-Context LLM Inference