Full Attention 反击：在百步训练内将 Full Attention 转化为 Sparse

TL;DR RTPurbo 挖掘全注意力 LLM 的内在稀疏性，仅需几百步训练，将仅少量检索头保留完整 KV 缓存，并结合轻量 16 维索引器实现动态稀疏注意力，在长上下文推理中达到最高 9.36 倍预填充加速与 2 倍解码加速，且精度近乎无损。

问题

问题背景

Long-context 推理已成为现代 LLM 的核心能力，但全注意力机制的计算成本随序列长度平方增长，严重制约推理效率，尤其在上万 token 乃至百万 token 场景下，传统全注意力几乎不可用。

现有方法局限

现有方案主要分两类：

• 原生稀疏训练：从零开始训练稀疏注意力，虽能保证效率，但需重新预训练，训练成本极高，且难以迁移到已有全注意力模型。
• 启发式 token 驱逐：基于固定 top-k 或局部性假设裁剪 KV cache，虽避免训练，但稀疏策略与查询无关，容易丢失关键长程信息，导致准确率显著下降。 两者无法兼顾效率、训练开销与精度，且常忽略注意力头的功能分化——并非所有头都需要全序列处理。

为什么难且重要

在保持精度前提下大幅降低长上下文推理成本是产业落地的关键，但稀疏化本质是信息瓶颈问题：如何在压缩 KV cache 的同时保留对当前查询最相关的 token。难点在于，不同查询关注的信息分布差异大，固定稀疏模式和全局阈值难以胜任。同时，全注意力模型的内在稀疏性未被充分利用，直接训练稀疏模型又代价过高。该问题兼具理论价值和工程挑战，直接影响 LLM 在长文档理解、多轮对话等场景的部署可行性。

行业类比

如同数据库系统中用自适应索引代替全表扫描加速查询，长上下文推理需要根据查询动态选择 token，而非固定规则。RTPurbo 相当于为 LLM 配备轻量索引器，用极少训练成本实现“按需检索”，类似向量数据库中的近似最近邻搜索，但完全融入自注意力机制。

核心洞察

• 全注意力 LLM 中仅少数检索头负责长程依赖，其余头天然稀疏，无需昂贵的原生稀疏预训练，通过几百步微调即可将完整模型高效稀疏化，打破效率与精度的固有取舍。传统高效注意力方法要么从头预训练稀疏架构（Longformer、BigBird），固定稀疏模式且训练成本高；要么在推理时启发式驱逐 token（StreamingLLM、H2O），以精度损失换速度。RTPurbo 首次揭示标准全注意力训练已孕育可迁移的稀疏结构，利用头专业化分析分离检索头，其余头仅需少量 token 参与运算，在 9.36 倍预填充加速下保持近乎无损精度，为稀疏推理开辟低成本适配路径。
• 利用 RoPE 的低维几何特性，检索头可通过 16 维索引器执行高效长程 token 检索，并采用查询相关的动态 top-p 选择替代固定 top-k，实现更精准的信息压缩。现有稀疏注意力常预设固定稀疏度或模式，无法适应不同查询对上下文的不同需求，导致关键 token 遗漏或计算浪费。RTPurbo 发现检索头的注意力分数矩阵可由低维投影近似，据此构建轻量索引器快速筛选候选 token；同时动态 top-p 门控根据每个查询自动调整保留 token 数量，确保召回相关 chunk 的同时避免固定 top-k 的刚性约束，使长上下文推理在保持精度下获得 2 倍解码加速。

方法

输入

RTPurbo 接收一个已完成标准全注意力预训练的 LLM，并假设该模型在长上下文推理时存在内在稀疏性。输入为长序列，目标是在保持精度的前提下大幅降低注意力计算量。

关键模块

1. 离线头级标定 (Offline Head-wise Calibration)
通过分析各注意力头对长上下文的依赖模式，将头分为两类：检索头 (retrieval heads) 负责长程信息聚合，需要访问完整 KV 缓存；其余 局部头 (local heads) 只需关注邻近 token，可进行稀疏化。

2. 低维令牌索引器 (Low-dimensional Token Indexer)
引入一个轻量级投影模块，将查询和键映射到 16 维子空间，在该低维空间中计算相关性分数以筛选相关 token。这使得长程检索不再依赖原始高维注意力，大幅降低计算开销。

3. 动态 top-p 选择 (Dynamic Top-p Selection)
摒弃固定 top-k 策略，根据每个查询的注意力分布动态调整保留的 token 数量（通过累积概率阈值 p 控制），更贴合查询的异构需求。

4. 两阶段低成本训练

• 阶段一 只训练低维投影层，冻结原模型，快速学习令牌检索能力。
• 阶段二 端到端自蒸馏训练，使用原全注意力模型作为教师，保证稀疏模型输出分布不偏离。整个过程仅需几百步训练，无需昂贵的额外预训练。

5. 硬件友好的 fast top-p 解码内核
针对动态 top-p 操作设计高效 CUDA 内核，实现预填充和解码阶段的实际加速。

输出

输出一个高度稀疏的推理模型：仅检索头保留完整 KV 缓存，局部头使用低维索引器与动态 token 筛选执行稀疏注意力。在长上下文基准测试中实现近无损精度，同时获得显著加速（如 1M 上下文中 9.36× 预填充加速、约 2.01× 解码加速）。

与同类方法的关键差异

相比原生稀疏训练（需从头修改预训练目标）或启发式 token 驱逐（基于固定规则，精度损失大），RTPurbo 挖掘了全注意力模型中已存在的内在稀疏性，以极小的微调代价（数百步）完成从稠密到稀疏的迁移，并首创性地将动态 top-p 应用于长上下文稀疏化，平衡了效率与准确性。

实验

实验设计

评估围绕两个维度展开：长文本基准（如多文档问答、摘要）和复杂推理任务（数学、代码）。对比对象包括 全注意力基线 和各类稀疏注意力方法（原生稀疏训练、启发式 token 驱逐）。RTPurbo 先离线校准识别 检索头，再训练低维 token indexer 和动态 top-p 选择，整个适配过程仅需数百步。效率测试在 1M 上下文下进行，覆盖 prefill 与 decode 阶段。

关键发现

• 全注意力 LLM 存在固有稀疏性，仅少部分头负责长程检索，且检索头仅依赖低维子空间（16 维），使得极轻量的 indexer 即可保持精度。
• 动态 top-p 选择比固定 top-k 更适应查询相关的 token 预算，避免信息丢失。
• 在几乎无精度损失的前提下，RTPurbo 实现最高 9.36 倍 prefill 加速、约 2.01 倍 decode 加速，且训练开销极小，收敛迅速。
• 该方法可直接应用于标准全注意力训练模型，无需昂贵的原生稀疏预训练。

与基线对比深度解读

传统高效推理方案面临效率、训练成本与精度之间的三角权衡：原生稀疏训练需从零开始，成本高且不兼容已有密集模型；启发式驱逐虽快捷但易丢失关键 token。RTPurbo 揭示密集模型已蕴含可迁移的稀疏结构，通过事后校准 + 低秩适配将“全注意力转入稀疏”，近乎零成本获得稀疏推理能力。这一发现意味着，稀疏推理的竞争壁垒不在是否从头预训练，而在能否挖掘模型内在稀疏性，为长上下文推理的部署优化提供了新范式。

行业影响

落地场景

RTPurbo 使得标准 full-attention 训练的大模型能以极低成本适配超长上下文推理，直接受益的产品形态包括：

• 长文档理解与对话：法律合同分析、财报问答、学术文献综述等 SaaS 工具；
• 代码辅助与自主代理：需要维护超长上下文的 IDE 插件或 Agent 记忆系统；
• 多轮对话与客服：历史对话越长响应越慢的问题可被缓解；
• 视频/多模态流分析：对长时序输入进行高效 token 检索。

商业价值

• 降本增效：prefill 阶段最高 9.36 倍加速、decode 约 2 倍加速，意味着同等硬件可支撑更多并发请求，直接降低 GPU 推理成本；同时避免从零做 native sparse pretraining 的巨额投入。
• 体验提升：长上下文下首 token 延迟大幅降低，流式输出更流畅，减少用户等待，提高付费转化与留存。
• 技术普惠：仅需数百步训练即可将现成 dense 模型转化为高稀疏推理模型，降低长上下文能力的准入门槛。

与现有产品/工作流的接口

RTPurbo 可作为一个推理引擎插件接入现有 vLLM / TensorRT-LLM 等服务框架：

• KV cache 策略：保留针对 retrieval heads 的完整 KV cache，其余头部使用动态 top-p 稀疏索引，与现有 paged attention 兼容；
• 集成路径：离线完成 head 校准与轻量 token indexer 训练 → 导出配置与权重 → 部署时替换标准 attention kernel 为作者提供的 hardware-aware fast top-p decoding kernel。
• 与量化/蒸馏协同：该方法本质是对注意力机制的稀疏化，可与 GPTQ / AWQ 等权重量化及 FlashAttention 等算子优化正交叠加。

具体落地用例

1. 企业级知识库问答：某法律科技平台部署 Llama-3.1-8B 处理数百页合同，采用 RTPurbo 后，1M 上下文的 prefill 延迟从约 20 秒降至 2 秒，首字响应更快，支撑更多法律助理同时在线查询。
2. 电商客服长对话：客服机器人需回顾整月聊天记录理解用户意图，利用 RTPurbo 动态选择真正关键的 retrieval heads 和 token，在 A100 节点上吞吐量提升近 2 倍，相同成本下服务更多店铺。

局限

• **模型与架构依赖性**：RTPurbo 基于注意力头专业化（retrieval heads）和 RoPE 低维检索子空间的假设，这些特性可能在不同训练范式或位置编码（如 ALiBi、NoPE）的模型中不成立。离线校准步骤需要为每个目标模型单独进行，当面对不同尺寸或微调变体时，工程部署复杂度较高。动态 top-p 选择引入的计算量具有输入依赖性，可能影响批量处理中的负载均衡与硬件调度。
• **泛化性与任务覆盖**：论文实验主要基于特定长上下文基准和推理任务，但对需要细粒度位置或精确计数能力的任务（如细粒度信息检索、长文档统计）评估不足。方法在非 Llama 类模型或不同模型规模（如 1B 以下或 100B+）上的效果尚不明确，且额外训练步骤（即使仅几百步）仍可能成为某些低资源场景的障碍。
• **信息保真度与稀疏性平衡**：虽然保留检索头的完整 KV 缓存可避免关键信息丢失，但低维投影索引（16 维）可能损失内容细节，导致在需要精确匹配或长距离依赖的任务上精度下降。此外，固定检索头比例（如 10%）的选择可能不是全局最优，不同下游任务的最优稀疏配置仍需手动调参。