KV-Fold 用简单的左折叠思路解决了长上下文推理的内存和精度痛点,做 LLM 推理优化或长文档处理的团队可以直接在现有模型上尝试,无需额外训练。
KV-Fold 是一种无需训练的长上下文推理协议,将键值(KV)缓存视为序列块上的左折叠累加器。模型在处理每个块时,基于累积的缓存进行条件处理,追加新生成的键和值,并将扩展后的缓存传递到下一步,重复这一单步更新过程。该方法在 Llama-3.1-8B 模型上的“大海捞针”基准测试中,在 152 次试验中实现了 100% 的精确匹配检索,覆盖 16K 到 128K 令牌的上下文和最多 511 层的链深度,且仅需单块 40GB GPU 内存。KV-Fold 的递归过程稳定,每步漂移短暂上升后饱和为平坦平台,对数值精度变化不敏感,跨块大小和模型家族表现一致。这项工作表明,冻结的预训练 Transformer 已经支持稳定的 KV 缓存递归形式,为无需架构更改或训练的长上下文推理提供了实用路径。
KV-Fold: One-Step KV-Cache Recurrence for Long-Context Inference
We introduce KV-Fold, a simple, training-free long-context inference protocol that treats the key-value (KV) cache as the accumulator in a left fold over sequence chunks. At each step, the model processes the next chunk conditioned on the accumulated cache, appends the newly produced keys and values, and passes the enlarged cache forward; the same one-step update is applied repeatedly, analogous to foldl in functional programming. Building on the KV cache concatenation primitive introduced for latent multi-agent communication, we repurpose it as a chunk-to-chunk recurrence for long-context inference. When processing chunk t, the model attends to the KV cache carried from earlier chunks as a prefix, reusing its internal state across segments without modifying or retraining the model. Despite its simplicity, the induced recurrence is stable: per-step drift rises briefly and then saturates into a flat plateau that persists across deep chains. This plateau is insensitive to a 10,000x change in numerical precision, robust across chunk sizes, and consistent across model families. At the task level, KV-Fold preserves exact information over long distances. On a needle-in-a-haystack benchmark, it achieves 100% exact-match retrieval across 152 trials spanning contexts from 16K to 128K tokens and chain depths up to 511 on Llama-3.1-8B, while remaining within the memory limits of a single 40GB GPU. Compared to streaming methods, which trade fidelity for bounded memory, KV-Fold maintains long-range retrieval while operating as a sequence of tractable forward passes. Overall, our results show that frozen pretrained transformers already support a stable form of KV-cache recurrence, providing a practical route to long-context inference without architectural changes or training.