摒弃锚点：基于Pulsar Attention的分布式系统统计上下文摘要

# 基于Pulsar Attention的分布式系统统计上下文摘要 来源：https://arxiv.org/html/2607.20457 ###### 摘要 使用大语言模型\(LLMs\)对长序列进行推理因自注意力的二次复杂度而计算成本高昂。分布式块级方法（如Star Attention）通过将上下文分片到多个主机来降低这一成本，但依赖于在每个主机前附加一个静态的、内容盲目的第一个块副本。我们提出Pulsar Attention，它用两个轻量级、内容感知组件替代静态锚点：一个稳定softmax的小型注意力汇聚前缀，以及通过Max-IDF启发式方法构建的紧凑跨块摘要（该启发式选择包含全局罕见标记的块组）。这使得第一阶段每GPU FLOPs相比Star Attention最多降低3.3倍，同时保持相同的KV缓存占用。在使用Llama-3.1-8B的RULER和BABILong评估中，Pulsar Attention在序列长度达128K标记时均优于Star Attention和密集注意力，相比密集基线绝对提升最高达4.7%。机器学习, ICML, 深度学习, 关键词3, 关键词4 ## 1引言 大语言模型\(LLMs\)在需要长上下文窗口的任务上表现出色，但扩展推理仍然成本高昂：自注意力在序列长度上是二次的，KV缓存线性增长，迅速耗尽单个GPU的设备内存。分布式方法通过将上下文分片到多个GPU来解决这一问题。Ring Attention\(Liu等，2023 (https://arxiv.org/html/2607.20457#bib.bib3)\)通过以环状方式循环KV块来计算精确的全局注意力，但需要在每一层进行协调通信。Star Attention\(Acharya等，2025 (https://arxiv.org/html/2607.20457#bib.bib2)\)在上下文编码期间消除了这种通信，通过并行处理独立块，在查询时通过轻量级softmax合并恢复全局注意力。为防止softmax崩溃，每个主机在其输入前附加一个静态的第一个块副本（称为锚点块）。锚点有效，但它无论相关性如何都复制相同的标记，使上下文编码期间每个主机的序列长度和FLOPs翻倍。由于它固定到序列起始，块i无法获得中间块的信息，随着序列变长，这一问题更加严重。我们提出Pulsar Attention，它用两个轻量级、内容感知组件替代静态锚点：(i) 一个小型注意力汇聚前缀（64个标记），以锚点成本的一小部分稳定softmax；(ii) 紧凑的Max-IDF块摘要，从每个前面的块中选择包含全局罕见标记的连续块组。摘要因果传播并适应每个块的内容，优先选择最可能携带任务关键信息的标记。我们的主要贡献包括： - •我们用内容感知摘要和注意力汇聚替代静态锚点，使第一阶段每GPU FLOPs相比Star Attention最多降低3.3倍，同时保持相同的KV缓存占用。 - •我们引入Max-IDF块组评分，并在BABILong上针对三种替代启发式方法进行验证，表明稀有标记选择能最有效地捕获跨块上下文。 - •我们在RULER上展示，Pulsar Attention在上下文长度高达128K时优于Star Attention和密集注意力，相比密集基线增益高达+4.7%。 ## 2背景 基于Transformer的LLM在自回归解码期间缓存所有先前标记的键和值向量。虽然这避免了冗余计算，但KV缓存内存随序列长度线性增长，在长上下文时迅速耗尽设备内存。Liu等人\(2023 (https://arxiv.org/html/2607.20457#bib.bib3)\)通过将KV缓存分布到多个GPU并以环状方式循环块来计算精确的全局注意力来解决这个问题。这允许任意长的序列，但需要在每一层进行跨主机通信。Acharya等人\(2025 (https://arxiv.org/html/2607.20457#bib.bib2)\)通过两阶段设计在上下文编码期间消除了这种通信。在第一阶段，上下文被分割成连续的块，在并行主机上独立处理；每个主机的输入前附加一个静态的第一个块副本（锚点块）以防止softmax崩溃。在第二阶段，查询被广播到所有主机，主机计算局部注意力分数并通过在线softmax合并聚合以恢复全局注意力。虽然有效，但锚点无论相关性如何都复制相同的标记，使第一阶段每个主机的序列长度和FLOPs翻倍。我们的方法解决了这个局限性。 ## 3方法论 参照图注 (a)第一阶段：统计上下文编码。主机i组装输入\[Ssink‖Σ0‖⋯‖Σi−1‖Bi\]，运行前向传播，然后丢弃汇聚和摘要KV条目，仅保留KV\(Bi\)。参照图注 (b)第二阶段：全局查询编码。每个主机对其保留的KV分片进行注意力计算；通过all_gather和在线softmax合并\(Milakov and Gimelshein，2018 (https://arxiv.org/html/2607.20457#bib.bib17)\)恢复精确的全局注意力，无需传输完整的KV缓存。 图1:Pulsar Attention概览。(左)第一阶段用内容感知前缀替代静态锚点块：一个小型注意力汇聚和所有因果前驱块的Max-IDF统计摘要。(右)第二阶段将查询广播到所有主机，并通过在线softmax将局部注意力分数合并为精确的全局输出。### 3.1第一阶段：统计上下文编码 与Star Attention类似，Pulsar Attention将长度为L的输入序列划分为n个连续块\{B0, B1, ..., Bn-1\}，每个块分配给一个并行主机（图1(a) (https://arxiv.org/html/2607.20457#S3.F1.sf1)）。我们不复制整个第一个块作为静态锚点，而是用两个轻量级、内容感知组件替代：一个注意力汇聚前缀和Max-IDF块摘要。上下文编码分四步进行。 **注意力汇聚。** 从B0中提取前ssink个令牌作为固定前缀Ssink（默认ssink=64）。遵循Xiao等人\(2023 (https://arxiv.org/html/2607.20457#bib.bib10)\)，这些令牌在块级注意力期间以完整锚点块成本的一小部分稳定softmax分布。 **摘要生成。** 在任何神经网络计算之前，每个主机从完整序列的令牌ID构建一个相同的语料级IDF表，这是一个O(L)的纯整数操作，无需主机间通信或GPU时间。使用该表，每个主机通过Max-IDF启发式方法对每个块内大小为m个令牌（默认m=32）的非重叠连续块组进行评分： score\(C\) = maxt∈C IDF\(t\) (1) IDF\(t\) = log\(n / max(df\(t\), 1)\) (2) 其中df\(t\)是包含令牌t的块数量。每个块选择top-k个块组，按位置顺序拼接形成块摘要Σi，每个块预算为σ个令牌（默认σ = 0.125 × |Bi|）。我们选择连续块组而非单个令牌，以在生成的键值表示中保留局部语法结构。Max-IDF优于平均启发式方法（TF-IDF, BM25），因为检索关键信息（如UUID、实体名称、数值）通常集中在每个块组中的单个稀有令牌上；平均会将该信号稀释1/m倍（参见第5.1节 (https://arxiv.org/html/2607.20457#S5.SS1)的实证比较）。 **因果组装。** 每个主机i只从因果前驱块中附加摘要以及Ssink，形成增强的第一阶段输入： Inputi = { \[B0\] i=0 ; \[Ssink ‖ Σ0 ‖ ... ‖ Σi-1 ‖ Bi\] i>0 } (3) 每个令牌保留其原始全局位置ID，保留由RoPE编码的令牌间距离（参见附录E (https://arxiv.org/html/2607.20457#A5)与连续重新编号的比较）。 **KV丢弃。** 前向传播后，Ssink和{Σj}j<|start_header_id|>system<|end_header_id|> You are a helpful assistant.<|eot_id|> <|start_header_id|>user<|end_header_id|> {task_instruction} {in_context_examples} {post_prompt} {long_context} Question: {query}<|eot_id|> <|start_header_id|>assistant<|end_header_id|> `` #### 上下文/查询分割。 Pulsar Attention（以及Star Attention）在分词前将输入分为两个不同的字符串： - •prompt_context：从<|begin_of_text|>到闭合标签（包括闭合标签）之间的所有内容。该字符串被分词、划分为大小为B的块，并由第一阶段处理。 - •prompt_query：Question: {query}片段及助手标题令牌。该字符串被单独分词并由第二阶段处理。 #### 停止词。 生成在以下任何一个处停止：<|end_of_text|>、<|eom_id|>、<|eot_id|>，或达到max_new_tokens（BABILong为128；RULER随任务变化，见表9 (https://arxiv.org/html/2607.20457#A4.T9)）。 #### 分词。 使用add_special_tokens=False生成令牌；<|begin_of_text|> BOS令牌作为模板字符串的一部分包含在内，不得再次添加。 ## 附录C评分启发式方法详情 所有启发式方法对从每个块中切出的非重叠连续块组（大小为chunk_size个令牌）进行操作。块组保留局部词序；根据分数选择top-num_chunks个块组，并按位置顺序返回，以保持Transformer的自然阅读顺序。 ### C.1语料级IDF表 在评分之前，对所有块进行一次扫描构建文档频率表。设N为总块数，df(t)为包含令牌t的块数： IDF(t) = log(N / max(df(t), 1))。 (8) 这存储为一个密集的整数索引张量idf_table，形状为(vocab_size,)。所有四种基于IDF的启发式方法都从该表进行一次收集；无需遍历词表的Python循环。 ### C.2TF-IDF分数 score_TF-IDF(C) = (1/|C|) * Σ_{p∈C} IDF(t_p)。 (9) 这相当于经典的类型平均TF-IDF，允许一次向量化查找：idf_table[chunk_tokens].sum() / len(chunk)。 **优势。** 平衡、通用；提升具有罕见但重复词汇的块组。 **弱点。** 一个稀有实体混杂在许多常见令牌中会被平均值稀释。 ### C.3BM25分数 BM25(C) = Σ_{t∈types(C)} IDF_BM25(t) * (tf(t) * (k1+1)) / (tf(t) + k1 * (1 - b + b * |C| / l̄))， (10) 其中k1=1.2, b=0.75, l̄是整个语料库中块组的平均长度，IDF_BM25(t) = log((N - df(t) + 0.5) / (df(t) + 0.5) + 1)。 **优势。** 当块组长度变化时比TF-IDF更鲁棒。 **弱点。** 计算量略高；当所有块组长度相同时收益减少。 ### C.4熵（类型-令牌比）分数 Entropy(C) = |{t: t∈C}| / |C|。 (11) 实现为 (bincount(C) > 0).sum() / len(C)。 **优势。** 与语料无关；无需IDF表。填充令牌或重复停用词的块组分数接近0；富含实体的块组分数接近1。 **弱点。** 忽略全局稀有性；一个具有32个不同常见停用词的块组与一个具有32个稀有实体的块组分数相同。 ### C.5Max-IDF分数 Max-IDF(C) = max_{p∈C} IDF(t_p)。 (12) 实现为 idf_table[chunk_tokens].max()。 **优势。** 充当纯探针检测器：如果块组中任何令牌仅出现在一个块中，则该块组保证被选中。这正是多跳任务和NIAH所需要的。 **弱点。** 忽略整体块组质量；一个稀有令牌混杂在31个无信息令牌中仍会胜出。 ### C.6均匀间隔（基线） 使用torch.linspace在块上均匀间隔选择恰好num_chunks × chunk_size个令牌： indices = floor(linspace(0, |B|-1, s)), s = num_chunks × chunk_size。 (13) 保留原始（稀疏）位置ID。 **优势。** 无需语料统计；确定性。 **弱点。** 随着预算增长，将预算浪费在低信息跨度上（在消融实验中确认：当s从128增加到1024时，qa3从36下降到25）。 ### C.7启发式方法总结表 表8：所有评分启发式方法总结。 ## 附录D超参数配置 表9 (https://arxiv.org/html/2607.20457#A4.T9)列出了我们主要实验中使用的所有超参数。除非另有说明，所有值在RULER和BABILong上固定。 表9：完整超参数配置。 ## 附录ERoPE位置ID分配 当一个块的第一阶段输入被组装为[SINK | S1, ..., Si-1 | Bi]时，插入的汇聚令牌和摘要令牌的位置ID可以以两种方式分配。 #### 稀疏（默认）。 每个令牌保留其在完整上下文中的原始全局位置。来自块j的摘要令牌保留其源位置的位置ID。因此组装的序列包含非选定令牌被丢弃而产生的间隙。这保留了RoPE训练编码的真实令牌间距离：来自块0的摘要令牌仍然被视为远离块3中的令牌，与模型的预训练几何结构相匹配。 #### 连续。 对于块i的第一阶段前向传播，组装的序列被重新编号为0,1,...,Li-1（原始Star Attention风格）。在此模式下，第二阶段查询位置ID偏移到刚超过最长的组装第一阶段长度之后，确保从查询到每个块KV的RoPE差值保持较小且在分布内。连续模式有时可能通过确保摘要令牌和本地块令牌以短相对距离出现而有所帮助，但会扭曲全局几何结构。本文报告的所有实验均使用稀疏模式，这与模型的预训练分布一致。 ## 附录FKV缓存丢弃 在每个第一阶段前向传播之后，模型的KV缓存包含组装输入[SINK | S1, ..., Si-1 | Bi]中所有令牌的条目。汇聚和摘要KV状态立即丢弃： Ki(layer) ← Ki(layer)[:,:,is_bi]（并对其他层进行相同操作）。