PaperlayerView exact revision on arXiv

Source-linked AI summary

CommunityKV: Efficient Long-Context Decoding via Graph Partitioning

Joe McKenna, Anastasios Alexandridis, Nathan Susanj, Jing Liu

arXiv:2610.00418v1cs.LG

TL;DR

Long-context attention has quadratic compute and linearly growing KV-cache transfer, while existing training-free sparse approaches rely on coarse heuristics or costly-to-update preprocessing. CommunityKV builds semantic communities from attention-score graphs and locally assigns new tokens during decoding. It reports up to 1.25× per-query-head throughput and 1.71× with query-group aggregation, with query-group accuracy nearly matching per-head accuracy.

  • Problem

    Long-context Transformers face quadratic self-attention cost and linearly growing KV-cache memory transfer, while training-free sparse methods trade accuracy or efficient updating.

    Source: §1 · p. 1

  • Method

    CommunityKV partitions a token graph induced from attention scores into semantic communities and locally assigns generated tokens to communities during decoding.

    Source: §1 · p. 2, §5 · p. 9

  • Results

    End-to-end throughput reaches 1.25× per query head and 1.71× with query-group aggregation, nearly matching per-head accuracy while auxiliary storage falls from 52–82% to 12–15% of KV-cache size.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Local updates support streaming decoding without global re-partitioning, while query-group aggregation nearly matches per-head accuracy with less auxiliary storage.

    Source: §5 · pp. 9–10

    Greedy assignment freezes prior communities and may cause partition drift over long horizons; retrieval-restricted edge approximation can miss true top-κ keys outside retrieved communities.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

CommunityKV addresses long-context attention costs and limitations in training-free sparse retrieval by grouping tokens into communities from attention scores. Its local update rule supports streaming decoding, and the reported results show higher throughput with query-group aggregation than with one graph per query head.

1 Introduction

  • Long-context Transformers face quadratic self-attention cost and linearly growing key-value cache memory transfer.The paper motivates sparse attention as a way to load only a portion of the KV cache at each generation step.Source: §1 · p. 1
  • A local update rule assigns generated tokens to communities in constant time, enabling streaming decoding without global re-partitioning.The method’s complexity analysis also reports sub-linear community retrieval and O(d) token updates independent of sequence length.Source: §1 · p. 2
  • CommunityKV formulates training-free sparse attention as community detection on a graph induced from attention scores, capturing direct dependencies and second-order co-attention without auxiliary embeddings or vector-space clustering.The combined graph supports retrieval of semantically coherent token groups.Source: §1 · pp. 1–2
  • 1.08–1.25× throughput with 17–29% additional peak memory per query head, versus 1.39–1.71× throughput with 0.5–8.5% additional peak memory and comparable accuracy using query-group aggregation.Evaluations span two model families and three long-context benchmarks.Source: §1 · p. 2

2 Related work

Related approaches either integrate sparsity into training or use training-free cache compression and selective retrieval. Their retrieval patterns range from fixed patterns and coarse heuristics to clustering, while CommunityKV differs from GraphKV by retrieving communities without permanently deleting context.

Architectural vs. training-free sparse attention

  • Training-integrated sparse attention methods achieve throughput gains but have patterns fixed at training time, limiting retrieval quality to the training distribution.Their upfront training cost amortizes only at very high inference volumes.Source: §2 · p. 2
  • Training-free methods either compress the KV cache through lossy eviction or merging, or preserve it and dynamically load selected tokens at each decoding step.Selective retrieval effectiveness depends on the sparsity mask’s fixed or dynamic pattern.Source: §2 · p. 2

Fixed-pattern sparse attention

  • Fixed-pattern methods use query-independent structures, including strided, sliding-window, dilated, and global-token patterns, which can miss relevant tokens outside their predefined structure.Sparse Transformer has O(n√n d) complexity, while Longformer and BigBird achieve O(nd) complexity.Source: §2 · p. 2

Dynamic sparse attention

  • Dynamic methods include block heuristics, data-independent locality-sensitive hashing, and k-means clustering, each with retrieval granularity, overhead, or updating trade-offs.Contiguous-block retrieval may include unrelated neighbors, while clustering faces stale centroids or expensive re-clustering as tokens arrive.Source: §2 · pp. 2–3

Comparison to GraphKV

  • Unlike GraphKV’s eviction of semantically redundant tokens, CommunityKV retains the full KV cache and dynamically retrieves communities for each query.This gives the model access to locally dense semantic structure without permanently deleting context.Source: §2 · p. 3

3 Method

CommunityKV builds a token graph from attention scores, partitions it into communities, and extends the partition locally during decoding.

3.1 Graph construction

  • The token graph combines direct query-key links with co-attention between keys sharing a query, producing a symmetric sparse graph with O(n) edges for constant κ.Initial sink positions are excluded because their attention can be high regardless of semantic relevance.Source: §3.1 · p. 4

3.2 Initial partitioning

  • Leiden maximizes modularity, with γ setting community granularity; partitioning plus centroid computation costs O(n log n + nd), versus O(n^2d) k-means clustering.Leiden guarantees connected communities.Source: §3.1 · p. 4, §3.2 · pp. 4–5

3.3 Incremental update

  • Retrieval ranks community centroids to form a bounded sparse context, while each new token updates retrieved edges and joins a neighboring or singleton community in O(d).Across a 4096-token generation horizon, modularity remains stable and attention mass stays above 85% at a 4096-token budget. Refreshing every 512 steps improves final-window top-1 agreement with dense attention from 68.05% to 82.81%.Source: §3.3 · pp. 5–6

3.4 Implementation

  • Top-κ selection is fused into FlashAttention’s QK loop, while Leiden partitioning overlaps later forward-pass layers through asynchronous workers.Fused selection adds 6–8% to prefill wall time; only the final layer’s partition latency is exposed, at 0.37–0.45 seconds across Qwen3.Source: §3.4 · p. 6

4 Experiments

Experiments assess long-context accuracy, end-to-end systems cost, and design ablations under controlled decoding and retrieval-granularity settings.

4.1.1 Datasets

  • The evaluation covers LongBench v2, BABILong QA1–QA5, and RULER, spanning long-context understanding, multi-hop retrieval, and configurable long-context tasks.BABILong requires retrieving 1–5 supporting facts from long inputs.Source: §4.1.1 · p. 6

4.1.2 Models

  • Experiments use Qwen3-4B, 8B, and 14B and Llama-3.1-8B-Instruct.Source: §4.1.2 · p. 6

4.1.3 Baselines

  • Baselines include dense FlashAttention-2, graph-based GraphKV, Quest, Multipole, direct retrieval methods, and eviction or compression methods.GraphKV tests dynamic retrieval against permanent eviction; Quest and Multipole span fast, coarse retrieval and higher-fidelity clustering.Source: §4.1.3 · p. 7

4.2 Long context understanding

  • 31.1% accuracy with 4k tokens for CommunityKV contrasts with GraphKV’s 22.5% at 32k tokens on LongBench v2 with Qwen3-8B.GraphKV requires an 8× larger budget to reach 22.5% accuracy; CommunityKV nearly matches the dense baseline with 4k tokens.Source: §4.2 · p. 7, Table 2 · p. 7
  • On LongBench v2 at cluster size 16 and a 4,096-token budget, CommunityKV scores 32.0% on Qwen3-14B, matching the dense upper bound exactly.Source: Table 3 · p. 7, §4.2 · p. 8

4.3 Computational efficiency

  • 1.39–1.71× full-request speedup comes from query-group aggregation, with 0.5–8.5% additional peak memory and accuracy of 31.8% versus 32.0% per query head.Per-query-head graphs achieve 1.08–1.25× speedup with 17–29% additional peak memory; throughput gains increase with context length for every model size.Source: §4.3 · p. 8, Table 5 · p. 9

4.4 Ablations

  • 32.0% accuracy at κ = 8 falls to 27.8% at κ = 32, while λ = 0.5 outperforms pure direct or co-attention on LongBench v2 with Qwen3-14B.Accuracy is 30.3% at κ = 2, 28.5% at λ = 1, and 28.8% at λ = 0.Source: §4.4 · p. 8
  • Greedy full-token retrieval scores 32.0% versus 28.6% for 256 centroid summaries, while query-group graphs score 31.8%, per-head graphs 32.0%, and layer-wise graphs 21.5%.Excluding sink tokens and always retrieving them scores 32.0%, versus 31.3% when they are partitioned; Leiden outperforms both controls at every budget and model scale.Source: §4.4 · p. 9

5 Conclusion

CommunityKV uses Leiden-partitioned token communities for sub-linear retrieval and constant-time local updates. Its conclusion reports higher throughput with query-group aggregation, alongside limitations and scaling results.

5 Conclusion

  • Leiden partitions the attention-score graph during prefill; retrieval costs O(√n d) per step, with constant-time local updates for generated tokens.The update avoids global re-partitioning.Source: §5 · p. 9
  • The default configuration leads every accuracy category, and Leiden beats connected-components and top-k-neighborhood controls at every budget and model scale.The controls use the same weighted graph and are calibrated to match CommunityKV’s mean group size.Source: Table 6 · p. 10, Table 7 · p. 10
  • 1.25× per-query-head and 1.71× query-group end-to-end throughput accompany auxiliary storage falling from 52–82% to 12–15% of KV-cache size.Query-group aggregation nearly matches per-head accuracy.Source: §5 · p. 10
  • Frozen greedy assignments may drift over long horizons; retrieval can miss true top-κ keys outside retrieved communities, especially for unstructured graphs or isolated-span tasks.Periodic re-partitioning every √n steps mitigates drift but adds latency in practice.Source: §5 · p. 10
  • Sharing community structure across adjacent layers could reduce graph memory and pipelining depth.Source: §5 · p. 10

A.1 Resolution parameter γ

  • Across a 64× range of resolution settings, predicted answers changed on at most 2 of 30 examples despite a 2.5× change in mean community size.The comparison used Qwen3-8B with a 4,096-token budget.Source: §A.1 · p. 15

A.2 Sub-linear community scaling

  • Community counts empirically follow O(√n); at 100k tokens, the graph forms approximately 50 communities, enabling sub-linear retrieval.Scoring community centroids costs O(√n d), versus O(nd) for all tokens.Source: §A.2 · p. 15

B.1 Incremental partition drift

The appendix tracks partition and retrieval stability during incremental decoding, then describes a fixed-continuation evaluation protocol.

B.1 Incremental partition drift

  • The stress test tracks graph modularity and attention-mass recall over a 4,096-token generation horizon using the default configuration.Recall is the fraction of the full attention distribution captured by retrieved tokens.Source: §B.1 · p. 15
  • 0.44–0.50 modularity and attention-mass recall above 85% persist through 4,096 generated tokens at a 4,096-token retrieval budget.The measurements are reported without global re-partitioning during generation.Source: §B.1 · p. 15, Figure 3 · p. 16

B.2 Fixed continuations

  • The fixed-continuation experiment teacher-forces a shared 16,384-token continuation and measures terminal-window NLL, KL divergence from dense, and dense top-1 agreement.Results average five prompts per domain.Source: §B.2 · p. 16, Table 8 · p. 16

B.3 Sampled generation

The sampled-generation appendix compares dense attention with approximate methods on the same generated histories.

B.3 Sampled generation

  • Generation uses 8,192 tokens at temperature 0.7 and top-p 0.95, with 10 prompts per domain and two sampling seeds.Source: §B.3 · p. 16
  • Because approximate methods produce different histories, each history is replayed through dense attention and its generating method for distribution comparisons.The comparison uses terminal 512-token windows at 2k, 4k, and 8k generated tokens.Source: §B.3 · p. 16
  • NLL is reported only for the fixed-continuation experiment, where methods share the same target tokens.Source: §B.3 · p. 16

C Implementation and systems analysis

The appendix covers fused top-κ complexity, masking Leiden latency through pipelining, memory accounting, and sampled-generation stability.

C Implementation and systems analysis

C.1 Fused top-κ selection

CommunityKV selects top-κ keys during prefill and pipelines graph partitioning to limit exposed latency. With fixed κ, graph construction matches the asymptotic cost of causal prefill, while measured overheads remain.

C.1 Fused top-κ selection

  • Graph construction fuses top-κ selection into the attention-score computation, with fixed κ yielding O(n²) complexity and no asymptotic overhead.GPU threads maintain sorted candidate queues as score tiles are computed; block-level candidates are merged to produce each row's top-κ indices and weights.Source: §C.1 · p. 17
  • 7.2–7.8% prefill overhead was measured on 131k-token Qwen3 inputs, despite the fixed-κ construction matching prefill's asymptotic O(n²) cost.The reported overheads are 7.8% for Qwen3-4B, 7.7% for Qwen3-8B, and 7.2% for Qwen3-14B.Source: §C.1 · p. 17

C.2 Pipelined graph partitioning

  • 0.37–0.45 s of Leiden partitioning latency remains exposed, corresponding to 2–3% of prefill wall time on 130k-token inputs.Asynchronous partitioning overlaps with same-layer FFN and subsequent-layer attention, masking the partition cost for all but the last layer.Source: §C.2 · p. 17

C.4 Memory footprint

CommunityKV maintains graph indexes alongside the full KV cache, so their aggregate memory depends on the number of independently maintained graphs. Query-group aggregation reduces that graph footprint substantially.

C.4 Memory footprint

  • At n = 100k tokens and κ = 8, each graph has an observed footprint of ∼9 MB, with observed per-graph usage below 20 MB across evaluated sequence lengths.Each graph stores edge weights in bfloat16 with 32-bit indices.Source: §C.4 · p. 18
  • At 100k tokens with κ = 8, query-group aggregation uses <3 GB of graph memory, versus <15 GB without aggregation for Qwen3-14B.The model has 40 layers and 40 query heads, yielding 1,600 per-head graphs or 320 graphs under query-group aggregation.Source: §C.4 · p. 18
  • With query-group aggregation, graph memory is roughly an order of magnitude smaller than the ∼16 GB dense KV cache at 100k context, but remains additive because the full cache is preserved.The graph footprint is <3 GB under query-group aggregation.Source: §C.4 · p. 18

D Expanded baseline results

The expanded baseline results cover LongBench v2 and RULER accuracy under specified token budgets. The supplied table descriptions establish evaluation scope, not comparative outcomes.

D.1 LongBench v2

  • Table 11 compares LongBench v2 accuracy across direct retrieval, eviction/compression, and CommunityKV methods at matched token budgets.The caption provides no accuracy values or winning method.Source: Table 11 · p. 19

D.2 RULER

  • Table 12 reports RULER accuracy for three direct-retrieval tasks at a matched 4,096-token budget, with task-specific context lengths for Qwen3 and Llama-3.1.The average covers only the three reported tasks; the caption names MK-2, MQ, and VT as Multikey-2, Multiquery, and variable tracking.Source: Table 12 · p. 19
  • Table 13 reports unweighted average accuracy over all 13 RULER tasks at a matched 4,096-token active-token budget, with dense attention as a reference.The caption provides no accuracy values or comparative outcome.Source: Table 13 · p. 20
Loading 2610.00418v1…