[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의 고성능 추론 엔진으로서의 정체성을 강화하는 변경입니다.

결론

이번 PR은 Denoising 루프마다 반복 호출되던 aten::nonzero·aten::cumsum 연산을 사전 연산(Precomputation)으로 옮겨, 루프 내부의 불필요한 호스트-디바이스 동기화와 연산 오버헤드를 제거했다. 절대적인 시간 단축은 작지만, 반복되는 연산을 식별해 루프 밖으로 빼고 Optional 인자로 기존 코드와의 호환성을 유지한 설계는 고성능 추론 엔진에서 의미 있는 아키텍처 개선의 모범을 보여준다.

참고 자료

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

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