新浪财经 4小时前
大模型推理加速全链路:内存管理、编译优化、量化与并行策略
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_font3.html

 

作者 | 金煜阳博士,清华大学助理研究员

策划 | QCon 全球软件开发大会

大模型推理成本已成为 AI 产业落地的核心瓶颈。随着模型参数规模向万亿级迈进、多模态智能体应用场景爆发,推理引擎需要应对算子实现、内存管理、量化压缩、异构调度与并行策略等多维度的技术挑战。本文整理自清华大学助理研究员金煜阳博士在 QCon 全球软件开发大会 2026 北京站的分享《大模型推理加速全链路:内存管理、编译优化、量化与并行策略》。

金煜阳系统介绍了其所在实验室在大模型推理加速全链路方面的探索,覆盖从底层算子优化到上层并行策略的完整技术栈。本文将逐一展开算子级性能调优中的细粒度图层优化方法、面向异构模型结构的 KV Cache 内存管理、编译与运行时协同的混合精度量化、基于负载特性的 CPU-GPU 异构调度,以及面向动态请求场景的自适应并行策略,最后介绍开源推理引擎赤兔在国产芯片生态中的实践成果。

以下是演讲实录(经 InfoQ 进行不改变原意的编辑整理)。

人工智能的市场规模正在以超出预期的速度扩张。去年来自 Precedence Research 和艾瑞咨询的数据预测,到 2030 年全球 AI 市场规模将达到 11 万亿。但今年年初 OpenAI 的持续爆发,以及智能体应用的集中涌现,让这个预测数字显得保守了许多,我相信它一定会再往上调整不少。与此同时,中国 AI 产业也处于快速发展期。从企业数量、人才培养、论文专利等维度与美国对比,我们在产业界和顶尖人才方面仍有差距,但追赶速度极快,可以说正在从跟跑进入并跑的阶段。

智能体的兴起正在改变大模型的应用范式。基座大模型提供基础能力,智能体将其与人类已有的所有软件工具连接起来,从而赋能各行各业。但在这一片繁荣背后,推理成本构成了整个产业链中最重的开销。能够训练基座大模型的厂家屈指可数,但几乎每一个 AI 应用背后都在调用推理 API。推理,以及推理所需的算力,就是这个产业背后的主要成本。

我们团队做的其中一件事情,就是构建一个大模型推理引擎,让它接在各类 AI 芯片之上、支撑各种模型、最终服务于上层的 AI 应用。底层芯片方面,除了英伟达的 GPU,国内也有华为昇腾、摩尔线程、燧原、沐曦、天数智芯等厂商在奋力建设国产芯片生态。上层,则有 VLM 和 SGLang 等优秀的开源推理引擎。我们清华大学实验室也开源了一个推理引擎,名为赤兔——这个在后面会详细展开。

理解推理引擎的优化,需要先理解大模型推理的两个基本阶段:Prefill 和 Decode。当用户输入 " 今天中饭吃什么 " 这句话时,系统会把整句话一次性处理,这个阶段叫做 Prefill,是计算密集型操作。之后,模型开始一个字一个字地输出回复,就像我现在说话一样,一个字接一个字地往外吐,这就是 Decode 阶段,也叫自回归生成。Decode 阶段每一次生成新 token,都需要访问之前所有的上下文数据,计算量不大,但访存量极大,因此是访存密集型操作。这两个阶段的特性截然不同,对底层资源的需求也完全不同,这为后续的异构调度和并行策略设计埋下了重要伏笔。

在这个基础之上,大模型的发展趋势给推理引擎带来了四个维度的挑战。首先是模型规模越来越大,DeepSeek 宣称的 V4 参数已达万亿级别,实际规模可能更大。其次是模型架构日益复杂。智能体场景下,用户与模型长期对话,上下文越来越长,还会从工具调用中产生视频、图片等多模态信息。架构层面,除了主流的 MoE(混合专家)结构,千问 Next 等模型还采用了混合注意力模式:一部分是传统的 Full Attention,另一部分是 Mamba 这样的线性注意力。这些异构架构对算力的需求分布完全不同,推理引擎必须具备足够的灵活性来应对。

从高性能计算的基本原理出发,推理系统的性能取决于五个关键维度:算子优化、内存管理、模型量化、异构调度和并行优化。算子优化是最底层的能力,就像一个公司里每个人的单兵作战能力。如果每个算子的执行效率都不高,再好的系统调度也无济于事。

去年 DeepSeek 的开源周活动很好地印证了这一点。他们陆续发布的 FlashMLA 和 DeepGEMM 等子模块中,有两个直接与算子实现相关,另外两个与并行通信相关。这说明,底层算子的性能是整个系统的基石。问题在于,现代 AI 硬件的算力增长迅速,但新功能的复杂性也在急剧上升。以 Tensor Core 为例,GPU 原本只有通用的 CUDA Core 做简单运算,但为了加速深度学习中的张量计算,英伟达专门加入了 Tensor Core。如何将 Tensor Core 与 CUDA Core 协同调度,如何利用 Distributed Shared Memory 等新硬件特性来缓解访存瓶颈,就成了算子实现中的关键难题。硬件为大模型负载而设计新特性,软件必须跟上才能把硬件性能榨干——这是一个相辅相成、持续迭代的过程。

手写算子可以针对特定形状达到极致性能,但问题在于,模型在并发度变化时会产生不同形状的算子切分,而每种形状的最佳实现都不同。FlashMLA 是一个典型例子——如果针对特定形状手工优化,性能远优于自动编译器 Triton 生成的代码。但为每一种形状都手写算子,开发成本极其高昂。而用 Triton 自动编译虽然省力,性能又大打折扣。这中间的根本原因,就是编译器难以自动适配不同硬件架构的底层特性。

编译过程包含了多个重要环节,我在这里做一个高度简化的说明。深度学习代码首先被抽象为计算图(这个概念由 TensorFlow 首倡),然后在图层进行优化:算子融合、计算等价或不等价变换等。融合后的算子再调用算子库或进行自动代码生成,最终才能在芯片上运行。在图层面,不同芯片架构需要不同的算子聚合或拆分策略;在算子层面,又涉及计算与访存的掩盖、异步流水线调度等细节,每种芯片设计各不相同。即使是英伟达自家的 H 系列,H20 和 H100 的计算访存能力不同、显存大小不同,所需的调度策略也不同。把这一切全部交给编译器自动完成,难度极大。

我们团队在图层变换和算子代码生成这两个方向都做了工作,这里重点介绍其中一项成果—— FlashTensor,这项基于张量属性的细粒度图层优化工作发表在去年的领域顶会上。这个工作的出发点是一个观察:模型参数量和上下文长度都在增长,但增长速度不同。更大参数量带来更高精度,这是 Scaling Law 的核心结论;更长上下文带来更好的记忆能力,这一点在代码智能体中体现得尤为明显。如果模型只能处理 200K 上下文,当对话内容超出时,它只能将 200K 压缩成 10K 作为背景知识,再进行下一轮处理。而如果上下文窗口达到百万 token 级别,模型就能持续保留最无损的信息,记忆能力显著增强。但问题是,上下文长度的增长速度快于参数规模的增长,这导致计算过程中中间结果的变化幅度不一致。举个例子:第一个算子算出的中间结果写入显存,第二个算子再读取。如果上下文的维度增长更快,通过不同维度的乘法运算,中间结果会爆炸式增大——将中间结果写入显存的过程反而成为瓶颈,前后计算本身可能并不大。

现有的粗粒度融合处理方式有两种。第一种是把第一个和第二个算子合并,消除了中间结果的写入读出开销。但合并后的计算图依赖关系复杂,并行度受限,无法让 GPU 上几千个线程同时工作,硬件效率难以发挥。第二种情况是,如果编译器识别出算子无法合并,就拆分成很多小算子。这时中间结果的访存开销极大——虽然每个小算子计算时间短,但不间断地写入读出中间张量会让性能骤降。我们在 A100 上测试过一个 Sparse Attention 算子,计算效率仅有 3.45%,而正常情况下应该达到 80% 到 90%。

我们提出的方法是更细粒度地审视这些依赖关系,定义出细粒度的属性来描述算子间复杂的依赖结构,再据此进行变换和切分,在最大化并行度的条件下重组计算。在这个场景下,我们在 A100 和 H100 上进行了多模型测试,与 TensorRT、PyTorch 原生实现、TVM、CUTLASS 等多个基线进行对比。衡量的是端到端推理时间,而非单个算子的加速比——整体推理时间提升了 1.5 到 1.86 倍。

大模型推理过程中会产生一类关键中间结果—— KV Cache。可以将其直观理解为模型对上下文的记忆:模型把用户的输入编码成自己能理解的形式存储起来。对话越长,KV Cache 就越大。它带来的核心问题是显存占用:单张 GPU 可能只有 80GB 显存,而 KV Cache 会随着对话持续增长,很快超出物理限制。硬件层面当然可以堆更大的显存,但从软件层面做更智能的调度显然是更经济的方案。基本思路是:当 KV Cache 需要被访问时,放在 GPU 显存中;不需要时,转移到 CPU 内存;当用户似乎要再次对话时,再预取回 GPU。用户在感觉上以为数据一直留在显存里。

目前这个方向上主要有两类工作。第一类是 KV Cache 压缩:分析哪些历史 token 对后续对话的相关性更高,只保留最有价值的部分。如果我说了一大段话,可能 80% 是废话,只提炼 20% 的关键信息就足够了。这类方法通过量化分析做出筛选。第二类是更高效的系统管理,不删减任何信息,而是借鉴操作系统的页表机制进行动态内存分配——这就是 VLM 的 Paged Attention 方法。基于页表的管理方式能大幅减少内存碎片。

我们的研究聚焦于一个更新的场景:异构大模型。这类模型内部混合了不同的注意力机制,Full Attention 做压缩滑动窗口来筛选相关信息,Mamba 做线性注意力,未来可能成为模型架构的一种主流形态。关键发现是,不同的注意力层需要不同的内存分配机制。如果用统一的大页表管理所有层,仍然存在大量浪费。我们提出了结合大页和定制化小页的多粒度内存管理方案,把粒度划分得更细致,尽可能避免碎片。在 H100 平台上针对新型异构模型进行测试,与当前最优的 VLM 方案相比,整体吞吐量提升了 4.92 倍。

模型参数规模的飞速增长与 GPU 显存的有限增加之间存在巨大的剪刀差。量化——用更低精度的数值表示来存储参数和执行计算,是缓解这一矛盾的核心技术手段。计算机中用不同精度表示数字,精度越高,有效数字越长,但需要的计算资源也越多;精度越低,有效数字短,但计算更快。这是一个典型的权衡问题。

一种常见的量化方法是在推理时将参数以低精度存储,但计算仍保持高精度。这能节省大量显存,却无法利用硬件中的低精度计算单元。如果能把权重和激活值同时量化,既量化存储又量化计算,就既节约显存,又提升吞吐率。但问题随之而来:全部降精度后,模型开始 " 不说人话 ",输出质量严重下降。

当前主流方法叫做混合精度量化。基本思想是:参数和计算中只有极少数关键点需要保持高精度,就像公司里有少数几个人需要特别关照。这少数 " 离群值 " 用高精度处理,剩余 99% 用低精度处理,理论上既节约内存又没有精度损失。然而在实际硬件上运行时,效率会骤降到只有原来的 30%,这是完全不能接受的。因为单独处理这些离群值就像单独给每个人发私信,调度开销极大。

我们的优化从编译和运行时两个层面入手,对离群值数据的重构和调度流程进行一系列改进,大幅降低了专门调度离群值的开销。测试工作不只在英伟达平台上进行,还在燧原的芯片上做了验证。在英伟达平台,算子性能相比当前最优方案提升 1.5 到 1.6 倍,相比 AWQ 提升 6 倍。在燧原 S60 加速卡上进行的端到端精度验证表明,混合精度量化后模型精度与原始 FP16 版本基本持平,而性能普遍提升 2 倍以上。

回到推理的两个阶段:Prefill 和 Decode。过去可能更关注 Decode,因为聊天场景下用户对逐字输出的速度极为敏感——如果一个字一秒钟才吐出来,等待感会很强。但在智能体工作模式下,情况发生了变化。以写代码为例,用户不仅关注代码输出速度,还关注多轮对话的效率。每一轮对话都会形成新的上下文,需要对整个上下文重新进行 Prefill。这意味着 Prefill 和 Decode 可能都很慢,都是影响用户体验的瓶颈。

Prefill 是计算密集型,Decode 是访存密集型,两者负载特征完全不同。传统做法是把它们都放在 GPU 上处理,但既然特征不同,是否应该考虑用不同的硬件资源来承接?进一步拆解 Decode 阶段,Attention 层和 MLP 层在不同 Batch Size 下的算术密度变化规律不同—— Attention 的可扩展性较强,MLP 扩展性较弱。这带来一个思路:现代计算系统中 CPU 有大的内存资源,虽然计算能力弱,但能不能让它处理那些 Batch Size 增加但算术密度变化不大、对计算需求小但对内存需求大的部分?而将对计算需求极大的 MLP 继续留在 GPU 上执行。

这个想法的本质是把显存占用大而计算密度低的任务卸载到 CPU,把计算密集任务留在 GPU。在特定场景下,比如 Batch Size 能拉得非常大的时候,这个方案的吞吐量提升非常显著,可以达到 7 倍左右。我们将这个工作命名为 FastDecode,在此基础上又做了进一步优化,称为 FastDecode Max。

当前大模型的参数量已经大到难以放在单台机器上,即使经过量化压缩也很难塞进一张 GPU。我们算过一个数:全量参数需要 9 张 H100 才能放下。因此,分离式架构和各种并行策略的组合成为必然选择。

DeepSeek 的部署配置非常能说明问题。在 PD 分离策略下,Prefill 阶段用 32 张 GPU 作为一个部署单元,Decode 阶段用 144 张 GPU。而且每个阶段内部采用的并行策略也不同:Attention 部分使用张量并行加数据并行的 4 乘 8 模式,MoE 部分则直接使用 32 个专家并行。这说明,不同子结构需要的并行策略是不同的,不存在一劳永逸的方案。

更进一步,不同应用场景带来的负载特征差异很大。简单的客服问答,上下文长度很短;科学研究或财务分析,上下文可能极长。用户是否开启深度思考模式,也会导致上下文长度剧烈变化。这些输入差异导致推理负载在计算密集型和访存密集型之间漂移,需要不同的并行策略来应对。另一个关键特征是推理场景的动态性。用户通常是白天集中请求、晚上几乎停歇,后端压力呈现明显的潮汐现象,我们从部分 Trace 数据中能清晰看到这种波动。在不同时段,并行策略应当不同:负载压力大时以吞吐量为优先,压力小时可以切换策略来降低时延,提供更高质量的推理服务。

这种动态变化要求推理引擎能够根据负载情况切换并行策略,实现自适应的调度,这是我们并行优化的核心方向。

开源大模型推理引擎——赤兔

最后介绍我们实验室的开源推理引擎——赤兔。它的核心使命是在国产芯片上提供高质量的推理性能。VLM 和 SGLang 已经是很优秀的开源引擎,但它们在国产芯片上适配度有限。尤其是在 H100 等高端 GPU 禁运的背景下,国产算力必须 " 支棱起来 "。算子也许已经写好了,但把推理系统整个跑起来、跑得好,仍然是很大的挑战。赤兔的目标就是成为一个支持国产算力、也支持国产大模型的开源推理引擎。

赤兔集成了前述的一部分技术。一些偏向于探索未来模型结构的前沿技术,比如面向异构模型架构的内存管理,暂时还没有完全融入,但我们在工程和研究上双线投入了大量人力去做适配,目前在国产硬件上已经能跑出不错的性能。同时在英伟达平台上,赤兔也能做到与主流引擎性能持平。

在 A800 集群上,赤兔相比 VLM 更快。我们实现了一个 FP8 量化方案——在硬件不支持 FP8 量化单元的情况下,通过软件手段实现了等效加速。在精度持平的前提下,速度明显优于 VLM。在华为昇腾芯片上,结合 FP4 量化能力,可以在更少的资源上提供较高性能。沐曦芯片也得到了支持。另一项工作是异构调度:在一台仅配备一个 CPU 和一个 GPU 的机器上,成功装入并运行了 617B 参数的完整模型。相比 VLM 的 8 张 H20 配置,虽然赤兔这套单机方案的吞吐量大约只有其一半,但以如此有限的资源撬动如此庞大的模型,本身已是不小的突破。

我们的开源计划持续推进,中间不断有新技术融入。当前版本是 0.5.4,即将发布的 0.6 版本会在智能体推理方面做进一步工作。多说一句,科研圈的主流工作集中在前述那些面向新型模型结构的技术探索上,而我们实验室额外投入了大量精力去做国产硬件的适配,这中间有相当多的工程工作。大模型推理成本对整个产业落地至关重要,我们尝试做的,包括算法层面的量化、KV Cache 压缩,以及软硬件协同优化,都是朝着降低这个成本的方向努力。

我非常认同 DeepSeek 团队持续在做创新,而我们要做的,就是通过技术手段把大模型对底层硬件算力的需求打破、压低,让推理变得更高效、更普惠。

金煜阳,工学博士,清华大学计算机系助理研究员。CCF 高性能计算专委会执行委员。主要研究领域包括并行应用性能分析与优化、高性能人工智能系统。在并行计算与系统领域顶级会议和期刊 SC、PPoPP、ATC、EuroSys、IEEE TPDS 等发表论文 16 篇,出版专著 2 部,研究成果获 IEEE TPDS 2022 最佳论文 Runner-up 奖。主持或参与国家级项目 4 项、省部级项目 1 项。担任清华大学学生超算团队指导老师之一,指导团队 3 次获得世界冠军。获 CCF 体系结构优秀博士学位论文激励计划、清华大学优秀博士学位论文、博士后创新人才支持计划等。

宙世代

宙世代

ZAKER旗下Web3.0元宇宙平台

一起剪

一起剪

ZAKER旗下免费视频剪辑工具

相关文章
评论
没有更多评论了
取消

登录后才可以发布评论哦

打开小程序可以发布评论哦

12 我来说两句…
打开 ZAKER 参与讨论