[sglang] SGLang LTX-2 VAE 디코딩 성능 최적화: channels_last_3d 도입으로 4.5배 속도 향상

PR 링크: sgl-project/sglang#27431 상태: Merged | 변경: +223 / -49

들어가며

최근 생성형 AI 모델, 특히 비디오 생성 모델에서 VAE(Variational Autoencoder) 디코딩 단계는 전체 파이프라인의 병목 구간이 되곤 합니다. SGLang의 LTX-2 모델 역시 Conv3d 연산이 주를 이루는 디코딩 단계에서 성능 최적화가 필요했습니다. 기존 구현은 channels_last_3d 레이아웃을 활용하지 못하고, 레이아웃 호환성 문제로 인해 잦은 메모리 복사와 변환이 발생하여 성능 저하와 높은 메모리 점유율을 야기했습니다. 본 PR은 이 문제를 해결하기 위해 레이아웃 보존형 패딩(Layout-preserving padding)을 도입하고 로더 정책을 개선했습니다.

코드 분석

1. vae_loader.py: 자동 레이아웃 정책 적용

기존에는 특정 모델에만 제한적으로 적용되던 channels_last_3d 최적화 정책을 LTX-2 모델로 확장했습니다.

# Before
if pipeline_name.startswith("QwenImage"):
    return True

# After
if isinstance(pipeline_config, QwenImagePipelineConfig):
    return True
if (isinstance(pipeline_config, (WanT2V480PConfig, LTX2PipelineConfig)) and server_args.num_gpus == 1):
    return True

단일 GPU 환경에서 LTX-2 모델이 자동으로 최적화된 레이아웃을 사용하도록 설정하여, 사용자가 별도의 환경 변수 없이도 즉각적인 성능 향상을 경험할 수 있게 했습니다.

2. ltx_2_vae.py: 레이아웃 보존형 패딩 구현

가장 핵심적인 변경은 LTX2VideoCausalConv3d 클래스 내의 패딩 로직입니다. 기존에는 repeat()와 concatenate()를 사용하여 텐서를 생성했는데, 이 과정에서 메모리 레이아웃이 NCDHW로 고정되어 channels_last_3d의 이점을 누릴 수 없었습니다.

# Before (Layout-hostile)
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, left, 1, 1))
hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)

# After (Layout-preserving)
out = torch.empty((b, c, t + left + right, h, w), memory_format=torch.channels_last_3d)
out[:, :, left : left + t].copy_(x)
# ... copy logic ...

새로운 구현은 torch.empty를 통해 channels_last_3d 메모리 포맷으로 할당한 뒤 copy_를 수행합니다. 이를 통해 레이아웃 변환 비용을 제거하고, cuDNN의 고속 NDHWC 커널을 그대로 활용할 수 있게 되었습니다.

왜 이게 좋은가

이번 최적화는 단순한 코드 수정을 넘어 하드웨어 가속기를 최대한 활용하는 방향으로 이루어졌습니다.

• 성능 향상: H100 GPU 환경에서 Conv3d 단일 연산은 약 3.67배, 전체 LTX2VideoDecoder3d.forward는 약 4.58배의 속도 향상을 보였습니다.
• 메모리 효율: 불필요한 중간 텐서 생성을 방지하여 피크 메모리 사용량을 약 9.7 GiB(13.5%) 절감했습니다.
• 안정성: _is_channels_last_3d_stride를 통해 레이아웃 상태를 확인하고, 기존 로직과 병행할 수 있도록 설계하여 하위 호환성을 완벽히 보장했습니다.

리뷰 과정에서 lru_cache를 도입하여 레이아웃 확인 로직의 오버헤드를 최소화한 점도 주목할 만한 기술적 디테일입니다. 이 사례는 고성능 딥러닝 모델 개발 시 메모리 레이아웃(Memory Layout) 최적화가 전체 시스템 성능에 얼마나 큰 영향을 미치는지 잘 보여줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html
• https://pytorch.org/docs/stable/generated/torch.empty.html

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