Source-linked AI summary
GroupKV: Hierarchical KV Cache Management for Long-Context Diffusion LLM Inference
Jinhao Wang, Zhexin Hu, Kangjie Zhou, Xin Zhou, Fangfang Liu
TL;DR
Long-context dLLM inference is constrained by KV memory growth, transfer overhead, and dynamic updates. GroupKV uses hierarchical group selection with predictive prefetching, staleness correction, and streaming prefill, achieving substantially longer serviceable contexts and faster offloaded inference while retaining competitive accuracy.
Problem
Long-context dLLM inference faces KV-cache memory growth, host-device transfer bottlenecks, and dynamic updates that complicate cache management and prefetch scheduling.
Method
GroupKV partitions context into contiguous groups for coarse-to-fine KV recovery, predicts next-layer groups for prefetching, corrects stale cache data, and streams prefill.
Results
GroupKV extends maximum serviceable context length by up to 48.00× and improves end-to-end performance by up to 3.73× in offload-based long-context settings while maintaining competitive accuracy.
Takeaways & Limitations
GroupKV is most beneficial in offload-heavy long-context regimes where host-device KV transfer dominates inference latency.
Takeaways & Limitations
Streaming prefill reduces temporary activation peaks but does not reduce decoding KV footprint or host-device transfer volume, and GroupKV's transfer compression is structurally bounded at approximately 2× when complete keys remain available.
Abstract
from arXiv · showhide
Diffusion large language models (dLLMs) are emerging as a promising generative paradigm that complements autoregressive decoding. In long-context settings, KV cache bloat and offloading transfer overhead have become primary bottlenecks in inference systems. Meanwhile, the periodic full-sequence recomputation and localized token updates in dLLMs make the KV lifecycle substantially more dynamic, complicating cache management and prefetch scheduling while making heavyweight token-level indexing or clustering schemes harder to amortize effectively during decoding. To address these challenges, we present \textsc{GroupKV}, a lightweight hierarchical KV cache management system for long-context dLLM inference. We observe that under block-wise decoding, tokens within the same generation block tend to access highly overlapping and spatially concentrated context regions, making group-level sparse selection effective. Building on this observation, \textsc{GroupKV} partitions the context into contiguous groups and performs coarse-to-fine sparse selection. \textsc{GroupKV} further exploits cross-layer consistency to enable predictive prefetching, and incorporates a staleness correction mechanism to maintain cache coherence under dynamic KV updates. Additionally, \textsc{GroupKV} adopts streaming prefill to reduce peak memory consumption during prefilling. Experiments show that \textsc{GroupKV} extends the maximum serviceable context length by up to $48.00\times$ under constrained GPU memory, improves end-to-end inference performance by up to $3.73\times$ in offload-based long-context settings, and maintains competitive task accuracy.
1 Introduction
Long-context dLLM inference is limited by KV memory growth, host-device transfer overhead, and a dynamic cache lifecycle. GroupKV addresses these challenges with hierarchical group selection, predictive prefetching, staleness correction, and streaming prefill.
- 1 Introduction: KV offloading extends serviceable context beyond GPU capacity but makes host-device data movement the primary bottleneck.Long-context KV caches grow with sequence length and batch size, quickly exceeding GPU memory.
- 1 Introduction: Block-wise dLLM decoding yields highly overlapping, spatially concentrated context access within each generation block, motivating contiguous groups as the selection unit.This locality reduces the need to union disparate per-token selections across the full context.
- 1 Introduction: GroupKV partitions the prompt into groups, selects important groups coarsely, and recovers token-level KV only within them for fine-grained computation.The approach uses lightweight group summaries to avoid the high amortized cost of complex clustering under periodic refresh.
- 1 Introduction: Cross-layer group-saliency consistency enables predictive prefetching that overlaps next-layer KV transfers with current-layer computation, while staleness correction handles local KV updates.The dynamic dLLM cache lifecycle makes prefetched data potentially stale before attention consumes it.
- 1 Introduction: Streaming prefill reduces temporary activation peaks during prefilling, while group sparsity, prefetching, and staleness correction remain the core decode-side mechanisms.Streaming prefill is orthogonal to decode-side KV management and does not reduce decoding KV footprint or transfer volume.
- 1 Introduction: 48.0× maximum serviceable context length and 3.73× end-to-end inference speedup are reported under constrained-memory and offload-based long-context settings, respectively.The system maintains competitive accuracy on LongBench, RULER, and NIAH benchmarks.
2 Related Work
Related work spans diffusion-language-model development, long-context extensions, autoregressive KV optimization, and dLLM-specific cache acceleration. GroupKV targets dynamic dLLM KV management with group summaries and selective repair rather than assuming autoregressive append-only caching.
- 2 Related Work: dLLMs support parallel token generation and bidirectional attention, while long-context variants extend context through methods such as NTK-based RoPE extrapolation and diffusion-aware NTK.These works establish long-context diffusion modeling as a distinct application area.
- 2 Related Work: Autoregressive KV methods use stable attention patterns, retrieval, compression, or hierarchical selection under an append-only decoding lifecycle.Their token-level structures can be reused throughout unidirectional decoding, unlike the dynamic updates in dLLMs.
- 2 Related Work: GroupKV differs from autoregressive optimizations by using lightweight group summaries and selective repair to support group-level prefetching under local token updates.This design targets the dynamic KV lifecycle created by bidirectional attention and iterative denoising.
- 2 Related Work: dLLM-specific methods exploit similarity across denoising steps, but existing approaches variously emphasize computation, GPU-memory eviction, or cache acceleration rather than offloaded long-context management.MaskKV remains bounded by GPU memory because it does not incorporate KV offloading.
3 Observations
The analysis finds that block-wise dLLM decoding exhibits overlapping, spatially concentrated attention and strong adjacent-layer consistency, motivating hierarchical group-level KV retrieval and predictive prefetching.
- 3.1 Intra-block Locality Enables Group-level Retrieval: Tokens in the same generation block attend to overlapping context regions concentrated in a few continuous segments.With top-k set to 20%, the observed overlap and spatial concentration make block-level context access substantially more structured than arbitrary token selection.
- 3.1 Intra-block Locality Enables Group-level Retrieval: GroupKV partitions context into contiguous groups, summarizes each group with a mean-pooled key, and ranks groups before token-level refinement.The design reduces the selection space while preserving semantic and positional information in the group-level representation.
- 3.1 Intra-block Locality Enables Group-level Retrieval: Group retrieval remains efficient when salient tokens occupy contiguous segments, because boundary expansion adds only limited extra tokens.The required group count is bounded by segment lengths plus a boundary-misalignment term; added boundary tokens generally receive low attention weights.
- 3.2 Inter-layer Consistency in Attention Saliency: Group-level selections are more stable across adjacent layers than token-level selections, supporting cross-layer predictive prefetching.Table 3 quantifies adjacent-layer recall at both granularities, while the accompanying analysis reports high consistency across consecutive layers.
- 3.2 Inter-layer Consistency in Attention Saliency: Adjacent layers largely preserve focal context regions, with roughly 11% of top-k tokens persistently salient across multiple layers.Attention changes are described primarily as intensity refinements around stable segments, making the current layer a useful prior for the next layer.
- 3.2 Inter-layer Consistency in Attention Saliency: Figure 4 contrasts full-KV prefetching, critical-path on-demand fetching, predictive prefetching, and GroupKV’s added intra-layer refinement.The comparison highlights overlap of predictive transfer with computation and further computational reduction from refinement.
4 GroupKV Design
GroupKV manages dynamic dLLM KV caches with hierarchical group-level selection, predictive cross-layer prefetching, and selective staleness repair. Streaming prefill separately reduces activation peaks during prefilling.
- 4.1 Group-Level Selection and Hierarchical KV Access: GroupKV selects salient contiguous groups before recovering token-level K/V, avoiding full token scans and reducing cross-device data movement.The method maintains lightweight per-group summaries and performs fine-grained refinement only within selected groups.
- 4.1 Group-Level Selection and Hierarchical KV Access: 3.81× theoretical data-movement reduction follows under group size g=8 and recovered-token fraction f=0.2.Per-layer transfer combines approximately Nctx/g group keys with 2fNctx token-level K/V, instead of 2Nctx entries.
- 4.2 Cross-Layer Predictive Prefetching: GroupKV overlaps current-layer computation with next-layer KV prefetching using cross-layer group-selection consistency and alternating compact buffers.The compute stream consumes prepared sparse KV while the prefetch stream asynchronously loads the next layer’s token-level KV.
- 4.1 Group-Level Selection and Hierarchical KV Access: Intra-layer refinement retains top-contributing recovered tokens for value-side attention, reducing matrix computation without changing coarse-grained transfer.This refinement is applied after candidate KV recovery within selected groups.
- 4.3 Staleness Correction: Selective staleness correction refreshes affected CPU group summaries and repairs only stale compact-buffer entries after generation-slot updates.Repair cost scales with stale hits rather than full context length, while stable prompt-only groups need no GPU-side repair.
- 4.4 Throughput-Oriented Streaming Prefilling: Streaming prefill reduces temporary activation peaks caused by full-context intermediate tensors but does not reduce decode-side KV footprint or transfer volume.It is orthogonal to the core decode-side mechanisms: group sparsity, predictive prefetching, and staleness correction.
5 Evaluation
GroupKV is evaluated for accuracy, throughput, context-length scalability, and component sensitivity across long-context dLLM workloads. It remains competitive in accuracy while improving serviceability and offload-based performance under constrained memory.
- 5.2 Accuracy Evaluation: GroupKV scores 32.73 on LLaDA and 37.43 on UltraLLaDA on LongBench while retaining 20% of the KV budget.On UltraLLaDA, it exceeds the full-attention baseline by +0.82 points and outperforms the InfiniGen-style baseline by +2.51 points.
- 5.2 Accuracy Evaluation: 70.22 and 62.81 RULER scores at 16K and 32K slightly outperform full-attention UltraLLaDA while Fast-dLLM fails with OOM.GroupKV remains robust across the evaluated 4K–32K context range.
- 5.2 Accuracy Evaluation: 98.17% average NIAH retrieval accuracy across contexts up to 128K, including 95% at 128K.Retrieval reaches 100% at 32K, indicating that aggressive group-level filtering does not catastrophically discard needle information in this stress test.
- 5.3 Efficiency Evaluation: 2.0–3.7× higher throughput than FlexGen-style offload baselines across 4K–128K contexts on A100-80G, while remaining operational where Fast-dLLM runs OOM.GroupKV reaches 10.64 tok/s at 4K and 0.39 tok/s at 128K; it also exceeds InfiniGen-style throughput by 1.25–1.35×.
- 5.4 Ablation Studies: Predictive prefetching raises 32K decoding throughput from 1.08 to 1.79 tok/s with peak VRAM increasing only from 15.87 to 16.00 GiB.At 64K, it improves throughput from 0.60 to 0.99 tok/s with VRAM increasing from 16.28 to 16.53 GiB.
- 5.4 Ablation Studies: The default g=8 and top-k=20% configuration provides a favorable accuracy–throughput trade-off, while k2=0.50 is the stable refinement operating point.Reducing k2 to 0.25 improves throughput by 6.4% and 8.9% at 32K and 64K but lowers LongBench by 0.54 points; k2=0.75 reduces throughput by 20–29%.
6 Conclusion
GroupKV is a memory-efficient long-context dLLM inference system built on intra-block locality and inter-layer consistency. It extends serviceable context length and improves offload-based performance while preserving a tunable accuracy–efficiency trade-off.
- 6 Conclusion: GroupKV extends maximum serviceable context length by up to 48.00× under constrained GPU memory.The system targets long-context diffusion LLM inference with memory-efficient KV management.
- 6 Conclusion: GroupKV improves end-to-end performance by up to 3.73× in offload-based long-context settings.Its benefits are most pronounced when host-device KV transfer dominates inference latency.
- 6 Conclusion: GroupKV provides a tunable accuracy–efficiency trade-off for memory-constrained long-context dLLM inference.Future work concerns adaptive budget selection across context lengths, batch sizes, hardware platforms, and dLLM architectures.