[sglang] SGLang Diffusion 모델의 NVFP4 추론 성능 최적화: CUTLASS 도입

PR 링크: sgl-project/sglang#22091 상태: Merged | 변경: +None / -None

들어가며

최근 SGLang 프로젝트에서는 Blackwell 아키텍처 기반의 Diffusion 모델 추론 성능을 극대화하기 위해 NVFP4(4-bit Floating Point) 연산 백엔드를 재검토했습니다. 기존의 comfy-kitchen 백엔드는 특정 모델(FLUX)에서는 준수했으나, 다양한 Diffusion 모델의 워크로드 전반에서는 최적의 성능을 내지 못하는 문제가 있었습니다. 본 PR은 기본 백엔드를 범용적인 JIT 기반의 CUTLASS 경로로 전환하고, 이를 검증하기 위한 포괄적인 벤치마크를 도입하여 성능을 최적화했습니다.

코드 분석

1. 벤치마크 도구 추가 (bench_diffusion_nvfp4_scaled_mm.py)

새로운 벤치마크 스크립트는 다양한 모델의 실제 런타임 형상(Shape)을 캡처하고, cutlass, flashinfer_auto, flashinfer_cudnn 백엔드 간의 성능을 비교합니다.

# Before/After: 기존에는 고정된 백엔드를 사용했으나, 이제는 여러 백엔드를 동적으로 비교합니다.
providers: dict[str, Callable[[], torch.Tensor]] = {
    "cutlass": lambda: sgl_kernel.cutlass_scaled_fp4_mm(
        quantized["x_fp4"], quantized["w_fp4"], quantized["x_sf"], 
        quantized["w_sf"], quantized["alpha"], DTYPE,
    ),
    "flashinfer_auto": lambda: flashinfer.mm_fp4(
        quantized["x_fp4"], quantized["w_fp4"].T, quantized["x_sf"], 
        quantized["w_sf"].T, quantized["alpha"], DTYPE, backend="auto",
    ),
    # ... (생략)
}

2. 백엔드 전환 로직

기존의 comfy-kitchen을 기본값에서 제외하고, 환경 변수를 통해 선택적으로 사용할 수 있도록 변경하였습니다. 이는 CUTLASS가 265개 테스트 케이스 중 252개에서 승리하는 압도적인 성능을 보였기 때문입니다.

왜 이게 좋은가

이번 최적화의 핵심은 '범용성'과 '성능'의 균형입니다. 벤치마크 결과에 따르면, FLOP 가중치를 적용한 상대적 시간 비교에서 CUTLASS는 1.013x로 가장 우수한 성능을 보였으며, comfy는 1.095x, flashinfer 계열은 1.4배 이상의 지연 시간을 보였습니다.

• 성능 수치: 265개 케이스 중 252개에서 CUTLASS가 가장 빠른 실행 속도를 기록했습니다.
• 교훈: 특정 모델에 최적화된 커스텀 커널보다, 잘 튜닝된 범용 라이브러리(CUTLASS)가 다양한 모델 형상을 처리할 때 더 안정적이고 빠른 성능을 제공할 수 있음을 보여줍니다.

리뷰어 피드백

리뷰어 BBuf는 벤치마크 결과를 통해 CUTLASS가 263/265 케이스에서 승리함을 확인하며, 기존 comfy-kitchen 제거가 타당함을 뒷받침했습니다. 또한, CI 환경에서의 안정적인 실행을 위해 벤치마크 파일의 독립성을 강조했습니다.

참고 자료

• CUTLASS — 고성능 CUDA 커널 라이브러리
• FlashInfer — LLM 및 Diffusion 모델을 위한 고성능 커널 라이브러리

참고 자료

• https://github.com/NVIDIA/cutlass
• https://flashinfer.ai/

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