[sglang] SGLang의 Ideogram4 추론 성능 최적화: Denoising 루프 내 오버헤드 제거

PR 링크: sgl-project/sglang#27443 상태: Merged | 변경: +26 / -10

들어가며

최신 Diffusion 모델, 특히 Ideogram4와 같은 대규모 모델을 추론할 때 가장 큰 병목 중 하나는 Denoising 루프 내부의 반복적인 연산입니다. 기존 SGLang 구현에서는 매 Denoising 스텝마다 Attention Mask 메타데이터를 생성하고, 스케줄러 값을 계산하는 과정이 포함되어 있었습니다. 이는 GPU 연산 효율을 저해하는 불필요한 호스트 측 오버헤드를 발생시킵니다. 본 PR은 이러한 반복적인 연산을 루프 진입 전으로 옮겨(Precompute) 추론 효율을 극대화했습니다.

코드 분석

1. ideogram.py: 모델 Forward 패스 최적화

기존에는 forward 함수가 호출될 때마다 build_varlen_mask_meta를 통해 마스크 메타데이터를 매번 생성했습니다. 이를 선택적 인자로 변경하여, 외부에서 미리 계산된 메타데이터를 전달받을 수 있도록 수정했습니다.

# Before
attn_mask = segment_ids > 0
attn_mask_meta = build_varlen_mask_meta(attn_mask)

# After
if attn_mask is None:
    attn_mask = segment_ids > 0
if attn_mask_meta is None:
    attn_mask_meta = build_varlen_mask_meta(attn_mask)

2. ideogram.py (Pipeline): 사전 연산(Precomputation) 도입

_prepare_denoising_loop 단계에서 스케줄러 값과 Attention Mask 메타데이터를 미리 계산하여 ctx.extra에 저장합니다. 이를 통해 _run_denoising_step 내부에서는 별도의 연산 없이 즉시 값을 참조할 수 있게 되었습니다.

# Precomputing values
schedule_values = schedule(step_intervals)
schedule_deltas = schedule_values[:-1] - schedule_values[1:]

# Storing in context
ctx.extra.update({
    "ideogram4_schedule_values": schedule_values,
    "ideogram4_attn_mask_meta": build_varlen_mask_meta(attn_mask),
    # ...
})

왜 이게 좋은가

성능 향상 포인트

PyTorch Profiler 분석 결과, 기존에는 aten::nonzero와 aten::cumsum이 Denoising 스텝마다 호출되어 각각 수백 밀리초의 오버헤드를 발생시켰습니다. 이번 최적화를 통해 해당 연산들이 프로파일링 결과에서 완전히 사라졌습니다. 비록 전체 추론 시간 대비 절대적인 시간 단축은 작지만(약 0.08% ~ 0.12%), 루프 내부의 불필요한 호스트-디바이스 간 동기화 및 연산 오버헤드를 제거했다는 점에서 매우 중요한 아키텍처 개선입니다.

교훈

1. 반복되는 연산의 식별: 루프 내부에서 매번 동일한 결과를 반환하는 함수 호출은 반드시 사전 연산(Precomputation)을 고려해야 합니다.
2. 유연한 API 설계: 기존 코드와의 호환성을 유지하면서 최적화된 경로를 제공하기 위해 Optional 인자를 활용하는 패턴은 라이브러리 설계 시 매우 유용합니다.

리뷰어 피드백

리뷰어들은 CI 테스트 통과를 확인하며, 특히 varlen(가변 길이) 어텐션 마스크 처리에 대한 효율적인 접근 방식을 긍정적으로 평가했습니다. 이는 SGLang의 고성능 추론 엔진으로서의 정체성을 강화하는 변경입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.expand.html
• https://github.com/sgl-project/sglang

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