vLLM 的 `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 函式）
模型的 forward 函式及其呼叫的相關函式（見下文）

考慮所有這些因素後，通常我們可以保證快取是安全的，不會導致任何意外行為。因此，快取預設啟用。如果您想除錯編譯過程，或者懷疑快取導致某些問題，可以透過設定環境變數 VLLM_DISABLE_COMPILE_CACHE=1 來停用它。

vLLM 的 torch.compile 整合的一個獨特之處在於，我們保證所有編譯都在提供任何請求之前完成。沒有請求會觸發新的編譯。否則，引擎將因該請求而被阻塞，並且響應時間將出現意外的峰值。

Python 程式碼編譯¶

在非常詳細的日誌中，我們可以看到

日誌

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 函式。在 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、模型中的權重和緩衝區，輸出是最終的隱藏狀態。請注意，lm_head 投影和取樣操作不包含在圖中。

計算圖的大多數輸入都具有靜態形狀，因為它們是模型權重和緩衝區，並且在模型的生命週期內不會改變。只有輸入 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 可以比較不同的配置，並且通常能找到一些更好的配置來執行核心。例如，我們可以看到以下日誌

日誌

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 庫）快得多。

不幸的是，由於自動調優需要相當長的時間（從幾秒到幾分鐘，取決於模型大小和批次大小），即使可以快取以供後續使用，為了使用者友好性，我們預設將其關閉。如果您想獲得最大效能，建議透過編譯特定形狀來嘗試。

Cudagraph 捕獲¶

vLLM 的 V1 架構使用分段 Cudagraph。完整的計算圖如上所述被分割，我們只捕獲注意力操作之間（包括任何注意力操作之前的第一個圖，以及所有注意力操作之後的最後一個圖）的圖段的 Cudagraph。這是基於一個常見的觀察：注意力之間的計算通常是逐令牌的，易於 Cudagraph 處理；而注意力操作對於 Cudagraph 相容性來說並非易事。因此，透過在 eager 模式下執行注意力操作，而其餘操作在 Cudagraph 中執行，我們保持了注意力操作的靈活性。

分段 cudagraph 還具有細粒度的記憶體管理。目的是僅將注意力核心排除在 cudagraph 之外，同時將所有其餘模組和記憶體分配操作保留在 cudagraph 中。這就是為什麼 V1 中的注意力操作將輸出張量作為注意力的輸入。

Cudagraph 由編譯器後端捕獲和管理，並在批次大小有相應 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 相容，並且僅在停用級聯注意力時相容。

vLLM 的 torch.compile 整合¶