🤖 AI Summary
This work addresses the limitation of block-diffusion large language models in long-context reasoning, where linear growth of the KV cache and PCIe bandwidth constraints hinder efficient offloading and retrieval of critical information. The authors propose a CPU-GPU协同 KV cache offloading mechanism that exploits the consistent correlation of KV entries across denoising steps within a block, enabling one-time selection of key KV entries for reuse throughout the entire block while overlapping selection with denoising computation. Through two key techniques, the method reduces KV selection overhead to the inverse of the original block size and achieves extremely low cache occupancy (5–10%) with near-lossless accuracy. Evaluated on three block-diffusion LLMs and five long-context tasks, the approach reduces per-block latency by up to 59% and improves throughput by 2.47× over GPU-only inference, with acceleration gains amplifying as context length increases.
📝 Abstract
Diffusion LLMs (dLLMs) improve GPU utilization over autoregressive decoding by generating multiple tokens per forward pass, but their KV cache still grows linearly with context, limiting throughput at long contexts. KV cache offloading to host DRAM alleviates this memory pressure, but the limited PCIe bandwidth necessitates recalling only a sparse subset of KV entries. In block dLLMs, the relevant KV entries remain consistent across denoising steps within a block, enabling high-accuracy selection by identifying the top-k entries once and reusing them throughout all denoising steps. This property appears attractive for offloading as it amortizes the selection overhead across the entire block, but it requires exact attention over the full KV cache, which is too expensive under offloading. We present HERALD, a KV offloading system for block dLLMs that resolves this through two opportunities that reduce the required selection compute by a factor of the block size and enable selection to be overlapped with denoising. Across three block dLLMs and five long context tasks, HERALD achieves near-lossless accuracy at 5-10% KV budget and up to 1.59x lower per block latency and 2.47x higher throughput over GPU-only inference, with speedups growing with context length.