CUDA 그래프 모드

CUDA 그래프 모드 (CUDA Graphs)

CUDA 그래프(CUDA Graphs)는 GPU 커널 실행 전체를 한 번에 기록해 두고 재사용하는 기법이에요. 커널을 하나하나 CPU에서 디스패치하면 그 사이사이 지연이 쌓이는데, CUDA 그래프로 묶어 두면 반복되는 작업을 짧은 시간에 재생할 수 있어요. vLLM의 V1 아키텍처에서는 torch.compile과 함께 이 CUDA 그래프가 추론 성능의 핵심 축을 담당해요.

이 문서에서는 vLLM V1의 새로운 CUDA 그래프 모드가 왜 생겼는지, 어떤 모드들이 있는지, 어떻게 동작하는지, 그리고 실전에서 어떻게 쓰는지를 순서대로 살펴볼게요.

왜 이번에 CUDA 그래프 모드를 다시 설계했을까요

처음에는 조각별(piecewise) 컴파일 방식을 써서 CUDA 그래프와 호환되지 않는 연산(주로 어텐션)을 제외하고 캡처했어요. 이 방식은 모든 어텐션 백엔드와 호환되면서도 어느 정도 가속을 얻을 수 있었죠. 이후 어텐션이 CUDA 그래프를 지원하는 경우를 위해 "전체(full) CUDA 그래프"를 추가했는데, 문제는 컴파일과 캡처가 너무 단단하게 붙어 있다 보니 전부 아니면 전무(all-or-nothing) 식의 유연하지 못한 구조가 되었어요.

게다가 많은 어텐션 백엔드가 통합된 "전체" 캡처를 아직 지원하지 않았고(그때 기준 FlashAttention 3만 지원), 일부는 순수 디코드 배치에서만 CUDA 그래프를 지원했어요(FlashInfer, FlashMLA, Mamba 등). 결과적으로 지원 여부가 백엔드마다 제각각이라 성능·호환성 트레이드오프가 헷갈리고 코드도 점점 복잡해졌죠.

그래서 vLLM은 다음과 같은 특징을 가진 더 세밀한 CUDA 그래프 솔루션을 만들기로 했어요.

프리필/혼합 배치와 (균일) 디코드 배치를 각각 구분해서 CUDA 그래프를 캡처
캡처 로직을 컴파일에서 최대한 분리 (같은 컴파일 결과로 조각별·전체 그래프 모두 캡처, 컴파일 없이도 전체 캡처)
배치 구성에 따라 런타임에 전체/조각별 그래프를 자동 선택
CUDA 그래프 동작을 중앙에서 관리해 코드 복잡도를 줄이고 확장성을 높임

이렇게 하면 시작 시간과 성능, 기능 지원 사이에서 훨씬 유연하게 트레이드오프를 고를 수 있어요.

CUDA 그래프 모드 (CudagraphModes)

CompilationConfig.cudagraph_mode가 바로 조절하는 단일 노브(knob)예요. vLLM은 이 값으로 CUDA 그래프 동작을 통째로 바꿔요.

NONE — CUDA 그래프를 완전히 끔. 디버깅할 때 유용해요.
PIECEWISE — 단일 모드 전략이자 예전 기본값이에요. 어텐션이나 CUDA 그래프와 호환되지 않는 연산은 그대로(eager) 두고, 나머지만 CUDA 그래프로 넣어요. 조각별 컴파일이 필요해요.
FULL — 단일 모드 전략이에요. 균일하지 않은 배치(프리필·혼합)에 대해서만 전체 CUDA 그래프를 캡처하고, 균일 디코드 배치는 같은 배치 크기의 균일하지 않은 배치용 그래프를 재사용해요. 작은 모델이나 프롬프트가 짧은 워크로드에 좋아요.
FULL_DECODE_ONLY — 균일 디코드에만 전체 CUDA 그래프를 쓰고 프리필·혼합에는 안 써요. 프리필이 덜 중요한 P/D(프리필/디코드 분리) 구성의 디코드 인스턴스에 적합해요. 덕분에 PIECEWISE 그래프용 메모리를 아낄 수 있어요.
FULL_AND_PIECEWISE — 기본값이예요. 균일 디코드는 전체 CUDA 그래프, 나머지는 조각별 그래프를 써요. 일반적으로 가장 성능이 좋은 설정인데, 특히 작은 모델이나 MoE에서 낮은 지연을 만들 때 효과적이에요. 다만 메모리를 가장 많이 쓰고 캡처 시간도 가장 길어요.

기본 동작은 이렇게 정해져요. V1에서 조각별 컴파일을 쓰면 성능을 위해 FULL_AND_PIECEWISE가 기본이고, (pooling 모델은 여전히 PIECEWISE) 조각별 컴파일이 불가능한 환경이라면 NONE이 기본이 되죠.

여기서 NONE, PIECEWISE, FULL은 단일 모드 구성이라 과거의 eager 실행, 조각별 CUDA 그래프, 전체 CUDA 그래프와 각각 같아요. 반면 FULL_DECODE_ONLY와 FULL_AND_PIECEWISE는 새로 추가된 이중 모드 구성이라, 런타임 배치에 따라 실제 런타임 모드를 골라주는 디스패처가 필요해요.

캐스케이드(cascade) 어텐션은 CUDA 그래프와 호환되지 않지만, 모든 CUDA 그래프 모드 구성과는 호환되도록 처리돼요. 캐스케이드 어텐션을 쓰는 배치는 항상(가능하면) PIECEWISE 모드로 보내지고, 아니면 NONE으로 처리돼요.

주의할 점이 하나 있어요. 모든 CUDA 그래프 모드가 모든 어텐션 백엔드와 호환되는 건 아니에요. 그래서 vLLM은 자동으로 가장 가까운 지원 모드로 "다운그레이드"해요. 예를 들어 백엔드가 순수 디코드(균일) 배치에만 CUDA 그래프를 지원한다면, 조각별 컴파일이 켜져 있을 때 FULL을 FULL_AND_PIECEWISE로, 아니면 FULL_DECODE_ONLY로 바꿔요.

상세 설계

새 CUDA 그래프 로직은 조각별 컴파일 위에 올라가 있고, 이중 CUDA 그래프 런타임 모드 전환을 지원해요. 핵심 구성 요소는 네 가지예요.

CUDAGraphWrapper — 감싼 호출 가능 객체에 CUDA 그래프 캡처·재생을 붙여주는 래퍼
CudagraphDispatcher — CUDA 그래프에 대한 단일 정보 원천(source of truth)이자, 그래프 간 디스패치를 처리하는 중앙 컨트롤러
CUDAGraphMode — 위에서 소개한 지원·런타임 모드를 표현하는 enum
BatchDescriptor — 디스패치에 쓰이는 런타임 배치의 고유한 표현

예전에는 CUDA 그래프 로직과 컴파일 로직이 PiecewiseBackend에 단단히 묶여 있었고 CUDA 그래프가 batch_size 묵시적으로 디스패치됐어요. 지금은 CUDA 그래프 로직이 CUDAGraphWrapper 클래스로 분리되어 전체·조각별 능력을 모두 담고, 디스패치는 런타임 모드 + BatchDescriptor(디스패치 키) 를 통해 CudagraphDispatcher가 명시적으로 수행해요.

BatchDescriptor

BatchDescriptor는 ForwardContext 안에 있는, 런타임 디스패치 키의 핵심 구조예요. 프로토타입은 이렇게 생겼어요.

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

여기서 num_tokens은 패딩된 토큰 길이이고, uniform은 모든 요청의 쿼리 길이가 같은지 여부예요. 많은 어텐션 백엔드는 배치가 균일할 때만 전체 CUDA 그래프를 지원해요. 순수 디코드 배치는 균일하지만 쿼리 길이가 꼭 1은 아닐 수 있는데(즉 num_tokens == num_reqs), 이건 스펙 디코드(spec-decode)의 검증 단계에서 "디코드" 배치가 쿼리 길이 1+num_spec_tokens를 가질 때 생기는 경우예요.

이 구조의 목표는 (패딩된) 배치를 CUDA 그래프 항목에 대응하는 최소한의 정보로 고유하게 식별하는 거예요.

CudagraphDispatcher

CudagraphDispatcher는 FULL 런타임 모드용 디스패치 키 집합과 PIECEWISE 런타임 모드용 키 집합, 두 벌을 유지하고, 모델 forward 실행 전에 올바른 런타임 모드와 디스패치 키를 골라줘요. 초기 키(패딩 입력에 대한 대략적 batch descriptor)를 받아 선택된 런타임 모드와 최종 batch descriptor를 반환하고, 그 결정을 forward 컨텍스트를 통해 CUDAGraphWrapper 인스턴스들에 알려줘요.

핵심은 CudagraphDispatcher가 가용 CUDA 그래프 키의 유일한 정보 원천이라는 점이에요. 그래서 CUDAGraphWrapper 인스턴스들은 forward 컨텍스트가 알려주는 그래프 선택을 그대로 신뢰할 수 있고, 덕분에 래퍼 코드는 단순해지고 로직은 디스패처에 모일 수 있어요.

디스패치 키는 initialize_cudagraph_keys 메서드로 초기화되는데, 이 메서드는 가능한 모든 어텐션 백엔드가 초기화된 뒤 gpu_model_runner가 호출해요. 디스패치 코드는 대략 이렇게 생겼어요.

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() 내부에서는 우선순위 FULL > PIECEWISE > None 순서로 기존 키를 찾아 반환해요. 디스패치 키가 없으면 eager 실행용 NONE 모드를 반환하죠.

CUDAGraphWrapper

CUDAGraphWrapper 인스턴스는 실행 가능한 객체를 감싸면서 CUDA 그래프 기능만 덧붙인 클래스예요. 각 래퍼 인스턴스는 특정 runtime_mode(PIECEWISE나 FULL로 제한)에 묶여 있고, 캡처/재생과 통과(directly calling)를 책임져요. 런타임에 각 래퍼는:

전역 forward 컨텍스트에서 runtime_mode와 batch_descriptor(디스패치 키)를 확인
runtime_mode가 NONE이거나 래퍼의 모드와 다르면 그냥 직접 호출
그렇지 않으면(래퍼 모드와 일치하면) CUDA 그래프 캡처(키가 없으면 새로 만들어 캐시) 또는 재생(키가 캐시에 있으면)

이 단계들은 래퍼가 forward 컨텍스트(디스패처가 제어)를 그대로 신뢰한다는 가정 위에 있어요. 덕분에 로직이 단순해지고 중앙화되며 래퍼와 디스패처 사이 상태 불일치 위험도 줄어들고, 같은 래퍼 클래스를 FULL과 PIECEWISE 두 모드에 모두 재사용할 수 있어요.

중첩 래퍼 설계

전체 CUDA 그래프와 조각별 CUDA 그래프가 공존하면서 서로 호환되도록 만드는 핵심 메커니즘이 중첩(nested) 래퍼 설계예요. 단일한 조각별 FX 그래프 위에 올라가 있는데요, 전체 모델 바깥에는 FULL 모드 래퍼를 두고 전체 CUDA 그래프 기능을 담당하게 해요. 그러면서 컴파일 안의 각 조각별 백엔드는 PIECEWISE 모드 래퍼로 감싸요.

그래서 FULL 런타임 모드에서는 조각별 래퍼가 활성화되지 않으므로 전체 CUDA 그래프를 캡처/재생해도 안전해요. PIECEWISE 모드에서도 FULL 래퍼와 PIECEWISE 래퍼가 충돌하지 않죠. NONE 런타임 모드에서는 두 래퍼 모두 활성화되지 않아 그냥 eager 실행으로 빠져요.

전체 CUDA 그래프 캡처와 워밍업

CUDA 그래프 캡처는 러너가 NONE이 아닌 런타임 모드로 모델 forward를 처음 호출할 때(_dummy_run) 일어나요. 전체 CUDA 그래프 캡처에서는 프리필/혼합 배치나 균일 디코드 배치를 구분하기 위해 어텐션 메타데이터를 적절히 설정해서, 기반 어텐션 백엔드가 원하는 커널 루틴을 실행하게 만들어요. 배치 유형을 구분하는 가장 중요한 속성은 attn_metadata의 max_query_len이에요(대부분의 어텐션 백엔드가 그렇죠). 균일 디코드용으로는 원하는 uniform_query_len을, 그 외 배치에서는 num_tokens을 설정해요.

워밍업 로직은 이제 래퍼가 관리하지 않고 GPU 모델 러너가 직접 제어해요. 워밍업 때는 NONE 런타임 모드를 할당해 eager 실행을 수행하죠. 전체 CUDA 그래프를 워밍업할 때는 워밍업 dummy_run 호출에서 어텐션도 명시적으로 실행하는 게 중요해요.

어텐션 백엔드의 CUDA 그래프 호환성

어텐션 백엔드의 CUDA 그래프 지원 수준을 나타내기 위해 AttentionCGSupport라는 새 enum 타입을 도입했어요. 값은 지원 능력 순서대로 정렬돼요. 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 that 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 that only contain query_len==1 decodes"""
    NEVER = 0
    """NO CUDA Graphs support"""

하이브리드 어텐션 백엔드(mamba mixer 모델 같은)를 쓰는 경우에는 모든 백엔드의 최소 능력을 기준으로 모델의 최종 능력을 정하고, 호환되지 않는 CUDA 그래프 모드는 그에 맞는 최적 모드로 다운그레이드해요. 예를 들어 최소 능력이 UNIFORM_BATCH면 FULL을 FULL_AND_PIECEWISE로, NEVER면 PIECEWISE로 바꾸는 식이죠.

이 글을 쓰는 시점에 전체 CUDA 그래프를 지원하는 백엔드는 다음과 같아요.

Attention Backend	cudagraph_support	Comments
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 어텐션을 쓰면 `UNIFORM_BATCH`로 설정
FlashMLA	`UNIFORM_BATCH`
FlashInferMLA	`UNIFORM_BATCH`
FlashInferMLASparse	`UNIFORM_BATCH`
AITER MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
CUTLASS MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
Mamba attention	`UNIFORM_SINGLE_TOKEN_DECODE`

표에 없는 백엔드는 모두 NEVER로 선언돼요.

사용 가이드

CLI에서는 compilation_config에 cudagraph_mode의 대문자 문자열을 바로 써요. --compilation-config '{"cudagraph_mode": "..."}'에서 ...는 NONE, PIECEWISE, FULL, FULL_DECODE_ONLY, FULL_AND_PIECEWISE 중 하나예요. 모든 PIECEWISE 관련 모드는 조각별 컴파일이 필요하고, 모든 FULL 관련 모드는 어텐션 백엔드의 CUDA 그래프 지원이 필요하다는 점만 기억하세요. 예를 들면:

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

파이썬 예시

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 fusion, sequence parallelism)

아쉽게도 일부 커스텀 컴파일 패스는 전체 그래프를 봐야 효과가 있어서 조각별 컴파일과 호환되지 않아요. 여기에는 AttnQuantFusionPass와 SequenceParallelismPass가 포함돼요. 단기 해결책으로 attention fusion이 켜지면(splitting_ops=[] 설정으로) 조각별 컴파일을 자동으로 끄고, CUDA 그래프 모드는 FULL이나 FULL_DECODE_ONLY(백엔드 지원 여부에 따라)를 사용해요. 다만 이것도 또 다른 최적화 비호환성과 헷갈리는 성능 트레이드오프를 만들 수 있어요.

장기적으로는 Dynamo 직후가 아니라 Inductor에서 그래프를 분할하는 기능을 추가했어요. CompilationConfig.use_inductor_graph_partition=True로 켤 수 있는데, 현재는 실험적이고 torch>=2.9에서만 동작해요. 이 방식은 전체 그래프를 컴파일해야 해서 조각별 컴파일 산출물을 재사용하지 못하므로 컴파일 시간도 늘어나요. vLLM이 2.9를 지원하게 되면 이 방식을 기본으로 할 계획이에요. 조각별 CUDA 그래프 캡처도 빨라지거든요.

성능에 관해

성능 비교 예시는 공식 문서에서 PR 20059의 댓글 링크들을 참고하면 돼요. 모델 종류와 배치 구성에 따라 어느 모드가 좋은지가 크게 달라지니, 실제 워크로드로 벤치마크해 보는 걸 권해요.

더 알아보기 (Learn more)

vLLM 공식 문서: CUDA Graphs
관련 문서: torch.compile 통합, 멀티모달 ViT CUDA 그래프