FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion

Abstract: Generating long-form content, such as minute-long videos and extended texts, is increasingly important for modern generative models. Block diffusion improves inference efficiency via KV caching and block-wise causal inference and has been widely adopted in diffusion LLMs and video generation. However, in long-context settings, block diffusion still incurs substantial overhead from repeatedly computing attention over a growing KV cache. We identify an underexplored property of block diffusion: cross-step redundancy of attention within a block. Our analysis shows that attention outputs from tokens outside the current block remain largely stable across diffusion steps, while block-internal attention varies significantly. Based on this observation, we propose FlashBlock, a cached block-external attention mechanism that reuses stable attention output, reducing attention computation and KV cache access without modifying the diffusion process. Moreover, FlashBlock is orthogonal to sparse attention and can be combined as a complementary residual reuse strategy, substantially improving model accuracy under aggressive sparsification. Experiments on diffusion LLMs and video generation demonstrate up to 1.44$\times$ higher token throughput and up to 1.6$\times$ reduction in attention time, with negligible impact on generation quality. Project page: https://caesarhhh.github.io/FlashBlock/.

• The paper introduces FlashBlock, which caches cross-step-stable block-external attention and reduces per-step attention and KV-cache work from O(BN) to O(B²), with O(B) composition overhead.
• FlashBlock improves Trado-8B-Thinking throughput by up to 1.44× at contexts reaching 800k tokens while keeping benchmark accuracy nearly unchanged at the recommended reuse threshold τ=2.
• The method extends to sparse attention and video diffusion, improving sparse-model accuracy and reducing LongLive-1.3B attention time by about 1.6×, while remaining limited by kernel integration, stability failures, and model-specific validation.

Motivation and problem setting

Block diffusion combines autoregressive block-wise causal inference with iterative denoising, enabling KV caching and variable-length generation in diffusion LLMs (dLLMs) and long-horizon video diffusion. Despite per-step efficiency gains, block diffusion still recomputes attention over the entire growing KV cache at every diffusion step within a block, so in long-context regimes the combined cost of attention computation and KV-cache memory traffic dominates inference latency. Prior KV-reuse methods for dLLMs target standard (non-block) diffusion and do not exploit the structural properties of block-causal attention. The paper addresses this gap by identifying and exploiting a redundancy specific to block diffusion.

Empirical insight: cross-step stability of block-external attention

The central empirical finding is an asymmetry between block-internal and block-external attention across diffusion steps. Measuring the similarity of attention outputs between adjacent diffusion steps (on Trado-8B-Thinking, layer 3), the authors observe that attention over tokens outside the current block ($A_{\text{out}}$) remains highly stable across steps, with coherent structure in the cross-step similarity heatmaps, whereas block-internal attention ($A_{\text{in}}$) varies substantially because tokens in the block are actively refined or unmasked. The appendix extends this analysis to LongLive-1.3B for video diffusion, confirming the same separation for the majority of layers and heads, while noting that a small subset of heads exhibits lower block-external stability — a caveat that motivates the head-wise selective reuse strategy used for video models.

The implication is direct: recomputing block-external attention at each diffusion step is largely redundant, so caching it eliminates both attention computation and KV-cache reads that scale with context length $N$.

Method

FlashBlock decomposes attention into block-internal and block-external components via the standard softmax statistics (weighted sums $U$ and normalizers $Z$). At a given diffusion step, the block-external output $A_{\text{out}} = U_{\text{out}}/Z_{\text{out}}$ and its log-normalizer $L_{\text{out}} = \log Z_{\text{out}}$ are cached. At subsequent steps, only block-internal attention is recomputed, and the full output is recovered by log-space composition with max-subtraction for numerical stability in half precision, following the FlashAttention streaming formulation. The cached tensors depend only on the block size $B$, not on $N$; the implementation copies the running accumulators inside the FlashAttention kernel when the key stream crosses the block boundary, incurring constant overhead per query. Per-step KV access and accumulation drop from $O(BN)$ to $A_{\text{in}}$0, with $A_{\text{in}}$1 composition cost.

Two safeguards handle cases where the stability assumption fails. First, a threshold $A_{\text{in}}$2 on the number of updated tokens $A_{\text{in}}$3 governs reuse: if $A_{\text{in}}$4, block-external attention is recomputed to preserve correctness — relevant for discrete dLLMs where unmasking causes abrupt representational changes. Second, for video diffusion, head-wise selective reuse is applied: heads whose empirically estimated cross-step similarity exceeds a threshold $A_{\text{in}}$5 (set to 0.9) reuse the cache, while the rest recompute.

A notable property of the decomposition is that it generalizes to sparse attention: the "internal" set can be reinterpreted as the keys selected by a sparse policy, and the cached block-external term becomes a residual contribution from unselected keys. Under this reading, FlashBlock is orthogonal to sparsification and can compensate for the information loss that aggressive sparsification introduces.

For discrete dLLMs, where abrupt unmasking can violate the reuse assumption, the authors add a reuse-aware distillation stage: a frozen dense-attention teacher supervises a student with LoRA adapters (rank 32) on the Q/K/V projections, trained for 5,000 iterations on DAPO-Math-17K with rollouts of 1,000–4,000 tokens. The objective combines a KL loss on the reused-attention student output with a regularization KL term ($A_{\text{in}}$6) that anchors the student's dense-attention output to the teacher, stabilizing training.

Results on diffusion LLMs

Evaluated on Trado-8B-Thinking under identical decoding settings (batch size 128, 800k context for throughput), FlashBlock with $A_{\text{in}}$7 yields a 1.44× throughput increase (312 → 451 tokens/s at block size 4; 532 → 674 at block size 8) with minimal quality loss. At block size 4, accuracy deltas on GSM8K, MATH500, and AIME are within 0.2%; at block size 8, they remain within 1.6%. Coding pass@1 on MBPP, HumanEval, LiveCodeBench-V2, and LiveBench stays comparable to the dense baseline, with MBPP even improving in both configurations.

Latency analysis shows the gap versus the dense baseline widens with context length; from 100k to 800k context under $A_{\text{in}}$8, the latency growth of FlashBlock is roughly half that of the baseline, implying a theoretical speedup upper bound approaching 2× in extreme long-context settings. The reduction in effective attention token count tracks the latency gains, confirming that the improvement comes from avoided computation and KV-cache traffic.

The ablation on $A_{\text{in}}$9 and distillation is instructive. Without distillation, quality degrades sharply as reuse becomes more aggressive: at $N$0, AIME falls from 33.33 to 6.67 and MATH500 from 86.00 to 48.40. Distillation recovers a large portion of this loss (AIME 13.33, MATH500 66.20 at $N$1) and nearly closes the gap at $N$2 (e.g., AIME 33.33 and 30.00 respectively). Since $N$3 and $N$4 yield nearly identical latency, the authors adopt $N$5 as the default. A limitation worth noting is that the efficiency–quality trade-off is characterized only up to $N$6, and the distillation is trained on math data, so its generalization to other domains is not directly evaluated.

Combination with sparse attention

When applied on top of SparseD (adapted to block diffusion by recomputing the sparse mask once per block rather than once per sequence, since block diffusion invalidates the fixed-context assumption underlying the original method), FlashBlock consistently improves accuracy across densities. At density $N$7 on Trado-4B-Instruct, it improves GSM8K by +7.96, MATH500 by +7.40, and HumanEval pass@1 by +9.76; gains remain positive but shrink at higher densities (+3.64/+5.40/+6.71 at 30%; +2.65/+3.20/+5.49 at 40%). Output-level analysis supports a mechanistic explanation: the L1 distance between sparse and full attention outputs, which grows rapidly for SparseD alone as density decreases (0.0014 at 50% to 0.0031 at 10%), is reduced by roughly 2–3.5× when the cached residual is included (e.g., 0.0005 at 50% and 0.0008 at 10%).

The authors concede a practical caveat: SparseD is implemented only at the PyTorch level without kernels compatible with paged KV cache, so no end-to-end latency comparison under vLLM-style inference is possible for the combined method; accuracy is the only reported metric for this setting.

Results on video generation

On LongLive-1.3B with VBench2, FlashBlock with head-wise selective reuse ($N$8) at 55% attention density reduces per-step attention time from 23.02 s to 14.43 s — approximately 1.6× — and end-to-end latency from 93.23 s to 84.64 s, while maintaining quality close to the dense baseline across Human Fidelity (0.7861 vs. 0.8188), Creativity, Controllability (0.3026, actually exceeding dense), Physics, and Commonsense dimensions. Compared to SpargeAttn at 30–50% density, the method achieves better scores on most macro categories; the paper notes that SpargeAttn's per-step mask evaluation introduces overhead that limits its practical acceleration, and that SpargeAttn alone can produce temporal inconsistencies (visible in the Human Fidelity breakdown, where identity preservation drops from 0.7663 dense to 0.3133 at 30% density, versus 0.7593 for FlashBlock).

Two constraints bound the video results and are stated plainly: LongLive's fixed 12-frame temporal attention window limits the effective KV cache size and reduces attention's share of total runtime, capping end-to-end speedup; and the head-wise threshold $N$9 is calibrated on a small sample set before inference, an additional tuning step not required for the dLLM setting.

Limitations and open questions

The paper identifies several boundaries of its results. The sparse-attention combination is evaluated only in terms of accuracy, not end-to-end latency, because existing sparse attention kernels are incompatible with block diffusion and paged KV caches; developing block-diffusion-aware sparse kernels is left open. The stability assumption underlying reuse is empirically validated but not guaranteed — it can fail under aggressive unmasking in discrete dLLMs, which is why distillation and the $U$0 threshold are needed, and quality still degrades measurably at $U$1 even with distillation. For video models, a minority of attention heads remain unstable across steps, and the head-wise selection threshold $U$2 introduces a hyperparameter whose sensitivity is not systematically studied. Finally, the throughput claims are tied to the Trado architecture and the nano-vllm implementation; whether the observed cross-step stability holds for other dLLM families at different scales remains unverified.

Conclusion

FlashBlock exploits a previously unexamined structural property of block diffusion — the cross-step stability of block-external attention — to cache and reuse attention outputs, reducing per-step attention computation from $U$3 to $U$4 without altering the diffusion process. It delivers up to 1.44× throughput gains on diffusion LLMs and 1.6× attention-time reductions on video diffusion with negligible quality loss, and doubles as a residual-reuse mechanism that substantially mitigates the degradation caused by sparse attention. The main open issues are operator-level integration with sparse attention under paged KV caches and the generalization of the stability property across model families and training regimes.