`torch.compile` 整合¶

在 vLLM 的 V1 架構中，torch.compile 預設啟用，並且是該框架的關鍵組成部分。本文件提供了一個簡單的演練示例，展示如何理解 torch.compile 的用法。

在整個示例中，我們將執行一個通用的 Llama 模型，並開啟除錯級別日誌以顯示所有詳細資訊。要使用的命令是 VLLM_LOGGING_LEVEL=DEBUG vllm serve meta-llama/Llama-3.2-1B。

注意

有關 torch.compile 整合的更多資訊和最新進展，請參閱這篇部落格文章。

編譯快取¶

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

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 例項的啟動時間。

考慮的因素包括

所有相關的配置（請參閱配置資料夾）中 compute_hash 函式
PyTorch 配置（請參閱 compiler_interface.py）中 compute_hash 函式
模型的 forward 函式以及 forward 函式呼叫的相關函式（見下文）

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

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

預設情況下，快取將編譯的產物儲存為二進位制檔案。如果您想為了除錯目的與生成的程式碼進行互動，請在編譯配置中將欄位 compile_cache_save_format 設定為 unpacked，或者省略該欄位並設定環境變數 VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked。

動態形狀和 vllm 守衛移除¶

torch.compile 被設計為在需要時毫不猶豫地守護動態形狀。這與 vLLM torch.compile 移除守衛的方法相矛盾，因為許多守衛可能是實質性的。

torch.compile 提供了兩種動態形狀：backed 和 unbacked。torch.compile 會守護 backed 動態形狀，並且不保證不會向其新增任何守衛。使用者程式碼、dynamo、inductor 和 autograd 都可以新增守衛。此外，對於 0/1 特化，即使在這些範圍內沒有遇到分支，backed 符號也會無條件地特化為 0、1 或 >=2。

相反，unbacked 動態形狀保證不會被守護，也不會進行 0/1 特化。然而，當遇到需要其值的分支並且沒有定義顯式的 unbacked 處理時，可能會丟擲資料依賴錯誤。框架正在趨向於一個不會丟擲 DDE 而是選擇通用路徑的狀態。使用 unbacked 的一個缺點是由於效能錯誤或選擇通用路徑而錯失了最佳化機會，並且使用了基於固定非示例輸入的提示（這將很快透過 override_hint API 修復）。選擇通用路徑的一個示例是，在不能透過引入 clone 來符號化證明的情況下，假設輸入不連續，在函式呼叫 contiguous() 和 reshape() 時。

backed_size_oblivious 是一個標誌，它允許在定義了顯式的 unbacked 處理的地方將 backed 符號視為 unbacked。在此模式下，框架程式碼中幾乎避免了 0/1 特化，並且預設的 0/1 特化不會發生。然而，仍然不能保證 torch.compile 不會發生守護，尤其是由於使用者程式碼或自定義 pass。backed_size_oblivious 在 PyTorch compile 中是實驗性的，並且可能被棄用。儘管如此，它比 backed 更安全的選擇，並且效能下降的可能性比 unbacked 更低。

配置動態形狀¶

DynamicShapesConfig 允許您透過設定 type 欄位來控制動態形狀的行為。您可以在三種模式之間進行選擇：BACKED（預設）、UNBACKED 和 BACKED_SIZE_OBLIVIOUS。

離線推理示例（使用 LLM 類）¶

在使用 LLM 類進行離線推理時，您可以透過 compilation_config 引數配置動態形狀

from vllm import LLM, SamplingParams
from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType

# Example: Using backed_size_oblivious (experimental, safer than backed)
llm = LLM(
    model="meta-llama/Llama-3.2-1B",
    compilation_config=CompilationConfig(
        dynamic_shapes_config=DynamicShapesConfig(
            type=DynamicShapesType.BACKED_SIZE_OBLIVIOUS
        )
    )
)

# Example: Using unbacked (strongest guarantee against guards)
llm = LLM(
    model="meta-llama/Llama-3.2-1B",
    compilation_config=CompilationConfig(
        dynamic_shapes_config=DynamicShapesConfig(
            type=DynamicShapesType.UNBACKED
        )
    )
)

# Generate outputs
prompts = ["Hello, my name is", "The future of AI is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
outputs = llm.generate(prompts, sampling_params)

線上服務示例（使用 vllm serve）¶

在使用 vllm serve 進行線上服務時，您可以透過 --compilation-config 標誌配置動態形狀

# Example: Using unbacked
vllm serve meta-llama/Llama-3.2-1B \
  --compilation-config '{"dynamic_shapes_config": {"type": "unbacked"}}'


# Alternative: Using dot notation (simpler for single values)
vllm serve meta-llama/Llama-3.2-1B -cc.dynamic_shapes_config.type=unbacked

選擇合適的模式¶

BACKED（預設）：當您願意為了最大效能而接受潛在的不安全守衛移除時使用。守衛可能會被不安全地新增然後忽略。
UNBACKED：當您需要最強的反守衛保證時使用。這是最保守的選擇，但可能會錯過一些最佳化機會。
BACKED_SIZE_OBLIVIOUS：當您希望在避免守衛和效能之間取得平衡時使用。這種實驗模式比 BACKED 更安全，但仍不如 UNBACKED 保守。

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 頭投影和取樣操作不包含在圖中。

計算圖的大部分輸入都具有靜態形狀，因為它們是模型權重和緩衝區，在模型生命週期內不會改變。只有輸入 ID 和位置 ID 具有符號形狀，即形狀可以從批次到批次改變。然而，它們將共享相同的符號形狀。也就是說，計算圖唯一改變的大小是批次大小（當前 forward 傳遞中處理的 token 數量）。

注意力操作很複雜，它需要與具有複雜形狀的 kv 快取進行互動。幸運的是，注意力操作的輸出與注意力操作的輸入查詢共享相同的形狀。因此，我們將整個注意力操作封裝到一個 PyTorch 自定義 op 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 個圖不同。這是預期的，因為我們按注意力 op 分割圖，得到 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（包括任何注意力操作之前的第一個圖，以及所有注意力操作之後的最後一個圖）。這是基於一個常見的觀察：注意力之間的計算通常是 token-wise 的，並且易於為 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 中。這可以在某些情況下（如小型模型或 MOE 的解碼速度）提高效能。有關更多詳細資訊，請參閱 CUDA Graphs。

torch.compile 整合¶