LLM 推理框架运行时优化：CUDA Graph、分块预填充与投机解码

分析主流 LLM 推理框架的运行时开销来源（Python GIL 与多进程通信），并梳理 CUDA Graph、多步调度、Chunked Prefill、投机解码、基于哈希的自动前缀缓存与 vLLM Engine V2 等优化手段。

主流 LLM 推理框架的运行时开销大致来自：

Python 性能：考虑用户易用性和开发效率，业界主流框架都采用 Python 为主要开发语言、C++实现模型和算子的方式。Python 一直存在让人诟病的 GIL 问题，框架中很多的代码无法有效并行。这导致在高并发场景，Python 执行速度成为了显著的瓶颈。
通信开销：为了规避 Python GIL 的限制来提升框架的并行度，即使是在单机内，主流框架通常都会引入多进程的设计。此外大模型推理也逐渐扩展至分布式推理，系统模块可能存在多个 Prefill、Decode 实例。系统在运行过程中，不同的模块、实例间需要频繁进行交互。而进程间的通信、消息的序列化和反序列化也不可避免的成为了瓶颈。
同步执行逻辑：主流的推理框架都会采用分页的 KV Cache 管理，执行流程是严格按照 KV Cache Block 分配、模型计算、KV Cache Block 更新的迭代式的过程。目前业界主流框架都是基本按照同步的设计，各个组件的开销不能被掩盖，当并发变大后累计的开销可能会达到模型计算的量级。
近期，大模型推理社区（vLLM，SGLang 等）普遍开始关注框架运行时开销，提出了多步调度、异步输出处理、独立 API Server 进程等工作，来分摊或掩盖部分开销。在实际业务场景中，也观察到高额的框架开销严重限制了系统吞吐，特别是在高并发（>1k）场景下，运行时开销已经接近或高于 GPU 运行时间，导致资源严重浪费和性能下降。

vLLM社区致力于提升性能，主要通过以下几个方面：

矩阵乘法优化：使用cuda kernel和cutlass库优化矩阵乘法，实现kernel fusion等功能。
层级优化：针对特定层（如ffn层）优化Triton内核，提升性能。
通信性能优化：开发定制的all-reduce CUDA内核，优化数据传输和内存访问，提高并行度。
注意力机制优化：采用基于CUDA的内核和其他注意力内核，如Flash Attention和FlashInfer。
量化支持：开发通用量化方法抽象层，实现并优化多种量化方法（FP8、IN8、GPTQ、AWQ等），引入高效量化内核如Marlin。
模型压缩：LLM Compressor工具支持多种量化方法，高效地将模型量化为vLLM可理解的格式，以提升性能。

CUDA Graph

Cuda Graph对vLLM的性能提升很大，毕竟vLLM是采用pytorch原生的op配合拓展op搭建的，有很多额外的消耗：user-written logic, PyTorch dispatcher logic, memory allocation overhead, and GPU driver/kernel overhead。使用了cuda graph可以避免掉这些开销。

We find that without CUDA graphs, LLaMA-7B inference executes at 30 tokens/sec, but with CUDA graphs enabled it executes at 69 tokens/sec for a 2.3x speedup.

唯一可能的缺点就是因为要适配dynamic shape，占用的显存相比之前大些，同时启动时间也因为要捕获graph慢了些。

Multi-step Scheduling

一个典型的 LLM 推理引擎可以分为以下组件：

API Server：负责请求接收和响应；
Scheduler：负责请求调度、KV Cache Block 分配等；
Model Runner：主要负责模型计算和采样；
Decoder：负责将采样得到的 Token ID 转化为字符形式输出。
以下是一个完全同步的推理引擎的 timeline 示意图：

针对运行时开销逐渐变大的问题，社区也在逐渐优化中，提出了包括异步输出处理、多步调度等方案来尽可能掩盖或分摊运行时开销。以目前最新的 vLLM 0.6.0 为例，开启 MultiStepScheduling 和异步输出处理优化后，整体的 timeline 大致可以表示为：

Multi-step Scheduling提出的原因很简单，之前vllm被诟病python/cpu overhead一直很大，也就是每次decoder时候，CPU与GPU的同步操作导致了较高的性能开销（你可以亲自打下vLLM运行时候的timeline，中间间隙比TRT-LLM要大些）。

具体就是GPU在每次从CPU接收下一步的输入，以及从GPU传输采样后的token，到CPU生成用户响应时，都会产生一定的“GPU bubble”（GPU等待CPU处理的时间，5-13ms of GPU bubble），具体是这四个：

Transfer of the sampled token from GPU to CPU for de-tokenization and response to client
Generation of output for user - Pythonization of tensors into python objects
CPU preparation and generation of next step’s input metadata
vLLM scheduler
多步解码指的是在调用 vLLM 调度器并处理采样的 tokens 之前，执行多次解码。

之前每次解码步骤中 GPU->CPU 的内存传输都是同步进行的，导致 GPU 出现空闲时间。多步解码后，这种内存传输可以在一个独立的 CUDA 流中进行，而此时 CPU 已经领先于 GPU，因此几乎没有额外开销。

Multi-step decoding will be able to amortize all these overheads over n-steps at a time.

下方两个图是官方提供的对比图，这两个时长大约为 200ms。第一行是 CUDA 核心操作，底部是 Python 跟踪。每个红框大约代表 4ms 的 GPU 空闲时间，两张图均如此。

首先是baseline，每个decode步骤都会导致4毫秒的GPU空泡（这里使用Pytorch profiler打印）：

使用Multi-Step-8之后，4毫秒的开销仅在每8次解码中出现一次：

通过使用Multi-Step，基于H100的8B Llama模型的请求吞吐量从20.66请求/秒提升至40.06请求/秒，性能显著提升。

具体提升指标：

不过需要注意这个只针对decoder阶段，prefill阶段还和之前一样

更多的细节可以看这两个链接:

多步解码缺点

这样的调度方式存在以下问题：

连续 N 个 step 结束之后，Model Runner 还是会处于等待状态，直到 Scheduler 调度新的一轮请求。所以 MultiStepScheduling 只能平摊、但不能完全消除 Scheduler 的开销。而随着并发增大，Scheduler 和 Decoder 等环节开销变大，Model Runner 的等待时间仍然会增加。
Model Runner 需要多个 step 后再进行返回，而这中间,用户是不能接收到这些新生成的 token。所以从用户视角，TTFT 和 ITL(词元间时延) 会显著上升，影响用户体验。
Chunked prefill

https://docs.vllm.ai/en/stable/models/performance.html

Chunked prefil这个功能的作用在于，可以将非常长的输入prompt拆分成更小的块进行处理，以避免长prompt阻塞其他请求，从而减少延迟并提高整体的吞吐。相比不用的情况，用户的性能延迟有了 2 倍的提升，吞吐量更高，延迟更低。

默认情况下，vLLM 调度程序优先考虑预填充，并且不会将预填充和解码批处理到同一个批次。此策略优化了 TTFT（第一个令牌的时间），但会导致 ITL（令牌间延迟）变慢，并且 GPU 利用率低下。

启用分块预填充后，策略会更改为优先处理解码请求。它会将所有待处理的解码请求批处理到批次中，然后再调度任何预填充。当有可用的 token_budget（max_num_batched_tokens）时，它会调度待处理的预填充。如果最后一个待处理的预填充请求无法放入 max_num_batched_tokens 中，它会将其分块。

前面提到，vLLM 将请求的处理分为prefill阶段和decode阶段，同一批被调度的请求要么都处于prefill阶段，要么都处于decode阶段。而 Chunked prefill 将填充请求的 prompt 分成多个块，并与解码请求一起处理。Prefill是计算密集型（compute-bound），而decode则是内存密集型（memory-bound），通过重叠这两种请求可以大大提高系统效率，正如上图（右上角 prefill 进行数学上等价的切分；在 prefill 切成的这些 chunks 的气泡处捎带其他 request 的 decode pass，可以看到R2的Prefill和R1的Decode在一起计算）。

此策略有两个优点：

它通过优先处理解码请求来提高 ITL 和生成解码。
它通过将计算密集型（预填充）和内存密集型（解码）请求定位到同一个批次来帮助实现更好的 GPU 利用率。
不过也需要注意，做了 chunked prefill 后，prefill 的开销会略微增大。因为计算后续 chunk 的 KV 时需要不断地从 GPU memory 中里读出当前 chunk 的 KV 到 kernel 里面；而不做 chunked prefill 时，最开端的那些 KV Cache 可以不用反复从 GPU memory 中反复读取进入 kernel，毕竟他们一直在 kernel 里面。

可以通过更改 max_num_batched_tokens 来调整性能。默认情况下，它设置为 512，在初始基准测试（llama 70B 和 mixtral 8x22B）中，它在 A100 上具有最佳 ITL。较小的 max_num_batched_tokens 可以实现更好的 ITL，因为有更少的预填充会中断解码。较高的 max_num_batched_tokens 可以实现更好的 TTFT，因为你可以将更多预填充放入批次中。

如果 max_num_batched_tokens 与 max_model_len 相同，这几乎等同于默认调度策略（除了它仍然优先处理解码）。
请注意，max_num_batched_tokens 的默认值（512）针对 ITL 进行了优化，它的吞吐量可能低于默认调度程序。
建议你将 max_num_batched_tokens 设置为大于 2048 以获得更高的吞吐量。

关于 Chunked prefill 细节可以阅读：

DeepSpeed-FastGen
SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills
官方RFC：
[RFC] Upstream Chunked Prefill · Issue #3130 · vllm-project/vllm
Speculative Decoding

Speculative Decoding也就是俗称的投机推理或者叫预测推理。

vLLM 这里讲了两个，分别是speculative decoding (draft-verification) && dynamic speculative deccoding（根据系统负载和推测精度动态调整长度）

Speculative Decoding，其目标是利用小模型来预测大型模型的行为，从而加速大型模型的执行。例如，我们可以让小模型生成候选序列，然后将此候选序列提供给大模型，允许大模型并行处理以验证这些候选序列的正确性。如果验证通过，我们可以一次性生成多个token，大大提升生成速度：

再详细些，看下面的例子。下面是一个关于随机采样工作方式的示例。图中每一行表示算法的一次迭代。绿色标记表示近似模型生成的token建议，目标模型则负责判断是否接受这些建议。红色和蓝色标记分别表示被目标模型拒绝的token及其修正。

例如，在第一行中，近似模型生成了5个token，生成的句子是“[START] japan’s benchmark bond”。目标模型将这些token与前缀拼接成一句话进行验证。在这次推理中，目标模型接受了前四个token（“japan”，“’s”，“benchmark”），但拒绝了最后一个token “bond”，并重新采样生成“n”。

通过这种方式，虽然每次都对近似模型的输出进行一次性验证，但目标模型在仅9次推理后就生成了37个token。尽管目标模型的总体计算量没有变化，但一次性验证多个token的延迟与逐个生成token相似（并行总是比一个一个吐字快），从而显著提高了生成速度。

不过还有另一种情况，当系统负载很高时，vLLM已经可以很好地处理批量请求，因此在这种情况下使用预测解码可能并不会提升性能。然而，在系统负载较低时，利用预测解码可以减少延迟。因此，Dynamic Speculative Decoding的核心是如何动态调整预测解码量，以在系统负载不同的情况下获得最佳性能。

Hash-Based Automatic Prefix Caching

vLLM还实现了基于哈希的自动prefix caching。

一般我们的LLM服务会有一个较长的系统提示来描述模型的行为。若两个请求共享同样的系统提示，则可以缓存并共享该系统提示的 KV 缓存，从而减少重复计算。此外，多轮对话也可以受益于这种缓存机制。vLLM实现了基于哈希的前缀缓存自动化，通过为每个缓存块分配哈希值，使其在不同请求之间实现共享。

基于hash的自动前缀缓存，缓存system prompt及多轮对话，减少请求数据量，这与sglang中采用的radix树维护kv cache不同，但都提高了缓存重用：

vLLM也支持多个lora adapter及热加载，同时支持结构化输出等；也支持了基于prometheus/grafana的实时监控，包括gpu内存使用/kvcache重用率/访问延迟吞吐等指标。

在多任务服务方面，vLLM现在支持同时加载多个 LoRA 适配器，并能将不同适配器的请求批处理。

vLLM Engine V2

vLLM正在进行一次重大的架构重构，但我们保证这一过程将会很快完成。我们正在开发新的 vLLM Engine V2，这将包括异步调度和基于Prefill Cache的设计。已经实现了对 Torch Compile 的全面支持，以便用户可以优化他们的模型。此外还新增了一个内存管理器，以便更好地支持多模态模型。

异步调度是重构的一个重点。通过在当前步骤执行的同时调度下一个步骤，会大幅减少了 GPU 的空闲时间。这样可以让 GPU 在完成当前步骤后立即开始执行下一个步骤。与之前相比，这种新的异步架构让调度和执行可以并行进行，显著降低了 GPU 的空闲时间。

通过Prefill Cache简化了设计。在当前的架构中需要复杂的逻辑来处理并行采样或抢占的情况。在未来，vLLM将使用Prefill Cache来自动识别并行采样和抢占场景中的可重用 KV 缓存，从而减少对手动处理的需求。之前，调度器直接与并行采样和前缀缓存模块进行通信。未来，调度器将只需与前缀缓存模块通信，自动识别重用机会。

重构内存分配器，以更好地适应不同类型的 KV 缓存。由于 KV 缓存大小不兼容，导致 GPU 内存浪费严重。以新的 LLaMA 模型为例，文本部分有全层的 KV 缓存，而图像部分只有 1/4 层的 KV 缓存，我们实际上为所有层分配了内存，导致浪费。我们计划通过新的内存分配器来优化这种情况。我们的方案是将 GPU 内存首先分区成较大的页面，然后进一步将每个大页面划分为不同的大小，以适应不同的 KV 缓存大小。

开启对 Torch Compile 的全面支持。目前，vLLM已经手动优化了最流行的模型，例如 LLaMA 模型的 CUDA 内核和缓存重用优化。我们的工作正在进行中，我们将利用 Torch Compile 来优化所有模型，从而使用户的自定义模型和架构也能高效运行。

优化预测解码。当前的预测解码功能已经不错，但在QPS很高时可能会影响性能。我们即将引入动态预测解码机制，以确保预测解码始终让vLLM运行得更快。与之前的机制不同，新的解码机制将无论在何种负载下都能带来更好的性能。

还将支持将更多的 KV 缓存存储到 CPU 上，从而增加多轮对话和系统提示的缓存空间。目前的vLLM只在 GPU 中存储 KV 缓存，未来将支持自动将缓存转移到 CPU，甚至可以与远程缓存数据库（如 LM Cache 和 Redis）配合使用，进一步扩展 KV 缓存存储。之前的vLLM只能在 GPU 中存储 KV 缓存供未来重用，而未来版本的vLLM可以将 KV 缓存存储在 GPU、CPU 甚至远程数据库中。