본문으로 건너뛰기

[vLLM] torch.compile 통합: PyTorch 컴파일러

들어가며

vLLM은 PyTorch의 torch.compile을 깊이 통합하여 모델 실행을 최적화한다. 단순히 torch.compile(model)을 호출하는 수준이 아니라, 커스텀 백엔드(VllmBackend), 그래프 분할(Piecewise compilation), 캐싱 전략을 자체 구현하여 LLM 서빙에 특화된 컴파일 파이프라인을 갖추고 있다.

소스 경로: vllm/compilation/

공식 문서

vLLM 공식 문서: Torch Compile

핵심 구조/코드 분석

컴파일러 인터페이스

# vllm/compilation/backends.py
def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
    if compilation_config.backend == "inductor":
        if envs.VLLM_USE_STANDALONE_COMPILE and hasattr(
            torch._inductor, "standalone_compile"
        ):
            return InductorStandaloneAdaptor(...)
        else:
            return InductorAdaptor()
    elif compilation_config.backend == "eager":
        return EagerAdaptor()
    else:
        compiler = resolve_obj_by_qualname(
            current_platform.get_compile_backend()
        )()
        return compiler

vLLM은 CompilerInterface 추상 클래스를 통해 여러 백엔드를 지원한다:

  • InductorAdaptor: PyTorch의 기본 Inductor 컴파일러
  • InductorStandaloneAdaptor: 독립 실행 모드 Inductor (더 빠른 AOT 컴파일)
  • EagerAdaptor: 컴파일 없이 즉시 실행 (디버깅용)
  • 커스텀 플랫폼별 백엔드

CompilerManager: 캐시 관리

class CompilerManager:
    def __init__(self, compilation_config):
        self.cache: dict[tuple[Range, int, str], Any] = dict()
        self.compiler = make_compiler(compilation_config)

    def initialize_cache(self, cache_dir, disable_cache=False, ...):
        self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
        if not disable_cache and os.path.exists(self.cache_file_path):
            with open(self.cache_file_path) as f:
                cache = ast.literal_eval(f.read())

컴파일 결과를 (compile_range, graph_index, backend_name) 키로 캐싱한다. 캐시는 Python 파일로 직렬화되며, ast.literal_eval()로 안전하게 파싱한다. 이를 통해 서버 재시작 시 컴파일 시간을 절약한다.

Piecewise 컴파일

# vllm/compilation/piecewise_backend.py
def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
    """Symbolic dim을 구체적 size로 치환한 Fake 입력 생성"""
    def concretize(sym_val):
        if not is_symbolic(sym_val):
            return int(sym_val)
        expr = sym_val.node.expr
        return int(expr.subs({s: size for s in expr.free_symbols}))

vLLM의 Piecewise 컴파일은 전체 모델 그래프를 여러 조각으로 분할하여 각각 독립적으로 컴파일한다. 이는 CUDA 그래프와 결합할 때 중요한데, 각 조각 사이에 Python 코드(KV 커넥터의 레이어별 로드/세이브 등)를 실행할 수 있기 때문이다.

입력 버퍼 복사 최적화

def make_copy_and_call(
    sym_tensor_indices: list[int],
    input_buffers: list[torch.Tensor | None],
    callable_fn: Callable[..., Any],
) -> Callable[..., Any]:
    def copy_and_call(*args):
        list_args = list(args)
        for i, index in enumerate(sym_tensor_indices):
            runtime_tensor = list_args[index]
            runtime_shape = runtime_tensor.shape[0]
            if input_buffers[i] is None:
                input_buffers[i] = runtime_tensor.clone()
            static_tensor = input_buffers[i][:runtime_shape]
            static_tensor.copy_(runtime_tensor)
            list_args[index] = static_tensor
        return callable_fn(*list_args)
    return copy_and_call

CUDA 그래프는 정적 메모리 주소를 요구하므로, 런타임에 동적 텐서를 정적 버퍼에 복사한 뒤 컴파일된 함수를 호출하는 래퍼를 생성한다. 지연 초기화(lazy init)로 첫 호출 시에만 버퍼를 할당한다.

왜 이 설계인가

  1. 서빙 특화 컴파일: 학습과 달리 서빙에서는 배치 크기가 매 스텝 변한다. vLLM은 미리 정해진 크기 범위(Range)별로 각각 컴파일하여 이 문제를 해결한다.

  2. Piecewise + CUDA Graph 조합: FULL 모드에서는 전체 모델을 하나의 CUDA 그래프로 캡처하지만, KV Connector 같은 비동기 연산이 필요하면 PIECEWISE 모드로 전환하여 조각 사이에 Python 코드를 삽입한다.

  3. AOT 캐싱: compile_cache_save_format 옵션으로 컴파일 결과를 디스크에 저장하고 재사용할 수 있다. 대규모 모델의 경우 컴파일 시간이 수 분이 될 수 있으므로 이 캐싱은 프로덕션 환경에서 필수적이다.

참고

댓글

관련 포스트

vLLM 의 다른글