作者：周琰轲，南京大学二年级硕士研究生，阿里巴巴实习生，RTP-LLM项目核心贡献者
-qTDs.png)
长上下文能力已成为现代大语言模型的基础要求，但全注意力机制（Full Attention）随序列长度呈平方级增长的计算复杂度，构成了推理阶段的核心瓶颈。为了缓解这一问题，目前业界主要探索了两条技术路线：
1.Training-free 稀疏注意力： 基于启发式规则（如 Token Eviction, KV cache compress）在推理时直接裁剪注意力矩阵。然而，这种“训推不一致”的策略往往会破坏模型在预训练阶段习得的注意力模式，导致在复杂任务中出现显著的精度退化。
2.原生稀疏架构： 替换Full Attn 为全新的高效Attention模块（如 Kimi Delta Attention, DeepSeek Sparse Attention），并在预训练或 Continue Pre-train 阶段将其引入。尽管精度有保障，但代价是极高的训练成本以及潜在的模型收敛稳定性风险。
结合近期的可解释性研究，我们观察到一个关键现象：经过 Full Attention 训练的模型，在 Head 和 Token 级别已经具备了高度的“原生稀疏性”。这自然引出了我们的核心研究问题：
能否以最小的适配代价(Minimal Surgery)，将 Full Attention 模型转换为高效的稀疏模型，同时严格保留其原始能力？
为了回答这个问题，我们需要一种尽可能适应模型原生稀疏性质的稀疏策略，我们认为这里主要存在三大挑战：
1.如何分辨哪些attention head是依赖局部信息的？哪些是依赖全文信息的？
2.对于依赖全文信息的head，如何高效判断哪些token是重要的？
3.如何确保对于任何query在保证精度的同时尽可能的获得高的稀疏度？
为此，我们提出了RTPurbo。通过轻量级训练（仅需约600 步，消耗约1M Label Tokens），RTPurbo在无需原生稀疏预训练的前提下，实现了高达9.36倍的Prefill加速与2.01倍的Decode加速，且在主流长文本和推理Benchmark上保持了近乎无损的精度。
-VKnS.png)
核心洞见：顺应 LLM 原生稀疏性的深度剖析
RTPurbo 的核心思想并非通过强行剪裁来改变模型，而是通过“外科手术式”的微调，激活 Full-attention 模型在训练阶段就已经形成的原生稀疏特性。这种体质在以下三个维度呈现出极强的规律性：
Attention Head 的功能分化：天然的异构先验
-fwEH.png)
类似人类的阅读与写作过程，LLM 在处理长文本时，往往会先从长上下文中“召回”相关信息，再在相对局部的上下文范围内完成推理并输出答案。
研究表明，绝大部分Attention Head实际上主要处理局部上下文信息，对远端信息的感知极弱。相比之下，只有极少数特定功能的 Head 承担了跨越长距离检索语义关联信息的任务，即我们定义的Retrieval Head。
这种功能上的高度分化意味着模型长程检索的能力并非均匀分布，而是集中在深层的特定通道中。为了验证这一点，我们选取一些长文本，在这段文本的前后各插入一段needle文本，通过测试后一段needle文本对前一段needle的attention召回程度，我们给每一个head打分。
图3结果支持了这一结论。这种功能特化为我们提供了天然的稀疏先验：对于大多数的局部头，我们可以安全地压缩kv cache；而将有限的计算资源精准投放到那 15% 的召回头上。
-ASLA.png)
RoPE 编码下的几何压缩性：低维检索空间的理论依据
现代的LLM通常会配置RoPE。然而，召回头的召回行为面临一个几何悖论：无论语义相关的 Token 距离多远，召回头都必须赋予其高分，但这与 RoPE 自带的“长程衰减（Long-term decay）”属性相悖 。
通过对注意力机制的频域分解，我们揭示了其中的奥秘：
公式表明，当两个token的距离\Delta增大时，\theta_i比较大的高频分量变化非常快，导致大部分信息无法被很好的保留下来；而低频分量受影响较小，展现出平滑的距离不变性，构成了长程检索的骨干。
据此，我们提出第二个核心观察：我们可以在很小的维度空间中精准重建召回头的attention分布。无需在原始的高维空间（如128维）进行暴力搜索，只需提取低频结构构建一个极低维子空间，就能以极低的算力成本精准重建召回头的注意力图谱 。
动态的token选取策略：打破固定 Budget 的魔咒
在确定了检索空间后，如何划定保留Token的边界？我们发现，模型所需的Token预算与具体任务的性质深度耦合，固定的 Top-k 策略在长文本推理中会面临严重的鲁棒性挑战 。
在复杂长文本理解或多跳逻辑推理任务中：如回答关于“加拉帕戈斯群岛的治理模式”这类开放式问题，模型需要从长达数十万字的背景中搜集大量语义相关的线索 。此时注意力会呈现出广阔的分布规律，可能需要召回 8k 个以上的 Token 才能还原 90% 的注意力 Mass 。如果此时死守较小的 Top-k 预算，将不可避免地导致关键语义信息的丢失。
在“大海捞针”（NIAH）等高指向性检索任务中：模型只需要定位极个别关键的事实锚点。此时注意力的分布表现出极高的确定性，往往仅需保留 2 个 Token 就能覆盖超过 96% 的注意力质量 。在这种场景下，任何过大的固定 Budget（如 Top-4k）都会引入 99% 以上的计算冗余。
基于这一观察，我们应该抛弃静态稀疏策略，转而采用动态的token选取策略：根据当前 Query 的难度和任务上下文，自动调节 Token 留存数量。
RTPurbo
-gDGp.png)
在方法论部分，我们将重点转向RTPurbo的具体实现路径。RTPurbo的核心架构由三部分算法组件及一个底层工程优化算子组成，旨在通过软硬件协同的设计，将前述的理论洞见转化为高效的推理能力。
离线head-wise校准
为了准确识别不同head的功能角色，我们通过一次性的离线校准过程对query head进行划分。具体而言，我们在长文本文档的起始和末尾分别插入完全相同的“Needle”序列，并计算后部序列对前部序列的注意力分数综合，以此来作为Head的检索得分。
利用这个分数对query head从大到小排序，根据setting选择多少比例的head作为召回头\mathcal{H}_{ret}，其它的作为局部头\mathcal{H}_{cal}。
动态sparse attention
在推理阶段，RTPurbo为两类head应用不同的计算模式，以兼顾计算效率和kv cache留存：
局部头 (\mathcal{H}_{loc})：在整个推理周期内，始终采用带Attention Sink的滑动窗口(SWA)策略。
召回头 (\mathcal{H}_{ret})：在prefill阶段，执行完整的causal attention计算，以构建并维护全量的KV cache；在decode阶段，引入query-aware的动态稀疏选择，利用训练好的低维投影矩阵W_h^Q, W_h^K对RoPE前的Q, K进行压缩，并重构query的attention分布s_h(m,\cdot)。随后应用动态Top-p策略从全局KV cache中选择关键的Active Token Set \mathcal{S}_h(m)进行稀疏注意力计算。
轻量化两阶段训练
为了使模型在稀疏机制下完美恢复性能，我们设计了极低成本的training pipeline：
Stage1--低维投影对齐：冻结大模型主干，仅对召回头训练低维投影参数。通过最小化投影分布与原始全量分布间的KL散度，将长程检索信号压缩至子空间。
Stage2 --端到端自蒸馏：启用稀疏模式，以原始未稀疏权重为“教师”，执行端到端的自蒸馏训练。出于效率考量，在训练过程中仅需对齐teacher model的top-10 Logits，极大降低了训练开销。整个过程通常仅需百步训练即可完美收敛。
Kernel优化：高效top-p decode
-GxTn.png)
朴素的 Top-p 实现，通常需要先对全部分数进行排序，然后再进行累加操作以寻找截断阈值。如果将这种常规做法直接代入 RTPurbo 的设计中，整个 decode 过程将被迫拆分为四个离散的阶段：低维 GEMV（特征打分）->Sort（全局排序）->Cumsum（累加寻址）->Sparse GEMV（稀疏注意力计算）。
特别是其中的排序过程将引入极其高昂的开销，甚至完全抵消稀疏计算带来的理论收益。
为了彻底打破这一工程瓶颈，我们正对性的优化了这一过程：
基于histogram的免排序top-p: 我们将压缩后的 Key 序列划分为多个 Block，每个 CTA（Compute Thread Array）在完成块级低维打分后，不进行任何排序，而是通过原子操作直接将其累加到一个 256-bin 的全局直方图中。
该直方图在每个 Head 下仅需 1 KB 的极小显存。最后一个完成计算的 CTA 仅需顺手扫描该直方图，即可在O(1)的内存额外开销下，直接锁定 Top-p 阈值并生成块掩码（Block Mask）。这一精巧设计将打分与阈值筛选完美融合在了单次 Kernel 启动中。
带宽优化的稀疏计算： 针对选定后跨度依然很大的长序列，我们设计了无共享内存的单 Warp CTA 结构。通过将全部中间状态驻留在寄存器中，极限提升了 SM 级别的并发内存请求数，并通过向量化加载（如 half2 指令）让稀疏注意力的计算与访存延迟深度重叠。
.png)
单kernel的测试结果表明，我们优化后的算子不仅彻底碾压了基于PyTorch实现的4-stage朴素实现，更在与高度优化的FlashAttention 2的对比中展现出显著优势，获得稳定的稀疏收益。
实验评估
为了全面验证RTPurbo的有效性，我们在Qwen3-Coder-30B-A3B（用于长文本理解）和推理模型Qwen3-30B-A3B-Think（用于长CoT推理任务）上进行了广泛的测评，并与其它SOTA工作进行了严格实验对比。
精度测评：长文本理解与长推理链推理
长上下文理解（LongBench & RULER）
.png)
.png)
传统的基于启发式规则或局部块稀疏的方法在面临复杂任务时暴露出明显的短板。例如，在 Multi-Q、Multi-K 以及 Multi-V等任务中，局部上下文往往偏离全局关键信息，导致 MInference、FlexPrefill 和 Quest 等方法出现严重的精度退化。
此外，采用静态 Top-k 预算的基线在 64K 长度下即因召回不足而表现不佳。 相比之下，RTPurbo 的动态 Top-p 选择器能够精准适应不同查询的复杂度，不仅在基础设置下保持领先，更在向超长上下文（高达 512K）外推时展现出绝对的鲁棒性。当多数基线在极端长度下遭遇灾难性崩溃时，RTPurbo 依然维持了高水准的准确率。
长CoT推理（AIME24/25 & MMLU-PRO）
.png)
现代数学与逻辑推理任务（如 AIME）呈现出极端的“生成不对称性”：输入提示往往不足 300 个 Token，但模型生成的推理思维链（CoT）却可能长达 32K 个 Token，将计算瓶颈完全推向了 Decode 阶段。
在如此严苛的长解码场景下，RTPurbo 实现了近乎无损的精度保留。在 AIME 基准上，RTPurbo 取得了与 Full-attention 教师模型完全一致的满血成绩（86.67）。
在 MMLU-PRO 的各个子学科测试中，其表现也与Full Attn基线紧紧咬合，充分证明经过稀疏化后，模型并未丧失其深度的通用推理能力。
稀疏度与加速比测评：动态策略释放极致性能
除了精度，RTPurbo 在稀疏度与加速效果上也展现出了巨大的优势：
高度自适应的动态稀疏度
.png)
RTPurbo的核心优势在于彻底抛弃了静态Top-k策略。对Decode阶段的某召回头特征分析显示，最佳的Token Budget是高度query-aware的：在32K上下文下，针对大海捞针任务（niah-S），RTPurbo平均仅需保留468.8个Active Tokens；而对于较为复杂的Multi-K任务，则会自动扩展至2462.1个，动态跨度高达5倍。
.png)
此外，这种按需分配的动态机制，使得当上下文长度被拉长至512K时，RTPurbo在保证高精度的同时，将稀疏度向上推升至97.1%以上。
端到端高效提速
我们将RTPurbo的单层执行效率与FA2以及其它稀疏策略进行了全面对比：
.png)
Prefill阶段： 得益于head-wise稀疏策略，RTPurbo对比FA2实现了从32K长度下的2.83x到1M长度下9.36x的超高加速比。
Decode阶段： 配合我们优化过的top-p decode算子，RTPurbo对比FA2同样实现了稳定且显著的收益，加速比从32K下的1.47x稳步提升至1M上下文的2.01x。
总结
我们的结果传达了一个重要的启示：高昂的原生稀疏预训练并非通往高效长上下文推理的唯一路径。通过轻量级的稀疏化适配，Full Attention模型完全能够胜任并展现出卓越的稀疏执行能力。这为未来在资源受限的环境下，以极低成本优化与部署顶尖的大型预训练模型，提供了一条极具潜力和理论支撑的实用化新路径。
团队介绍
RTP-LLM 是阿里巴巴智能引擎团队自研的高性能大模型推理引擎，支持了淘宝、天猫、高德、饿了么等核心业务的大模型推理需求。智能引擎源自阿里巴巴搜索、推荐和广告技术，是阿里 AI 工程领域的先行者和深耕者。团队专注于 AI 工程系统的建设，主导建立了大数据 AI 工程体系 AI・OS，持续为阿里集团各业务提供高质量的 AI 工程服务。
RTP-LLM 项目已开源，欢迎交流共建：
https://github.com/alibaba/rtp-llm
合作团队南京大学计算机软件所围绕智能化软件与系统研究主题，以多领域融合创新为特色，不断寻求原创性突破，在大模型支撑技术与应用技术方面取得了一系列进展，承担了一系列重要项目。