业界前沿 | 一念 LLM 大语言模型推理加速

文摘 2024-09-13 12:00 北京

导读本文介绍了一念 LLM 大语言模型推理加速。

本次分享的主要内容包括：

1. 大语言模型概要介绍

2. 一念 LLM(KsanaLLM)基本框架

3. 一念 LLM 框架调度

4. 一念 LLM 在 GR 模型的应用

5. 未来规划

6. Q&A

大语言模型概要介绍

首先来看一下大语言模型的结构。在 Transformer 结构下的大语言模型推理的过程中，一个 token 或者一个字的生成的过程大致上可以分成两步：

Step 1: 根据已有信息，也就是 input 的已知信息，估计下一个 token 的概率分布；

Step 2: 根据采样的策略，从概率分布里面挑出最有可能的下一个 token。

这个过程有可能是以概率最大的，偏 greedy 的方式来做，要考虑到后期生成的 token 的概率，从总体上去做采样。这是跟传统深度学习推理不太一样的地方。这两步是一个循环的过程，当生成了下一个 token 之后，这个 token 会进到下一步 Step 1 里面去再生成再下一个 token，这就是推理的基本逻辑。

这里引出一个经常提到的概念，KVCache。刚才提到的在 step 1 的时候，是根据已有信息，这里的已有信息包含两个意思，一个是原始的输入，另一个是之后生成的 token。如果我们把一个生成的过程拆开，前面部分是最原始的输入，生成第一个 token A。第二步从概率分布的逻辑上来说其实是要把前面的部分再加上 A 去估计下一个 token，依此循环。这会导致一个问题，计算量是在不停增长的，而且会与前面已生成的部分和 input 部分成正比，可以想象到这样的逻辑一定会越跑越慢。

在 Transformer 结构里面，存在一个计算的特性，当前 token 的结果只与前面的 token 有关，可以把前面 token 的计算结果进行缓存，形成两个阶段：

Prefill 阶段：输入后走一遍全部的过程，这是全量的走模型的过程，走完之后，会产生一些中间结果。这些中间结果被缓存起来，放入到图中标红的下一步的过程中，KVCache 在进入 attention 之前，跟现有的新生成的 token 的结果做一个 concat，然后再做计算。之后又是一个 token 生成的过程。
Decoding 阶段：通过 KVCache 的优化，decoding 阶段的计算量和前面的 token 数就变得无关了。这里其实是一个近似的无关。因为在其他主要的部分都是无关的，但是在 attention 计算的地方，是被恢复成了一个全长的 token，然后进行 attention。

这就是 KVcache 的存在的意义，让 decoding 阶段的计算尽量复用以前已经计算过的结果，这样对前部分的数量就没有依赖，从而提高整体的推理速度。

这里存在的问题是，首先，开始输入很多 input，之后每次只会输入一个，比如这里输进去的是一本书，可能是成百上千万的输入，之后计算的数全部都是1。所以在推理的时候就会出现这样的一种现象，比如一张 A100 的卡，它能够并行推理的 token 数跟 GPU 的 TFLOPS 的关系是，当 token 数增长的时候，GPU 会被更充分地利用起来，到达一个阈值之后，可能跟算值本身的时限有关，基本上到达了该 GPU 卡能够提供的最大 TFLOPS。

我们发现在 prefill 的运转区间，可以把 GPU 压满。但是因为后续 decoding 的并行每一次只预测了一个 token 数，也就是并行度非常小，在生成过程中，GPU 都处于一种不饱和的工作状态。对于不饱和，最简单的处理方式就是做 batch，把 batch 加大。这里面存在一个问题，如何才能把 batch size 加大？由于 KVCache 的存在，而且由于在大模型的情况下，KV cache 占用显存非常厉害，batch size 会受到显存的限制。

我们需要看一看，显存到底是被怎么消耗掉的。一个正常的执行过程，包括了 prefill 和 decoding 两个阶段。一个模型加载之后会占用一部分显存，之后第一把执行 input token 的时候，显存会有个快速的消耗，之后随着 token 的逐步生成，显存的消耗在慢慢变大，这和常规的深度学习的推理非常不一样。因为它有 prefill 过程和后面的生成的过程。这个消耗过程其一是跟长度有关系，其二是跟我们加的输入的长度也有关系，当我们 batch size 扩大的时候，这里的显存就会成倍地上涨。

从显存角度来看待的话，可以列个很简单的公式，首先是模型占用了多少参数，然后在模型的推理过程中有很多的中间变量其实也会占用一部分参数，另外有多少 token 的 KVCache 缓存也会占一部分，最终是要小于显存的大小。

正常而言，比如一个 llama-13B 的模型，只要这种模型一上去，显存基本就固定了。如果我们想优化的是把 batch size 做大，就可以通过优化 β，将显存使用变小。这里主要涉及到 KVCache 的量化技术。

占用了 KV cache 的这些 token 的数量跟什么相关呢？可以这样去列一个公式，实际上是 batch 内平均的 token 的数量。比如现在已经输入了多少 token，以及生成了多少 token 的一个总数乘以 batch size。这里还存在一个 γ 系数，其实就是 batch 之中不同的请求之间的 token 可复用的 KVCache。

一念 LLM 框架的名字是取自一念三千，现在的大模型转瞬之间就会生成各种不同的结果，而一念在佛教里面就是一刹那的意思。一念三千也代表我们的目标是希望大模型在一瞬之间生成世间万象。

因此这里会对 latency 和 throughput 两个方向重点优化。

一念 LLM 基本框架

一念 LLM 的基本框架如上图所示。从上往下分别为：

最上层就是咱们经常实现的如 Llama、Baichuan、QWen 等模型，与常规的深度学习模型的推理方案不一样，我们采用的是手写模型。我们抛弃了计算图，计算图以前都是用角度算值，拼成模型，这最大的好处是有很多的灵活性，便于算法人员去实现。而这样会带来另外一个问题，即深度优化困难，以至于最后大家会走向一个方向，去做融合算子，然后把融合算子放到图里面去，如果符合图的 pattern，用融合算子去替换。

对于 Transformer 结构，因为结构非常简单，并且现在结构基本已开始收敛，大家不再卷模型结构，更多的是卷模型效果。也就是说，这一块我们需要实现的模型的类型是比较少的，这就让手写模型和手写算子变得有利可图。

英伟达、Facebook 等各个大厂为大语言模型场景写的算子都非常的大。甚至像 attention 这种，按现在基本就是一个标准的，用 Flash Attention 做一个大算子来做。

我们为什么要彻底丢掉计算图呢？因为我们还希望去优化整个模型推理过程中的显存，只要已经规定好了模型，就可以盯着这个模型，仔仔细细地把里面显存的使用全部去调好，使得整个显存使用最小化。省出来的显存都可以拿去做其他事，去做 KVCache 相关的事。

另外一个是高效调度以提高吞吐，后文中还会详细介绍。

下面就是算子择优，我们的底层算子，很多时候会是开源的封装。因为其实当模型结构相对固定之后，硬件厂商会专门针对这些大的算子进行优化，提高性能，最终目的是卖掉他们的卡。现在各大厂商甚至为了不同的 tensor 的大小，去写不同的算子，从而获得收益的。

多硬件的支持方面，现在业界的主流框架，基本上都是英伟达，当然也有支持英特尔的 CPU。然而像支持 GPU 和华为昇腾等国产卡，还并不完善。从国内厂商的角度来看都会面临一个问题就是高性能的 GPU 卡进不来。从业务安全的角度，我们也必须要去支持不同的硬件。

但是为什么不是按不同厂商使用不同框架，比如英伟达用一个框架，华为用华为的框架，两边的都能用到一个好的收益？实际上，当你在做这一层优化的时候，你会发现面临不同的框架，需要重复做。另外各个框架本身对后期业务逻辑的适配也会不同。

我们通过上面统一框架，下面支持多种硬件，这样就相当于做到调度这层一次优化在所有硬件上都可以用，这样更有利于整个平台的运营。

一念 LLM 框架调度

第一个问题是 ContinuousBatching 和 PageAttention，从最原始的公式来看它其实就是在有效地优化 batch size。

正常情况下，我们会把不同的请求打成一个 batch 去做 GPU 的推理，这个过程往往是一个 batch 一个 batch 地进去，要等到这个 batch 里面的请求全部处理完才退出。从 GPU 的任务调度上来说这是最简单的使用方式，但这种使用方式最大的问题在于它的有效 batch 会越来越低，因为每个请求的输入长度不一样，输出长度不一样。比如有可能第一个请求是一个摘要的任务，会丢入一个 1,000 字的文章让大模型用一句话总结出来；第二个请求可能是一个扩写的任务，给了一小段话，让模型把它扩写成一篇长文。也就说输入输出的不匹配，会导致有些请求很快就结束了，有些请求还要跑很久。这样的话，要等到所有的请求都完成，这个 batch 才会退出，这会导致有效的 batch 到后面越来越小。另外一个问题是，后面很多请求结束了，GPU 算力就会闲置。

因此很容易想到的是能否当一个请求处理完成后，就把另外新的请求输入进去，这就是 ContinuousBatching 的基本思想。这个想法其实早在 22 年就已经被微软提出来，但是这么好的一个想法一直都没爆发，是因为存在一个问题，就是它的 KVCache 的显存的操作成本比较高，比较麻烦。去年由伯克利提出的 PageAttention解决了这一问题，于是整个 ContinuousBatching+PageAttention 的机制就迎来了一个爆发期，现在成为了大语言模型推理框架的一个标配功能。

简单讲因为整块操作成本高，所以将它切小，它采用了常规操作系统里面的内存分页的机制来管理显存，把 KVCache 切成不同的块，用指针去引用的方式来进行计算。这样就显著降低了显存操作的粒度。这样 ContinuousBatching 的整体操作效率就得到了提升。

这里面还遗留了一个问题，就是 input 在请求与请求之间的共享问题，这在以前的深度学习的推理里面很少关注到，但是在推荐里面从很早的时候就已经有类似的工作。比如一个用户多个 item 的 rank 操作，因为用户都是一样的，所以这一部分的推理只需要做一次。然后不同的 item 的推理多推多次，再融合到一起，再去做后面的推理工作。

请求是有共性的，这些共性的请求其实可以只算一次，然后把计算结果缓存，就可以把 KVCache 存起来，等到下一个请求，只需要处理后面的。Batch 之间也可以继续复用，这样整个后期请求的推理响应就可以一下提上去。

KVCache 机制看起来非常美好，但也是有成本的，因为 KVCache 会占用显存，而且一旦缓存大，就会面临 cache 换入换出和命中问题。比如放两个前序在显存里面，就得占两份显存。最终在具体的执行节点上面，命中率就决定了这个机制最终的收益。所以 prefix catching 更多就相当于上面公式里面的 γ。

我们在服务节点的前置加了一个叫 prefix-token 的路由器。在这里，路由需要平衡两件事情，命中率以及传统路由的问题，包括负载平衡、容灾等问题。例如，如果前序都是一样的，就直接打到一个节点上去，它一定会被命中。但是实际上还是会面临负载平衡和容灾的事情怎么解决的问题。所以我们构建了一个路由表，例如经常看到的一些角色扮演的服务，比如正在跟宋江聊，首先需要告诉模型你现在正在扮演的是宋江，宋江是一个什么样的人、之前的生平、有什么样的能力、他的性格等一些重要的事件，就像用户简历一样。然后再说你跟宋江之前聊过什么，下面请生成你准备作为宋江回复给用户的信息。这个过程里面有大量的信息其实是跟用户 profile 一样。还有另外一种，比如作为爱因斯坦或者牛顿这样的角色，对不同的角色，在进到模型之前，我们就会把前面的这一段做一个 prefix-token。对相同的一段给他指定一个具体的路由表，在有限的机器里面去选中一个机器集合。

另外刚才也提到一个问题，我们需要解决不能有太多份的 cache，cache 份数多了，显存里就全是 cache，没有显存去计算当前要执行的请求了。我们通过另一个维度的管理，对于单一的节点，相当于是一个 server-set 其实是有限的。从这张图可以感觉到，最后每个节点只命中了两个 set，从而达到对于单个节点而言 cache 份数是相对可控的，同时又能够在 set 的维度完成负载平衡和容灾。

最后再讲一下，CPU 跟 GPU 的混合推理，其实就是优化 M。

前面提到把模型从 FP16 变成 INT8 或者 INT4，这种是量化操作，效果是有损的。从框架层面需要支持，待业务评估是否使用。这里讲的是无损的一些方式，计算密度的问题。我们一般讲的大语言模型是计算强密集的，通常指的是 Transformer 结构部分。实际上它还有一个 token embedding 部分，这部分就是查表，类似推荐里面的 sparse 操作。往往这部分又是放在 GPU 上做的，这个表一旦大了就会占显存。而且可以看到一般在业务应用的时候，通常都会扩 token，扩词表。因为原始的模型里面的那些词并不能完全覆盖到我们的业务里。业务里面有可能会有一些特殊场景的词。开源模型如标准的 Llama-13B 在 v2 的时候是 3.2 万的词表。词表会随业务扩展，Llama-13B 原生的词表可能只占整个参数的百分之一。但是当把它扩到三十万的词表时，它对整个模型参数量的占比就会到 11.8% 的显存。而它又是个 sparse 操作，很直接的一个想法，就是把它丢到 CPU 上去改。我们自己在测试 30 万词表的 Llama-13B 时，能够有 10% 的性能提升。但这 10% 也不一定是能拿到的收益，跟词表大小有直接相关。因为这会涉及到 CPU 跟 GPU 的联合推理的问题，意味着 CPU 执行完的结果要拷贝到 GPU 上，多了一个成本。如果节省的显存不足以去 cover 成本，收益就可能是负的。所以需要根据实际业务所用的模型去调整是否选用这个机制。

一念 LLM 在 GR 模型的应用

目前大家都在考虑大语言模型的 Transformer，生成式的结构，能否用在推荐系统中。推荐系统的推理因为量非常大，耗时要求非常高，推理成本比常规的 AI 推理要高很多。

当我们把常规的模型结构，变成一个大语言模型的结构，采用一个长序列输入，计算量可能是成千上万倍，成本也可能随之增加。

生成式推荐其实是基于历史序列去预测候选 item 的 action。这里单个用户会有大量的 item 要预测，因为 rank 往往是千级别的。对于同一用户的历史序列，其实是一样的，这是很好的符合 prefix catching 的场景。可以对这个用户所有的 item 预测请求，做一次 prefix-cache，之后只做 item 部分的推理。这可以实现原来整个计算量是 prefix-token 加上一个待预测的 item，需要推理的 token 量乘以 item 的数量。相当于每一个 token 的计算量是乘的关系，所以会导致我们的计算量成千上万倍的增长，因为 token 是成千上万的。通过这个功能，就可以实现 item 变成了 item 的数量加上 token 数量，也就是把乘变成了加。这样最后的计算量就是跟 item 数量线性相关，就跟现在正常推荐的推理相似。

这里只是解决了一个计算量的问题，还有另一个问题是 latency 的问题。如果是以万级的 token 输入，想要最后控制在 10 毫秒以下也是非常困难的，哪怕业务能够接受更长时间的 latency，也不是将阈值从 10 放松到 50 毫秒这样的一个状态，而是要放松到秒级。这里对于 item 的预测是可以分开的。在不知道 item 的情况下，就已经知道了用户的序列，可以提前计算。比如在用户请求刚刚过来时，就可以把序列发给用户了，然后等到把 item 做了召回初排之后，再去执行这一部分。这部分就只有一个 item 的 token 耗时，这也是我们传统意义上讲的 rank 的请求耗时。这个耗时就可以做到只跟 token 的最后 item 有关而跟前面的 prefix-token 的数量无关。这样的话就可以把整个系统跟现有推荐的推理系统基本对齐。

未来规划

未来的规划就是围绕整个架构的几层：

1. 对模型的支持

常用的大语言模型；业务的 GR 的推荐模型的支持。

2. 调度层面的优化

计算/显存的流水线，包括现在热门的投机解码等业界先进技术，会持续跟进。

3. 硬件的支持

不只是在硬件上跑起来，更重要的是在硬件上定制化算子的开发。例如华为等国内其他公司也在做加速类的硬件。要单独提一下 CPU，因为现在最新的英特尔芯片，已经在 CPU 盒里面加进了矩阵计算的硬件单元。这种情况下，CPU 其实是可以去承担一定程度的高密度的矩阵计算的，性能会好很多。

Q&A

Q1：之前在做 CTR 推理的时候做了很多类似显存分配、动态 batch、多流并行、kernel launch 等工作，未来有哪些 CTR 推理的能力和经验是可以借鉴到 LLM 推理上的？CTR 和 LLM 之间的区别和优化侧重点有哪些共性和不同点？

A1：这里没有提到传统的深度学习推理里面的那些优化，准确来说这些优化全部都有效。只是在大语言模型推理场景下面，因为长序列，序列输入是以 token 方式去做并行这样一个特殊性，引入了一些特殊的优化方法，包括动态 batch，其实很大程度上也会是跟刚才提到的 continuous batching 结构类似，实际上 continuous batching 就是更细化的动态 batch 操作。

多流的并形其实可以用在单个请求单个 batch 前向推理的过程优化里面。这部分相当于已经有一定的 batch 了，要去生成一个 token，就要经过图的过程。所有以前用的优化都可以继续使用。

只是可能有一部分，比如手写算子，图优化，因为没有图了，所以也就不需要图优化了。

简而言之，目前的 GR 模型其实并不是一个连续生成的模型，其实对 KVCache 连续生成这一点上的依赖没那么重。就像在现有体系下是计算图去实现。走一次前项，再走一次前项，只不过多了一个 KVCache 的输入。然后图存在一些变化的部分，用传统计算图的一些优化方式去实现也都是没问题的。

这个理论可能会涉及到极致优化的问题，因为手写算子极致优化这件事情本身就是相当于损失易用性来实现的。这张图上的应用，是跟咱们算法同学配合上线的。在大语言模型的场景下面，如果有一个新模型，就又得去支持它。从工程层面上来说，需要重新把从训练到推理的这一套，全部的给重新弄一遍，这其实就是取决于模型结构本身的稳定性的问题。因为在大语言模型场景，结构已经非常稳定，跟推荐场景比起来简单太多了。推荐场景会有很多的 sparse，小算子来回拼来拼去的事情。在大语言模型这种情况下全都没有了，只剩下大算子。

像 kernel launch 这种类型的优化，在现有的大语言模型场景下？全部都可以用。现在的主流的大语言模型推理框架基本上都不是用图。除了 tensorRT LLM 算是用图的，但其实它里面也是很多大算子，只是用图大致串了一下几个大算子。

Q2：一念现在用的是连续的 batch，在大 batch 中间，模型的 forward 部分和解码采样是会有一个串形的执行流水线，这样是不是可能会出现 GPU 的空隙？

A2：实际上在解码采样这件事情上，目前主要的优化手段还是把解码采样的时间段尽量放短，因为现在主要的方式仍然是 token 依次生成。当前大语言模型，因为显存的问题，显存占用太大了，比较难像以前一样。这个 batch 跟那个 batch 交换着做一个流水线。

现在尽可能将一个 batch 打大，打大之后显存就已经没了，没法再切到另一个 batch 的显存来做下一步的推理，所以现在这一块基本上是串行。现在主要的优化是想办法把解码采样这一环节压缩，把它压的更短，不要就那个地方的解码采样，现在很多都在最后全部采用 GPU 算子来做，而不是在 python 层面去写，写了再用。其实都是为了极致的压缩，因为它这整个过程现在就是个串行的，比以前的优化的空间小很多，不能够让流水线让下一个 batch 的计算能够进来，还需等前一个解码采样算完，希望尽可能的用 GPU 去算解码采样。

现在主流的方案基本都是这样的，因为中间提到的 decoding 是有状态的，其中的 KVCache 不可能在这么短的时间内（一个采样的时间）把它换出去，再换另外一个 batch 进来。这样的操作会得不偿失。

Q3：一念有没有做过跟 TensorRT-llm 的对比，是否有跟 A800 或者 4090 做推理。它的性价比如何？

A3：A800 或者 A100 的有测，其他普通的非线上卡没有测。跟 tensorrt-llm 的对比的话，在具体场景，有 10%-20% 的收益。这个收益主要源于我们在下面对开源算子进行了封装，包括 FastTransformer 的算子，vllm 的算子和 TensortRT 的算子。我们集成了开源的大语言模型推理框架，以及为业务定制的算子。这里面存在一个问题，因为做了大量的融合算子都用 FP16 在推理，在这种情况下 FP16 有精度损失。大家都知道在业务实际应用的时候，有些业务可能会有非常强的效果一致性的需求，有时候就需要把算子退化到比如 pytorch 的算子来跟训练做对齐。

对于 TensorLRT-LLM 来说，它其实完全用 FP16 的性能并不是那么好，只有开启 INT8 和 FP16 的混合量化的机制之后性能才上来。

Q4：如果是对于 GR 生成推理的话是更适合做一个 multi stream 的并发，还是也像 llm 那样做一个异步大 batch 之后，再去做一个串行执行？

A4：在大语言模型情况下为什么没有做这件事情，是因为显存不够。但是在 GR 的场景，模型大小不是很大，像现在做的大语言模型，13B 的规模那就是 26G 的显存没了。但如果模型可能只有 G 这种级别，剩下的显存就是足够大的，就有空间去做多 batch。对于现在大语言模型的场景，很多问题都暴露在了显存大小这件事情上。

以上就是本次分享的内容，谢谢大家。

分享嘉宾

INTRODUCTION

袁镱博士

腾讯

专家工程师

袁镱博士，腾讯公司专家工程师，负责无量系统和一念 LLM 等机器学习训练和推理框架研发。

数据空间技术与系统

数据空间技术与系统全国重点实验室面向国家数据空间建设的中长期战略需求和重大任务，开展数联网基础软件与数据空间操作系统的技术体系、标准规范、核心系统、试验环境、应用示范与开源生态等重点任务研究。