本文作者袁镱博士是腾讯公司专家工程师，负责无量系统和一念 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 等，但是也存在两方面的问题：
为了便于算法人员使用和尝试新技术，vLLM 采用了 Python 作为主要调度管理功能的实现语言，导致显存管理和调度效率较低。
主要支持 Nvidia 的 GPU 等国外主流厂商的硬件，对国产硬件没有支持。国产硬件配套的推理框架，缺乏对业界最新的推理加速技术的支持。
一念 LLM 通过对异构硬件的底层抽象，构建统一的高性能调度管理，实现了：
应用业界最新的推理加速技术，推理单价相比业界开源框架低 20%+。结合业务场景进行针对性优化，单价可以降低 60%+。
将最新的技术同时应用到国外主流的 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 主要由以下模块组成：
图 3 一念 LLM 模块图
内存 / 显存统一管理模块负责分配和管理内存和显存，其中最重要的功能是分配和管理 PagedAttention 机制所需的 Block 和推理过程中所需的临时显存。在后期，与调度配合，实现更细化的调度能力。
请求调度器模块负责调度多个请求执行，协调内存 / 显存统一管理，实现推理过程的流水线化。
kv-cache 缓存调度负责 kv-cache 在请求之间的缓存复用。
加速算子包括不同硬件的模型并行，计算量化，算子融合等功能相关的算子。包括了自研的高性能算子，经过框架适配的开源框架优秀算子。相关算子随着业界发展迭代。
统一抽象接口负责将不同硬件的算子以相同的操作方式对接到执行流水线。采用是类似 Nvidia Cuda API 的接口。
LLM 模型是基于统一抽象接口和加速算子库来支持的 LLM 模型。
模型请求调度模块用于调度不同的请求到后端推理节点。与传统机器学习推理不同，LLM 模型具有推理时间长，状态数据大的特点，请求调度模块更加请求的特点和后端服务节点的状态进行调度，优化系统性能。
在大语言模型推理过程中，一般会使用到 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 主要的内存和显存操作。
在国外高端卡进口受限的局面下，形成的新问题。目前业界最新的推理框架（比如：最早提出 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，有时就会为了对齐效果，将对效果影响最大的算子替换为性能更低的算子。
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%+。
在 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，筚路褴褛，以启山林。
