CommunityKV: Efficient Long-Context Decoding via Graph Partitioning
Joe McKenna ⋅ Anastasios Alexandridis ⋅ Nathan Susanj ⋅ Jing Liu
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 open weight models and long context benchmarks, demonstrating that our approach achieves up to $1.86\times$ higher decoding throughput than exact FlashAttention-2 with negligible accuracy degradation.
Chat is not available.
Successful Page Load