开源!快手提出面向长上下文推理的序列压缩注意力机制KSA!

admin 2026-05-16 06:06:25 网络安全文章 来源:ZONE.CI 全球网 0 阅读模式

文章总结: 快手技术团队提出KwaiSummaryAttention(KSA)注意力机制,通过插入SummaryToken对长序列进行分段压缩,将KVCache显存占用降至1/8。该方案在RULER-128K任务中相比基线提升3.69-5.48分,支持完整历史回溯且不影响短序列性能。采用三阶段CPT训练流程可实现无缝接入现有模型,代码已开源供社区使用。 综合评分: 87 文章分类: 技术标准,解决方案,AI安全,安全工具,安全开发


cover_image

开源!快手提出面向长上下文推理的序列压缩注意力机制KSA!

原创

快手技术 快手技术

快手技术

2026年5月14日 11:45 北京

在小说阅读器读本章

去阅读

近日,快手推荐模型团队提出了Kwai Summary Attention(KSA),通过插入Summary Token的形式将序列信息进行分段提纯,大幅削减了长序列情况下的KV Cache存储和计算,推理成本优势明显。在序列8倍压缩的设置下,团队验证了Hybrid KSA架构在From Scratch和CPT两种训练范式下,对长上下文建模均带来收益。在RULER长序列检索任务中,Hybrid KSA在4K至128K全长度区间持续领先其他模型变体。在RULER-128K上,相比最佳基线模型,Hybrid KSA分别取得+3.69(From Scratch)和+5.48(CPT)的提升。

随着大模型时代OpenClaw、Code Agent等强Token量消耗的应用工具不断进化,行业的发展趋势逐渐向低成本的inference架构倾斜,同时,旗舰模型每单位成本下输出的Token量将会成为衡量模型能力的重要指标。近一段时间,大模型应用的蓬勃发展与基础模型的长上下文理解与记忆能力的提升密不可分,目前基础模型1M的有效序列窗口一定程度上缓解了“金鱼记忆”问题,支持长时间的深度思考与多轮对话。然而,继续扩长依然面临着资源成本问题,而标准Attention架构在超长序列下面临结构性瓶颈:一方面KV Cache显存随序列长度线性增长;另一方面,计算复杂度随序列长度二次方增长。这使得超长序列的训练与推理成本迅速恶化。

在KSA中,从推荐中常见的序列压缩角度出发,考虑在序列层面按照一定压缩比进行KV Cache的压缩提纯,这条路径追求以可接受的显存成本换取对远端历史的完整、可回指、可解释的信息保留,拥有完整的长上下文回溯能力。团队初步验证了8倍序列压缩是有效果保证的,在From Scratch和CPT上均有收益,在KV缩减、短序列Bench效果不掉的情况下,长序列Bench能够有显著提升,验证了序列级KV Cache压缩切实可行。KSA相关代码已开源,欢迎大家体验。

  • Kwai Summary Attention技术报告:

    https://arxiv.org/abs/2604.24432

  • 开源仓库:

    https://github.com/Kuaishou-OneRec/KSA

  • 模型参数:

    https://huggingface.co/OpenOneRec/KSA-4B-base

一、背景:上下文注意力的两难

过去几年,LLM领域围绕长序列建模这一问题,业界的技术路线主要集中在两个方向:

第一条路线是压缩每个token的KV表示。从MHA到GQA,再到DeepSeek系列提出的MLA,以及近期的NSA/DSA,通过head分组共享或低秩投影,逐步降低单个token的KV体积。这一方向能将压缩率做到百分之几的量级,但存在一个结构上限:KV Cache总量与序列长度依然保持1:1的线性关系。序列翻倍,cache也随之翻倍。

第二条路线是用高效注意力变体替换大多数层,代表是Hybrid-SWA与Hybrid-GDN/Linear。这些方法让多数层只看局部窗口,或将历史压缩进一个定长的recurrent state,将cache与序列长度彻底解耦。代价同样明显:SWA 完全丢弃窗口外的信息,线性注意力将历史压进固定容量的状态,长距依赖要么不可见,要么被模糊化,难以支撑需要精确回看历史的长链路推理与Agent轨迹。

二、方法

在这两者之间,我们认为存在一条有前景的中间路径:在序列层面按照一定压缩比进行KV Cache的压缩与提纯。

基于这一思路,我们提出了KSA。其核心并非追求最小KV Cache,而是在保持cache与序列长度线性关系的前提下,引入一个语义级压缩比k,将增长斜率从1降为1/k。默认k=8时,KV Cache的开销只有原来的1/8。这样既避开了GQA/MLA那种与序列1:1绑死的困境,也没有Linear/SWA把历史压成定长状态后细节糊掉的代价,在可控的显存成本下,把远端历史原原本本地保留下来,支持准确回查到具体位置。

2.1 模型设计

Summary Token 压缩

给定长度为n的输入序列,KSA 将其按chunk size k切分为n/k个chunk,在每个chunk末尾插入一个可学习的Summary token。两类token承担不同职责:

  • Summary token只对所属chunk内的text token可见,负责将该chunk的语义压缩为1个表征;
  • Text token的视野由两部分组成:近处可见最近若干chunk的原文token(sliding chunk),远处可见更早chunk的summary token。

这一设计将token空间与状态空间显式解耦:summary token专注短距语义压缩,text token通过历史summary中转访问长距离上下文。这样每一段历史要么被原文完整覆盖、要么被summary代理。

Sliding Chunk Attention

#

#

对于近期文本序列的感知选型,一个直接的方案是复用Sliding Window Attention (SWA),让text token关注最近W个token。然而,当SWA与summary token组合时,会出现一个关键问题:若窗口边界恰好切穿某个chunk中部,该chunk会同时失去两种覆盖,原文只暴露一半,而summary 又因chunk未完全越窗而未被采用,导致这段历史信息被丢失。

因此,我们设计了Sliding Chunk Attention (SCA):将滑动粒度从token级提升至chunk级,保证每个历史chunk要么整体落在窗口内(原文全部可见),要么整体落在窗口外(仅通过summary访问),不存在中间态。基于这一设计,KSA在前向过程中能够实现对前序信息的完整覆盖且避免冲突。

与现有KV Cache压缩技术组合

KV Cache总量可以沿三个维度分解:

数数维度

  • GQA压缩的是Head数;
  • MLA压缩的是Embedding维度;
  • KSA压缩的是Token数。

KSA与它们相互独立,压缩率可以直接相乘:

GQA、MLA做的是每个token的KV体积的压缩,KSA做的是参与计算的token数量的压缩,三者压缩维度正交,可以叠加使用,压缩率可观。

上图给出了组合使用的效果:在长序列下,KSA+GQA的KV Cache相对MHA压到 0.78%,KSA+MLA压到 0.22%,理论收益显著。

2.2 系统实现

#

连续KV Cache布局

#

解码阶段是典型的memory-bound场景,内存访问的连续性直接决定推理效率。KSA将KV Cache组织为三段物理相邻的buffer:

  • Current Chunk:当前正在填充的chunk;
  • Sliding Chunk Text:最近若干chunk的原文,以ring buffer循环写入;
  • Summary Token Buffer:所有历史summary token,按序追加。

这一布局下,text token解码所需的”近处原文+远处摘要”恰好对应buffer中的一段连续slice,每步attention仅需一次内存读取,无需concat、gather或动态mask构造。换言之,cache的物理布局本身即编码了可见性规则。这是KSA在大幅压缩cache的同时仍能保持解码吞吐的关键设计。

CPT三阶段Recipe

为了将KSA无缝接入成熟预训练模型(如Qwen3-4B-Base),我们设计了一套三阶段CPT训练流程:

  • 注意力蒸馏:为summary分支初始化独立的Q/K/V投影权重,从layer-wise的attention输出、distribution-wise的logits KL以及objective-wise的语言建模loss三个粒度对齐Full Attention teacher;
  • 参数退火:使用随训练步数线性衰减的插值系数\lambda,将summary分支的独立权重平滑并回主干权重,推理时不引入任何额外参数;
  • 渐进长度扩展:32K→64K→128K分阶段扩展序列长度,让summary机制逐步适应更长上下文。

三、实验结果

我们在两种设置下进行了完整评测:From Scratch(1.9B参数、400B tokens、128K)与Qwen3-4B-Base CPT(85B tokens、128K),对比方法覆盖Full Attention、Hybrid-SWA、Hybrid-SCA、Hybrid-GDN/Linear。

3.1 长上下文检索:Hybrid-KSA反超Full Attention

From Scratch长序列RULER-128K:

Hybrid-KSA超过Full Attention16.60 分,超过次优基线Hybrid-GDN5.48分。

CPT设置下的RULER-128K大海捞针测试Hybrid-KSA取得71.67,超过Full Attention5.81分。

3.2 通用能力:与Full Attention差距最小

Hybrid-KSA是所有Hybrid方案中与Full Attention能力差距最小的,甚至略高于Full,并在MBPP (62.20)、HumanEval (62.50) 等结构化生成任务上领跑。From Scratch设置下,Hybrid-KSA在MATH上取得36.92,相比Full Attention高出13.54分;GSM8K上高出10.85分。这表明summary机制对需要长链路依赖的数学与代码推理任务,反而提供了更好的归纳偏置。

四、总结与展望

长序列建模一直是推荐系统中的难题。然而在过去一段时间里,LLM在长序列能力上的演进,几乎没有从推荐系统借鉴到相关思路,其性能提升更多依赖于infra的进步,尤其是并行计算能力的持续增强,技术路径上有一定的“错位”。

在Agent时代,模型会向低推理成本的结构设计倾斜,比如DSA的Top-k indexer排序,V4的Sequence Compressor压缩,都可以看到推荐系统思想的影子。从趋势上看,推荐系统低成本serving的技术和思想与大模型预训练的底层技术交汇将会越来越紧密。许多在推荐场景中已被反复验证的问题,正在LLM体系中被重新定义与解决。快手推荐模型团队长期深耕这两个方向,致力于推动推荐系统与大模型技术的深度融合。欢迎更多同行加入,一起构建更高效、更智能的大模型与推荐系统。

五、关于我们

快手推荐模型团队的使命是打造下一代推荐系统。这里有超强的算力和业界top院系成员,我们致力于利用大模型的世界知识和学习范式为推荐系统带来新的增长曲线,寻找算力到用户价值的最优转化方案。目前我们在打造业内领先的生成式推荐大模型OneRec,包括但不限于生成式推荐、多模态理解、强化学习、大规模稀疏MoE、高效训练、蒸馏剪枝等前沿技术探索等方向,探索算力到业务效果的最优转化路径。

六、招聘

  • 招聘岗位

  • 推荐大模型算法专家-【OneRec】

  • 推荐算法工程师-【Ranking方向】

  • 推荐算法工程师-【Matching方向】

  • 职位描述

  • 应用大模型技术优化亿级用户视频推荐系统,显著提升用户停留时长、点击率、留存率等核心业务指标;

  • 基于生成式推荐模型重构推荐系统架构与底层逻辑,推动生成式推荐范式在多种业务场景中的规模化落地;

  • 持续提升生成式推荐模型性能,涵盖生成式推荐模型扩展(Scaling)、强化学习偏好对齐、大语言模型知识融入与推荐系统协同等关键方向。

  • 任职要求

  • 对AI技术具备强烈热情,认同技术驱动业务与社会价值提升;

  • 高度关注推荐大模型/大语言模型技术前沿,对行业动态有持续追踪和理解;

  • 在推荐系统、计算广告、信息检索、NLP、LLM、强化学习、多模态、高效训练/推理等至少一个领域具备扎实技术积累与业务落地经验;

  • 具备优秀的工程实现能力,熟练掌握常用数据结构和算法;

  • 能够快速阅读和复现前沿文献,具有良好的问题分析、解决及团队协作能力。

  • 加分项

  • 在SIGKDD、ICML、NeurIPS、ACL、WWW、RecSys等国际顶级会议发表过论文;

  • 有ACM/ICPC、TopCoder等算法竞赛获奖经历;

  • 具备大语言模型与强化学习结合(RL)在真实业务场景的落地经验者优先。

  • 招聘岗位面向

  • 校园招聘:面向2026届毕业生

  • 日常实习:面向2027届及之后的在校同学

  • 社会招聘:诚邀算法领域专家

  • 我们希望您具备

  • 对人工智能技术充满热情,坚信技术可推动业务发展;关注LLM/推荐大模型的前沿动态,具备扎实的算法基础;

  • 拥有优秀的工程实现能力,能够快速复现学术界最新成果;

  • 如有在ICML、NeurIPS、ACL等顶级会议发表的论文,或获得ACM/ICPC等竞赛奖项,将成为重要加分项。

简历投递方式:

扫描下方二维码投递或将简历发送至邮箱:[email protected]

【推荐阅读】

点击【阅读原文】,了解KSA详情!


免责声明:

本文所载程序、技术方法仅面向合法合规的安全研究与教学场景,旨在提升网络安全防护能力,具有明确的技术研究属性。

任何单位或个人未经授权,将本文内容用于攻击、破坏等非法用途的,由此引发的全部法律责任、民事赔偿及连带责任,均由行为人独立承担,本站不承担任何连带责任。

本站内容均为技术交流与知识分享目的发布,若存在版权侵权或其他异议,请通过邮件联系处理,具体联系方式可点击页面上方的联系我

本文转载自:快手技术 快手技术 快手技术《开源!快手提出面向长上下文推理的序列压缩注意力机制KSA!》

评论:0   参与:  0