技巧精选

FlexBudget:统一 FlexAttention 内核公平对比 KV 缓存淘汰策略

精选理由

AMD 研究员做了个开源方案,把 KV 缓存淘汰策略和内核实现解耦,跑长上下文想省显存的可以看看怎么公平测不同策略。

长上下文推理时 KV 缓存占用过大,需要丢弃部分条目,现有方法各绑一套自研内核,难以判断效果来自策略还是代码速度。AMD 的 Huanran Wang 将在 PyTorchCon 北美展示 FlexBudget,把内核和内存预算固定,只替换重要性规则,实现公平比较。该方法无需训练,完全基于 PyTorch 核心 FlexAttention 实现,将于 10 月 20-21 日在圣何塞发布。

原文 · PyTorch

When you run a model over a long context, the KV cache gets huge, so you have to throw some of it away. There are competing rules for choosing what to drop: keep the most recent tokens, keep the ones that get the most attention, and so on. The problem is that every method ships its rule welded to its own kernel. So if one method beats another, you can't tell whether the rule was smarter or the code was just faster.

At #PyTorchCon North America, Huanran Wang of @AMD will present the poster: FlexBudget: Interchangeable KV Importance Signals under One Compiled FlexAttention Budget We hold the kernel and the memory budget fixed, and swap only the rule. Same code, same allowance, honest comparison. It's training-free, and it's all core PyTorch.

Join us in San Jose, October 20-21 to learn more: https://t.co/jBApW8nESi