HERALD: High-Throughput Block Diffusion LLM Inference via CPU-GPU Cooperative KV Cache Retrieval
Omin Kwon, Doyeon Kim, Jongseok Park, Seung Yul Lee, Sehoon Kim, Ion Stoica, Jae W. Lee
Block diffusion LLMs (block dLLMs) improve decode throughput by generating tokens in parallel, and are emerging as a next-generation alternative to autoregressive LLMs. At long contexts, however, the KV cache remains a major throughput bottleneck in both memory bandwidth and capacity. For autoregressive LLMs, dynamic sparse attention is often combined with KV offloading to alleviate both bottlenecks, keeping the full KV cache in CPU DRAM and fetching only a critical subset to the GPU. However, KV subset selection and retrieval over bandwidth-limited PCIe incur substantial offloading overheads that often stall the GPU. We present HERALD, the first dynamic sparse attention method for block dLLMs that supports KV offloading. HERALD exploits a property of block dLLMs: a critical KV subset selected once per block is reused across all of its denoising steps, so the next block's selection can be hidden behind the current block's denoising. HERALD realizes this with a two-stage selection algorithm co-designed with its execution system. While the GPU denoises the current block, the CPU speculatively searches the full KV cache for the next block with a single query merged from the current block's B queries, cutting the selection compute by Bx; the GPU then refines the candidates into the final top-K subset with the next block's real queries. At K=1024, HERALD achieves near-lossless accuracy relative to Dense on LongBench and reasoning tasks across four block dLLMs, and improves decode throughput by up to 3.70x over Dense GPU-only inference, 2.59x over MAGE, a state-of-the-art block-dLLM sparse attention method, and 8.38x over MAGE-Offload, its direct KV-offloading extension.
尚无 AI 解读，展示原文摘要