动态回收LLM推理中的预填充缓存,提升内存利用率。
Elastic KV Cache for LLM Serving:A Working Reclamation Mechanism, and Why Chunked Prefill Already Closes the Gap
- 通过用户态虚拟内存管理,动态借还预填充缓存区。
- 实测小块预填充对延迟影响仅1%,收益有限。
- 适合高并发、大模型推理场景下的内存优化需求。
LLM推理引擎在启动时固定分配关键值(KV)缓存,为最坏情况的预填充激活预留空间。解码阶段该预留空间闲置却无法释放给缓存池,因大预填充仍需它。本文提出弹性KV缓存机制,在解码时将预留空间借出,并在预填充前归还,依赖调度器的前瞻视图控制。该机制基于CUDA虚拟内存路径实现,每层使用两个物理句柄映射到连续虚拟地址,保持注意力核函数不变,无需驱动修改。解分配耗时数毫秒,重分配数十毫秒,兼容CUDA图和前缀缓存,且永不触发显存溢出。静态分配相同内存存在预填充突发崩溃风险,故动态切换必要。进一步实验发现,若小块预填充显著增加延迟才值得回收。在真实解码负载中注入长提示的对比测试显示,8192与32768分块下首令牌时间差异仅约1%,因预填充受计算瓶颈限制,解码每步仅消耗约1个词元/序列。降低max_num_batched_tokens可获得更多缓存,延迟几乎相当。在张量并行下,预留比例从TP1的16%降至TP4的2.7%。本文明确指出回收机制仍有效的边界条件,并开源了可复用的用户态弹性虚拟内存分配器。
原文摘要 · Abstract (English)
An LLM serving engine sizes its key-value (KV) cache once, at startup, permanently setting aside a reserve for the worst-case prefill activation. During decode-dominant phases that reserve sits idle, yet it cannot be handed to the KV pool because it is exactly the memory a large prefill needs. We ask whether this reserve is reclaimable, and build a mechanism to test it. Our elastic KV cache lends the reserve to the KV pool during decode and returns it before prefill, driven by the scheduler's one-step-ahead view of the next batch. It is pure userspace on the CUDA virtual-memory path: two physical handles mapped into one contiguous virtual range per layer, so the attention kernel is unchanged and no driver patch is required. It decommits in a few milliseconds and recommits in tens of milliseconds, works with CUDA graphs and prefix caching, and never triggers an out-of-memory event. A static commit of the same memory is unsafe, crashing on prefill bursts, which makes the dynamic toggle necessary. Having built the mechanism, we test the premise it rests on and report an honest negative result. It only pays off if a small prefill chunk size badly hurts prefill latency. In a controlled experiment injecting long prompts into a live decode load, that penalty is small (median time-to-first-token differs by about 1% between chunk sizes of 8192 and 32768 tokens), because prefill is compute bound and decode consumes only about one token per sequence per step. Simply lowering max_num_batched_tokens recovers more KV than the controller does, at nearly equal latency. The reserve also dilutes under tensor parallelism, from 16% of KV at TP1 to 2.7% at TP4. We state precisely when reclaiming the reserve could still help, and release the mechanism as a reusable userspace elastic-VMM allocator.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。