[sglang] SGLang 성능 최적화: RTX 4090에서 MXFP4 MoE 추론 속도 6배 향상시키기

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

들어가며

최근 SGLang 프로젝트에서 흥미로운 성능 최적화 PR이 병합되었습니다. RTX 4090(SM89)과 같은 Hopper 이전 아키텍처에서 MXFP4 MoE(Mixture of Experts) 모델을 실행할 때, 추론 속도가 예상보다 6배가량 느린 문제가 발견되었습니다. 프로파일링 결과, 전체 추론 시간의 대부분이 MXFP4 행렬 곱셈 커널에서 발생하고 있었으며, 원인은 Triton 커널이 사용하는 warp 개수가 너무 적게 할당되는 것이었습니다. 본 글에서는 이 문제를 어떻게 해결했는지 코드 수준에서 분석합니다.

코드 분석

문제의 핵심은 mxfp4.py 파일 내의 _swizzle_mxfp4 함수에 있었습니다. Hopper 아키텍처에서는 스케일 레이아웃이 warp 개수를 결정하지만, 그 이전 세대에서는 적절한 제약 조건이 없어 기본값인 '1 warp'로 동작하고 있었습니다.

python/sglang/srt/layers/quantization/mxfp4.py

변경 전에는 Hopper 아키텍처를 위한 로직만 존재했으나, 변경 후에는 CUDA capability가 9.0 미만인 경우 명시적으로 num_warps를 제약 조건으로 추가합니다.

# Before
    elif get_platform().is_cuda and get_platform().device_capability < (9, 0):
        # No scale layout carries the warp count below Hopper, and the tile
        # formula gives the batch-1 decode tile a single warp.
        opt_flags.update_opt_flags_constraints({"num_warps": num_warps})

이 변경을 통해, 기존에 1 warp만 사용하던 커널이 8 warps를 사용하도록 강제되어 GPU 연산 유닛의 활용도를 극대화하게 되었습니다.

왜 이게 좋은가

이번 최적화는 단순히 코드 몇 줄을 추가한 것이지만, 성능에 미치는 영향은 매우 큽니다. RTX 4090 환경에서 측정된 결과는 다음과 같습니다.

• 성능 향상: 기존 약 2024 tok/s에서 최적화 후 약 114135 tok/s로, 약 5~6배의 성능 향상을 기록했습니다.
• 병목 해소: 프로파일링 결과 42ms가 소요되던 커널 시간이 획기적으로 단축되었습니다. 이는 GPU가 180W 수준의 낮은 전력 소모에서 벗어나 더 높은 연산 효율을 낼 수 있게 되었음을 의미합니다.

일반적 교훈

1. Warp Occupancy 확인: Triton 커널 작성 시 num_warps는 성능에 직결되는 핵심 파라미터입니다. 특히 작은 타일 크기에서는 기본 설정이 최적이 아닐 수 있습니다.
2. 아키텍처별 분기: 하드웨어 세대(Hopper vs Pre-Hopper)에 따라 커널의 동작 방식이 달라질 수 있으므로, device_capability를 활용한 세밀한 제어가 필요합니다.
3. 프로파일링의 중요성: 막연한 짐작보다 torch-profiler를 통한 병목 지점 파악이 최적화의 첫걸음임을 다시 한번 확인했습니다.

결론

이번 PR은 SGLang이 다양한 GPU 아키텍처에서 최상의 성능을 낼 수 있도록 하는 중요한 개선입니다. 특히 MXFP4와 같은 최신 양자화 기법을 사용할 때, 하드웨어 특성에 맞는 커널 튜닝이 얼마나 중요한지 잘 보여주는 사례입니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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