`torch.compile` 集成¶

在vLLM的V1架构中，torch.compile默认启用，是该框架的关键组成部分。本文档通过一个简单的示例演示如何理解torch.compile的用法。

在本示例中，我们将使用v1版本运行一个常见的Llama模型，并开启调试级别日志以显示所有详细信息。使用的命令是VLLM_USE_V1=1 VLLM_LOGGING_LEVEL=DEBUG vllm serve meta-llama/Llama-3.2-1B。

编译缓存¶

在非常详细的日志中，我们可以看到：

INFO 03-07 03:06:55 [backends.py:409] Using cache directory: ~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0 for vLLM's torch.compile

vLLM会综合考虑所有可用因素，并决定一个目录来存储所有编译产物。这意味着，您可以直接复制整个~/.cache/vllm/torch_compile_cache目录到您的部署场景中，从而节省大量编译时间，进而加速vLLM实例的启动时间。

考虑的因素包括：

所有相关配置（参见 config.py中的compute_hash函数）
PyTorch配置（请参阅 compiler_interface.py中的compute_hash函数）
模型的前向函数及其调用的相关函数（见下文）

综合考虑所有这些因素后，通常我们可以确保缓存的使用是安全的，不会引发任何意外行为。因此，默认情况下缓存是启用的。如果您想调试编译过程，或者怀疑缓存导致了某些问题，可以通过设置环境变量VLLM_DISABLE_COMPILE_CACHE=1来禁用缓存。

vLLM的torch.compile集成有一个独特之处，我们保证在开始处理任何请求之前完成所有编译。不会有任何请求触发新的编译。否则，引擎会被该请求阻塞，导致响应时间出现意外峰值。

Python代码编译¶

在非常详细的日志中，我们可以看到：

Logs

DEBUG 03-07 03:06:52 [decorators.py:203] Start compiling function <code object forward at 0x7f08acf40c90, file "xxx/vllm/model_executor/models/llama.py", line 339>

DEBUG 03-07 03:06:54 [backends.py:370] Traced files (to be considered for compilation cache):
DEBUG 03-07 03:06:54 [backends.py:370] xxx/torch/_dynamo/polyfills/builtins.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/torch/nn/modules/container.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/torch/nn/modules/module.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/attention/layer.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/distributed/communication_op.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/distributed/parallel_state.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/custom_op.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/layers/activation.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/layers/layernorm.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/layers/linear.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/layers/rotary_embedding.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/layers/vocab_parallel_embedding.py
DEBUG 03-07 03:06:54 [backends.py:370] xxx/vllm/model_executor/models/llama.py

DEBUG 03-07 03:07:07 [backends.py:462] Computation graph saved to ~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/computation_graph.py
DEBUG 03-07 03:07:07 [wrapper.py:105] Dynamo transformed code saved to ~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/transformed_code.py

这是关于Python代码编译的内容，即通过Dynamo进行图捕获。它尝试追踪代码xxx/vllm/model_executor/models/llama.py:339处的函数，该函数是我们编译模型的forward前向传播函数。在前向传播过程中，根据日志显示，Dynamo还会内联调用其他函数，包括来自xxx/torch/nn/modules/module.py的一些PyTorch函数（由PyTorch的nn.Module使用，因为模块属性访问会触发函数调用），以及来自vLLM的一些通信/注意力/激活函数。在决定使用哪个缓存目录时，所有被追踪的文件都会被考虑进去。这样一来，上述文件中的任何代码更改都会触发编译缓存未命中，从而导致重新编译。

Dynamo编译的结果是一个存储在~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/transformed_code.py中的新函数。通常，该函数会从模块中解包张量，然后将其传递给追踪的计算图。计算图存储在~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/computation_graph.py中。

计算图处理¶

计算图中的每个张量都有形状标注。输入包括输入ID、位置ID、模型权重和缓冲区，输出则是最终的隐藏状态。请注意，图中未考虑语言模型头部的投影和采样操作。

计算图的大多数输入具有静态形状，因为它们属于模型权重和缓冲区，在模型生命周期内不会改变。只有输入ID和位置ID具有符号化形状，即每批次的形状可能不同。不过，它们会共享相同的符号化形状。也就是说，计算图中唯一变化的尺寸是批次大小（当前前向传播中处理的令牌数量）。

注意力操作较为复杂，需要与形状多变的kv缓存进行交互。幸运的是，注意力操作的输出形状正好与其输入查询的形状相同。因此，我们将整个注意力操作封装到PyTorch自定义算子torch.ops.vllm.unified_attention_with_output中，这样Dynamo就不会尝试检查任何内部操作。如此一来，尽管注意力操作很复杂，但从Dynamo的角度来看，我们仍然可以将模型的计算图捕获为一个完整图。

计算图会进一步被splitting_ops（通常是注意力运算操作）分割成多个片段。因此，在~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/computation_graph.py文件中，我们可以看到许多子模块，每个子模块都是分割后的一个图片段：

注意力操作本身是一个子模块。
计算图中从一个注意力操作到下一个注意力操作的部分是一个子模块。

每个子模块都可以通过其索引进行标识，并将被单独处理。

计算图编译¶

在非常详细的日志中，我们还可以看到：

DEBUG 03-07 03:52:37 [backends.py:134] store the 0-th graph for shape None from inductor via handle ('fpegyiq3v3wzjzphd45wkflpabggdbjpylgr7tta4hj6uplstsiw', '~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/iw/ciwzrk3ittdqatuzwonnajywvno3llvjcs2vfdldzwzozn3zi3iy.py')
DEBUG 03-07 03:52:39 [backends.py:134] store the 1-th graph for shape None from inductor via handle ('f7fmlodmf3h3by5iiu2c4zarwoxbg4eytwr3ujdd2jphl4pospfd', '~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/ly/clyfzxldfsj7ehaluis2mca2omqka4r7mgcedlf6xfjh645nw6k2.py')
...
DEBUG 03-07 03:52:45 [backends.py:134] store the 15-th graph for shape None from inductor via handle ('f7fmlodmf3h3by5iiu2c4zarwoxbg4eytwr3ujdd2jphl4pospfd', '~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/ly/clyfzxldfsj7ehaluis2mca2omqka4r7mgcedlf6xfjh645nw6k2.py')
DEBUG 03-07 03:52:45 [backends.py:134] store the 16-th graph for shape None from inductor via handle ('fvj3ccoi7m34f3dnr4itmu55mmun44l5xymwhrjlwisylsk7q6jy', '~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/tf/ctfftkglj7b4lcttq5cymx6cew372uoauupqn6ldsvpiucavqcjc.py')

这意味着第一段计算图（形状为None表示符号形状）由Inductor编译（使用密钥fpegyiq3v3wzjzphd45wkflpabggdbjpylgr7tta4hj6uplstsiw）。编译后的内核存储在~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/iw/ciwzrk3ittdqatuzwonnajywvno3llvjcs2vfdldzwzozn3zi3iy.py中。您可以打开该文件查看Inductor最终运行的代码。

还有一个细节：你可以看到第1个图和第15个图具有相同的键，而第0个图和第16个图则不同。这是预期的，因为我们通过注意力操作分割图，得到了3个独特的子图：

注意力机制前的第一层
每一中间层，从一个注意力操作到下一个注意力操作
注意力机制后的最终层

如果我们已经有缓存目录（例如第二次运行相同的代码），我们将看到以下日志：

DEBUG 03-07 04:00:45 [backends.py:86] Directly load the 0-th graph for shape None from inductor via handle ('fpegyiq3v3wzjzphd45wkflpabggdbjpylgr7tta4hj6uplstsiw', '~/.cache/vllm/torch_compile_cache/1517964802/rank_0_0/inductor_cache/iw/ciwzrk3ittdqatuzwonnajywvno3llvjcs2vfdldzwzozn3zi3iy.py')

这次完全绕过了Inductor编译过程，我们将直接从磁盘加载读取上次获得的编译产物。

上面的示例仅使用Inductor为通用形状(即符号形状)进行编译。我们也可以使用Inductor为某些特定形状进行编译，例如：

vllm serve meta-llama/Llama-3.2-1B \
  --compilation_config '{"compile_sizes": [1, 2, 4, 8]}'

然后它还会专门为批处理大小1, 2, 4, 8编译一个特定的内核。此时，计算图中的所有形状都是静态且已知的，我们将开启自动调优以获得最佳性能。首次运行时可能会较慢，但下次运行时就可以直接跳过调优阶段，直接运行已调优的内核。

当所有形状都已知时，torch.compile可以比较不同的配置，并经常能找到一些更好的配置来运行内核。例如，我们可以看到以下日志：

Logs

AUTOTUNE mm(8x2048, 2048x3072)
  triton_mm_4 0.0130 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=2
  triton_mm_8 0.0134 ms 97.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=64, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=4
  triton_mm_12 0.0148 ms 87.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=128, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=4, num_warps=4
  mm 0.0160 ms 81.6% 
  triton_mm_16 0.0165 ms 78.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=16, BLOCK_N=128, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=8
  triton_mm_3 0.0199 ms 65.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=2
  triton_mm_1 0.0203 ms 64.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=16, BLOCK_N=32, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=2, num_warps=2
  triton_mm_7 0.0203 ms 64.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=16, BLOCK_N=64, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=3, num_warps=4
  triton_mm_2 0.0208 ms 62.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=16, BLOCK_N=64, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=5, num_warps=4
  triton_mm_11 0.0215 ms 60.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=16, BLOCK_N=128, B_PROLOGUE_CAST_TYPE=None, EVEN_K=True, GROUP_M=8, num_stages=3, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 2.0428 seconds and 7.5727 seconds precompiling

这意味着，对于一个形状为8x2048x3072的矩阵乘法运算，torch.compile会尝试使用不同配置的triton模板，其执行速度远快于默认代码（默认代码会调用cublas库）。

遗憾的是，由于自动调优耗时较长（从几秒到几分钟不等，具体取决于模型大小和批次大小），尽管可以缓存供后续使用，但为了用户体验，我们默认将其关闭。如果您希望获得最佳性能，建议通过编译特定形状来尝试启用该功能。

CUDA图捕获¶

vLLM的V1架构采用了分段式cudagraph技术。如上所述，完整计算图被分割处理，我们仅捕获注意力操作之间的子图部分（包括首个注意力操作前的初始子图，以及所有注意力操作后的最终子图）。这一设计基于常见现象：注意力操作之间的计算通常具有token粒度的特性，易于适配cudagraph；而注意力操作本身要实现cudagraph兼容则较为复杂。因此，通过在eager模式下执行注意力操作，其余操作在cudagraph中运行，我们保持了注意力操作的灵活性。

分段式cudagraph还具备细粒度的内存管理功能。其目的是仅将注意力内核排除在cudagraph之外，同时保留所有其他模块和内存分配操作在cudagraph中。这就是为什么V1版本中的注意力操作会将输出张量作为注意力层的输入。

cudagraphs由编译器后端捕获和管理，并在批次大小有对应的cudagraph捕获时重放。模型的调用者（模型运行器）只需确保正确管理输入缓冲区即可。所有中间缓冲区都由编译器后端自动管理。

默认情况下，vLLM会尝试确定一组尺寸来捕获cudagraph。您也可以通过配置cudagraph_capture_sizes来覆盖它：

vllm serve meta-llama/Llama-3.2-1B \
  --compilation-config '{"cudagraph_capture_sizes": [1, 2, 4, 8]}'

这样它将仅捕获指定尺寸的cudagraph。对cudagraph捕获进行细粒度控制会很有帮助。

完整捕获Cudagraph¶

如果使用与cudagraph兼容的注意力后端，可以将注意力机制作为cudagraph的一部分包含进来。这可以在某些情况下提高性能，例如较小模型的解码速度。通过--compilation-config '{"full_cuda_graph": true}'启用此功能。

目前仅FlashAttention 3兼容，且仅在禁用级联注意力时可用。

torch.compile 集成¶