FlashBlock: 효율적인 장문 컨텍스트 블록 확산을 위한 어텐션 캐싱

FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion

최신 생성 모델에서 분 단위 동영상이나 장문의 텍스트와 같은 장문 콘텐츠 생성은 점점 더 중요해지고 있습니다. 블록 확산(block diffusion)은 KV 캐싱과 블록 단위 인과 추론을 통해 추론 효율성을 향상시키며, 확산 언어 모델 및 비디오 생성 분야에서 널리 사용됩니다. 그러나 장문 컨텍스트 환경에서 블록 확산은 여전히 증가하는 KV 캐시에 대해 반복적으로 어텐션을 계산하면서 상당한 오버헤드를 발생시킵니다. 본 연구에서는 블록 확산의 간과된 특성인, 블록 내의 어텐션 단계 간의 중복성을 발견했습니다. 분석 결과, 현재 블록 외부의 토큰에서 생성되는 어텐션 출력은 확산 단계에 걸쳐 대체로 안정적인 반면, 블록 내부의 어텐션은 크게 변화합니다. 이러한 관찰을 바탕으로, 우리는 확산 과정을 수정하지 않고 안정적인 어텐션 출력을 재사용하여 어텐션 계산과 KV 캐시 접근을 줄이는 캐시된 블록 외부 어텐션 메커니즘인 FlashBlock을 제안합니다. 또한, FlashBlock은 희소 어텐션과 상호 보완적으로 작동하여, 공격적인 희소화 환경에서도 모델 정확도를 크게 향상시킬 수 있습니다. 확산 언어 모델 및 비디오 생성 실험 결과, FlashBlock은 최대 1.44배 더 높은 토큰 처리량을 제공하고, 어텐션 시간을 최대 1.6배 단축하며, 생성 품질에는 거의 영향을 미치지 않는 것으로 나타났습니다. 프로젝트 페이지: https://caesarhhh.github.io/FlashBlock/.

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 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 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. Project page: https://caesarhhh.github.io/FlashBlock/.