Technical Summary: BlockBatch

Problem Statement

Diffusion Language Models (dLLMs) generate text by iteratively denoising multiple token positions in parallel, offering an alternative to strictly autoregressive decoding. However, efficient inference faces a critical granularity trade-off. Existing block-wise inference methods (e.g., Fast-dLLM) select a single block size per request. Small blocks preserve local conditioning but require many denoising steps (high NFE), while large blocks increase parallelism but risk premature commitments and accumulating cache errors (KV-cache drift). Current acceleration methods treat block size as a fixed hyperparameter, failing to exploit the complementary nature of different block sizes for different parts of the same generation.

Methodology: BlockBatch

The authors propose BlockBatch, a training-free online inference framework that treats block size as a branching dimension rather than a fixed parameter. Instead of selecting one block size, BlockBatch executes multiple branches with different block sizes (e.g., 4, 8, 16, 32, 64, 128) for the same request within a single batched forward pass.

The framework relies on three core mechanisms to coordinate these branches:

1. Confidence-Gated Merge: Branches exchange token proposals only when they are compatible (i.e., their decoded prefixes match exactly). A destination branch accepts a token proposed by a source branch only if the destination branch's own probability distribution assigns that token a confidence score above a fixed threshold. This allows lagging branches to make progress without forcing inconsistent token commitments.
2. Leader-Based Synchronization: To prevent lagging branches from wasting computation in stale regions of the KV-cache space, a synchronization rule is applied. If a leading branch advances beyond a lagging branch by a specific threshold number of tokens, the lagging branch copies the leader's sequence state and KV cache, effectively resetting its trajectory to the most advanced reliable path.
3. Periodic Full-Sequence Refreshes: Local block denoising updates only a subset of the sequence, which can lead to accumulated drift from the globally consistent KV state. BlockBatch periodically recomputes the KV cache for the entire sequence of active branches. This re-anchors the branches to a globally consistent state, correcting errors accumulated during local updates.

System Optimization: To ensure efficiency, BlockBatch packs the unmasked query positions from all active branches into a single variable-length buffer. This allows the model to process multiple branches in one forward pass, utilizing FlashAttention's varlen kernel, ensuring that the latency of multi-branch denoising is comparable to a single branch.

Key Characterizations

The paper provides theoretical and empirical characterizations of why this approach works:

• KV-Cache Trajectory: Different block sizes induce different trajectories in the KV-cache space. Small blocks follow conservative paths with fine-grained conditioning, while large blocks make aggressive local progress. Full-sequence refreshes act as global corrections that prevent these trajectories from drifting too far apart.
• Token-Level Dynamics: Branches often share an initial prefix, bifurcate at semantically decisive positions, and later re-agree on syntactically lightweight tokens (e.g., whitespace, operators, digits). The framework leverages this by merging branches when they agree on high-confidence tokens and synchronizing them when they diverge significantly.

Experimental Results

BlockBatch was evaluated on three representative dLLMs (LLaDA-1.5-8B, LLaDA-Instruct-8B, Dream-Base-7B) across four benchmarks (GSM8K, MATH, HumanEval, MBPP).

• Efficiency: BlockBatch reduces the average number of denoising function evaluations (NFEs) by 26.6% compared to Fast-dLLM (the state-of-the-art baseline).
• Speed: It achieves an average 1.33× end-to-end speedup over Fast-dLLM, with speedups reaching up to 5.8× in specific configurations (e.g., LLaDA-Instruct-8B on MBPP).
• Accuracy: Despite the aggressive reduction in NFEs and latency, BlockBatch preserves accuracy comparable to or better than the baselines across all tasks.
• Comparison: It consistently outperforms Vanilla diffusion decoding, Fast-dLLM, and LocalLeap in terms of NFE efficiency and latency across all twelve model-task settings tested.

Significance and Claims

The paper claims that block-size diversity is a practical and previously underexplored axis for branch-parallel dLLM inference. By demonstrating that different block sizes provide complementary accuracy-efficiency tradeoffs, the authors argue that treating block size as a dynamic branching dimension allows for more efficient decoding than any single fixed configuration.

The significance of BlockBatch lies in its ability to:

1. Exploit the inherent structure of dLLM inference where different granularities are useful for different parts of a generation.
2. Coordinate multiple decoding trajectories without training, using only inference-time operations (merge, sync, refresh).
3. Provide a general framework that improves efficiency while maintaining the quality of generated text, identifying block-size diversity as a viable strategy for accelerating diffusion-based language models.