一、引言

在线服务场景中，如聊天机器人，大语言模型（Large Language Model，LLM）推理一般都采用流式输出（streaming）的形式，LLM推理的首token时延就是用户感受到的LLM推理服务的响应时间，直接影响用户体验。对于在线服务，为了提升用户体验，所以都希望首token时延要小，一般在一秒左右比较好。LLM推理的首token时延（time to first token, TTFT [1]）与模型参数规模、Prompt长度、Batch Size、GPU资源等因素有关。本文就聊聊首token时延的优化，特别是通过System Prompt Caching [2]优化首token时延。

二、首token时延

在LLM推理过程中，生成首token是计算密集型任务，生成首token阶段也称为prefill phase或context phase，生成首token的时间与处理输入给大模型的Prompt的计算量有关，与Prompt长度直接相关。例如，在Prompt长度相对较长的情况下（Prompt计算时间显著超过模型参数IO时间）[3]，再考虑到FlashAttention2等技术优化，生成首token的时间与输入Prompt的长度近似成线性关系。

在我们通常的应用中，例如，个人助理聊天机器人、RAG客服系统等，输入给大模型的Prompt一般包含System Prompt和User Prompt两部分。

System Prompt，是对大模型的一些任务指示或者角色定义，是与Prompt Template相关的，且一般内容比较长，例如，"You are a Blockchain Development Tutor. Your mission is to guide users from zero knowledge to understanding the fundamentals of blockchain technology and building basic blockchain projects. Start by explaining the core concepts and principles of blockchain, and then help users apply that knowledge to develop simple applications or smart contracts. Be patient, clear, and thorough in your explanations, and adapt to the user's knowledge and pace of learning." [4]。网上甚至有传言ChatGPT的 System Prompt 可能包含1700 tokens [5]。

User Prompt，也就是用户输入的问题或者诉求，一般是相对较短的，例如，"I'm new to blockchain technology. Can you help me understand what it is and how it works?" [4]。

总之，我们看到 System Prompt 是相对较长的输入，且我们对LLM的每一次请求都可能带着相同的System Prompt作为输入（作为Prompt的一部分）。这样就导致，对于用户多次请求，LLM推理需要重复计算System Prompt，造成GPU资源浪费，特别是增加了不必要的首token时延。

所以，对于首token时延优化，如果能省去对于System Prompt的重复计算，那将会显著提升首token生成速度。System Prompt Caching方法的提出[2]，就是为了避免重复计算System Prompt，从而提高首token生成速度。

三、System Prompt Caching

System Prompt Caching [2]，也称为 Prefix Sharing [6, 7, 8]，其基本思想是对System Prompt部分进行一次计算，并缓存其对应的Key和Value值（例如，存放在GPU显存中），当LLM推理再次遇到相同的（甚至部分相同的）System Prompt时，则可以直接利用已经缓存的System Prompt对应的Key和Value值，这样就避免了对于System Prompt的重复计算。

System Prompt Caching主要分两种形式。

第一种形式，Prefix Sharing [6, 7, 8]，适用于 “Prompt = System Prompt + User Prompt”这样的场景，其中System Prompt就是Prefix。例如，给大模型输入的翻译指令，具有相同的System Prompt (Shared Prefix)。

如果用户每次请求都带有相同的System Prompt (Shared Prefix)，那么我们可以一次处理System Prompt，计算其对应的Key和Value值，并缓存在GPU显存中。在处理用户其他请求时，对于相同的System Prompt，就不用再重复计算，而是直接从显存中加载System Prompt对应的Key和Value值，这样就可以显著减少Prefill Phase (Context Phase)的计算时间，可以显著减少生成首token的时间。

第二种形式，Prompt Cache，属于相对高级的用法，是对整个输入Prompt对应的Key和Value值进行Caching操作，不局限于shared prefix。这种方式需要使用Prompt Cache模板[2]，可以针对Prompt的不同部分分别执行KV Cache。

特别地，对于多轮对话场景，上述第二种方式，即Prompt Cache，可以支持Session Prompt Cache。在一个多轮对话session里，输入给LLM的Prompt，会携带多轮对话历史，涉及到很多重复计算。通过 Session Prompt Cache 可以显著减少不必要的重复计算，节省GPU资源，提高对话响应速度和用户体验。

四、TRT-LLM中的实现

Nvidia开源的TensorRT-LLM（TRT-LLM）推理引擎已经支持了System Prompt Caching（Prefix Sharing）功能[9]。我们实测，当System Prompt在Prompt中占比较大时（即System Prompt比User Prompt长），System Prompt Caching功能可以带来较大的性能提升，可以显著减少生成首token的时延。

不过，截止发稿时间，我们发现在TensorRT-LLM里，System Prompt Caching与FP8 KV Cache、INT8 KV Cache并不兼容。我们期待TensorRT-LLM的下一个版本可以修复这些问题，这样我们就可以同时获得 System Prompt Caching 与 FP8/INT8 KV Cache 带来的推理加速效果。

最后我想说，选择LLM推理解决方案，我推荐 Triton & TRT-LLM。其中，需要Triton支持RESTFul API流式输出可以通过增加一个HTTP2gRPC模块来实现（过渡方案），可以实现兼容OpenAI接口协议。Triton未来也会直接支持基于RESTFul API的流式输出。TRT-LLM的功能日臻完善，性能优越，易于运维监控，并较好的支持了FP8推理计算，相对于INT8推理，FP8推理具有更好的模型效果[11]。FP4计算都要来了，FP8计算要尽快用起来啦。