AIxiv专栏是机器之心发布学术、手艺内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
本文作者袁镱博士是腾讯公司专家工程师，负责无量系统和一念LLM等机器学习训练和推理框架研发。
以 OpenAI 的 GPT 系列模型为代表的大语言模型（LLM）掀起了新一轮 AI 应用浪潮，但是 LLM 推理的高昂成本一直困扰着业务团队。
腾讯 PCG 机器学习平台中心自研了高功能 LLM 推理引擎：一念 LLM。在传统的算子融合，ContinousBatching 等推理加快手艺的基础上，通过显存优化，异步调理和盘算复用等手艺，在相同精度的推理中，一念 LLM 相比 vLLM，TensorRT-LLM 等著名开源框架的推理单价低 20%+。
另外，为了应对国外高端 GPU 卡供应不足的问题，一念 LLM 在高功能 LLM 推理框架领域第一次同时支持 Nvidia GPU 卡和华为 NPU 卡。目前一念 LLM 已在 QQ 智能体等 PCG 主要的 LLM 业务场景上线。
本文以一个简单的公式，逐步分析 LLM 推理的功能瓶颈，介绍当前 LLM 推理的相关手艺，以及一念 LLM 的设计决策逻辑。
“夫一心具十法界，一法界又具十法界、百法界；一界具三十種世間，百法界即具三千種世間。此三千在一念心，若無心而已，介爾有心即具三千。”
— 出自：佛教天台宗摩訶止觀卷五上（大四六・五四上）
“一念亦稱一心，指心念活動之最短時刻；三千表示世間與出世間一切善惡、性相等人、物差別之總和。一念三千即謂，於凡夫當下一念之中，具足三千世間之諸法性相。”
— 出自：佛光大辭典 (慈怡法師主編) 词条 “一念三千”
一念 LLM，取 “一念三千” 之意，寓意 “一念之间，用大模型生成世间万象”。
简介
以 OpenAI 的 ChatGPT 为代表的大语言模型（LLM）掀起了新一轮 AI 应用浪潮，业务团队都在探索基于 LLM 重构现有应用或者构建新的 APP。大语言模型的大参数量导致了巨大的盘算和显存需求，使得 LLM 的要求推理成本高昂。LLM 推理框架成为 2023 年以来的业界研究热点。当前有多个著名的开源项目，比如：UCBerkeley 的 vLLM 和 Nvidia 的 TensorRT-LLM。
这些框架完成了诸多业界先进的推理加快手艺，比如：ContinousBatching、PagedAttention 等，但是也存在两方面的问题：
1. 为了便于算法人员使用和尝试新手艺，vLLM 采用了 Python 作为主要调理管理功能的完成语言，导致显存管理和调理效率较低。
2. 主要支持 Nvidia 的 GPU 等国外主流厂商的硬件，对国产硬件没有支持。国产硬件配套的推理框架，缺乏对业界最新的推理加快手艺的支持。
一念 LLM 通过对异构硬件的底层抽象，构建统一的高功能调理管理，完成了：
1. 应用业界最新的推理加快手艺，推理单价相比业界开源框架低 20%+。结合业务场景进行针对性优化，单价可以降低 60%+。
2. 将最新的手艺同时应用到国外主流的 Nvidia GPU 和国产的华为 NPU 上，避免软件手艺被硬件供应能力影响。
一念 LLM 已开源，欢迎共建。代码地址：https://github.com/pcg-mlp/KsanaLLM
问题分析
为了构造一个高功能的 LLM 推理框架，我们需要从源头上分析推理功能的瓶颈。我们从两个方面来分析：1）调理和显存管理；2）盘算功能优化，类似算子融合，量化等盘算优化等。
调理与显存管理
在一个 LLM 推理系统中，我们希望 token 的生成速度能够最大化。由于 GPU 硬件的特性，将多个要求组合成一个大 batch 并行盘算是 LLM 推理主要的盘算速度提高手段。从 A100 的推理压测结果图可以给我们不少启示：
图 1 并行盘算 token 数与 GPU 实际盘算效率关系。图片来源：https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-fastgen/README.md
当前向推理的 token 数增大时，GPU 有效的盘算量吞吐逐步增大，当到达 200Tflops 之后，趋于稳定。这个稳定的上限与硬件的最大 Tflops 参数（A100 的 float16 的标称 flops 为 312Tflops）以及 LLM 推理算子的完成效率有关。
图 2 GPT-2 模型推理流程示例。图片来源：https://medium.com/@joaolages/kv-caching-explained-276520203249
在 LLM 模型的推理流程大致分为 prefill 和 decoding 两个阶段。在 prefill 阶段（图 2 中 '$' 之前部分输出，生成 'A' 的流程），prompt 的多个 token 被一起输出给模型，输出的 token 数量可能几百或者几千。在 decoding 阶段（图 2 中红色输出部分），模型每次输出上一次生成的 token，生成下一个 token，循环这个逻辑直到回答结束。
从图 1 中，我们可以标出两个阶段所处的工作区间。在输出的 token 数超过 300 的情况下，prefill 阶段可以处于 GPU 全力工作的区间，而由于 decoding 阶段一个要求每次输出的 token 只有一个，则处在 GPU 极大浪费的状态。decoding 阶段如果要达到 GPU 盘算资源的充分利用，batch size 应该增大到 300 左右。然而，实际情况下由于 GPU 显存的限制，batch size 远远小于这个数。
我们需要进一步分析显存是如何影响 batch size 的。我们可以列一个简化的公式来帮助分析：
其中 M 是模型参数占用的显存，α 是每个要求推理流程中的显存占用，BS 是 batch size，β 是每个 token 对应的 kv cache 所需的显存，TN 是缓存 kv cache 的 token 数量，Mem 是 GPU 的显存大小。在不使用量化等手段的情况下，选定模型和 GPU 硬件后，β 和 Mem 是固定。如果要增大 batch size，就需要降低 M，α 和 TN。
M 主要由放在显存上的参数量决定的。α 主要是由模型盘算逻辑中的中间变量占用的显存空间决定的，而 TN 与 BS 相关，一个简单的关系是如果 batch 内要求的 token 平均数量为 TA，那么
。γ 表示 batch 中差别要求之间 token 不能复用 kv-cache 的比例。所以，从显存节省的角度优化系统的吞吐，就有下面两条路径：
优化 M：在对 latency 影响较小的前提下，将参数 offload 到内存或者其他存储上。
优化 α：优化推理盘算逻辑中的中间变量显存占用。
优化 γ：优化 batch 中差别要求之间复用的 kv-cache 比例。
盘算功能优化
LLM 模型由于模型结构非常固定，尤其 ChatGPT 的成功让比传统 transformer 结构更简单的 decoder-only transformer 结构模型成为当前的主流。这种稳定而简单的结构让手撸大算子成为了 LLM 推理加快框架的首选，类似 FlashAttention 等针对单个大结构深度优化的算子库深受大家追捧。各个开源框架都纷纷推出自己的定制算子，Nvidia 等硬件厂商也都提供了与自身硬件高度适配的算子库，甚至不惜为同一结构的差别参数大小模型单独开发算子。
对开源算子的支持能力，决定了框架是否能有一个持平业界的基线功能。
方案设计
下面我们先简单介绍一下一念的主要模块，稍后按照从前面提到的多个功能优化角度介绍一念 LLM 的设计。
一念 LLM 的基本结构
一念 LLM 主要由以下模块组成：
图 3 一念 LLM 模块图
内存 / 显存统一管理模块负责分配和管理内存和显存，其中最重要的功能是分配和管理 PagedAttention 机制所需的 Block 和推理流程中所需的临时显存。在后期，与调理配合，完成更细化的调理能力。
要求调理器模块负责调理多个要求执行，协调内存 / 显存统一管理，完成推理流程的流水线化。
kv-cache 缓存调理负责 kv-cache 在要求之间的缓存复用。
加快算子包括差别硬件的模型并行，盘算量化，算子融合等功能相关的算子。包括了自研的高功能算子，经过框架适配的开源框架优秀算子。相关算子随着业界发展迭代。
统一抽象接口负责将差别硬件的算子以相同的操作方式对接到执行流水线。采用是类似 Nvidia Cuda API 的接口。
LLM 模型是基于统一抽象接口和加快算子库来支持的 LLM 模型。
模型要求调理模块用于调理差别的要求到后端推理节点。与传统机器学习推理差别，LLM 模型具有推理时间长，状态数据大的特点，要求调理模块更加要求的特点和后端服务节点的状态进行调理，优化系统功能。
ContinousBatching&&PagedAttention（优化有效 batch size）
在大语言模型推理流程中，一般会使用到 GPU 进行加快。在一个要求的依次生成 token 的流程中会有大量使用 kv-cache 来降低盘算量，但是 kv-cache 本身会占用 GPU 显存资源。目前 LLM 推理的功能瓶颈主要是因为 LLM 参数量大导致的显存带宽瓶颈，为了提高服务吞吐，需要尽量使用多个要求组成一个大 batch 来推理，以充分利用 GPU 的盘算能力。
在通常情况下，由于大语言模型盘算流程中用到了一个自增长的 kv-cache，为了加快 GPU 的盘算流程，通用方案（图 4 (a)，典型代表为 FasterTransformer）都是按 batch 为单位调理执行。由于 batch 中差别的要求 token 数量差异大，batch 粒度的调理方式会导致 GPU 盘算能力的浪费，后续的要求不能得到及时处理，影响推理服务的吞吐能力。在 batch 执行的后期，可以理解为有效输出 token 的 batch size 在逐渐变小。
以图 4 (a) 为例，到第 5 步时，只有两个要求还在推理，到第 6 步，有效 batch size 就只有 1 个了。
图 4 差别调理方案示意图。
为了充分利用 GPU 的盘算能力， 需要细化要求调理的粒度。于是有了按要求粒度调理的 ContinousBatching 手艺，如图 4 (b) 所示，第二个 batch 的第一条和第三条要求在第一个 batch 最后一条要求执行完之前就开始了执行，GPU 盘算资源的利用效率得到了提升。Batch 越大，要求长度的差异也越大，ContinousBatching 对系统吞吐的提升就越大。
ContinousBatching 的手艺被提出后，并没有引起推理加快框架的爆发增长。其中最大障碍是原有的 GPU 盘算中对 kv-cache 连续空间访问方式，导致 ContinousBatching 在 token 生成后调理要求有很大的显存操作开销。为了解决这个问题，PagedAttention 手艺提出了类似操作系统虚拟页的显存管理机制，将 kv-cache 的整个连续空间切分为多个连续块（Block），使得按要求粒度的调理变得高效，让 ContinousBatching 手艺被广泛应用。
为了完成 GPU 盘算资源的充分利用，大语言模型推理框架必须要完成 ContinousBatching 功能，一念 LLM 有了要求管理器。在前面问题分析的时候提到过，在差别的场景下，调理器的优化逻辑差别，甚至需要设计比 ContinousBatching 更小粒度的调理策略。我们抽象出了调理策略接口，用于完成差别的调理策略。纯 C++ 的完成让调理的异步逻辑更高效。
为了完成 PagedAttention 的功能，一念 LLM 设计了显存 / 内存统一管理模块，同时为了便于后期完成多模型，Multi-LoRA，状态缓存等功能，显存 / 内存统一管理模块收拢了一念 LLM 主要的内存和显存操作。
多硬件算子抽象（硬件和算子决定系统的 TFlops 上限）
在国外高端卡进口受限的局面下，形成的新问题。目前业界最新的推理框架（比如：最早提出 PagedAttention 的 vLLM 和 Nvidia 的 TensorRT-LLM）主要支持 Nvidia 的 GPU 等国外主流厂商的硬件，对国产硬件缺乏支持。国产硬件配套的推理框架，缺乏对业界最新的推理加快手艺的支持。目前相关的新手艺主要集中在调理或者更高的算法层面，与硬件关系不大，所以最合理的方式是使用统一的算子抽象来屏蔽下层硬件差异，从而完成一次优化，所有硬件上可用。
在 Nvidia GPU 生态下，一念 LLM 的算子库包含了来自 FasterTransformer，vLLM，TensorRT-LLM，pytorch 的开源项目算子，以及部分自研算子。
在华为生态，推荐的使用方式是用华为生态软件，使用图优化等方式来加快，但是图盘算存在优化控制粒度的问题。在 LLM 这种相对稳定的模型结构上，也不能发挥盘算图优化的优势。一念选择了相对底层的 AscendC 接口来完成自定义算子的方案。这套接口与 Nvidia Cuda 的接口类似，有 device，stream 等常用的对象接口。AscendC 接口当前在成熟度和功能方面与 Nvidia Cuda 还有不少差距。通过与华为共建和华为卡的广泛使用，我们相信 AscendC 这层接口完成的 LLM 算子功能会越来越好。
在算子使用上，通过功能和效果两个维度来选择算子。从功能方面，根据差别算子在差别硬件上的功能特点，择优选择。与功能相对，有的业务场景会希望推理的结果与训练的结果严格对齐，从而降低评估和效果调优成本。比如：要求生成的长文内容对齐。导致推理服务和训练框架在长文本生成内容上不对齐的主要原因是推理流程普遍使用的 float16 的表示精度有限，差别算子完成在数学上等价，但是实际精度误差差别，而且误差会随着推理长度增长而累积，于是出现了差别算子的推理结果在前几十个 token 相同，然后结果差异越来越大的情况。
当出现这种情况时，框架需要在功能与效果之间进行 tradeoff，有时就会为了对齐效果，将对效果影响最大的算子替换为功能更低的算子。
Prefix Caching，基于 prefix-token 的 kv-cache 缓存调理（优化 γ）
ContinousBatching 方案的调理仍然是要求粒度，并未对要求输出内容进行针对性优化。如图 1 (a,b) 所示，所有要求的前三个 token 都是 (1，2，3)，我们称要求的相同输出部分为 prefix-tokens。在当前的调理逻辑下，prefix-tokens 的盘算在每个要求的盘算中都会被执行，而在当前主流的 decoder-only 的模型结构下，prefix-tokens 的盘算结果以及对后续 token 盘算结果的影响是相同的，也就是说对 prefix-tokens 的盘算只需要进行一次。
当前要求粒度的调理导致了重复盘算，从而导致了 GPU 盘算资源的浪费。而且这个浪费的盘算比例会随着 prefix-tokens 的占比和 batch size 增大而增大。在类似角色扮演等应用场景中，prefix-tokens 的占比可能达到 50% 以上，而 batch size 会超过 30，那么将会有超过 50% 的盘算被浪费掉。
要完成高效的 prefix-token 的显存和盘算复用，面临以下问题：
常规的盘算流程是以矩阵方式盘算的，相关算子的优化也都是基于矩阵的规则大小盘算。如果要完成不规则的盘算，需要重新开发和优化算子，手艺通用性和开发成本都非常高，可能新完成的算子最终的功能与现有算子差异巨大，得不偿失。
调理上 batch 内的差别要求的相同输出长度短，导致盘算节省的收益小。
所以要完成收益的最大化，我们需要完成下面的优化：
调理要求到差别的 batch，完成 batch 内要求的相同 prefix-token 长度最大化。
在 batch 的输出处理逻辑中，需要基于开源算子，以最小的代价去掉相同输出的盘算。
调理上需要匹配显存与盘算的复用逻辑，让盘算的调理与显存的管理协调一致。
基于这样的需求，我们设计了以下的架构框架：
图 5 Prefix Caching 功能模块关系图。
总体上，为了提升 batch 内要求的相同 prefix-token 长度，增加了基于 prefix-token 分析的要求路由器模块。在调理器上改造为 prefix-token 与剩余部分的两阶段调理，同时调理策略上针对 prefix-token 和 kv-cache 缓存情况进行了优化。为了配合调理策略的执行，增加了 kv-cache 缓存管理器模块，后期可以完成 kv-cache 缓存在显存 / 内存 / 外部存储的三级调理管理能力。
在传统的分布式系统中，要求路由器主要考虑后端服务节点的要求响应和负载情况进行要求分发。为了提高后端服务节点 prefix-tokens 缓存的命中率，在路由模块上增加根据 prefix-tokens 路由的策略，构造了下图所示的路由表。其中 PT 代表 prefix-tokens，用 PTi 表示第 i 个 prefix-tokens。同时我们用 S 表示要求的服务节点 Server，用 Sj 表示第 j 个节点。用 SS 表示多个服务节点的集合 ServerSet，用 SSk 表示第 k 个 ServerSet。
图 6 Prefix token 路由表示例。
通过将差别 PT 映射到 SS 上，完成差别 PT 要求的服务扩缩容。同时通过控制单个服务节点所归属的 SS 数量来控制需要处理的 PT 种类，从而提高缓存命中率。
在特征批量处理的场景中，prefix-token 在输出中的占比 80%+。一念开启 KV-cache 缓存功能后的吞吐率提升 60%+，等效于单价下降 40%+。
CPU/GPU 混合推理（优化 M）
在 transformer 模型的执行流程中，业界常规的做法是将所有的算子放到 GPU 上执行，相应的模型参数也被放到了显存中。由于 LLM 模型参数量大，导致用于并行推理的显存空间被挤压。在业务常用的 7B，13B 模型中，模型的词表有变大的趋势，导致 token embedding 的参数量占比较大。以 llama-13B 为例，原始词表大小为 3.2 万，token embedding 参数占比为 1.2%。如果词表大小扩展到 30 万，embedding 参数占总参数量的 11.8%。但是我们发现 token embedding 的操作并不是盘算密集型的，而是一个典型的 sparse 查表操作。于是一念将 token embedding 参数放到内存中，用 CPU 执行 token embedding 操作，完成 CPU/GPU 混合推理，如下图所示。在词表大小为 30 万的 llama-13B 模型上，提升吞吐率 10%+。
图 7 Cpu/GPU 混合推理示意图。
临时显存优化（优化 α）
在深度学习网络执行流程中，会使用到很多临时变量来存储中间结果。在差别的框架中，会有差别的临时变量回收策略，一般是基于盘算图来优化的，在 LLM 模型的推理流程中，很多变量大小都是动态增长的，会导致盘算图的显存优化失效。一念没有采用盘算图的方式来进行推理，而是采用算子拼接的方式直接描述模型，从而完成临时变量的自管理。通过预先分配显存然后重复使用的方式，最小化临时变量的显存消耗。
未来计划
现在一念算是有了第一阶段的起点，解决了调理优化和多硬件支持的基础问题。
在国产硬件支持方面，目前只支持华为 NPU，后期还要支持腾讯自研的紫霄以及其他国产芯片。
在调理 / 显存 / 算子层面还需要根据业务场景和硬件的特点持续优化。同时在算法手艺层面，还不断有一些新的方向出现。下面简单说一下两个有趣的方向：Speculative Decoding 和稀疏化。
Speculaitve Decoding：在前面的分析流程中，一直是以加大 batch size 的方式来提升服务整体的 throughput，其中的主要瓶颈点是显存。可能存在下面的场景：a）显存有剩余，但是有不足以增加一个要求到 batch 中；b）要求很少，batch size 不能放大。SpeculativeDecoding 可以通过猜测多个可能的 token 输出，然后并行验证的方式，降低 latency，让当前要求尽快结束，从而释放出显存空间来响应新的要求。这个流程中猜测准确率是关键。于是有多种预测方式，比如：UCBerkeley 的 Big Little Decoder 利用小模型来快速猜测，蚂蚁金服的 Lookahead 框架基于 Trie-based retrieval 来进行猜测等。
稀疏化：大语言模型的大参数量和 kv-cache 都会带来大的盘算量和显卡存储消耗，让人不禁会问，这些存储和盘算是否都是必须的。围绕各种问题，产生了很多稀疏化的尝试。大致可以分为模型参数稀疏化和 kv-cache 稀疏化两个方向。模型参数稀疏化在深度学习模型推理加快领域一直有比较多的研究，本文不再赘述。kv-cache 作为 transformer 结构引入的新变量，也有很多有意思的研究，比如：基于 kv-cache 内容压缩的 GEAR，基于 token 重要性压缩的 KeyFormer。
如果要让这些手艺成功落地，对显存和调理管理都提出了更严苛的要求。
结语
大语言模型的能力越来越强，但大语言模型在应用场景中 ROI 正向仍然是一个非常挑战的问题。LLM 推理在 LLM 应用成本中占比大，任何小小的进步都能获得不错的成本收益，切实帮助业务完成更好 ROI。
在国外高端硬件供应不足的当下，统一框架以及国产硬件支持的可控亦是完成业务安全的必要路径。在相关软件生态不成熟的背景之下，会有很多困难。相信随着国产硬件的成长，会越来越好。
一念 LLM，筚路褴褛，以启山林。