[sglang] SGLang, GPU 간 VAE 디코딩 최적화를 통한 이미지 생성 속도 향상

PR 링크: sgl-project/sglang#28071 상태: Merged | 변경: +2766 / -1415

들어가며

이미지 및 비디오 생성 모델에서 VAE(Variational Autoencoder) 디코딩은 전체 생성 과정에서 상당한 시간을 차지할 수 있습니다. 특히 고해상도 이미지나 비디오 시퀀스를 생성할 때 VAE 디코딩 병목 현상은 사용자 경험을 저하시키는 주요 원인이 됩니다. SGLang의 최신 PR은 이러한 VAE 디코딩 과정을 최적화하여 여러 GPU에 걸쳐 효율적으로 처리함으로써, 이미지 및 비디오 생성 속도를 크게 향상시키는 것을 목표로 합니다.

이 PR은 기존의 제한적인 병렬 디코딩 방식을 개선하여, 다양한 병렬 처리 설정(TP-only, SP-only, CFG-only, 혼합 DiT 등)에서도 VAE 디코딩을 효율적으로 수행할 수 있도록 새로운 공유 공간 분할(spatial-shard) 디코딩 경로를 도입했습니다. 또한, 모든 VAE 모델에 대해 무조건적인 최적화를 적용하는 대신, 실제 성능 향상이 발생하는 경우에만 공간 분할 디코딩을 활성화하는 'auto' 정책을 도입하여 범용성과 효율성을 높였습니다.

이번 글에서는 이 PR의 주요 변경 사항을 살펴보고, 각 변경이 왜 성능 향상에 기여하는지, 그리고 이 최적화가 가지는 일반적인 교훈은 무엇인지 코드 예시와 함께 자세히 알아보겠습니다.

코드 분석

1. runtime/layers/spatial_parallel.py (새로운 공통 유틸리티)

이 PR은 VAE 디코딩을 위한 새로운 공통 로직을 runtime/layers/spatial_parallel.py 파일에 구현했습니다. 이 파일은 다음과 같은 핵심 기능들을 제공합니다:

• Height Split/Gather Helpers: VAE의 잠재 표현(latent representation)을 높이(height) 기준으로 분할하고 다시 모으는 유틸리티 함수를 제공합니다. 이는 여러 GPU에 걸쳐 VAE 디코딩을 병렬로 수행할 때 필수적인 연산입니다.
• Variable-Height Gather: 분할된 데이터의 높이가 다를 경우에도 올바르게 모을 수 있는 기능을 지원합니다.
• Halo Exchange: 공간 분할 시 경계 영역의 데이터를 주고받는 'halo exchange' 메커니즘을 구현합니다. 이는 이미지 생성 시 경계 부분의 연속성을 보장하는 데 중요합니다.
• Spatial 2D/3D Conv Wrappers: 기존의 2D 및 3D 컨볼루션 연산을 공간 분할 환경에 맞게 래핑합니다. 이를 통해 각 GPU는 자신이 담당하는 부분의 잠재 표현에 대해서만 컨볼루션 연산을 수행하게 됩니다.
• Rank-Aware Padding: 병렬 처리 시 각 랭크(GPU)에서 발생하는 패딩을 올바르게 처리합니다.

이러한 공통 유틸리티의 도입은 모델별로 중복되는 공간 분할 로직을 제거하고, 코드 재사용성을 높여 유지보수를 용이하게 합니다.

2. diffusers_spatial.py (Diffusers 스타일 VAE 적응)

diffusers 라이브러리 스타일의 VAE 디코더를 사용하는 모델들을 위해, 이 PR은 diffusers_spatial.py 파일을 통해 기존 모듈을 새로운 공간 분할 경로에 맞게 조정합니다.

• Conv2d 모듈 교체: Conv2d 모듈을 공간 분할 환경에 맞는 새로운 래퍼로 교체합니다.
• GroupNorm/Attention 래핑: Group Normalization 및 Attention 모듈도 'gather-then-reshard' 의미론을 적용하여 공간 분할 디코딩 과정에 통합합니다.

이를 통해 diffusers 기반의 다양한 VAE 모델들이 새로운 최적화된 디코딩 경로를 활용할 수 있게 됩니다.

3. VAEConfig 및 관련 설정 변경

VAEConfig 클래스는 VAE 디코딩 관련 설정을 관리하며, 이번 PR에서 가장 많은 변경이 이루어진 부분 중 하나입니다.

기존 설정:

class VAEConfig(ModelConfig):
    # ...
    use_parallel_decode: bool = False
    parallel_decode_mode: str = "tiled"

변경 후 설정:

AUTO_PARALLEL_DECODE_MODE = "auto"
SPATIAL_SHARD_PARALLEL_DECODE_MODES = ("spatial_shard", "spatial")

# ...

class VAEConfig(ModelConfig):
    # ...
    use_parallel_decode: bool = True
    parallel_decode_mode: str = AUTO_PARALLEL_DECODE_MODE
    auto_parallel_decode_min_latent_elements_per_rank: int = 4096

    def auto_parallel_decode_prefers_spatial_shard(self) -> bool:
        return False

    def should_use_auto_spatial_shard_parallel_decode(
        self, z: torch.Tensor, world_size: int
    ) -> bool:
        return _should_use_auto_spatial_shard_parallel_decode(
            tuple(z.shape), 
            world_size,
            self.auto_parallel_decode_min_latent_elements_per_rank,
        )

주요 변경 사항은 다음과 같습니다:

• use_parallel_decode 기본값 변경: False에서 True로 변경되어 병렬 디코딩이 기본적으로 활성화됩니다.
• parallel_decode_mode 기본값 변경: `

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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