[sglang] SGLang Diffusion 모델의 Ulysses Attention 파이프라인 최적화: Copy Engine을 활용한 성능 향상

PR 링크: sgl-project/sglang#43521 상태: Merged | 변경: +1390 / -370

들어가며

최근 대규모 Diffusion 모델(Wan2.1, FLUX, HunyuanVideo 등)의 추론 성능을 최적화하는 과정에서, 분산 학습 및 추론의 핵심인 UlyssesAttention의 병목 현상이 주요 과제로 떠올랐습니다. 기존에는 모든 헤드(head)에 대한 All-to-All 교환(exchange)을 먼저 수행한 뒤 어텐션 연산을 수행하는 '순차적(sequential)' 방식을 사용했습니다. 이 PR은 이 과정을 파이프라이닝하여, 한 헤드 그룹이 어텐션을 수행하는 동안 다음 헤드 그룹의 데이터를 Copy Engine을 통해 이동시키는 방식으로 성능을 획기적으로 개선했습니다.

코드 분석

1. ipc_a2a_multi.py: 파이프라인 전송 로직 구현

핵심 변경 사항은 ulysses_pipelined_attention 함수와 이를 지원하는 _PipelinePool입니다. 기존에는 모든 데이터를 한 번에 교환했으나, 이제는 헤드 그룹 단위로 작업을 쪼개어 오버랩합니다.

Before (순차적 방식):

# 기존에는 모든 헤드에 대해 교환 후 어텐션 수행
exchange_all_heads(q, k, v)
output = attend(q, k, v)

After (파이프라인 방식):

# 헤드 그룹별로 교환과 어텐션을 오버랩
for group in head_groups:
    # Copy Engine을 통해 다음 그룹의 데이터를 비동기 이동
    move_group_data_async(group + 1)
    # 현재 그룹의 어텐션 연산 수행
    attend(group)

이 구조는 _PipelinePool을 통해 메모리 버퍼를 공유하며, 여러 해상도에서 동일한 버퍼 풀을 재사용함으로써 메모리 오버헤드를 최소화합니다.

2. USPAttention._pipelined_ulysses: 레이아웃 지원

다양한 모델(FLUX, Qwen-Image 등)의 복잡한 레이아웃(replicated prefix, suffix 등)을 지원하기 위해, 어텐션 코어와 교환 로직을 분리했습니다. 특히, 첫 번째 호출 시 순차적 방식과 파이프라인 방식의 결과가 동일한지 검증하는 'First-sight check'를 도입하여 수치적 정확성을 보장합니다.

왜 이게 좋은가

성능 향상

• 속도: B300 GPU 환경에서 Wan2.1 14B 모델 기준, Ulysses 2 설정에서 약 11.4%, Ulysses 4 설정에서 약 11.8%의 성능 향상을 보였습니다. 일부 모델(FLUX 2048²)에서는 최대 24%의 속도 개선이 측정되었습니다.
• 메모리 효율: 공유 버퍼 풀을 도입하여 기존 순차적 방식 대비 메모리 사용량을 약 50% 수준으로 절감했습니다(예: 3.6GB → 2.1GB).

일반적 교훈

1. 통신과 연산의 오버랩: GPU의 Copy Engine을 활용하여 데이터 이동 시간을 어텐션 연산 시간 뒤로 숨기는 것은 분산 시스템에서 가장 효과적인 최적화 전략 중 하나입니다.
2. 동적 그룹화: 고정된 그룹 개수 대신, GPU의 연산 능력(SM count)과 시퀀스 길이에 따라 최적의 그룹 개수를 자동으로 결정하는 로직(_AUTO_PIPELINE_GROUPS)은 범용적인 성능 최적화의 모범 사례입니다.
3. 정확성 검증: 성능 최적화가 수치적 결과(bit-identical)를 해치지 않도록 첫 호출 시 검증 로직을 포함하는 것은 프로덕션 환경에서 필수적인 안전장치입니다.

참고 자료

• SGLang Documentation
• PyTorch Distributed Communication

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