ICML 2026poster0 citations

FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion

Zhuokun Chen, Jianfei Cai, Bohan Zhuang

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 language models and video generation. However, in long-context settings, block diffusion still incurs substantial overhead from repeatedly computing attention over an ever-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, substantially 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. When integrated, it substantially improves model accuracy under aggressive sparsification by offsetting much of the performance loss induced by sparsity. Experiments on diffusion language models 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.

DiffusionTransformerCausalityVision
BibTeX
@inproceedings{
chen2026flashblock,
title={FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion},
author={Zhuokun Chen and Jianfei Cai and Bohan Zhuang},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=4jfuNNghPS}
}