一张RTX 4090,24GB显存,跑一个32B参数的大模型做agent任务。
不做任何KV压缩,显存直接爆掉,连模型都跑不起来。
换上TriAttention,模型稳稳跑起来,顺利读完6份文档,自动生成了一份完整周报。
这不是社区大神的魔改,而是一篇来自MIT、英伟达、浙大的联合论文。
https://arxiv.org/pdf/2604.04921
核心思路是在pre-RoPE空间里,用Q/K的三角集中度来估计每个KV token到底有多重要,然后只保留真正重要的那些。
打个比方来说,别的方法压KV cache像是把所有行李都塞进压缩袋,不管里面是羽绒服还是砖头一律压扁。
TriAttention是先翻一遍行李箱,把砖头扔掉,只给羽绒服打包。
TriAttention demo演示,展示单张RTX 4090上Qwen3-32B完成OpenClaw agent任务的完整过程。
作者之一Yukang Chen在X上发布了这组对比,左边不压缩,显存直接报错;右边开了TriAttention,agent一路读完6份文档,周报完整输出。
2.5倍吞吐,10.7倍内存缩减
效果怎么样?数字说话。
在AIME25数学推理任务上,TriAttention在匹配Full Attention准确率(40.8%)的前提下,吞吐量提升了2.5倍。
再看内存:KV cache内存缩减10.7倍。
在AIME25(Qwen3-8B)上的性能权衡。(A) 在相同准确率(40.8%)下,TriAttention的吞吐量比Full Attention高2.5倍。(B) TriAttention在保持与Full Attention相同准确率的同时,将KV缓存内存减少了10.7倍。
注意,这里说的是KV cache memory,不是整机显存,也不是模型参数占用的总内存。
但就算只是KV cache这一项,对长序列推理场景来说,KV cache往往就是压垮显存的最后一根稻草。
砍掉这一项,就是能跑和不能跑的分界线。
主实验是在Qwen3-8B上做的,覆盖AIME24、AIME25、MATH500等任务。
在32K token的生成长度条件下,TriAttention几乎没有牺牲精度,但把推理效率拉到了一个新台阶。
单张4090跑通32B大模型
这篇论文附录中提到了一个真实部署案例。
场景是OpenClaw,一个多轮agent工作流。任务是读6份markdown文档,生成一份周报。
模型是Qwen3-32B,用了AWQ INT4量化,跑在一张RTX 4090(24GB)上。
不压缩KV cache直接跑这个任务?显存当场爆掉。
长系统提示加上多轮文档读取,KV cache膨胀到显存根本兜不住。
TriAttention接管之后,agent顺利读完所有文档,生成了完整报告。
模型用的是Qwen3-32B AWQ INT4量化版,不是原始FP16满血版;跑的是OpenClaw agent工作流,不是通用长文本benchmark。
但它刚好证明了「一个完整的、有实际生产价值的agent任务,可以在消费级硬件上跑通」。
vLLM插件已就位,MLX实验性起步
TriAttention不只停在论文里。
作者已经在GitHub仓库中提供了vLLM集成,README明确写到TriAttention包含一个vLLM插件,并给出了OpenAI兼容API的server mode、Python API以及OpenClaw接入说明。
相比论文中的实验结果,这属于仓库层面的工程化扩展。
这意味着,你不需要改模型架构,不需要重新训练,只需要挂上这个插件,就能在现有的vLLM推理管线上获得KV压缩收益。
在Apple Silicon方向上,官方仓库里单独放了一份docs/mlx.md,覆盖M1到M4全系芯片,基于MLX框架和mlx-lm运行,附带示例代码和硬件benchmark。
TriAttention官方仓库已提供MLX实验性支持文档,覆盖M1-M4芯片https://github.com/WeianMao/triattention/blob/main/docs/mlx.md
不过,官方文档标题中也标注了这还是实验性支持,这说明他们已经在早期试水MLX了,但离成熟的Mac本地部署还有距离。
KV压缩赛道的两条路线
KV cache压缩赛道存在两条路线。
一条是量化派。
Google Research在3月24日发布了TurboQuant,官方博客中的定位是「在零精度损失下实现极致压缩」的方案,主打把KV cache和向量搜索的bit数压到极低。
Google Research官方博客中LongBench基准测试图,TurboQuant在LongBench基准测试中,相较于多种压缩方法,在Llama-3.1-8B-Instruct模型上展现出稳健的KV缓存压缩性能
社区已经有人在Apple Silicon上用TurboQuant跑通了Gemma 4 31B。
另一条是选择性保留派。
TriAttention就是这条路线的新代表,不压bit,而是直接判断哪些token的KV值得留、哪些可以扔。
两条路线的终点其实一样:让大模型跑在消费级硬件上,显存不炸,精度不掉。
但方法论完全不同。
量化是把每个行李都压扁,选择性保留是直接减少行李数量。
理论上,两者甚至可以叠加使用。
目前还没有严格的同模型、同硬件、同任务的head-to-head对比,所以「谁碾压谁」还说不了。
但可以确定的是,这两条路线正在加速向消费级部署推进。
一年前,「本地跑大模型」还是极客圈的行为艺术,跑个7B都要折腾半天。
现在,32B模型在单张消费级卡上完成agent任务,Apple Silicon上的MLX生态一周一个新仓库,vLLM插件让KV压缩变成「挂上就用」的一键方案。
KV cache压缩这条赛道,正在从论文里的消融实验,变成每个开发者都能触碰到的工程现实。
作者简介
Weian Mao
Weian Mao
Weian Mao现为MIT CSAIL博士后研究员,博士毕业于阿德莱德大学AIML,师从沈春华教授。其当前研究聚焦大语言模型,尤其关注推理效率与长上下文推理中的KV cache压缩;此前也从事过计算机视觉与蛋白质设计等方向研究。
Xi Lin
Xi Lin
Xi Lin是浙江大学计算机科学与技术专业高年级本科生,研究兴趣集中在高效AI的算法—系统协同设计,尤其关注面向硬件友好的稀疏与量化模块设计,以及高效推理策略。其工作与高性能计算、机器学习系统等方向密切相关。
Wei Huang
Wei Huang
Wei Huang现为香港大学博士生,研究聚焦Efficient AI与大型视觉/语言模型。
目前,他在NVIDIA Research实习,与Yukang Chen等研究者合作,并在Song Han 指导下开展相关研究,参与了QeRL、LongLive等工作。
参考资料:
https://arxiv.org/abs/2604.04921
https://x.com/yukangchen_/status/2041366586423165152
https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/
本文来自微信公众号“新智元”,编辑:元宇 ,36氪经授权发布。