一句话总结

来自KAIST AI和Google DeepMind的研究人员提出了声明式注意力（Declarative Attention, DA），这是一种内禀协议，允许语言模型在思维链中通过、和模式声明关注位置，使推理引擎能够跳过大部分KV缓存读取。在15个长上下文任务的零样本评估中，DA使Gemma-4-31B的参与token数减少52.0%，Qwen-3.6-27B减少31.1%，准确率分别下降1.27个百分点和2.75个百分点。

核心贡献

• 声明式注意力是一种零样本协议，引导语言模型在其思维链中声明其关注目标（全局、聚焦、局部），使推理引擎能够跳过大部分KV缓存读取，无需外部评分。
• 在15个长上下文任务中，使用声明式注意力的现成模型在解码期间的总参与token数分别减少52.0%（Gemma-4-31B）和31.1%（Qwen-3.6-27B），准确率仅下降1.27个百分点和2.75个百分点，且随着模型规模增大，准确率差距缩小。
• 提供了一种高效的vLLM集成，具有与FlashAttention兼容的块对齐、原地KV缓存掩码，并且基于屋顶线的墙钟时间分析预计在相应模型上解码成本分别降至原始版本的0.71倍和0.77倍。

引言

长上下文Transformer推理的主要瓶颈是在每个解码步骤加载键值（KV）缓存的内存带宽成本，这使得全注意力变得极其昂贵。先前的工作尝试通过预测哪些token是相关的来稀疏化注意力，但静态启发式方法无法适应特定查询的需求，而动态方法仍然会产生每步O(N)的扫描成本。作者引入了声明式注意力，这是一种零样本提示协议，引导模型在其思维链中明确声明将关注何处。推理引擎解析这些声明的模式转换（全局、聚焦、局部），直接从生成的文本构建注意力掩码，消除了每步的选择开销，并实现了高效的长上下文解码，准确率损失极小。

方法

作者提出了声明式注意力（DA），这是一种协议，引导大型语言模型重构其思维链推理，使其注意力计划显式且可读。DA不依赖辅助评分器或预测注意力权重，而是要求模型将其推理组织成连续的跨度，在这些跨度内注意力范围保持稳定，并使用预定义的标签语法声明该范围。

为此，作者设计了一个特定的提示结构，由持久支架和可变上下文区域组成。支架包括系统指令、用户问题和详细的模式说明，确保协议的持久基础。长输入上下文被划分为约2048个token的可寻址段，称为“魔术块”。这些段以模拟工具使用记录的形式呈现给模型，其中助手似乎为每个块调用检索工具。这种格式将段边界与分隔用户、助手和工具消息的特殊token对齐，模型由于其后训练而能流畅地跟踪这些边界。

模型在三种不同的推理模式之间自由切换，每种模式服务于特定目的，并决定上下文的可见程度。如下图所示：

在<global>模式下，模型关注所有上下文段，用于导航和浏览完整上下文以定位相关信息。在<focus>模式下，模型仅关注标签中指定的特定上下文段，从而能够针对特定区域进行推理，而无需处理整个上下文的计算成本。最后，在<local>模式下，模型不关注任何上下文段，仅依赖其响应中已积累的信息进行自包含推理和答案合成。在所有模式下，模型持续关注支架及其自身生成的token。

为了实现相应的注意力掩码，作者引入了一个与推理引擎并行运行的DA状态机。状态机从默认的<global>模式开始，并解析输出流，根据开标签的闭合字符检测模式转换。由于vLLM等现代推理引擎以固定大小的块存储键值缓存，状态机以块粒度应用掩码。它将保留的token跨度向外舍入到块边界，以确保不会丢弃任何声明的token，从而允许FlashAttention等现有内核保持不变。该系统通过注意力元数据构建器上的钩子与vLLM集成，在每个解码步骤重写请求的KV缓存块表，使得只有保留的块对注意力内核可见。这种设计显著减少了每步读取的KV块数量，以更多的解码步骤换取更低的每步注意力成本。

实验

本文在15个长上下文检索和推理任务上评估了差分注意力（DA），使用了Gemma和Qwen家族的六个模型，与原始全注意力和无掩码变体进行比较。DA实现了显著的token节省（最多减少52%的参与token），准确率下降适中，且注意力掩码是节省和准确率成本的主要来源。相对于原始方法，准确率随模型规模提高，token节省随上下文长度增加，转化为在优化硬件上理论解码时间加速1.3-1.4倍。效率源于廉价的聚焦和局部模式，并且对于较大模型，协议遵循是稳健的。

基准测试套件涵盖15个长上下文来源，从短文档到百万token代码仓库，单跨度检索任务平均长度为9K–11.5K token。在差分注意力实验中，廉价的聚焦和局部模式主导生成，其节省随上下文长度增长，而全局注意力仅占少数token。对于大型模型，协议遵循近乎完美，聚焦成功率达到99%，但小型模型难以遵循，协议遵循失败与准确率下降相关。基准测试中的单跨度检索任务平均上下文长度为9K至11.5K token，标准差高达7.1K。在Gemma-4-31B上，聚焦和局部注意力模式共同生成73%的token，并相对于原始步骤实现76–99%的每token注意力节省。聚焦成功率随模型规模急剧上升，从Gemma-4-E4B的58%上升到最大模型的99%，使得协议遵循仅在小型模型上成为瓶颈。

在长上下文检索和推理任务中，DA协议使Gemma-4-31B的总参与token减少52%，Qwen-3.6-27B减少31%，平均准确率仅下降1.3和2.8个百分点。自定义注意力掩码至关重要：没有它，token减少幅度缩小，参与token甚至可能超过原始版本。在Gemma上，这种权衡略为有利，更大的token节省伴随着更小的准确率惩罚。DA在两个模型上均持续减少参与token，Gemma在所有任务上大约减半。移除自定义掩码（DA^nm）会侵蚀节省效果，并且在Qwen上，某些单跨度任务的参与token高于原始版本。DA下的准确率仍然很高，仅在少数任务上有1–4个百分点的孤立下降，其他任务无损失。与Gemma（1.27个百分点，52%）相比，Qwen的平均准确率下降更大（2.75个百分点），token减少更少（31%）。

DA协议将Gemma-4-31B的估计解码墙钟时间降至原始版本的71%，Qwen-3.6-27B降至77%，在单个B200加速器上。节省完全来自削减全局注意力KV读取，这是原始版本的主要成本，而矩阵乘法和本地内存成本略有增加，因为DA生成更多解码步骤。Qwen的收益更大，因为其本地内存（GDN状态）非常小，而Gemma的大型滑动窗口注意力本地读取限制了总体增益。DA在两个模型上将全局内存读取时间大约减半，但总加速受限于掩码无法减少的固定本地内存成本。在DA下，Gemma的本地SWA读取占注意力时间的42%，而Qwen的GDN状态仅占5%，这解释了为什么Qwen的总解码时间相对减少更大。

评估使用长上下文检索基准，平均上下文长度为9K–11.5K token，模型如Gemma-4-31B和Qwen-3.6-27B。差分注意力协议将大部分计算转移到廉价的聚焦和局部模式，从而大幅减少参与token，准确率损失极小，而小型模型的协议遵循度较低，与准确率损失相关。自定义注意力掩码对于这些节省至关重要，尽管解码墙钟时间有所改善，但总体加速受限于协议无法减少的固定本地内存成本。