CUDA 圖¶

本文件介紹了 vLLM v1 中新增的 CUDA 圖模式，這是對之前 torch.compile 整合的改進。總結來說，我們

添加了靈活的 cudagraph_mode 配置
使完整的 CUDA 圖支援與編譯解耦
引入了 CUDA 圖排程器作為中央控制器，自動為每個批次選擇所需的執行時模式和 CUDA 圖

本文件將討論

動機
CUDA 圖模式
詳細設計
不同 CUDA 圖模式的示例用法

注意

本文件中，我們將純粹的解碼（max_query_len=1）或投機解碼（max_query_len =1+num_spec_tokens）稱為統一解碼批次，反之則為非統一批次（例如，預填充或預填充-解碼混合批次）。

注意

以下內容主要基於 Pull Request #20059 的最後一個提交為基礎。

動機¶

最初的分段編譯是為了允許分段 CUDA 圖捕獲，排除了不支援 CUDA 圖的操作（主要是 Attention）。這使得 CUDA 圖能夠提供一些加速，同時保持與所有 Attention 後端的相容性。後來我們透過不進行分段編譯來增加了對“完整 CUDA 圖”的支援，以便在 Attention 支援 CUDA 圖的情況下進一步降低延遲。然而，編譯和 CUDA 圖捕獲之間的緊密耦合導致了一種“全有或全無”的體驗，靈活性很低。許多 Attention 後端也還沒有準備好統一的“完整” CUDA 圖捕獲（例如，目前只有 FlashAttention 3 支援），或者只支援純解碼批次的 CUDA 圖（例如 Flashinfer, FlashMLA, Mamba 等）。這導致了令人困惑的效能/相容性權衡、不一致的 CUDA 圖支援以及日益複雜的程式碼結構。

這促使我們尋求一種更細粒度的 CUDA 圖解決方案，具有以下特性

明確區分預填充/混合批次或（統一）解碼批次，並分別捕獲。
將 CUDA 圖捕獲邏輯與編譯分離（儘可能可行），以實現功能正交性，這表明
- 使用相同的編譯圖捕獲分段和完整 CUDA 圖，並且
- 在沒有編譯的情況下進行完整 CUDA 圖捕獲。
在執行時根據批次組成在完整和分段 CUDA 圖之間進行排程。
集中控制 CUDA 圖行為，以減少程式碼複雜性並提供更大的擴充套件性。

這些特性為各種啟動/效能權衡和功能支援提供了最大的 CUDA 圖捕獲和編譯靈活性。

`CudagraphModes`¶

CUDAGraphMode 是您在 CompilationConfig.cudagraph_mode 中調整的唯一引數。

NONE — 關閉 CUDA 圖。適用於除錯。
PIECEWISE — 單一模式策略（也是過去的預設設定）。這是最靈活的：Attention 或其他不支援 CUDA 圖的操作保持即時執行，所有其他操作都進入 CUDA 圖。需要分段編譯。
FULL — 單一模式策略，僅捕獲非統一批次的完整 CUDA 圖，然後具有相同 batch_size 的統一解碼批次可以重用非統一批次的 CUDA 圖，因為它們是相容的；可能對小型模型或具有小提示的工作負載有利。
FULL_DECODE_ONLY — 統一解碼的完整 CUDA 圖，不對預填充/混合等進行 CUDA 圖捕獲；適用於 P/D 設定中的解碼例項，其中預填充不太重要。這樣我們可以節省 PIECEWISE CUDA 圖所需的記憶體。
FULL_AND_PIECEWISE —（預設模式）統一解碼的完整 CUDA 圖，其他批次的已分段 CUDA 圖；通常是最具效能的設定，尤其是在小型模型或 MoE 的低延遲方面，但同時需要最多的記憶體且捕獲時間最長。

預設值：如果您使用的是 v1 且啟用了分段編譯，我們預設設定為 FULL_AND_PIECEWISE 以獲得更好的效能（對於池化模型，仍然是 PIECEWISE）。否則，例如，如果分段編譯不可用，我們預設設定為 NONE。

雖然 NONE、PIECEWISE 和 FULL 是單模式配置，並且分別相當於過去即時執行、分段 CUDA 圖和完整 CUDA 圖的實現，但 FULL_DECODE_ONLY 和 FULL_AND_PIECEWISE 是新新增的雙模式配置，它們需要透過排程根據執行時批次動態切換到具體的執行時模式。

注意

在此，單模式 NONE、PIECEWISE 和 FULL 被視為 CUDA 圖排程的執行時模式。如果使用雙模式，排程器將根據批次組成，始終排程到其成員模式之一（如果不存在合適的 CUDA 圖，則可能加上 NONE）。

雖然級聯 Attention 不相容 CUDA 圖，但它現在與所有可能的 CUDA 圖模式配置相容。如果批次使用級聯 Attention，它總是會被排程到 PIECEWISE 模式（如果可用）（否則是 NONE）。

注意

並非所有 CUDA 圖模式都與所有 Attention 後端相容。我們會自動將模式“降級”到最接近的支援模式。例如，如果一個後端只支援純解碼/統一批次的 CUDA 圖，我們會將 FULL 模式轉換為 FULL_AND_PIECEWISE 模式（如果啟用了分段編譯），否則轉換為 FULL_DECODE_ONLY 模式。

詳細設計¶

概述¶

新的 CUDA 圖邏輯建立在分段編譯之上，並支援雙 CUDA 圖執行時模式切換。該系統包含以下核心元件

CUDAGraphWrapper：處理對包裝後的可呼叫物件進行 CUDA 圖捕獲和重放的包裝器。
CudagraphDispatcher：中央控制器，包含 CUDA 圖的唯一事實來源，並負責在它們之間進行排程。
CUDAGraphMode：描述支援和執行時模式的列舉（如上所述）。
BatchDescriptor，作為執行時批次的唯一表示，用於排程。

參見下圖，快速比較 CUDA 圖與 Inductor 編譯之前的當前設計模式。我們可以看到，之前 CUDA 圖邏輯和編譯邏輯被緊密耦合到 vllm 的 PiecewiseBackend 中，並且 CUDA 圖是透過 batch_size 被動排程的。現在，CUDA 圖邏輯已分離到 CUDAGraphWrapper 類中，該類負責完整和分段 CUDA 圖的功能，並且透過 CudagraphDispatcher 使用執行時模式加上 BatchDescriptor 作為排程鍵進行顯式排程。

之前

之後

`BatchDescriptor`¶

BatchDescriptor 是 ForwardContext 中的一個元件，與 CUDA 圖執行時模式一起，作為執行時排程鍵的核心結構。原型如下：

class BatchDescriptor(NamedTuple):
    num_tokens: int
    num_reqs: int
    uniform: bool = False
    has_lora: bool = False

其中 num_tokens 可以是填充後的 token 長度，uniform 表示所有請求是否具有相同的查詢長度。許多 Attention 後端只支援批次統一時的完整 CUDA 圖；純解碼批次是統一的，但可能不是查詢長度為 1（即 num_tokens == num_reqs），這在投機解碼的驗證階段發生，其中“解碼”批次的查詢長度為 1+num_spec_tokens。

此結構的目標是使用最少的項唯一標識一個（填充後的）批次，這些項對應於一個 CUDA 圖項。

注意

未來，BatchDescriptor 的原型可能會擴充套件以適應更通用的情況，例如，包含更多項，如 uniform_query_len 來支援多種不同的統一解碼長度設定（ Pull Request #23679），或支援 CUDA 圖的非 token 長度感知輸入模型所需的其他修改（例如，某些多模態輸入）。

`CudagraphDispatcher`¶

CudagraphDispatcher 負責維護兩組有效的排程鍵，一組用於 FULL 執行時模式，另一組用於 PIECEWISE 執行時模式，並在執行模型的前向傳播之前排程正確的執行時模式和排程鍵。它將接收初始鍵（一個粗略的批次描述符，用於填充後的輸入），並返回選定的執行時模式和最終的批次描述符，然後透過前向上下文將此決策告知 CUDAGraphWrapper 例項。請注意，CudagraphDispatcher 是可用 CUDA 圖鍵的唯一事實來源，CUDAGraphWrapper 例項可以毫無顧慮地信任前向上下文關於要排程到哪個 CUDA 圖的資訊。這使我們能夠簡化 Wrapper 程式碼並將邏輯集中在排程器中。

排程鍵透過排程器的 initialize_cudagraph_keys 方法進行初始化，該方法在所有可能的 Attention 後端初始化完成後由 gpu_model_runner 呼叫。這是我們將來可以做得更花哨的地方，並“準備”所有可能的 CUDA 圖組合。目前，我們只是根據編譯配置中 decode_mode/mixed_mode 的 cudagraph_mode 和 cudagraph_capture_sizes 的有效組合追加可用的鍵。

排程程式碼如下：

batch_descriptor=BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=...)
runtime_mode, batch_descriptor = cudagraphdispatcher.dispatch(batch_descriptor)
# execution
with set_forward_context(
    ..., 
    cudagraph_runtime_mode=runtime_mode, 
    batch_descriptor=batch_descriptor,
):
     output = self.model(...)

在 dispatch() 方法內部，排程器將搜尋合適的 CUDA 圖執行時模式和現有的排程鍵以返回。我們基本上按照優先順序搜尋現有鍵：FULL>PIECEWISE>None。如果排程鍵不存在，則預設返回 NONE 模式以進行即時執行。實現可以在此處找到。

下面是在模型執行器中執行時的工作流程的簡化圖示：

`CUDAGraphWrapper`¶

CUDAGraphWrapper 例項包裝了一個可執行物件，並簡單地模仿了可執行物件，附加了 CUDA 圖功能。每個 Wrapper 例項都繫結到一個特定的 runtime_mode，該模式被限制為 PIECEWISE 和 FULL 模式，並負責捕獲/重放和傳遞（直接呼叫）可執行物件。執行時，每個 Wrapper 會

檢查全域性前向上下文中的 runtime_mode 和 batch_descriptor（排程鍵）。
如果 runtime_mode 是 NONE 或 runtime_mode 與 Wrapper 的模式不匹配，則直接呼叫可執行物件。
否則，即 runtime_mode 與 Wrapper 的模式匹配，Wrapper 將執行 CUDA 圖捕獲（如果鍵不存在，則建立新條目並快取它）或重放（如果鍵存在於快取中）。

以上步驟基於 CUDA 圖 Wrapper 將直接信任前向上下文中的內容的假設（由排程器控制）。這使我們能夠簡化和集中邏輯，減少複雜性以及 Wrapper 和排程器之間狀態不匹配的風險。它還允許 Wrapper 類同時用於 FULL 和 PIECEWISE 執行時模式。實現可以在此處找到。

巢狀 Wrapper 設計¶

使完整 CUDA 圖和分段 CUDA 圖共存併兼容的核心機制是巢狀 CUDA 圖 Wrapper 設計，它建立在只有一個分段 FX 圖的分段編譯之上。我們用一個 FULL 模式 Wrapper 包裹整個模型以實現完整的 CUDA 圖功能；同時，每個分段後端都透過一個 PIECEWISE 模式 Wrapper 在編譯內部進行包裝。

下面的流程圖應清晰地描述其工作原理。

因此，對於 FULL 執行時模式，捕獲/重放完整 CUDA 圖是安全的，因為分段 Wrapper 未被啟用。PIECEWISE 模式的情況也類似，因為 FULL 模式 Wrapper 和 PIECEWISE 模式 Wrapper 之間沒有衝突。對於 NONE 執行時模式，FULL 和 PIECEWISE Wrapper 都不會被啟用，因此我們直接回退到即時執行。

完整 CUDA 圖捕獲 & 熱身¶

CUDA 圖捕獲發生在 Runner 第一次使用非 NONE 執行時模式呼叫模型前向傳播時（使用 _dummy_run）。對於完整的 CUDA 圖捕獲，我們透過正確設定 Attention 元資料來顯式捕獲不同的情況（例如，預填充/混合批次或統一解碼批次），以確保底層的 Attention 後端啟動所需的核心例程。區分預填充/混合批次或統一解碼批次，最重要的屬性是 attn_metadata 中的 max_query_len（對大多數 Attention 後端而言）。我們將其設定為所需的 uniform_query_len 以用於統一解碼，否則將其設定為非統一解碼批次的 num_tokens。

CUDA 圖 Wrapper 不再管理熱身邏輯。熱身過程現在由 GPU 模型 Runner 直接控制，其中 NONE 執行時模式被分配用於熱身的即時執行。在為完整 CUDA 圖進行熱身時，在熱身 dummy_run 呼叫期間顯式執行 Attention 也很重要。

Attention 後端對 CUDA 圖的相容性¶

為了表示 Attention 後端對 CUDA 圖的相容性，我們引入了一種新的列舉型別 AttentionCGSupport，它是一種列舉型別，用於跟蹤 Attention 後端支援 CUDA 圖的能力。該值按能力順序排序，即 ALWAYS> UNIFORM_BATCH> UNIFORM_SINGLE_TOKEN_DECODE> NEVER。

class AttentionCGSupport(enum.Enum):
    """ Constants for the CUDA Graphs support of the attention backend
    Here we do not consider the cascade attention, as currently
    it is never CUDA Graphs supported."""

    ALWAYS = 3
    """CUDA Graphs always supported; supports mixed-prefill-decode"""
    UNIFORM_BATCH = 2
    """CUDA Graphs supported for batches the only contain query lengths that are
    the same, this can be used for spec-decode 
        i.e. "decodes" are 1 + num_speculative_tokens"""
    UNIFORM_SINGLE_TOKEN_DECODE = 1
    """CUDA Graphs supported for batches the only contain query_len==1 decodes"""
    NEVER = 0
    """NO CUDA Graphs support"""

假設我們有混合 Attention 後端（例如，在 Mamba 混合模型中）。在這種情況下，我們尋求所有後端的最低能力來確定模型的最終能力，並且我們可能會透過將模式降級到最適合的模式來解決不相容的 CUDA 圖模式。例如，如果最低能力是 UNIFORM_BATCH，則將 FULL 模式降級到 FULL_AND_PIECEWISE 模式；如果最低能力是 NEVER（對於 -O3 編譯模式），則降級到 PIECEWISE 模式。有關完整的回退策略，請參閱此處 _check_and_update_cudagraph_mode 的程式碼。

下表列出了撰寫本文時支援完整 CUDA 圖的後端。

Attention 後端	cudagraph_support	評論
FlashAttention v2	`UNIFORM_BATCH`	實際上是 `ALWAYS`，但出於效能考慮，回退到 `FULL_AND_PIECEWISE`
FlashAttention v3	`ALWAYS`	為批次提供了統一的例程，因此 `FULL` 模式很好。
Triton Attention	`ALWAYS`	推薦 `FULL_AND_PIECEWISE`，因為它為預填充/混合和純解碼批次提供了不同的核心。
AITER FlashAttention	`UNIFORM_BATCH`
FlashInfer	`UNIFORM_SINGLE_TOKEN_DECODE`	使用 Blackwell 上的 TRTLLM Attention 時將設定為 `UNIFORM_BATCH`
FlashMLA	`UNIFORM_BATCH`
FlashInferMLA	`UNIFORM_BATCH`
AITER MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
CUTLASS MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
Mamba attention	`UNIFORM_SINGLE_TOKEN_DECODE`

未列出的後端均宣告為 NEVER。

使用指南¶

現在 CLI 直接使用 cudagraph_mode 的大寫字串作為 compilation_config：--compilation-config '{"cudagraph_mode": "..."}'，其中 ... 應該是 NONE、PIECEWISE、FULL、FULL_DECODE_ONLY 和 FULL_AND_PIECEWISE 之一。請注意，所有 PIECEWISE 相關的模式都需要分段編譯，而所有 FULL 相關的模式都需要 Attention 後端對 CUDA 圖的支援。例如

vllm serve --model meta-llama/Llama-3.1-8B-Instruct --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}'

Python 示例¶

import os
os.environ.setdefault("VLLM_LOGGING_LEVEL", "DEBUG")

import vllm
from vllm.config import CUDAGraphMode

compilation_config = {"mode": 3, "cudagraph_mode": "FULL_AND_PIECEWISE"}
model = vllm.LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    dtype="auto",
    compilation_config=compilation_config,
)
sampling_params = vllm.SamplingParams(
    temperature=0,  # greedy decoding
    max_tokens=1024,
)
outputs = model.generate(
    ["My name is John and"],
    sampling_params=sampling_params,
)

分段編譯和完整圖自定義通道（Attention 融合，序列並行）¶

不幸的是，一些自定義編譯通道需要看到整個圖才能生效，因此與分段編譯不相容。這包括 AttnFusionPass 和 SequenceParallelismPass。作為短期解決方案，當啟用 Attention 融合時，我們自動停用分段編譯（透過設定 splitting_ops=[]）。我們使用 CUDA 圖模式 FULL 或 FULL_DECODE_ONLY（取決於後端支援）。然而，這會導致另一個最佳化不相容和令人困惑的效能權衡。

長期來看，我們增加了在 Inductor 中分割槽圖的能力，而不是在 Dynamo 之後立即分割槽。可以透過 CompilationConfig.use_inductor_graph_partition=True 啟用，但目前仍處於實驗階段，並且僅在 torch>=2.9 時可用。這也增加了編譯時間，因為它需要編譯整個圖，並且無法重用分段編譯的工件。一旦 vLLM 支援 2.9，我們計劃將其作為預設方法，因為它也將加速分段 CUDA 圖捕獲。

關於效能¶

有關示例，請參閱以下連結