💬 NLP

FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion

This paper proposes FlashBlock, a novel mechanism that accelerates long-context block diffusion by caching and reusing stable cross-step attention outputs from tokens outside the current block, thereby significantly reducing computational overhead and KV cache access while maintaining generation quality.

Original authors: Zhuokun Chen, Jianfei Cai, Bohan Zhuang

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhuokun Chen, Jianfei Cai, Bohan Zhuang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are directing a movie. Every scene (or "block") needs to be filmed, but to make the story flow, the director constantly looks back at everything that happened in previous scenes to keep the characters consistent.

In the world of AI video and text generation, this "looking back" is called attention. The AI has to re-read its entire history (the "KV cache") every single time it takes a tiny step to refine the next part of the story. As the story gets longer, this process becomes incredibly slow and exhausting, like a director who has to re-watch the entire 10-hour movie every time they want to adjust a single line of dialogue in the current scene.

The paper introduces FlashBlock, a clever trick to speed this up without ruining the movie. Here is how it works, using simple analogies:

1. The Problem: The "Re-Reading" Trap

In current "Block Diffusion" models, the AI generates content in chunks (blocks). Even though it only changes the current chunk, it still re-calculates how that chunk relates to everything that came before it.

The Analogy: Imagine you are writing a long book. Every time you write a new paragraph, you re-read the entire book from page 1 to the current page to make sure your new sentence fits. As the book gets longer, you spend 99% of your time re-reading old pages and only 1% actually writing new ones.

2. The Discovery: The "Stable Background"

The researchers noticed something interesting while watching the AI work.

The "Inside" vs. "Outside": When the AI works on a specific block (the current scene), the details inside that block are changing rapidly (actors moving, dialogue shifting). However, the influence of the old scenes (the background context) is surprisingly stable.
The Analogy: Think of a news anchor reading the news. The anchor's face and voice (the current block) might change expression or tone slightly every second. But the news ticker scrolling at the bottom of the screen (the old context) stays exactly the same for a long time.
The Insight: The AI was wasting energy re-calculating the "news ticker" (the old context) every single second, even though it hadn't changed.

3. The Solution: FlashBlock (The "Memo" System)

FlashBlock acts like a smart assistant who keeps a memo of the stable parts.

How it works: Instead of re-reading the whole history, the AI says, "Okay, the old stuff hasn't changed much. I'll just grab my cached memo of how the past relates to the present, and I'll only do the hard math for the new stuff I'm currently working on."
The Metaphor: It's like a chef cooking a stew. The broth (the old context) has been simmering and is stable. Instead of tasting the whole pot from scratch every time they add a new spice, the chef trusts the broth's flavor and only tastes the new spice they just added.

4. The Result: Faster and Smarter

By skipping the re-calculation of the stable "old stuff," FlashBlock achieves two main things:

Speed: It makes the AI up to 1.44 times faster at generating text and video. It's like cutting the time it takes to write a chapter in half because you stopped re-reading the whole library every time.
Quality: Because the "memo" is so accurate, the quality of the story doesn't drop. The AI still understands the context perfectly; it just does it more efficiently.

5. The Bonus: Teamwork with "Sparse Attention"

The paper also mentions that FlashBlock works great with another technique called "Sparse Attention" (which is like only reading the most important sentences).

The Analogy: If "Sparse Attention" is like a reader who only reads the headlines, they might miss some details. FlashBlock acts as a safety net, filling in the missing details from the "memo" of the skipped parts. This allows the AI to be even faster (by reading fewer words) without losing the plot.

Summary

FlashBlock is a smart caching system for AI. It realizes that when generating long stories or videos, the "old" parts of the story don't change much from one moment to the next. By remembering these stable parts instead of re-calculating them, it frees up the AI to focus its energy on the new, changing parts, making long-form generation much faster without losing quality.

Technical Summary: FlashBlock

1. Problem Statement

While block diffusion has improved inference efficiency in long-context generative models (e.g., diffusion language models and video generation) by combining KV caching with block-wise causal inference, it still suffers from significant overhead in long-context settings. In standard block diffusion, every diffusion step requires recomputing attention over the entire growing KV cache. As context length increases, the combined cost of attention computation and KV cache access becomes the dominant factor in inference latency.

The paper identifies an underexplored property of block diffusion: cross-step redundancy. Specifically, during the denoising process, tokens outside the current block being updated remain largely static, leading to highly stable attention contributions from these "block-external" tokens across adjacent diffusion steps. Conversely, attention within the current block ("block-internal") varies significantly as tokens are actively refined or unmasked. Current methods fail to exploit this stability, repeatedly recomputing redundant block-external attention.

2. Methodology: FlashBlock

The authors propose FlashBlock, a mechanism that caches and reuses block-external attention outputs to reduce computation and memory I/O without modifying the underlying diffusion process.

Core Mechanism

FlashBlock decomposes the full attention output into two components:

Block-Internal Attention ( $A_{in}$ ): Computed over tokens within the current block ( $J_{in}$ ), which are actively changing.
Block-External Attention ( $A_{out}$ ): Computed over tokens outside the current block ( $J_{out}$ ), which remain stable across steps.

The method caches the block-external attention output ( $A_{out}$ ) and its log-normalizer ( $L_{out} = \log Z_{out}$ ) from the previous diffusion step. In subsequent steps, instead of re-accessing the full KV cache for $J_{out}$ , the model reuses the cached values. The full attention output is reconstructed by combining the cached $A_{out}$ with the newly computed $A_{in}$ using log-space composition (similar to FlashAttention) to ensure numerical stability:
$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$
where $m = \max(L_{out}, L_{in})$ .

Adaptive Reuse Strategies

Threshold-based Reuse: In discrete diffusion models, token updates can be abrupt. The method introduces a threshold $\tau$ on the number of updated tokens ( $M^{s+1}$ ). If $M^{s+1} < \tau$ , cached attention is reused; otherwise, it is recomputed to ensure correctness.
Head-wise Selection (Video): For video diffusion, where cross-step variability is higher, the method estimates similarity per attention head. Only heads with similarity exceeding a threshold $\gamma$ reuse the cache.
Compatibility with Sparse Attention: FlashBlock is orthogonal to sparse attention. It treats the unselected keys in sparse attention as the "block-external" set, caching their residual contributions. This allows FlashBlock to compensate for information loss caused by aggressive sparsification.

Training Adaptation (Optional)

To address potential distributional mismatches in discrete diffusion language models (dLLMs) caused by abrupt token unmasking, the authors propose a reuse-aware distillation strategy. A student model employing FlashBlock is trained to match the output distribution of a frozen teacher model using dense attention, utilizing LoRA adapters to keep the adaptation lightweight.

3. Key Contributions

Empirical Insight: The paper provides an empirical study revealing strong cross-step stability in block-external attention, contrasting it with the high variability of block-internal attention.
FlashBlock Mechanism: A novel caching mechanism that exploits this redundancy to reduce attention computation and KV cache access, implemented via log-space aggregation.
Synergy with Sparsity: Demonstration that FlashBlock serves as a complementary residual reuse strategy for sparse attention, mitigating quality degradation from aggressive sparsification.
Efficiency Gains: Significant acceleration in inference throughput and latency reduction for both diffusion language models and video generation models with negligible impact on generation quality.

4. Experimental Results

Experiments were conducted on diffusion language models (Trado-8B/4B) and video diffusion models (LongLive-1.3B).

Throughput: FlashBlock achieved up to 1.44× higher token throughput on diffusion language models (e.g., increasing from 312 to 451 tokens/s for block size 4).
Latency: Attention time was reduced by up to 1.6×.
Quality: Generation quality remained comparable to dense baselines. On mathematical reasoning (GSM8K, MATH500) and coding benchmarks (HumanEval, MBPP), accuracy differences were minimal (within 0.2% to 1.6% depending on block size).
Long-Context Scaling: As context length increased from 100k to 800k, the latency gap between FlashBlock and the baseline widened, with FlashBlock showing approximately half the latency growth of the baseline under certain thresholds.
Sparse Attention Combination: When combined with SparseD (at 20%, 30%, and 40% density), FlashBlock consistently improved accuracy over SparseD alone (e.g., improving GSM8K accuracy from 34.72% to 42.68% at 20% density), effectively recovering information lost to sparsification.

5. Significance

The paper claims that FlashBlock addresses a critical bottleneck in long-context block diffusion: the redundant recomputation of stable attention contributions. By decoupling the stable "block-external" context from the dynamic "block-internal" updates, the method substantially reduces the computational and memory I/O costs associated with growing KV caches.

The authors emphasize that this approach is orthogonal to existing sparse attention techniques, offering a complementary strategy that enhances both efficiency and accuracy. The results suggest that FlashBlock provides a scalable solution for efficient long-context generation in both text and video domains, enabling faster inference without compromising the quality of the generated content.