CommunityKV: Efficient Long-Context Decoding via Graph Partitioning¶
会议: NeurIPS 2026
arXiv: 2610.00418
代码: https://github.com/amazon-science/community-kv
领域: LLM效率
关键词: 稀疏注意力, KV缓存优化, 图划分, 社区发现, 长上下文解码
一句话总结¶
本文将训练无关的稀疏注意力检索建模为注意力图上的社区发现问题，利用预填充阶段天然产出的注意力分数构建融合直接交互与共注意力的稀疏图并运行 Leiden 算法聚类，在解码时通过常数时间局部增量更新实现亚线性复杂度的语义连贯 KV 块动态检索。
研究背景与动机¶
随着大语言模型在代码库级推理、长程智能体规划与多模态理解等场景的广泛应用，处理长上下文输入已成为核心能力。然而，标准自注意力机制在长序列下面临严重的二次方计算复杂度与线性增长的 KV 缓存显存搬运瓶颈。尽管模型在生成每个新 token 时实际仅依赖历史上下文中的一小部分关键信息，但在免微调（training-free）的实用约束下，现有的稀疏注意力方案往往陷入两难困境。
现有方法主要分为静态驱逐与动态检索两派。静态驱逐或压缩方法（如 H2O、SnapKV）直接丢弃被判定为低贡献的 token，但这种永久删除具有不可逆性，极易损坏复杂多跳推理所需的细粒度上下文；而动态检索方法中，基于固定连续分块与统计极值的方法（如 Quest）粒度过于粗糙，单个高分 outlier 会迫使系统拉取整个物理块的不相关 token；基于局部敏感哈希（LSH）或 K-means 向量聚类的方法（如 Multipole Attention）则面临昂贵聚类开销以及自回归生成过程中中心点漂移（centroid drift）导致的频繁重新聚类瓶颈。
图论视角提供了一种将上下文内在拓扑结构显式化的途径，但已有图方法（如 GraphKV）仍停留在利用图衰减信号进行静态 KV 驱逐的思路。本文的核心切入点在于反向思考图结构的价值：与其用图相似度决定扔掉什么，不如将注意力图作为保留完整 KV 缓存基础上的动态检索路由索引。核心 idea：将训练无关的稀疏注意力建模为注意力得分诱导图上的社区发现问题，结合直接注意力与共注意力拓扑利用 Leiden 算法形成语义连贯的 token 社区，并在解码阶段通过常数时间局部贪心分配实现免全局重划分的自回归增量更新。
方法详解¶
整体框架¶
CommunityKV 包含三个紧密衔接的阶段： 1. 图构建（Graph Construction）：在 Prefill 阶段复用标准矩阵乘法产生的注意力分数，过滤初始注意力汇点（attention sinks）后，为每个 query 保留前 \(\kappa\) 个最大权重的键，构建包含一阶直接注意力和二阶共注意力的对称稀疏图； 2. 初始划分（Initial Partitioning）：在稀疏图上运行 Leiden 社区发现算法，最大化模块度指标，将全量 token 划分为若干语义紧凑的社区并计算各社区质心，该过程通过流水线异步重叠隐藏在后序前向传播中； 3. 增量更新与检索（Incremental Update & Retrieval）：在自回归解码的每一步，利用当前 query 与各社区质心计算相似度并检索 Top 社区拼接候选集，在候选子集上执行精确注意力计算，同时仅基于候选集合以 \(\mathcal{O}(d)\) 的常数时间局部贪心更新新 token 的边与社区归属，避免昂贵的全局重新聚类。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["长上下文输入序列 Prefill"] --> B["融合图构建：结合直接注意力与共注意力拓扑<br/>WarpSelect 内核融合 Top-κ 键并剔除 Attention Sink"]
B --> C["Leiden 社区划分与流水线调度<br/>最大化模块度形成语义连贯簇并计算质心"]
C --> D["自回归生成 Step t"]
D --> E["增量社区检索：Query 检索质心并合并 Sink 候选集<br/>计算局部候选注意力与 Top-κ 边近似"]
E --> F["常数时间局部贪心分配：基于 ΔQ 更新新 Token 归属<br/>保持全量 KV 缓存无损并完成生成"]
F -->|下一 Token| D关键设计¶
1. 融合图构建：结合直接注意力与共注意力拓扑 传统向量空间聚类直接在键向量上执行 K-means，既丢失了注意力上下文依赖，又带来高达 \(\mathcal{O}(n^2 d)\) 的初始化负担。本文直接利用 Prefill 阶段已计算的注意力权重矩阵 \(\alpha \in [0, 1]^{n \times n}\)。首先排除初始的注意力汇点集合 \(\mathcal{S}\)（因为它们天然具有不成比例的超高注意力分值，若纳入会形成连接无关语义块的超大度中介 hub）。对于每个 query \(i\)，选出前 \(\kappa\) 个最高分键的集合 \(\mathcal{K}_i\)。直接注意力图捕捉显式依赖关系： $\(w^{(1)}_{ij} = \alpha_{ij} \mathbf{1}(j \in \mathcal{K}_i)\)$ 仅有一阶直接注意力不足以刻画功能相似性。当两个键 \(j\) 和 \(k\) 同时被同一个查询 token 赋予高权重时，它们在语义或功能上往往高度协同。为此引入二阶共注意力图： $\(w^{(2)}_{ij} = \sum_{m=|\mathcal{S}|+1}^n \alpha_{mi} \alpha_{mj} \mathbf{1}(\{i, j\} \subseteq \mathcal{K}_m)\)$ 最终合成的无向对称加权图为： $\(w = \frac{\lambda}{2} \left[ w^{(1)} + (w^{(1)})^\top \right] + (1 - \lambda) w^{(2)}\)$ 其中超参数 \(\lambda \in [0, 1]\) 平衡直接注意力与共注意力。由于每个节点仅向 \(w^{(1)}\) 贡献 \(\kappa\) 条边、向 \(w^{(2)}\) 贡献至多 \(\kappa^2\) 条边，当 \(\kappa\) 为常数时整图边数严格为 \(\mathcal{O}(n)\)。从马尔可夫随机游走视角看，行随机的 \(\alpha\) 对应离散转移核，对称化后的 \(w\) 是时间可逆随机游走的转移核，其高连通子图天然对应游走不易逃逸的亚稳态语义块。
2. Leiden 社区划分与流水线调度 传统基于模块度（Modularity）的划分若采用 Louvain 算法可能产生断开的子社区。本文采用 Leiden 算法，通过局部移动、细化与聚合三阶段交替迭代优化模块度： $\(Q = \frac{1}{2m} \sum_{ij} \left( w_{ij} - \gamma \frac{w_i w_j}{2m} \right) \mathbf{1}(c_i = c_j)\)$ 其中 \(m = \frac{1}{2} \sum_{ij} w_{ij}\) 为总边权，\(w_i = \sum_j w_{ij}\) 为节点加权度，\(c_i\) 为社区标签，分辨率参数 \(\gamma\) 严格控制社区细粒度。受稀疏图模块度分辨率极限（Resolution Limit）约束，社区数量自动随序列长度呈 \(\mathcal{O}(\sqrt{n})\) 亚线性增长。结合 Leiden 算法在稀疏图上的准线性时间与极少迭代轮数，单层图划分复杂度仅为 \(\mathcal{O}(n \log n + nd)\)，远低于 K-means 的 \(\mathcal{O}(n^2 d)\)。为了避免划分引入额外 Prefill 延迟，系统在 GPU 底层将 Top-\(\kappa\) 选择融合进 FlashAttention 内核，各层计算完毕后异步提交 Leiden 任务，与后续层的 FFN 和注意力计算重叠掩盖，最终仅最后一层暴露 0.37–0.45 秒的时延。
3. 增量社区检索：Query 检索质心并合并 Sink 候选集 在自回归解码步，若像传统检索方法那样直接扫描全部历史 token，开销仍是 \(\mathcal{O}(nd)\)。由于前期已形成 \(\mathcal{O}(\sqrt{n})\) 个社区，系统维护各社区键向量的平均质心。当前 query \(q_i\) 仅需与 \(\mathcal{O}(\sqrt{n})\) 个质心计算点积，耗时压缩至 \(\mathcal{O}(\sqrt{n} d)\)。随后根据质心得分从高到低贪心累加对应社区的原始 token，直至填满预设的 token 预算，并合并保留的注意汇点 \(\mathcal{S}\) 与当前 token \(i\)，构成紧凑的局部候选集 \(\mathcal{R}_i\)。最终注意力计算仅在该稀疏子集上展开，彻底摆脱了全序列显存读取带宽限制。
4. 常数时间局部贪心分配：基于 ΔQ 更新新 Token 归属 自回归生成过程中若频繁全局重新聚类，系统吞吐将断崖式下跌；若完全不更新聚类，新生成的 token 将变为无组织碎片。本文设计了轻量局部贪心更新规则：在解码步 \(i\) 计算出局部候选集上的近似注意力分值 \(\tilde{\alpha}_{ij}\) 并筛选局部 Top-\(\kappa\) 键集合 \(\tilde{\mathcal{K}}_i\) 后，仅对新加入的节点 \(i\) 及其邻居引入增量边。保持已有 token 的社区标签绝对冻结，节点 \(i\) 仅在两个选项中决策：成立新的单节点孤立社区，或者加入其邻居所属的既有社区 \(C \in \{ c_j \mid j \in \tilde{\mathcal{K}}_i \}\)。其增量模块度收益为： $\(\Delta Q(i \to C) \propto \sum_{j \in C \cap \tilde{\mathcal{K}}_i} w_{ij} - \gamma \frac{w_i w_C}{2m}\)$ 因为 \(|\tilde{\mathcal{K}}_i| = \kappa\)，最多仅需比对 \(\kappa\) 个候选社区的 \(\Delta Q\)，并在选定后以 \(\mathcal{O}(d)\) 时间增量更新目标社区的质心。整个更新过程触碰边数至多为 \(\kappa(\kappa+1)\)，总耗时严格为 \(\mathcal{O}(d)\)，与当前序列总长度 \(n\) 完全无关，达成了即时流式解码下的常数级更新。
一个完整示例¶
假设解码进行到第 100,000 个 token（此时历史划分为约 50 个社区）： 1. 质心打分：当前 query 仅与 50 个质心进行内积评分，迅速定位到得分最高的 3 个语义社区（例如包含特定函数实现的变量名与核心逻辑块，共计 3,800 个 token）； 2. 候选集拼接：取上述 3,800 个 token，补齐固定的 10 个 Attention Sink token 以及若干局部滑动窗口 token，总数拼满 4,096 的预算； 3. 稀疏注意力与近似边获取：在 4,096 个候选 token 上做 FlashAttention 并获取当前 query 的 Top-8 邻近 token，其中 6 个属于社区 A，2 个属于社区 B； 4. 贪心归属决策：代入 \(\Delta Q\) 公式比对加入社区 A、社区 B 或自立新社区的增量收益，由于与社区 A 连边密集且权重高，收益最大，直接打标归入社区 A 并更新社区 A 质心，全过程耗时不足 0.1 毫秒。
实验关键数据¶
主实验¶
论文在 LongBench v2（涵盖代码、长对话与多文档长程理解）及 BABILong（64k 长度下的多跳事实检索）等基准上系统评测了 Qwen3 与 Llama-3.1 家族模型。
在 LongBench v2 上，Qwen3 家族在 4,096 token 预算下对比密集注意力（Dense FlashAttention-2）及代表性基线结果如下：
| 模型 | 方法 | 128 | 512 | 1024 | 2048 | 4096 (满预算) | 密集上限 (Dense) |
|---|---|---|---|---|---|---|---|
| Qwen3-4B | Quest | 0.2 | 3.0 | 6.6 | 12.5 | 14.1 | 27.4 |
| Qwen3-4B | Multipole | 25.8 | 25.6 | 25.8 | 25.6 | 25.8 | 27.4 |
| Qwen3-4B | CommunityKV (本文) | 25.4 | 25.8 | 26.6 | 26.2 | 27.4 | 27.4 |
| Qwen3-8B | Quest | 5.6 | 15.1 | 17.5 | 22.5 | 24.7 | 31.4 |
| Qwen3-8B | Multipole | 29.4 | 28.4 | 29.8 | 28.2 | 28.8 | 31.4 |
| Qwen3-8B | CommunityKV (本文) | 28.5 | 29.5 | 30.1 | 29.9 | 31.1 | 31.4 |
| Qwen3-14B | Quest | 9.3 | 21.5 | 24.7 | 29.0 | 30.6 | 32.0 |
| Qwen3-14B | Multipole | 30.4 | 29.2 | 28.0 | 28.8 | 28.8 | 32.0 |
| Qwen3-14B | CommunityKV (本文) | 30.2 | 30.5 | 30.7 | 29.7 | 32.0 | 32.0 |
在单卡 NVIDIA H200 上生成 32,768 token（分别测试 64k 与 128k 输入上下文）的端到端吞吐量（Tok/s）与显存开销对比如下：
| 模型 | 上下文 | 配置方案 | 吞吐量 (Tok/s) | 加速比 | 辅助显存/KV比例 | 峰值显存 (GiB) | 峰值显存增幅 |
|---|---|---|---|---|---|---|---|
| Qwen3-4B | 64k | Dense | 88.95 | 1.00× | - | 32.80 | - |
| Qwen3-4B | 64k | Per-query head | 99.73 | 1.12× | 67.2% | 39.69 | +21.0% |
| Qwen3-4B | 64k | Query group | 133.86 | 1.50× | 15.0% | 33.83 | +3.1% |
| Qwen3-4B | 128k | Dense | 73.13 | 1.00× | - | 47.19 | - |
| Qwen3-4B | 128k | Per-query head | 90.77 | 1.24× | 54.9% | 60.65 | +28.5% |
| Qwen3-4B | 128k | Query group | 125.37 | 1.71× | 12.0% | 51.22 | +8.5% |
| Qwen3-8B | 128k | Dense | 63.49 | 1.00× | - | 61.95 | - |
| Qwen3-8B | 128k | Per-query head | 79.08 | 1.25× | 51.7% | 72.74 | +17.4% |
| Qwen3-8B | 128k | Query group | 103.02 | 1.62× | 11.7% | 63.86 | +3.1% |
| Qwen3-14B | 128k | Dense | 48.60 | 1.00× | - | 79.95 | - |
| Qwen3-14B | 128k | Per-query head | 56.90 | 1.17× | 67.8% | 97.16 | +21.5% |
| Qwen3-14B | 128k | Query group | 75.50 | 1.55× | 12.6% | 84.49 | +5.7% |
消融实验¶
在 LongBench v2（Qwen3-14B，4,096 预算）上对核心设计要素的详细消融结果如下：
| 消融维度 | 具体配置 | 准确率 (%) | 说明 |
|---|---|---|---|
| 图稀疏度 \(\kappa\) | \(\kappa=2\) | 30.3 | 连边过稀疏导致图碎片化 |
| 图稀疏度 \(\kappa\) | \(\kappa=4\) | 31.1 | 接近最优 |
| 图稀疏度 \(\kappa\) | \(\kappa=8\) (默认) | 32.0 | 保持稀疏度与拓扑连通性的黄金平衡点 |
| 图稀疏度 \(\kappa\) | \(\kappa=16\) | 28.7 | 边数过多模糊了社区边界 |
| 图稀疏度 \(\kappa\) | \(\kappa=32\) | 27.8 | 噪声连边显著降低社区分辨率 |
| 边权重平衡 \(\lambda\) | \(\lambda=0.0\) (纯共注意力) | 28.8 | 缺少直接注意力引导 |
| 边权重平衡 \(\lambda\) | \(\lambda=0.5\) (默认对半平衡) | 32.0 | 显式依赖与二阶协同互补最佳 |
| 边权重平衡 \(\lambda\) | \(\lambda=1.0\) (纯直接注意力) | 28.5 | 忽视了同被 Attend 到的键协同模式 |
| 图聚合粒度 | No Aggregation (每 Head 一图) | 32.0 | 表达能力最高，但显存占用大 |
| 图聚合粒度 | Query Group (每 GQA 组一图) | 31.8 | 精度几乎无损，显存开销锐减至 1/5 |
| 图聚合粒度 | Layer-Wise (整层共享一图) | 21.5 | 跨 Head 语义混杂导致严重性能崩塌 |
| Attention Sink | Excluded (图外独立保留) | 32.0 | 消除超高度节点对社区划分的破坏 |
| Attention Sink | In-Graph (纳入图划分) | 31.3 | Sink 形成虚假全局 Hub 干扰拓扑划分 |
在相同图结构下替换划分算法的对照控制实验表明： - 采用简单的连通分量划分（Connected Components），4B / 8B / 14B 在 4096 预算下准确率仅有 20.3% / 21.8% / 27.0%； - 采用 Top-k 局部邻域扩展（Top-k Neighborhoods），准确率仅为 18.2% / 26.2% / 29.6%； - 相比之下，Leiden 算法产出的社区划分在所有规模和预算上均具有决定性优势（27.4% / 31.1% / 32.0%）。
关键发现¶
- 动态检索对静态驱逐的代际碾压：GraphKV 在同样构建注意力图的前提下，采用静态驱逐策略在 Qwen3-8B 上需要 32k 缓存预算才达到 22.5% 的准确率；而 CommunityKV 仅需 4k 预算即达到 31.1%，证明了“保留完整 KV、按需检索”远优于“不可逆丢弃”。
- 工程落地的黄金组合 Query-Group Aggregation：在 GQA 架构下，按 KV Head 组聚合注意力图，不仅把辅助显存占 KV 的比例从 52–82% 压缩至 12–15%，更将端到端生成加速比推高至 1.39×–1.71×，且在 LongBench v2 上的精度仅下降 0.2%（31.8% vs 32.0%）。
- 长程生成稳定性验证：在无全局重新聚类的纯增量模式下，连续生成 4,096 个 token，图模块度指标稳定在 [0.44, 0.50]，注意力质量召回率保持在 85% 以上；若在超长生成中引入每 512 步全局重划分（Refresh），在 16k 上下文续写下的 Top-1 密集一致性进一步从 68.05% 跃升至 82.81%。
亮点与洞察¶
- 训练无关稀疏注意力的方法论重塑：首次将图论中的社区发现与马尔可夫亚稳态理论引入 KV 缓存路由，避免了无监督向量空间聚类与语言模型注意力分布脱节的问题，直接“就地取材”复用计算好的注意力热度矩阵。
- \(\mathcal{O}(\sqrt{n})\) 检索与 \(\mathcal{O}(d)\) 增量更新的复杂度闭环：借助图划分分辨率极限自然实现的 \(\sqrt{n}\) 社区缩放，巧妙化解了长文本检索的线性瓶颈；再通过局部候选集近邻的贪心指派，将新生成 token 的图融入耗时缩短至与上下文总长度脱钩的常数时间。
- 底层系统与算法的极致软硬协同：将 Top-\(\kappa\) 选择用 WarpSelect 熔进 FlashAttention 内核，使构图额外开销控制在 7% 左右，且利用多层流水线异步并行把 Leiden 耗时完全掩盖在 FFN 算力之后，工程完成度极高。
局限与展望¶
- 解码增量更新的局部近似偏差：在生成步更新时，新 token 仅在检索出的局部候选集中评估 Top-\(\kappa\) 边，若由于上一轮检索失误漏掉了真正的全局高权 token，可能导致图拓扑发生误差累积；
- 极端跨层语义多样性尚未完全解耦：Layer-wise 聚合出现断崖式性能暴跌（21.5%），表明不同层之间的注意力图拓扑差异极大，目前尚无法跨层复用图结构，显存节省仍受限于 GQA 的 KV 头数量；
- 分层社区结构的潜在潜力：目前仅采用了单层扁平社区质心检索。未来结合 Leiden 算法天生具备的多级聚合分层树，有望探索多层级分级检索，将检索复杂度进一步压低至 \(\mathcal{O}(d \log n)\)。
相关工作与启发¶
- vs GraphKV: 两者均利用 Prefill 阶段注意力诱导稀疏图，但 GraphKV 基于信号衰减执行永久的静态 Token 剪枝/驱逐，丢失的信息无法找回；CommunityKV 保持 KV 缓存完整无损，仅将图作为动态索引，在 8 倍更小的预算下大幅超越其表现。
- vs Multipole Attention: Multipole 使用 K-means 在键向量空间聚类，带来了 \(\mathcal{O}(n^2 d)\) 的高昂构建成本与漂移重新聚类开销；CommunityKV 利用标量图拓扑与 Leiden 算法，初始化开销降至 \(\mathcal{O}(n \log n + nd)\)，解码更新仅需 \(\mathcal{O}(d)\)。
- vs Quest: Quest 依赖简单的连续物理分块和 Min-Max 统计近似，粒度粗糙且极易受块内 Outlier 干扰；CommunityKV 从全局拓扑语义聚合非连续的关联 token 块，在长程复杂推理任务中准确率优势极为显著。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立自注意力图拓扑、马尔可夫稳定性与长文本 KV 缓存社区动态检索的严密映射，视角独特且优雅。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4B 到 14B 两大模型家族、三套权威长文本 Benchmark、端到端 H200 真实吞吐/显存测量及详尽消融。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰、复杂度分析严谨、系统流水线图解生动。
- 价值: ⭐⭐⭐⭐⭐ 兼顾免训练即插即用、超低额外显存与实质端到端加速，是长文本大模型工程部署的实用重磅成果。