💬 NLP

FlashBlock: Attention Caching for Efficient Long-Context Block Diffusion

本文提出了 FlashBlock，一种通过缓存并复用当前块之外 token 的稳定跨步注意力输出，从而显著降低计算开销和 KV 缓存访问量，同时保持生成质量的新颖机制，旨在加速长上下文块扩散。

原作者： Zhuokun Chen, Jianfei Cai, Bohan Zhuang

发布于 2026-07-07

📖 1 分钟阅读☕ 轻松阅读

原作者： Zhuokun Chen, Jianfei Cai, Bohan Zhuang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在执导一部电影。每一场戏（或称“块/block”）都需要被拍摄，但为了让故事流畅，导演会不断回顾之前发生的一切，以保持角色的一致性。

在 AI 视频和文本生成的世界里，这种“向后看”的过程被称为注意力（attention）。AI 每当要通过微小的步骤来完善故事的下一个部分时，都必须重新阅读它的整个历史（即“KV cache”）。随着故事变得越来越长，这个过程会变得异常缓慢且令人精疲力竭，就像一位导演每想要调整当前场景的一句台词，就不得不把整部 10 小时的电影重看一遍。

这篇论文介绍了一种名为 FlashBlock 的巧妙技巧，它能在不破坏电影质量的前提下提高速度。以下是它的工作原理，使用了简单的类比：

1. 问题所在：“重读”陷阱

在目前的“块扩散”（Block Diffusion）模型中，AI 是以块（块/blocks）为单位生成内容的。尽管它只改变当前的这一块，但它仍然会重新计算这一块与之前发生的一切之间的关系。

类比： 想象你正在写一本长篇小说。每当你写一个新段落时，你都会从第一页开始重读整本书，直到当前页，以确保你的新句子符合逻辑。随着书的内容变长，你 99% 的时间都在重读旧页面，只有 1% 的时间在真正进行创作。

2. 发现：“稳定的背景”

研究人员在观察 AI 工作时发现了一些有趣的现象。

“内部”与“外部”： 当 AI 处理特定块（当前的场景）时，该块“内部”的细节变化很快（演员在移动，对话在切换）。然而，来自“旧”场景（背景上下文）的影响却出奇地稳定。
类比： 想象一位新闻主播正在播报新闻。主播的面部表情和声音（当前的块）可能每秒钟都在发生细微的变化。但屏幕下方滚动的新闻标题栏（旧的上下文）在很长一段时间内都是完全不变的。
洞察： AI 过去一直在浪费精力去重新计算这个“新闻标题栏”（旧的上下文），即便它并没有发生变化。

3. 解决方案：FlashBlock（“备忘录”系统）

FlashBlock 就像一个记录着稳定部分的智能助手。

工作原理： 它不再重新阅读整个历史，而是说：“好吧，旧的内容变化不大。我直接调用关于‘过去如何关联现在’的缓存备忘录，我只对目前正在处理的‘新内容’进行复杂的数学运算。”
隐喻： 这就像一位厨师在炖汤。汤底（旧的上下文）已经炖好了，非常稳定。厨师不需要每次加入新香料时都从头品尝整锅汤，而是信任汤底的味道，只品尝刚刚加入的那种新香料。

4. 结果：更快、更聪明

通过跳过对稳定的“旧内容”的重新计算，FlashBlock 实现了两个主要目标：

速度： 它使 AI 生成文本和视频的速度提升了高达 1.44 倍。这就像因为停止了每次写新章节都要重读整个图书馆，从而将写完一个章节所需的时间缩短了一半。
质量： 由于这个“备忘录”非常准确，故事的质量并不会下降。AI 仍然能完美理解上下文，只是处理得更加高效。

5. 加分项：与“稀疏注意力”的协作

论文还提到，FlashBlock 与另一种被称为“稀疏注意力”（Sparse Attention，类似于只阅读最重要的句子）的技术配合得非常好。

类比： 如果说“稀疏注意力”是一个只看标题的读者，他们可能会错过一些细节。而 FlashBlock 则充当了一个安全网，利用“备忘录”填补了被跳过的部分所缺失的细节。这使得 AI 即使在阅读更少文字的情况下也能更快，且不会跑题。

总结

FlashBlock 是一个智能的 AI 缓存系统。它意识到，在生成长篇故事或视频时，故事的“旧”部分在每一时刻之间变化并不大。通过记住这些稳定的部分而不是重新计算它们，它释放了 AI 的精力，使其能够专注于新的、变化的场景，从而在不损失质量的前提下，让长篇内容的生成变得更加快速。

技术摘要：FlashBlock

1. 问题陈述

虽然块扩散（block diffusion）通过结合 KV 缓存与块状因果推理，提高了长上下文生成模型（如扩散语言模型和视频生成模型）的推理效率，但在长上下文设置下仍面临显著的开销问题。在标准的块扩散中，每一个扩散步骤都需要对整个不断增长的 KV 缓存重新进行注意力计算。随着上下文长度的增加，注意力计算与 KV 缓存访问的综合成本成为了推理延迟的主导因素。

本文指出块扩散中一个尚未被充分探索的特性：跨步冗余性（cross-step redundancy）。具体而言，在去噪过程中，当前正在更新的块之外的 token 保持高度稳定，导致这些“块外部（block-external）”token 在相邻扩散步骤中的注意力贡献非常稳定。相反，当前块内的注意力（“块内部/block-internal”）则随着 token 的积极精炼或去掩码过程而发生显著变化。现有方法未能利用这种稳定性，反复进行冗余的块外部注意力重计算。

2. 方法论：FlashBlock

作者提出了 FlashBlock，一种通过缓存并复用块外部注意力输出，从而在不修改底层扩散过程的情况下减少计算量和内存 I/O 的机制。

核心机制

FlashBlock 将完整的注意力输出分解为两个部分：

块内部注意力 ( $A_{in}$ )：在当前块内的 token ( $J_{in}$ ) 上计算，这些 token 处于活跃变化状态。
块外部注意力 ( $A_{out}$ )：在当前块外的 token ( $J_{out}$ ) 上计算，这些 token 在不同步骤间保持稳定。

该方法缓存了前一个扩散步骤中的块外部注意力输出 ( $A_{out}$ ) 及其对数归一化因子 ( $L_{out} = \log Z_{out}$ )。在后续步骤中，模型不再重新访问用于 $J_{out}$ 的完整 KV 缓存，而是复用缓存的值。完整的注意力输出通过使用对数空间组合（log-space composition）（类似于 FlashAttention）将缓存的 $A_{out}$ 与新计算的 $A_{in}$ 进行重构，以确保数值稳定性：
$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$
其中 $m = \max(L_{out}, L_{in})$ 。

自适应复用策略

基于阈值的复用：在离散扩散模型中，token 的更新可能是突发的。该方法引入了一个针对更新 token 数量 ( $M^{s+1}$ ) 的阈值 $\tau$ 。如果 $M^{s+1} < \tau$ ，则复用缓存的注意力；否则，进行重新计算以确保正确性。
逐头选择（视频）：对于跨步变异性更高的视频扩散，该方法会估计每个注意力头的相似度。只有相似度超过阈值 $\gamma$ 的注意力头才会复用缓存。
与稀疏注意力的兼容性：FlashBlock 与稀疏注意力是正交的。它将稀疏注意力中未被选中的 key 视为“块外部”集合，并缓存它们的残差贡献。这使得 FlashBlock 能够补偿由激进稀疏化导致的的信息损失。

训练适配（可选）

为了解决离散扩散语言模型（dLLM）中由于突发 token 去掩码可能导致的分布失配问题，作者提出了一种**感知复用的蒸馏（reuse-aware distillation）**策略。使用 FlashBlock 的学生模型通过使用 LoRA 适配器进行轻量化训练，旨在匹配使用密集注意力的冻结教师模型的输出分布。

3. 核心贡献

经验洞察：本文通过实证研究揭示了块外部注意力具有极强的跨步稳定性，这与块内部注意力的高变异性形成了鲜明对比。
FlashBlock 机制：一种新型缓存机制，利用这种冗余性来减少注意力计算和 KV 缓存访问，并通过对数空间聚合实现。
与稀疏性的协同作用：证明了 FlashBlock 可以作为稀疏注意力的补充性残差复用策略，减轻因激进稀疏化导致的质量下降。
效率提升：在扩散语言模型和视频生成模型上均实现了显著的推理吞吐量提升和延迟降低，且对生成质量的影响微乎其微。

4. 实验结果

实验在扩散语言模型（Trado-8B/4B）和视频扩散模型（LongLive-1.3B）上进行。

吞吐量：FlashBlock 在扩散语言模型上实现了高达 1.44倍的 token 吞吐量提升（例如，对于块大小为 4 的情况，从 312 tokens/s 提升至 451 tokens/s）。
延迟：注意力时间减少了高达 1.6倍。
质量：生成质量与密集基准模型相当。在数学推理（GSM8K, MATH500）和编程基准测试（HumanEval, MBPP）上，准确率差异极小（取决于块大小，在 0.2% 到 1.6% 之间）。
长上下文扩展性：随着上下文长度从 100k 增加到 800k，FlashBlock 与基准模型之间的延迟差距扩大，在特定阈值下，FlashBlock 的延迟增长仅为基准模型的一半左右。
与稀疏注意力结合：当与 SparseD（在 20%、30% 和 40% 密度下）结合使用时，FlashBlock 始终比单独使用 SparseD 表现出更高的准确率（例如，在 20% 密度下将 GSM8K 准确率从 34.72% 提升至 42.68%），有效地恢复了因稀疏化丢失的信息。

5. 意义

本文声称 FlashBlock 解决了长上下文块扩散中的一个关键瓶颈：即对稳定的注意力贡献进行冗余重计算的问题。通过将稳定的“块外部”上下文与动态的“块内部”更新解耦，该方法大幅降低了与不断增长的 KV 缓存相关的计算和内存 I/O 开销。

作者强调，这种方法与现有的稀疏注意力技术是正交的，提供了一种能够同时增强效率和准确性的互补策略。结果表明，FlashBlock 为文本和视频领域的长上下文高效生成提供了一种可扩展的解决方案，能够在不牺牲生成内容质量的前提下实现更快的推理。