[sglang] SGLang의 Spectral Progressive Diffusion 도입: 추론 속도 최대 2.78배 향상

PR 링크: sgl-project/sglang#27524 상태: Merged | 변경: +2543 / -55

들어가며

Transformer 기반의 Diffusion 모델들은 시퀀스 길이에 대해 O(n²)의 복잡도를 가지는 Attention 연산으로 인해 고해상도 이미지 생성 시 막대한 연산 비용이 발생합니다. 특히 확산 과정의 초기 단계에서는 고주파 세부 정보가 거의 활성화되지 않음에도 불구하고, 전체 해상도에서 연산을 수행하는 것은 큰 낭비입니다. 본 PR은 Spectral Progressive Diffusion 기술을 SGLang에 도입하여, 초기 단계에는 낮은 해상도에서 연산하고 이후 고해상도로 스펙트럼 업샘플링을 수행함으로써 추론 속도를 최대 2.78배까지 개선했습니다.

코드 분석

1. 핵심 모듈: runtime/pipelines_core/stages/progressive_resolution/

이 PR은 새로운 파이프라인 스테이지 아키텍처를 도입했습니다. 핵심은 spectral_ops.py에서 구현된 GPU 기반 DCT(Discrete Cosine Transform)입니다.

# Before: 항상 풀 해상도에서 연산 수행
# After: ProgressiveDenoisingStage를 통한 단계적 해상도 확장
class ProgressiveDenoisingStage(DenoisingStage):
    def _on_resolution_change(self, ctx, batch, server_args, h_px, w_px):
        # 해상도 변경 시 RoPE 및 위치 임베딩 업데이트
        update_positional_embeddings(ctx, h_px, w_px)

2. 파이프라인 라우팅

기존 파이프라인에 _DenoisingStageRouter를 도입하여 progressive_mode 설정에 따라 동적으로 연산 경로를 분기합니다. fullres 모드일 경우 기존과 동일한 동작을 보장하여 하위 호환성을 유지합니다.

# pipelines/flux.py 내 라우터 예시
class _FluxDenoisingStageRouter(PipelineStage):
    def forward(self, ...):
        if self.progressive_mode == "fullres":
            return self.fullres_stage.forward(...)
        return self.progressive_stage.forward(...)

왜 이게 좋은가

이 최적화의 핵심은 Bayes-optimal한 해상도 전환입니다. 모델의 VAE 잠재 공간의 전력 스펙트럼을 분석하여, 노이즈보다 신호가 더 중요해지는 시점에 정확히 해상도를 높입니다.

• 성능 수치: RTX A6000 기준 Wan 2.1 T2V 모델에서 최대 2.78배의 속도 향상을 기록했습니다.
• Cache-DiT 호환성: 기존 Cache-DiT와도 결합 가능하며, 해상도 전환 시 캐시된 컨텍스트를 적절히 리프레시하여 일관성을 유지합니다.

일반적인 교훈으로, 모델의 연산이 모든 단계에서 동일한 해상도를 필요로 하지 않는다는 점을 활용하면, 연산 효율을 극대화하면서도 품질 저하를 최소화할 수 있음을 보여줍니다.

리뷰 피드백 반영

리뷰어들의 요청에 따라 분산되어 있던 단위 테스트들을 python/sglang/multimodal_gen/test/unit/test_progressive.py로 통합하여 관리 효율성을 높였으며, 119개의 테스트 케이스를 통해 DCT 연산의 정밀도와 파이프라인 전환 로직의 안정성을 검증했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.fft.html
• https://arxiv.org/abs/2605.18736

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