본문으로 건너뛰기

[sglang] [성능 최적화] Wan2.2 모델을 위한 최적의 torch.compile 모드 찾기: 왜 'default'가 더 빠를까?

PR 링크: sgl-project/sglang#28304 상태: Merged | 변경: +53 / -15

들어가며

딥러닝 추론 엔진에서 성능을 쥐어짜기 위해 가장 먼저 고려하는 기술 중 하나는 torch.compile입니다. PyTorch 2.0부터 도입된 이 기능은 모델의 그래프를 분석하여 Triton 커널로 컴파일함으로써 실행 속도를 비약적으로 향상시킵니다. 하지만 torch.compile에는 여러 가지 '모드(mode)'가 존재하며, 흔히 max-autotune이 가장 빠를 것이라고 오해하곤 합니다.

이번 글에서는 sgl-project/sglang 레포지토리에 반영된 [perf] Use default torch compile mode for Wan2.2 T2V A14B PR을 통해, 최신 비디오 생성 모델인 Wan2.2 T2V A14B에서 왜 max-autotune-no-cudagraphs보다 default 모드가 더 나은 성능을 보이는지, 그리고 이를 시스템적으로 어떻게 유연하게 관리하도록 개선했는지 분석해 보겠습니다.

코드 분석: 하드코딩에서 설정 기반으로

기존 SGLang의 DiT(Diffusion Transformer) 구현에서는 컴파일 모드가 환경 변수에 의존하거나 특정 값으로 고정되어 있었습니다. 이번 PR은 이를 모델별 설정(ModelConfig)에서 제어할 수 있도록 구조를 개선했습니다.

1. DiTConfig의 확장성 확보

먼저, 모든 DiT 모델의 기반이 되는 DiTConfig 클래스에 torch_compile_mode 필드를 추가했습니다.

Before:

class DiTConfig(ModelConfig):
    # ... 기존 필드들
    prefix: str = ""
    quant_config: QuantizationConfig | None = None

After:

class DiTConfig(ModelConfig):
    # ... 기존 필드들
    prefix: str = ""
    quant_config: QuantizationConfig | None = None
    # 기본값을 max-autotune-no-cudagraphs로 설정
    torch_compile_mode: str = "max-autotune-no-cudagraphs"

이렇게 기본값을 지정함으로써 기존 모델들과의 하위 호환성을 유지하면서도, 특정 모델에서만 이 값을 덮어쓸 수 있는 구조를 만들었습니다.

2. Wan2.2 모델에 최적화된 모드 적용

이번 PR의 핵심인 Wan2.2 T2V A14B 모델 설정입니다. 벤치마크 결과에 따라 이 모델은 default 모드를 사용하도록 명시적으로 변경되었습니다.

Before (Implicit): 별도의 설정이 없어 시스템 기본값인 max-autotune-no-cudagraphs를 사용함.

After:

# python/sglang/multimodal_gen/configs/pipeline_configs/wan.py
class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
    # ...
    def __post_init__(self) -> None:
        self.dit_config.boundary_ratio = self.boundary_ratio
        # Wan2.2 모델에 대해 default 모드 강제 적용
        self.dit_config.torch_compile_mode = "default"

3. 런타임에서의 동적 모드 선택 로직

설정값이 실제 torch.compile 호출 시점에 어떻게 반영되는지 보겠습니다. denoising.py와 같은 런타임 스테이지에서 환경 변수와 설정값을 조합하여 최종 모드를 결정합니다.

Before:

mode = os.environ.get(
    "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
)

After:

# python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
dit_config = getattr(self.server_args.pipeline_config, "dit_config", None)
mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE") or getattr(
    dit_config,
    "torch_compile_mode",
    "max-autotune-no-cudagraphs",
)
compile_kwargs["mode"] = mode

이 로직은 우선순위를 명확히 합니다:

  1. 사용자가 명시적으로 설정한 환경 변수(SGLANG_TORCH_COMPILE_MODE)
  2. 모델 설정 파일에 정의된 값(dit_config.torch_compile_mode)
  3. 최후의 보루인 기본값(max-autotune-no-cudagraphs)

왜 이게 좋은 최적화인가?

1. 놀라운 성능 향상 수치

PR 작성자가 제공한 벤치마크 결과는 매우 흥미롭습니다. H100과 B200 GPU 모두에서 default 모드가 더 우수한 결과를 보였습니다.

H100 (4 GPUs) 결과:

  • Denoise 속도: 217.8s → 206.0s (+5.70% 향상)
  • Peak Memory: 15.23 GB → 14.00 GB (1.23 GB 절감)

B200 (4 GPUs) 결과:

  • Denoise 속도: 110.3s → 105.0s (+4.96% 향상)
  • Peak Memory: 19.04 GB → 18.00 GB (1.04 GB 절감)

2. 'Max-Autotune'의 함정

일반적으로 max-autotune은 다양한 커널 구성을 시도하여 가장 빠른 것을 선택하므로 최상의 성능을 낼 것으로 기대됩니다. 하지만 Wan2.2와 같이 파라미터가 14B에 달하는 거대 모델에서는 다음과 같은 이유로 default가 유리할 수 있습니다.

  • Compilation Overhead: max-autotune은 컴파일 시간이 매우 길고 메모리 소모가 큽니다. 때로는 최적화된 커널을 찾는 과정에서 발생하는 오버헤드가 실제 추론 이득보다 클 수 있습니다.
  • Memory Efficiency: 벤치마크에서 보이듯 default 모드는 메모리 사용량이 약 1GB 더 적습니다. 이는 대규모 모델 추론 시 OOM(Out of Memory) 위험을 줄이고 더 큰 배치 사이즈를 확보할 수 있게 해줍니다.
  • Inductor의 한계: 특정 아키텍처(DiT)나 연산 패턴에서는 Inductor가 max-autotune 모드에서 생성하는 복잡한 커널보다 default 모드의 표준적인 최적화가 하드웨어(H100/B200)의 Tensor Core를 더 효율적으로 활용할 수도 있습니다.

3. 유연한 아키텍처 설계

단순히 값을 바꾼 것에 그치지 않고, DiTConfig를 통해 모델별로 최적의 컴파일 전략을 다르게 가져갈 수 있는 구조를 만든 점이 훌륭합니다. 예를 들어, PR 내에서 LTX-2 모델 역시 default 모드를 사용하도록 함께 업데이트되었습니다. 이는 향후 새로운 모델이 추가될 때마다 매번 런타임 코드를 수정하지 않고도 설정만으로 최적화를 수행할 수 있음을 의미합니다.

결론

이번 PR은 "무조건 가장 높은 옵션이 최고는 아니다"라는 엔지니어링의 기본 원칙을 다시 한번 일깨워줍니다. 실제 벤치마크를 통해 데이터에 기반한 결정을 내리고, 이를 시스템적으로 유연하게 수용할 수 있도록 리팩토링하는 과정이 인상적입니다.

SGLang과 같은 고성능 추론 엔진을 다루는 개발자라면, torch.compile 적용 시 반드시 다양한 모드에 대한 프로파일링을 병행해야 함을 시사합니다.

References

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글