Tair：基于KV缓存的推理加速服务
内容介绍
一、大模型推理服务碰到的挑战
二、Nvidia TensorRT-LLM 推理加速库
三、基于KVCache优化的推理加速服务
阿里云智能集团数据库产品事业部高级技术专家王正恒。
团队本次部分内容主要介绍云数据库Tair如何基于NVDIA TensorRT-LLM提供推理加速的服务。
分成三部分，第一部分介绍大模型推理服务团队碰到的挑战。第二部分，团队会请专家介绍TensorRT-LLM推理加速库进行的相关工作。第三部分团队基于Tair的KV-Cache的优化的推理加速服务的内容。
一、大模型推理服务碰到的挑战
1.大模型推理：性能与功能
第一个问题为训练里自带推理，为何需要单独拉出推理服务。最重要的原因为团队希望能够基于给定的负载，用到一个最小的GPU进行团队的服务。实现该功能团队需要去补足现有模型的一些特有的功能。第一个功能为模型量化，通过模型量化节省团队的算力，KV-Cache，提高一个团队的内存使用率。第二个功能为KV-Cache，通过团队的KV-Cache管理支持更长的上下文，支撑好团队的推理体验。第三个功能为团队的服务化，服务化内容包括如下。首先为团队要去提供一个IPC的一个交互界面，像chat gpt一个对话的一个形式。第二个团队如何基于用户各种完全访问python不一致的问题进行调度，调度的本质为如何能够根据用户的负载去提升团队整体的GPU效率，两点相辅相成。
二、Nvidia TensorRT-LLM 推理加速库
由季光专家进行相关讲解。
TensorRT-LLM是英伟达推出专门针对大语言模型推理的一个加速库。它的特点回答了部署在语言模型需要解决的问题。
第一个特点为高性能，用TensorRT-LLM可以获得SOTA-LLM的推理性能。第二个特点为功能丰富，可以进一步的降低延迟保证模型精度，提高模型推理的效果。第三特点为开箱即用，使用方便简单。
具体介绍高性能。该表格来自于今年8月份ml perf4.1发布出来的数据。该产品是用自选的一个硬件的一个平台，利用任意的一个软件跑它规定的benchmark，对于LLM workload有三个模型要跑，即表格展现的三种。
英伟达所采用的软件TensorRT-LLM在指定的硬件上能够取得最好的成绩。对比ml perf4.1其他的一些厂商提供的数据都没有团队提供的数据好。在TensorRT-LLM的实际运用当中团队发现在给一些特定的workload上，TensorRT-LLM相比于其他的一些解决的方案能提高20%以上的性能。
为了实现卓越的性能，团队主要采取了两项策略：优化时间线（timeline）和优化内核（kernel）。
1.timeline优化
首先介绍timeline优化。该图它是来自于优化GPU程序的工具：insight systems。它可以采集的GPU的程序跑出来的timeline，蓝色的线段是一段GPU程序，团队称为GPU的kernel。黄色的部分代表的是在时间线上可以调用一些API，对timeline进行一些标记，比如知道这一部分正好是一个transformer block，要把它给标记出来，橙色的部分跨过了很多个 GPU的kernel把它给标记起来。在图里面团队发现有一些地方它是空白的，即没有GPU kernel的执行。空白的地方是GPU资源的一种浪费。
举例：TensorRT-LLM在优化前跑流水线并行，流传定型中timeline PP等于4。团队在分析timeline的时候发现空白（空炮或bubble）是由于团队不必要的一些数据的同步导致。于是团队进行了分析发现末尾绿色的unaviodable bubbles是不可避免，红色的团队通过重新组织团队的数据的传输模式，消除掉同步可以把它优化。
此图为团队已经把非必要的空炮消掉的状态。通过这种方式，团队在PP应用场景下能够获得两倍甚至三倍以上的加速。
接下来介绍功能丰富。TensorRT-LLM提供了丰富的功能，一方面市面上常见的这些模型团队都已经支持上。另一方面涉及到一些模型的修改以及模型后面像采样的一些后处理，像投机采样等投机采样方法团队同样支持。介绍TensorRT-LLM支持量化的功能。量化指把高比特的数据用低比特来表示。在训练大语言模型的时候有可能用是BF16，在推理的时候要把BF16把它给变成fp8，BF16本身是16个比特来代表1个浮点数，把它给变成IP8之后只有8个比特，显存的占用量可以减少一半。此外 GPU上面的IP8的计算的吞吐是BF16或者是IP16这些16比特的数值的两倍，所以它的计算时间理论可减少一半。综合可得出大语言模型推理的量化是重要的功能。TensorRT-LLM可实现weight activation和cache的量化。Weight即训练好模型之后直接加载了模型到显存里边的训练好的权重数据。activation即中间计算结果。
重点介绍KV-Cache。一方面是在一个生成一个句子时反反复复要用的，也是生成不同的句子的时候如果有共同的前缀，可以反复使用的。如果能够完成将 W8 a8 C8三者都给量化成硬盘或IP8了，该情况下能够获得较好优化性能。但是团队大部分的人没有接受过量化的算法，Nvidia提供了一个工具： model optimizer（model opt）可以跟TensorRT-LLM进行无缝的衔接，用model opt对指定的模型进行完量化之后，它可直接导出TensorRT-LLM的checkpoint并加载到TensorRT-LLM。另外如果一个量化的专家想用一个第三方的量化工具也是完全可以成功。
TensorRT-LLM支持多种多样的部署和调度，PP是一种并行方式，另外更常见的一种用法 tensor parallel，它实现的是在batching较小情况也能够把计算给分担到不同的GPU上面去。团队正在开发的pd分离，把 the prefill跟decode这两个阶段可以部署在不同的GPU上面，可以把它认为是 pipeline并行的一种特例。
TensorRT-LLM支持服务化，除本地可以运行大语言模型之外，可以把它部署为可以被远程调用的服务直接跟Nvidia的triton server集成在一起。团队也提供了扩展的接口，可以自由的接入第三方的服务。
介绍inflammation的原理。batching即一次处理多个查询。举例将“昨天天气不错风和日丽”这句话翻译成英语，文言文，日语和朝鲜语，此时一共提供了4个prompt，大语言模型会根据的提示来吐出后续的一些要翻译出来的话。把这4个prompt打成1个batch送到大语言模型里面去，可以更高吞吐的方式把它给计算完。
在大语言模型进行计算时可把它理解为两个步骤。第一个步骤是将“昨天天气不错”翻译成英语，处理这句话的程序叫prefill，所花的时间较长，后面它开始吐英文单词，假设4个单词将此句翻译完成，它是token，但token不等价于单词，但可以粗略的把他们等同，比如英文四个单词可翻译完毕，但翻译成文言文，一个单词即可翻译完毕。程序吐出一个单词的步骤即为一次解码，同一个batching里不同的这些句子它可能不等长,于是它结束的位置不一样。团队如果用一种比较难的实践即static batch上面图所表达的意思。等最长的句子都已经生成完毕重新开始下一轮的prefill，即为Static Batching。Inflight Batching基本思想为把中间的空炮填上，一个句子结束后把下一个句子的prefill给填到末尾上。
但TensorRT-LLM并不满足于实现简单的功能，团队希望把prefill跟decode，既然他们在同一个batch里边了，他们计算能够真正的达成同一个batch进行计算的地方。比如团队把prefill跟decode的矩阵成合在一起用同一个batch进行计算。对于Iteration部分 prefill采用FM ha，decode采用是m ha，需要团队分开处理。团队尽量能够打batch的地方都把它给打通使得Inflight Batching在图所表示的思想下的不同的实现当中取得最好的性能。
2.开箱即用
TensorRT-LLM是一个能开箱即用的加速库。团队提供了一个简化的安装pip install即可将其安装。进行二次开发时可以进行一个自行的构建。同时它提供简化的API，可以使用团队的high level API满足对于lm输入和输出文本的简易接口需求。
对运行的过程要进行一个非常精细的控制时会提供底层API，而且TensorRT-LLM在运行的过程当中还提供了很多的性能调整参数。团队简化了参数的配置的过程实现一定的自动化，使得即便是使用默认的参数，也能够取得一个相当不错的性能。假如对性能还有更高的要求，也可以自由的设置里面的这些参数。
3.应用与落地
TensorRT-LLM是一个高性能高可靠的推理加速库，已经在很多头部的互联网公司得到落地。Tair作为阿里云重点产品团队希望与其强强联合实现二者合作共赢。
三、基于KVCache优化的推理加速服务
从TensorRT-LLM量化服务化调度batch 介绍TensorRT-LLM的推理加速服务。下面讲解基于团队云数据库的KVCache优化的推理加速服务。
1.推理两阶段和常见推理请求
介绍prefill跟decoding概念和推理两阶段以及常见的推理请求。推理两阶段在第一阶段叫 prefill。团队经常用的输入到大语言模型里并把问题输进去，处理的过程叫prefill。后面指标叫time to first token，把问题扔到大语言模型里面到它第一个返回回来这段的时间即为它的响应时间。第二阶段叫decoding，它的机制为inter-token latency。token一直在循环输出一个一个单词输入，所看到的总是一个单词一个单词在蹦，下一个阶段生成token的时间。可认识到团队第一个token的数据影响到团队推理服务的整体的服务时间。通过团队常见应用可认识到后面的数据生成速度只要能赶上人肉眼读文字的速度就可以。
具体了解两个阶段里进行的工作。第一个阶段prefill即把问题扔到大语言模型里面进去，团队对用户的prompt进行一个解析并去填充KVCache(在后续计算过程中避免后面反复的去计算token，采用KVCache来缓存团队之前的计算的结果，能够通过内存的读取来减少团队计算的频率）。
在第一阶段，prefill的阶段是GPU计算的过程。第二阶段decoding的过程是接二连三的token去进行的，每次计算token的时候团队需要读取整体的参数，从内存里面读取出来，通常也内存访问密集型。因此在推理两个阶段对GPU的需求不同。
推理常见的请求第一个为system prompt 。system prompt可理解为为团队经常看到的隐藏在一边的东西。假设是英文翻译模型，收到的是英文，要求输入的是中文，这段一般都是隐藏在边，都是在开头的反复的都会使用 system prompt。
第二阶段经常碰到的团队流行的RAG系统去一些团队的知识库的文档，把它取出来，或者是团队一些多人对话里边的一些历史记录。
第三阶段是真正扔到单元模型里面去的团队的用户请求或者问题，最后decoding段产生即为结果。了解根据推理的两阶段和团队常见的推理请求如何作为结合来优化团队整体的推理服务。
2.云数据库Tair概览
首先介绍团队云数据库Tair。它是一个能够从内存，持久内存包括磁盘型的能够覆盖到团队整体的各种各样的性能，数据量，成本的一个Kv型云数据库。它具有非常多的存储介质和丰富的数据接口，完全兼容radius，同时支持从256兆到数TB范围内的Qps集成。
团队在云数据库Tair上面积累了非常强大的kv缓存经验以及多级高吞吐异构池化的一个经验。
3.Tair：推理KVCache优化
了解数据库如何优化KVCache。通过该图首先了解 system prompt。第一阶段为固定，目前流行趋势为该系统越来越长，并且是非常固定的，这段固定如果团队提前把它需要计算的逻辑缓存下，就可以避免后续反复计算和GPU计算。
第二阶段为文档类或者是团队的聊天对话的历史， RAG一般是固定的。Chunk是固定的，当团队查阅一些文档，可通过频繁使用Chunk进行一些预计算，甚至无需等到每一次真正需要的时候再去进行计算。第三部分 prefill 和decoding两个阶段对GPU的需求是不一样的，有一个为计算密集，一个是存储密集型。
团队把两个放在一起时对调度带来非常大的挑战，团队可以选择直接去把prefill和decoding两个放在不同的GPU类型进行计算，通过团队的池化来进行KVCache的传输。
4.KVCache用法
介绍KVCache用法。举例：多人对话会有system prompt，团队的使用者问题和答案是团队第一轮。第一轮问完得不到答案，团队可能会追加团队第二个问题，实际上团队看到的时候在第一轮结束的时候，团队decoding的些user的question或者answer，团队也是去算了它的KVCache，当然团队如果要基于首轮对话的再去推导，第二轮对话的时候，团队要重新去计算一遍。在decoding的过程中团队已经计算出来了，团队相当于在直接去复用它不用再去进行计算.
这里提出一个挑战，随着团队的多人对话它token时间越来越长，它的KVCache开始会越来越大。此时团队无论是存储还是团队的传输，都带来非常大的挑战，团队把这部分的KVCache如果是放到远端，团队的memory disk上面去，大家可以不用去考虑它的一些真堵，它的容量，存在本地盘上的本地盘都真受到限制，本地容量受到限制或者内存，扔到远端去，团队提供了一个超大容量的超大吞吐的自动弹性，或者说容的一个远端磁化。
5. RAG KVCache
介绍RAG KVCache。左图为目前RAG目前涵盖的知识库，包括书，代码，视频，音频。通常情况会把他们先切成不同的小的Chunk作为团队知识库一小部分，这些Chunk会进入到团队经常用的一些向量数据库里面去进行embedding。正常的一个用户搜索过来之后，首先会去找到跟用户请求相关的RAG与Chunk，并把它读取出拼成一个新的一个用户的请求扔到团队的推理引擎里。
RAG与Chunk是团队预先切好的。在取Vector Database到生成prompt团队是可以并行走。团队存储的RAG KVCache，一个是团队可以直接去复用，第二如果团队Chunk是可以反复用的，可以直接进行离线计算全部计算。第三个可跟RAG的Vector Database获取去进行连线加速整个it的访问效率。
下面介绍基于内存池的prefill和decoding分离。此为目前整个业界大家都在进行的事情。团队在prefill和decoding的处理的模式不同的基础划分为两个异构的方式，二者都放到特定的GPU。
了解整个流程：在prefill时团队可以从内存池里面去拉到团队的context Caching 或RAGCache，此时去生成团队的user KVCache完全生成了团队prefill阶段需要生成的一个Cache。此时数据也已经存在了团队的远端内存池上面去，团队在decoding的阶段有预取或者异步的逻辑，可以快速的把KVCache从团队的prefill的节点数据传到decoding节点。KVCache在decoding的阶段可以避免重复的计算的逻辑达到整体的调度加速效果。
这是基于千万27b的模型上长下文上看到的实验数据。因为KVCache池化整个成本可以降低到20%。团队和NV合作使得prefill和decoding分离的基础上，加上池化的能力和额外的一些调度优化整体的吞吐提升了30%。团队使用NVDIA 的GPU一些和TensorRT-LLM的推理引擎，算子的优化，软硬的结合本身的time to first token能够降低30%。谢谢大家。
