@cwolferesearch 训练 LLMs 管理自己的 KV cache 这个想法很…
@cwolferesearch The idea of training LLMs to manage their own KV cache is super…
NGC 论文提出让 LLM 在 decoding 中管理自身 KV cache：每隔 δ 个 token 对固定大小 cache block 打分，按 eviction rate ϵ 保留部分 block，以稳定 peak cache size。方法复用现有 attention query/key 打分，并将 eviction decision 与 token prediction 一并纳入 RL/GRPO 的 verifiable loss 端到端训练。
训练 LLM 管理自身 KV cache 的想法让我很感兴趣。最近的 neural garbage collection（NGC）论文是关于这个主题的一篇不错的文章。
Reasoning model / agent 显然需要长序列来处理复杂推理、长周期任务、tool call 等。然而，KV 的大小会随序列长度线性增长，形成 KV cache 瓶颈。
为了解决这个问题，已有若干 heuristic 被提出；例如，只保留最近的 token、保留 attention score 高的 token 等。但这些 heuristic 往往会降低性能，并且是否有效取决于 domain / task。
与其使用 heuristic，不如尝试教 LLM 管理自己的 KV cache。具体来说，NGC 通过实现一种 eviction cadence 来做到这一点。在 decoding 过程中，每隔 δ 个 token，NGC 会为其所有 KV cache block 打分，并定义一个 eviction rate ϵ，使得只保留 (1 - ϵ) 的 KV cache block。通过这种方式，可以确保 peak cache size 保持稳定。
为了给 KV cache block 打分，NGC 不使用任何新的或专门的 model / module。相反，它复用了 LLM 现有的 attention 机制。model 取最近的 query vector，将 KV cache 划分为固定大小的 block，然后基于 query vector 给之前的 key 打分。
NGC 没有为管理 KV cache 进行专门训练，而是直接把 KV cache management 纳入使用 RL / GRPO 训练时的 verifiable loss。RL objective 同时包含：
- 用于普通 token prediction 的 component。
- 用于 KV cache eviction decision 的 component。
这样，我们就可以用 RL 对 model 进行端到端训练，使其正确地 evict KV cache block（类似于预测 token），同时仍然使用 outcome reward。