本技术揭示了一种基于智能感知的大规模数据前缀KV排序技术与系统,该技术涉及利用重要信息收集模块,依据历史查询数据和前缀计算结果,对基数树中每个节点的token进行重要性评估和排序。
背景技术
现有前缀感知的大模型推理系统往往直接使用大模型进行推理可能生成次优的结果。例如,当被问及模型训练数据中未包含的最近事件时,模型可能会给出错误的答案。此外,由于诸如幻觉等问题,模型的回答可能包含不准确或误导性的信息。为了提高回答的质量,应用程序通常会在用户查询(Query)前添加包含含丰富信息的上下文,即前缀(Prefix),以形成完整的请求(Request),然后将请求输入到大语言模型(LLM)中。例如,检索增强生成技术会从外部知识库中搜索与用户查询相关的文档。高级的GPT插件,如Chameleon,在系统提示中包括工具定义,并使用少样本示例来引导LLM执行复杂的推理任务。
大模型推理的过程中,会为每个token在模型的每层中都计算生成一个K张量和一个V张量。然而,研究人员发现不同的请求的前缀中往往部分或完全相同(例如,类似的查询可能会使用RAG检索部分或全部相同的相关文档;同一个GPT插件可以被多次使用,从而在不同请求中生成相同的系统提示)。因此,现有前缀感知的大模型推理系统通常会存储并重用这些共享前缀的K和V张量(称为前缀KV)以避免前缀部分的KV重新计算。当一个包含重复前缀的新请求到达时,系统会复用其前缀KV加载到GPU内存中,从而避免了前缀KV的重新生成,只需要对查询部分进行计算即可。因此,重用前缀KV会降低大模型生成第一个token的时间(time to first token,TTFT),提升用户体验。下面用一个例子来说明这个过程:
举个例子,假如有一个用户使用一个专门用于中文翻译成英文的大模型插件,用户依次想翻译:苹果、香蕉这两个词。当用户输入“苹果”这个查询(Query)时,系统会在前面添加前缀(Prefix):“请把用户的中文翻译成英文。”,因此送入大模型的推理请求(Request)实际上是“请把用户的中文翻译成英文。苹果”。大模型会为整个请求中的每个token都计算生成K张量和V张量,并且把整个前缀对应的K和V存储到本地磁盘中(由于数据总量很大CPU内存存不下,需要往本地磁盘存储)。假设前缀的token划分结果是“请把/用户/的/中文/翻译/成/英文。”由于前缀被划分为7个token,所以在模型的每层中就有7个K和7个V张量。当用户第二个查询“香蕉”到来的时候,系统送入大模型的请求是“请把用户的中文翻译成英文。香蕉”,这时,为了完成每层的推理计算,大模型不需要重新计算前缀的KV,因为可以直接复用之前存下来的7个K和7个V,然后只计算“香蕉”的K和V即可,从而避免了冗余计算,缩短了TTFT。
此外,现有前缀感知的大模型推理系统还会考虑token重要性,减少部分token的加载,例如在复用前缀KV时,仅加载“请把”、“中文”、“翻译”、“英文”这四个token的K和V就能获得和加载全部前缀KV“请把用户的中文翻译成英文”相同的大模型推理结果,因为其他的token其实并不重要,对推理结果没有帮助。
术语的说明:
Token:输入给大模型的句子被分词后所形成的词汇/短语等,是后续进入大模型被编码的基本单位。
查询(Query):用户给LLM输入的问题(自然语言描述);
前缀(Prefix):加在用户输入的查询问题前面的自然语言句子;
键(Key,K):大模型推理时会为每个token在模型的每层计算生成一个K张量;
值(Value,V):大模型推理时会为每个token在模型的每层计算生成一个V张量;
前缀KV:前缀在推理中所生成的KV张量。
块(Chunk)存储与基数树(Radix Tree):
现有系统为了提高K和V传输到GPU的效率,采用数据块(chunk)为粒度进行存储,即把连续的几个token对应的K或者V存储在一起。例如上面的例子中的“请把”的K和“用户”的K存储成一个块对象。
此外,为了尽可能增加重用的前缀KV比例,现有系统采用基数树(Radix Tree)的数据结构来记录已经计算过KV的前缀,方便新请求进行匹配和搜索。在基数树中,每个节点代表前缀中的几个token,并且从根到任意叶子节点的路径表示一个完整的前缀。如果两个前缀有相同的token,那么共同的token将只在树中表示一次。继续以上面的例子来解释这个过程:
假设有另一个专用用于中文翻译成日语的大模型插件。用户输入“梨”的时候,该插件实际提交给大模型的请求是:“请把用户的中文翻译成日语。梨”。由于这个插件的前缀与之前插件的前缀并不完全相同,因此仅可以复用前面部分token对应的KV。
现有前缀感知的大模型推理系统存在的问题:
虽然仅加载部分重要的前缀KV可以降低数据加载量,然而现有系统在加载重要前缀KV时会产生读放大的问题。例如上面的例子中,当加载重要token“请把”的K或者V的时候,由于不重要token“用户”的K或V是存储在同一个数据块中的,因此导致也会把它从本地磁盘加载到CPU内存中,造成了读放大的问题。总的来说,现有系统会读取不必要的数据从而浪费读取带宽,导致读延迟增加。
实现思路