Star Attention 是什么？英伟达加速长文本推理的块稀疏注意力

Star Attention 是一种加速大模型长文本推理的注意力机制，由英伟达在 2024 年提出，论文标题是《Star Attention：长序列上的高效大模型推理》。它专为提升基于 Transformer 的大语言模型在长序列上的推理效率而设计。

长文本推理的瓶颈

大模型处理长文本时，注意力计算量随文本长度平方增长。当输入有几十万甚至上百万个 token 时，光是读完这段输入（预填充阶段）就要花很长时间，单台机器的显存也装不下。常见的做法是把计算分布到多台机器上，但标准注意力要求每个 token 都能看到前面所有的 token，机器之间需要频繁地交换数据，通信开销很大。

两阶段的处理流程

Star Attention 把推理过程分成两个阶段，像星形的网络结构一样组织计算。

第一阶段，上下文编码：把很长的上下文切成若干块，分配给多台机器（主机）并行处理。每台主机只在自己负责的那一块内部计算注意力，也就是局部注意力，机器之间不需要通信。一个巧妙的设计是：每一块前面都会加上第一块的内容作为“锚点块”。研究者发现，这样做能让每一块的注意力分布更接近全局注意力的效果，避免因为只看局部而导致的偏差。

第二阶段，查询编码和生成：用户的问题和后续生成的 token，需要看到整个上下文。这时会把问题广播给所有主机，每台主机计算问题对自己那一块上下文的注意力，再把结果汇总到一台中心主机上合并，得到全局的注意力结果，用于生成回答。只有这一阶段需要通信，而且通信量很小。

效果

论文报告，Star Attention 能把推理速度提升最多 11 倍左右，同时保持 95% 到 100% 的准确率。它可以直接用于大多数用全局注意力训练的现有模型，不需要额外的微调。对于需要处理超长文档、长对话历史和大型代码库的应用，这类方法能显著降低推理时间和成本。