torch.compile 통합

torch.compile 통합 (torch.compile integration)

vLLM의 V1 아키텍처에서는 torch.compile이 기본으로 켜져 있고, 프레임워크의 중요한 축이에요. 이 문서는 torch.compile 사용법을 이해하기 위한 간단한 예제를 따라가면서, vLLM이 어떻게 파이썬 코드를 컴파일하는지 실제 로그와 함께 보여줄게요.

전체 예제에서는 흔히 쓰는 Llama 모델을 실행하면서 디버그 레벨 로깅을 켜서 모든 상세 내역을 볼 거예요. 사용할 명령은 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의 compute_hash 함수)
  • PyTorch 설정(compiler_interface.py의 compute_hash 함수)
  • 모델의 forward 함수와 forward가 호출하는 관련 함수들

이 모든 요소를 고려하면 캐시를 써도 안전하다고 보장할 수 있어서, 캐시는 기본으로 켜져 있어요. 컴파일 과정을 디버깅하거나 캐시가 문제를 일으킨다고 의심되면 환경 변수 VLLM_DISABLE_COMPILE_CACHE=1로 끌 수 있어요.

vLLM의 torch.compile 통합에서 독특한 점 하나는, 요청을 서빙하기 전에 모든 컴파일이 끝남을 보장한다는 거예요. 어떤 요청도 새로운 컴파일을 트리거하지 않아요. 그렇지 않으면 엔진이 그 요청에서 블로킹되고 응답 시간에 예상치 못한 스파이크가 생기거든요.

기본적으로 캐시는 컴파일 산출물을 바이너리 파일로 저장해요. 생성된 코드를 디버깅 목적으로 직접 보고 싶다면 compilation config에서 compile_cache_save_format=unpacked를 설정하거나, 환경 변수 VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked를 설정하면 돼요.

동적 shape와 vllm guard dropping

torch.compile은 필요할 때 동적 shape에 대해 거침없이 guard를 다는 방식으로 설계됐어요. 그런데 vLLM은 이런 guard가 실질적인 경우가 많아서 guard를 버리는(drop) 방식이라 서로 충돌하죠.

torch.compile은 backed와 unbacked 두 종류의 동적 shape를 제공해요. backed 동적 shape는 guard가 붙고, guard가 추가되지 않을 거란 보장은 없어요. 유저 코드, dynamo, inductor, autograd 모두 guard를 추가할 수 있어요. 게다가 0/1 특수화의 경우 backed 심볼은 그런 범위에서 분기하지 않아도 무조건 0, 1, 또는 >=2로 특수화돼요.

반대로 unbacked 동적 shape는 guard가 붙지 않음이 보장되고 0/1 특수화도 없어요. 다만 그 값을 필요로 하는 분기를 만났는데 명시적 unbacked 처리가 정의돼 있지 않으면 데이터 의존 오류(data dependent error, DDE)가 날 수 있어요. 프레임워크는 DDE를 던지기보다 일반 경로를 고르는 방향으로 수렴하고 있어요. unbacked를 쓰면 성능 버그나 일반 경로 선택 때문에 최적화 기회를 놓칠 수 있고, 고정된 비-예제 입력 기반 힌트를 쓴다는 단점도 있어요(이건 곧 override_hint API로 해결 예정). 일반 경로를 고르는 예로는, 함수가 contiguous()와 reshape()에서 입력이 연속적임을 상징적으로 증명할 수 없을 때 clone을 도입하는 경우가 있어요.

backed_size_oblivious는 unbacked 처리가 정의된 곳마다 backed 심볼을 unbacked처럼 취급하게 하는 플래그예요. 이 모드에서는 프레임워크 코드에서 0/1 특수화를 대부분 피하고 기본 0/1 특수화도 발생하지 않아요. 다만 특히 유저 코드나 커스텀 패스 때문에 torch.compile이 여전히 guard할 수 있다는 보장은 없어요. backed_size_oblivious는 PyTorch 컴파일에서 실험적이라 폐기될 수 있어요. 그래도 backed보다 안전한 선택이고, unbacked보다 성능 저하 확률이 낮아요.

동적 shape 설정

DynamicShapesConfig를 type 필드로 설정하면 동적 shape 동작을 제어할 수 있어요. 세 가지 모드 중 하나를 고를 수 있죠: BACKED(기본), UNBACKED, BACKED_SIZE_OBLIVIOUS.

오프라인 추론 예시 (LLM 클래스 사용)

오프라인 추론에서 LLM 클래스를 쓸 때는 compilation_config 파라미터로 동적 shape를 설정할 수 있어요.

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 플래그로 동적 shape를 설정해요.

# 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 (기본): 최대 성능을 위해 잠재적으로 불안전한 guard 버리기를 감수할 수 있을 때 써요. guard가 불건전하게 추가됐다가 무시될 수 있어요.
  • UNBACKED: guard에 대한 가장 강한 보장이 필요할 때 써요. 가장 보수적이지만 일부 최적화 기회를 놓칠 수 있어요.
  • BACKED_SIZE_OBLIVIOUS: guard를 피하는 것과 성능 사이의 균형을 원할 때 써요. 실험적이지만 BACKED보다 안전하고 UNBACKED만큼 보수적이진 않아요.

파이썬 코드 컴파일

이건 Dynamo가 그래프를 캡처하는 파이썬 코드 컴파일 단계예요. xxx/vllm/model_executor/models/llama.py:339 코드로 함수를 트레이스하는데, 이게 우리가 컴파일하는 모델의 forward 함수예요. forward 실행 중에는 Dynamo가 인라인하는 다른 함수들도 호출돼요. 로그에서 보이듯 PyTorch의 일부 함수(torch/nn/modules/module.py — 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에 저장돼요.

계산 그래프 처리

계산 그래프는 모든 텐서에 shape 주석이 붙어 있어요. 입력은 input IDs, position IDs, 모델의 가중치와 버퍼이고, 출력은 최종 hidden states예요. lm head 투영과 샘플링 연산은 그래프에 포함되지 않아요.

계산 그래프 입력 대부분은 모델 가중치와 버퍼라 모델 수명 동안 변하지 않아 정적 shape를 가져요. 입력 IDs와 position IDs만 symbolic shape를 갖고, 배치마다 변할 수 있어요. 다만 같은 symbolic shape를 공유하죠. 즉 계산 그래프에서 유일하게 변하는 크기는 배치 크기(현재 forward에서 처리하는 토큰 수)예요.

어텐션 연산은 복잡하고 KV 캐시와 상호작용해야 해서 shape가 복잡해요. 다행히 어텐션 출력은 어텐션 입력 쿼리와 같은 shape를 공유해요. 그래서 어텐션 연산 전체를 PyTorch 커스텀 op torch.ops.vllm.unified_attention_with_output로 감싸서, Dynamo가 내부 연산을 하나도 들여다보지 않게 해요. 이렇게 하면 어텐션 연산이 복잡해도 Dynamo 관점에서는 모델 계산 그래프를 전체 그래프(full-graph)로 캡처할 수 있어요.

계산 그래프는 splitting_ops(보통 어텐션 연산)에 의해 조각으로 더 나뉘어요. 그래서 computation_graph.py 파일에 서브모듈이 여럿 보이는데, 각 서브모듈이 분할 후 그래프의 한 조각이에요.

  • 어텐션 연산 자체가 하나의 서브모듈
  • 한 어텐션 연산에서 다음 어텐션 연산까지의 계산 그래프 부분이 하나의 서브모듈

각 서브모듈은 인덱스로 식별되고 개별적으로 처리돼요.

계산 그래프 컴파일

어텐션으로 그래프를 분할하면 3개의 고유 서브그래프가 생겨요.

  • 어텐션 이전의 첫 레이어
  • 한 어텐션 연산에서 다음 어텐션 연산까지의 각 중간 레이어
  • 어텐션 이후의 마지막 레이어

캐시 디렉터리가 이미 있으면(예: 같은 코드를 두 번째 실행) Inductor 컴파일을 완전히 건너뛰고 디스크에서 이전 컴파일 산출물을 로드해요.

위 예제는 일반 shape(symbolic shape)에 대해 Inductor로 컴파일한 거예요. 특정 shape에 대해서도 컴파일할 수 있어요.

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

이러면 배치 크기 1, 2, 4, 8에 특화된 커널도 컴파일해요. 이때는 계산 그래프의 모든 shape가 정적이고 알려져 있어서, 최대 성능을 위해 오토튜닝을 켜요. 처음 실행할 땐 느릴 수 있지만, 다음 실행부터는 튜닝을 건너뛰고 튜닝된 커널을 바로 실행해요.

모든 shape가 알려지면 torch.compile은 다른 설정들을 비교해 더 나은 커널 설정을 찾는 경우가 많아요. 예를 들어 shape 8x2048x3072의 행렬 곱셈에서 triton 템플릿을 여러 설정으로 시도하는데, 기본 코드(cublas 라이브러리로 디스패치)보다 훨씬 빠른 설정을 찾아요.

아쉽게도 오토튜닝은 꽤 오래 걸리기 때문에(모델 크기와 배치 크기에 따라 몇 초에서 몇 분), 나중에 캐시할 수 있음에도 사용자 친화성을 위해 기본은 꺼져 있어요. 최대 성능이 필요하면 특정 shape를 컴파일해서 오토튜닝을 시도하는 걸 추천해요.

Cudagraph 캡처

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의 일부로 포함시킬 수도 있어요. 이러면 작은 모델이나 MoE의 디코드 속도 같은 경우에 성능이 좋아질 수 있어요. 자세한 내용은 CUDA Graphs 문서를 참고하세요.

더 알아보기 (Learn more)