[sglang] SGLang의 Wan2.2-TI2V 최적화: Triton 커널을 통한 메모리 트래픽 병목 해결

PR 링크: sgl-project/sglang#34584 상태: Merged | 변경: +250 / -25

들어가며

최근 대규모 확산 모델(Diffusion Model)인 Wan2.2-TI2V를 추론할 때, eager 모드와 torch.compile 모드 사이의 성능 격차가 컸습니다. 특히 H100/H200 환경에서 per-token adaLN 경로가 전체 연산 시간의 약 14%를 불필요한 메모리 이동에 낭비하고 있었습니다. 본 PR은 이 병목을 Triton 커널로 해결하여 메모리 트래픽을 최적화하고, eager 모드에서도 torch.compile을 능가하는 성능을 달성했습니다.

코드 분석

1. Triton 커널을 통한 adaLN 슬라이스 융합 (wan_temb_table_slices.py)

기존 방식은 (scale_shift_table + temb).chunk(6) 연산 후 각 슬라이스에 대해 .contiguous()를 호출하여 불필요한 복사가 발생했습니다. 이를 Triton 커널로 융합하여 한 번의 패스로 연속적인 메모리 레이아웃을 생성합니다.

Before:

# 기존 방식: 중간 텐서 생성 및 strided 슬라이스 복사
(scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)

After:

# Triton 커널: 한 번의 패스로 6개의 슬라이스를 연속적으로 생성
@triton.jit
def _temb_table_slices_kernel(out_ptr, temb_ptr, table_ptr, ...):
    # ... 커널 로직 ...
    tl.store(out_ptr + (j * rows + row) * D + cols, w + t, mask=mask)

2. RoPE 캐시 호이스팅 (wanvideo.py)

매 블록마다 반복되던 cos_sin_cache 생성을 트랜스포머 레벨로 끌어올려(hoisting) 한 번만 수행하도록 변경했습니다.

Before:

# 매 블록마다 반복되는 연산
cos_sin_cache = torch.cat([cos.to(torch.float32).contiguous(), ...], dim=-1)

After:

# 트랜스포머에서 한 번만 생성하여 전달
rope_cos_sin_cache = torch.cat([...], dim=-1)
for block in self.blocks:
    hidden_states = block(..., rope_cos_sin_cache=rope_cos_sin_cache)

왜 이게 좋은가

이 최적화는 단순히 코드의 효율성을 높이는 것을 넘어, 메모리 대역폭(Memory Bandwidth) 병목을 정면으로 해결했습니다.

1. 성능 향상: H100 기준 Denoise 단계에서 약 13.1%의 속도 향상을 기록했습니다. 이는 불필요한 (B, S, D) 텐서 복사를 제거한 결과입니다.
2. 비트 단위 정확성(Bit-exact): float32 연산 과정에서 반올림 오차가 발생하지 않도록 설계되어, 기존 eager 구현과 동일한 결과를 보장합니다. BitExactFusionGate를 통해 런타임 검증을 수행하며, 실패 시 안전하게 기존 코드로 폴백합니다.
3. 교훈: 딥러닝 모델 추론에서 torch.compile에만 의존하기보다, 특정 연산 패턴(예: strided slice copy)을 Triton 커널로 직접 최적화하는 것이 하드웨어 가속기를 더 효율적으로 활용하는 방법임을 보여줍니다.

결론

이번 개선은 SGLang이 하드웨어 수준의 최적화를 통해 얼마나 큰 성능 이득을 얻을 수 있는지 보여주는 좋은 사례입니다. 특히 eager 모드에서 compile 모드보다 빠른 성능을 달성함으로써, 모델의 유연성과 속도를 모두 잡았습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.contiguous.html

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