[SGLang] torch.compile & Inductor: PyTorch 컴파일러 통합

2026년 4월 11일수정: 2026년 4월 11일

들어가며

SGLang v0.3에서 torch.compile 통합을 도입하며 Decode 처리량에서 최대 1.5배 성능 향상을 달성했다. PyTorch Inductor 백엔드를 통해 커널 퓨전, 메모리 접근 최적화, Triton 커널 자동 생성을 활용한다. SGLang은 여기에 Piecewise 컴파일 전략을 결합하여, 동적 형상의 Prefill에서도 컴파일 이점을 얻는다.

이 글에서는 python/sglang/srt/compilation/ 디렉터리의 compile.py, backend.py, compilation_config.py를 중심으로 torch.compile 통합을 분석한다.

컴파일 아키텍처 개요

SGLang의 컴파일 파이프라인은 4개의 핵심 컴포넌트로 구성된다.

┌──────────────────────────────────────────────────┐
│  install_torch_compiled()                         │
│  - 모델의 forward에 컴파일 래퍼 설치               │
└─────────────┬────────────────────────────────────┘
              │
              ▼
┌──────────────────────────────────────────────────┐
│  SGLangBackend                                    │
│  - torch.compile의 커스텀 백엔드                   │
│  - FX Graph를 받아 분할/최적화                     │
└─────────────┬────────────────────────────────────┘
              │
              ▼
┌──────────────────────────────────────────────────┐
│  CUDAPiecewiseBackend                             │
│  - 분할된 각 조각에 대해 CUDA Graph 캡처           │
│  - 배치 크기별 그래프 관리                         │
└─────────────┬────────────────────────────────────┘
              │
              ▼
┌──────────────────────────────────────────────────┐
│  CompilerManager                                  │
│  - Inductor/Eager 컴파일러 선택                   │
│  - 컴파일 캐시 관리                               │
└──────────────────────────────────────────────────┘

CompilationConfig: 분할 설정

CompilationConfig는 컴파일 대상 형상과 분할 지점을 정의한다.

class CompilationConfig:
    def __init__(self, capture_sizes, compiler="eager", enable_debug_mode=False):
        self.capture_sizes = capture_sizes
        self.compiler = compiler
        self.enable_debug_mode = enable_debug_mode
        self.split_ops = []
        self.split_ops.extend(SPLIT_OPS)

    def add_split_op(self, op: str):
        self.split_ops.append(op)

    def get_capture_sizes(self):
        return self.capture_sizes

split_ops는 FX Graph에서 분할 지점으로 사용될 연산을 등록한다. @register_split_op 데코레이터로 자동 등록할 수 있다.

SPLIT_OPS = []

def register_split_op(op_name=None):
    def decorator(op_func):
        name = op_name or op_func.__name__
        SPLIT_OPS.append(f"sglang.{name}")
        return op_func
    return decorator

install_torch_compiled: 컴파일 래퍼 설치

install_torch_compiled()는 모델의 forward 메서드를 컴파일된 버전으로 교체한다. 핵심은 "첫 호출 시 컴파일"하는 지연 전략이다.

def install_torch_compiled(module, *, dynamic_arg_dims=None,
                           backend_factory=None, compile_config=None,
                           fullgraph=True, graph_pool=None):
    unbound_fwd = module.__class__.forward
    original_code = unbound_fwd.__code__
    state = {"compiled": False, "compiled_callable": None}

    def _ensure_compiled(self, *args, **kwargs):
        if state["compiled"]:
            return
        # Dynamic dim 마킹
        sig = inspect.signature(unbound_fwd)
        ba = sig.bind(self, *args, **kwargs)
        for name, dims in (dyn_map or {}).items():
            if name in ba.arguments and ba.arguments[name] is not None:
                _mark_dynamic_on_value(ba.arguments[name], dims)

        # 컴파일 실행
        compiled_callable = torch.compile(
            bound, fullgraph=fullgraph, backend=backend_factory
        )
        compiled_callable(*args, **kwargs)  # Dynamo 트리거
        state["compiled"] = True
        state["compiled_callable"] = compiled_callable

트램폴린 패턴: 실행 시 Piecewise CUDA Graph 모드인지 확인하고, 컴파일된 코드와 원본 코드를 동적으로 전환한다.

    def trampoline(self, *args, **kwargs):
        use_compiled = is_in_piecewise_cuda_graph()
        if use_compiled:
            if not state["compiled"]:
                _ensure_compiled(self, *args, **kwargs)
            return state["compiled_callable"](*args, **kwargs)
        else:
            return unbound_fwd(self, *args, **kwargs)

    module.forward = types.MethodType(trampoline, module)

이 설계 덕분에 Piecewise CUDA Graph 캡처 중에만 컴파일된 코드가 실행되고, 일반 Prefill에서는 원본 코드가 실행된다.

SGLangBackend: 커스텀 백엔드

SGLangBackend는 torch.compile의 백엔드로, FX Graph를 받아 Piecewise 분할과 최적화를 수행한다.

def make_backend(graph, compile_config, inductor_config, graph_pool,
                 piecewise_compile_index, total_piecewise_compiles,
                 sym_shape_indices, compiled_graph_for_general_shape,
                 sglang_backend):
    backend_cls = CUDAPiecewiseBackend if not is_npu() else NPUPiecewiseBackend
    return backend_cls(
        graph, compile_config, inductor_config, graph_pool,
        piecewise_compile_index, total_piecewise_compiles,
        sym_shape_indices, compiled_graph_for_general_shape,
        sglang_backend,
    )

NPU(Ascend)와 CUDA를 백엔드 레벨에서 분리한다.

CompilerManager: Inductor vs Eager

CompilerManager는 컴파일러를 선택하고 캐시를 관리한다.

def make_compiler(config: CompilationConfig):
    if config.compiler == "eager":
        return EagerAdapter()
    elif config.compiler == "inductor":
        return InductorAdaptor()
    else:
        raise ValueError(f"Unknown compiler: {config.compiler}")

class CompilerManager:
    def __init__(self, config):
        self.cache = dict()
        self.compiler = make_compiler(config)

    def compile(self, graph, example_inputs, inductor_config,
                graph_index=0, num_graphs=1, runtime_shape=None):
        compilation_counter.num_backend_compilations += 1
        compiled_graph = self.compiler.compile(graph, example_inputs, ...)
        return compiled_graph

컴파일 캐시는 (runtime_shape, graph_index, compiler_name) 튜플로 키를 구성한다.

def load(self, graph, example_inputs, graph_index, runtime_shape=None):
    handle = self.cache[(runtime_shape, graph_index, self.compiler.name)]
    compiled_graph = self.compiler.load(handle, graph, example_inputs, ...)
    return compiled_graph

Dynamic Shape 마킹

torch.compile에서 동적 축을 올바르게 처리하기 위해, 함수 시그니처의 타입 어노테이션에서 동적 차원을 자동 추론한다.

def _infer_dynamic_arg_dims_from_annotations(forward_fn):
    sig = inspect.signature(forward_fn)
    dyn = {}
    for name, p in sig.parameters.items():
        ann = p.annotation
        if ann is torch.Tensor:
            dyn[name] = 0
        elif getattr(ann, "__name__", "") in ("IntermediateTensors",):
            dyn[name] = 0
    return dyn

Inductor 최적화 설정

set_torch_compile_config()에서 Inductor의 핵심 최적화를 활성화한다.

def set_torch_compile_config():
    import torch._inductor.config

    # 좌표 하강법 튜닝: 최적의 타일 크기 탐색
    torch._inductor.config.coordinate_descent_tuning = True
    # 고유 커널 이름: 프로파일링 용이
    torch._inductor.config.triton.unique_kernel_names = True
    # FX 그래프 캐시: 재컴파일 시간 단축
    torch._inductor.config.fx_graph_cache = True

    # Dynamo 캐시 크기 제한 완화
    torch._dynamo.config.accumulated_cache_size_limit = 1024

IntermediateTensors: PP 지원

Pipeline Parallel에서 스테이지 간 전달되는 텐서를 위한 래퍼이다.

@dataclass
class IntermediateTensors:
    tensors: dict[str, torch.Tensor]
    finished_sending: Optional[set[str]] = None
    finished_recving: Optional[set[str]] = None

    def __getitem__(self, key):
        if isinstance(key, str):
            return self.tensors[key]
        elif isinstance(key, slice):
            return self.__class__({k: v[key] for k, v in self.tensors.items()})

성능 효과

SGLang v0.3 블로그에 따르면, torch.compile + CUDA Graph 통합은 Llama-8B Decode에서 약 1.5배 처리량 향상을 달성했다. 핵심 요인은 다음과 같다.

최적화	효과
커널 퓨전	여러 소형 연산을 단일 커널로 합침
메모리 레이아웃	접근 패턴에 맞게 텐서 레이아웃 최적화
Triton 코드 생성	커스텀 CUDA 커널 대신 자동 생성
CUDA Graph 통합	컴파일된 커널의 런칭 오버헤드 제거

참고

SGLang 의 다른글

이전글 [SGLang] Model Loader: 가중치 로딩 인프라와 최적화
현재글 : [SGLang] torch.compile & Inductor: PyTorch 컴파일러 통합
다음글 [SGLang] Warmup: GPU 초기화와 JIT 사전 컴파일