分享好友 健康资讯首页 健康资讯分类 切换频道

英伟达MIT出手,华人团队重磅开源,大模型推理内存暴降10倍

2026-05-14 21:090036kr

一张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氪经授权发布。

举报
收藏 0
打赏 0
评论 0
“易中天”的狂飙,被英伟达们盯上了
5月12日,中际旭创突破万亿市值后,再次迎来新的“里程碑”——股价突破1000元,成为创业板历史上第二只千元股。上一个创业板千元股是国内医美赛道的龙头企业爱美客,靠玻尿酸等注射类产品撑起了千亿市值,不过到现在其股价已从巅峰跌去近八成。而接过这个位置的,是光模块赛道。2026年以来,在AI算力基础设施扩张的带动下,光模块成为A股市场最火热的板块之一。被市场戏称为“易中天”的三只龙头股——新易盛、中际

0评论2026-05-140