TideKV reduces memory use for longer AI language model prompts
TideKV: Delayed KV Eviction via Layer-Wise Chunked Prefill for Memory-Efficient Long-Context LLM Inference
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
When large language models read very long texts, they need a lot of memory to remember what they've seen so far. The researchers developed TideKV, a new way to manage this memory use so the model can handle much longer texts without running out of memory. TideKV waits to clear some memory until each layer of the model has fully processed the current chunks, which keeps performance high but uses less memory at peak times. Tests showed it works well on different computer platforms, letting models process longer inputs more efficiently.
What this means in practice
- •For ai inference engineers: Run language models with much longer input contexts using less peak memory on GPUs and edge devices.
- •For embedded system developers: Enable long-context LLM inference on resource-limited platforms like Jetson where memory is typically a bottleneck.
Authors
Yu Wang, Zhu Wang, Ning Hu, Yifan Guo, Xiu Ma, Bin Guo, Zhiwen Yu
Abstract
Long-context large language model (LLM) inference is important for modern intelligent applications but remains constrained by peak memory during prompt prefill. Existing KV cache compression reduces decoding memory but typically requires full-prompt prefill to preserve global context, retaining a high construction peak. Online chunk-level eviction lowers memory earlier but selects before each layer has observed the complete prompt, potentially degrading cache construction and decoding. We present TideKV, an on-GPU cache-construction architecture using layer-major delayed eviction. TideKV processes all prompt chunks within each Transformer layer, retains full-prompt KV only for the active layer, and evicts at layer end. This schedule preserves full-prompt-aware selection while reducing activation and temporary memory. Experiments across multiple computing platforms show that TideKV preserves the inference quality and final cache size of the evaluated compression methods. At 64K on A100, it reduces the construction-time transient footprint by 83.89%. On RTX 4090, it increases time to first token by only 1.06% at 64K and enables 128K prefill. TideKV also completes 128K on Jetson, whereas matched post-prefill compression runs out of memory on both resource-constrained platforms. These results show that delayed layer-end eviction substantially expands feasible context lengths without sacrificing cache utility.