[sglang] SGLang Diffusion 최적화: CFG Gating을 통한 추론 속도 20% 향상

PR 링크: sgl-project/sglang#25848 상태: Merged | 변경: +303 / -0

들어가며

Diffusion 모델의 추론 과정에서 필수적인 Classifier-free guidance(CFG)는 매 denoising step마다 조건부(conditional)와 비조건부(unconditional) 모델을 각각 한 번씩, 총 두 번의 forward pass를 수행해야 합니다. 이는 연산 비용을 두 배로 증가시키는 주된 요인입니다. SGLang 프로젝트에 도입된 이번 PR은 'Adaptive Guidance' 개념을 활용하여, denoising 스케줄의 후반부에서 비조건부 연산을 생략하고 이전 단계의 잔차(residual)를 재사용함으로써 추론 속도를 획기적으로 개선합니다.

코드 분석

1. 환경 설정 및 파라미터 추가 (python/sglang/multimodal_gen/envs.py)

새로운 환경 변수 SGLANG_DIFFUSION_CFG_GATE_STEP을 도입하여 사용자가 어느 시점부터 gating을 적용할지 제어할 수 있게 했습니다.

# Before: 없음
# After:
SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0

2. Denoising 로직 최적화 (python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py)

핵심 로직은 _predict_noise_with_cfg 함수 내부에 구현되었습니다. 특정 gate step 이후에는 비조건부 모델을 직접 실행하는 대신, 캐시된 delta (conditional - unconditional) 값을 재사용합니다.

# After: CFG Gating 로직 적용
if can_reuse:
    neg_pred = _unwrap(tuple(p - d for p, d in zip(pos_t, delta_t)))
    cfg_gate_state["reused"] += 1
else:
    neg_pred = predict_fn(cfg_policy.branches[1])
    neg_t = _wrap(neg_pred)
    cfg_gate_state["delta"] = tuple(p.detach() - n.detach() for p, n in zip(pos_t, neg_t))
    cfg_gate_state["fresh_uncond"] += 1

이 코드는 pos_pred를 먼저 계산한 뒤, 이전에 저장된 delta_t를 사용하여 neg_pred를 근사합니다. 모델이 변경될 경우 model_id를 체크하여 캐시를 무효화하는 안전장치도 포함되어 있습니다.

왜 이게 좋은가

이번 최적화는 단순히 연산량을 줄이는 것을 넘어, Diffusion 모델의 고질적인 병목인 '이중 forward pass' 문제를 해결합니다. 벤치마크 결과에 따르면, SGLANG_DIFFUSION_CFG_GATE_STEP=0.5 설정 시 E2E latency가 약 21.3% 감소하였으며, 특히 DenoisingStage에서는 24.9%의 성능 향상을 보였습니다.

일반적 교훈

1. Temporal Redundancy 활용: Denoising 과정의 후반부는 노이즈가 적어 모델의 예측값이 급격히 변하지 않습니다. 이러한 특성을 이용해 연산을 근사하는 것은 매우 효과적인 전략입니다.
2. 상태 관리의 중요성: 모델이 교체되거나 파라미터가 변경될 때 캐시된 값을 무효화하는 로직(model_id 체크)이 포함되어야 안정적인 서비스가 가능합니다.
3. 설정 가능한 최적화: 모든 상황에서 성능 향상이 보장되지 않을 수 있으므로, 사용자가 임계값을 조정할 수 있게 설계하는 것이 좋습니다.

리뷰어 피드백

리뷰어 BBuf는 CI 테스트의 안정성을 위해 태그 기반의 재실행 전략을 제안했습니다. 이는 복잡한 분산 환경에서의 테스트 신뢰성을 높이는 데 기여합니다.

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