Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs

Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Diffusion LLMs§

Overview§

Diffusion Large Language Models (dLLMs) generate text non-autoregressively, but inference is bottlenecked by the lack of effective KV caching and scalable parallel decoding. Prior work treats caching and parallel decoding separately, ignoring I/O bottlenecks that arise when both are combined. Flash-dLLM is a training-free acceleration framework that co-designs cache layout and decoding to maximize GPU memory throughput.

Core Problem: I/O-Bound KV Reuse in dLLMs§

In standard autoregressive decoding, each step appends one token and reads the KV cache sequentially. In dLLMs, a denoising step updates many token positions simultaneously, and draft-and-verify schemes issue sparse, position-dependent reads across the cache. This creates fragmented memory access patterns: instead of one contiguous scan, the GPU performs many small, scattered loads of $K,V$ tensors. The arithmetic intensity collapses, so the kernel becomes memory-bandwidth-bound rather than compute-bound.

Let cache size be $C = 2 \cdot L \cdot H \cdot d \cdot n_{\text{ctx}}$ bytes. The effective bandwidth $B_{\text{eff}}$ for scattered access is significantly lower than peak $B_{\text{peak}}$, giving inference time: $$T \approx \frac{C_{\text{read}} + C_{\text{write}}}{B_{\text{eff}}} \gg \frac{C}{B_{\text{peak}}}.$$ Flash-dLLM targets raising $B_{\text{eff}}$.

Method 1: I/O-Aware Fused KV-Cache Kernel§

Flash-dLLM fuses the gather, concatenation, and attention-read phases into a single kernel. Key techniques:

• Layout-aware tiling: KV cache is stored in a padded, position-mapped layout so draft tokens map to contiguous cache blocks, avoiding per-token gathers.
• Fused scatter-update: New $K,V$ for verified tokens are written in-place with vectorized stores, eliminating a separate copy.
• Register-level reuse: Attention over the draft window uses shared-memory staging so each cache line is loaded once and reused across the verification batch.

Symbolically, attention for a query block $Q$ becomes: $$\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V,$$ where $K,V$ are streamed from the fused cache with minimal redundant movement.

Method 2: Cache-Driven Draft-and-Verify§

Instead of an auxiliary draft model, the dLLM acts as both drafter and verifier. The draft step reuses the existing KV cache to propose $\gamma$ tokens in parallel; verification is a single forward pass over the concatenated draft, accepting tokens where the dLLM's own distribution agrees: $$\text{accept } x_i \iff \frac{p_{\text{dLLM}}(x_i \mid x_{<i})}{q_{\text{draft}}(x_i \mid x_{<i})} \ge u,\quad u \sim \mathcal{U}(0,1).$$ Because drafter and verifier share weights and cache, no extra memory is required, and acceptance is computed with the same fused kernel.

Implementation Sketch (PyTorch-style)§

import torch, torch.nn.functional as F

def flash_dllm_verify(cache_k, cache_v, draft_ids, model, gamma):
    # draft_ids: [B, gamma] proposed tokens from the dLLM itself
    q = model.q_proj(draft_ids)                      # [B, gamma, H, d]
    k_new = model.k_proj(draft_ids)
    v_new = model.v_proj(draft_ids)
    # Fused cache update: append contiguously (padded layout)
    cache_k = torch.cat([cache_k, k_new], dim=1)
    cache_v = torch.cat([cache_v, v_new], dim=1)
    # I/O-aware attention: one load per cache line
    attn = F.scaled_dot_product_attention(q, cache_k, cache_v)
    logits = model.lm_head(attn)                     # [B, gamma, vocab]
    # Self-verification: accept iff draft matches dLLM argmax/ratio test
    probs = logits.softmax(-1)
    accept = probs.gather(-1, draft_ids[..., None]).squeeze(-1) > 0.5
    return accept, cache_k, cache_v

Results§

On GSM8K and HumanEval, Flash-dLLM achieves 5.1$\times$ and 11.0$\times$ speedups over Elastic-Cache, the strongest prior baseline, while improving memory efficiency and scaling to longer contexts and larger batches. Because it is training-free, it can be dropped into existing dLLMs without retraining.