[sglang] SGLang: ROCm 환경에서 RMSNorm 최적화 - Triton에서 aiter 커널로 전환

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

들어가며

대규모 언어 모델(LLM) 및 확산 모델(Diffusion Model)과 같은 최신 AI 모델들은 엄청난 연산량을 요구합니다. 이러한 모델의 성능을 최적화하는 것은 사용자 경험을 향상시키고, 더 넓은 범위의 하드웨어에서 효율적으로 작동하도록 하는 데 필수적입니다. 특히, 모델의 핵심 구성 요소 중 하나인 Layer Normalization (LayerNorm)은 각 레이어의 출력을 정규화하여 학습 안정성과 성능에 기여하지만, 동시에 상당한 연산 오버헤드를 유발할 수 있습니다.

이번에 분석할 GitHub PR은 sgl-project/sglang 레포지토리에서 ROCm(AMD GPU) 환경의 확산 모델 경로에서 사용되는 RMSNorm(Root Mean Square Normalization) 연산의 성능을 최적화하는 것을 목표로 합니다. 기존의 Triton 기반 RMSNorm 구현을 aiter 라이브러리의 최적화된 HIP RMSNorm 커널로 교체하여, 단일 RMSNorm 커널 실행 시간을 약 30% (430 us → 290 us) 단축시키는 중요한 개선을 이루어냈습니다. 이는 전체 denoise 단계 시간의 약 0.2%에 해당하며, 작은 개선처럼 보일 수 있지만, 이러한 최적화가 모여 전체 시스템의 효율성을 크게 향상시킬 수 있습니다.

코드 분석: layernorm.py의 핵심 변경사항

이 PR의 핵심 변경사항은 python/sglang/multimodal_gen/runtime/layers/layernorm.py 파일에 집중되어 있습니다. 주요 변경점은 다음과 같습니다.

1. aiter 라이브러리 임포트 및 사용 조건 추가

aiter 라이브러리의 RMSNorm 커널을 사용하기 위한 조건부 임포트 로직이 추가되었습니다. SGLANG_USE_AITER 환경 변수가 설정되어 있고 현재 플랫폼이 HIP(ROCm)인 경우에만 aiter 커널을 사용하도록 합니다.

Before:

_is_cuda = current_platform.is_cuda()
_is_npu = current_platform.is_npu()
_is_musa = current_platform.is_musa()
_is_cpu = current_platform.is_cpu()
_is_xpu = current_platform.is_xpu()

if _is_cuda or _is_xpu:
    from sgl_kernel import fused_add_rmsnorm, rmsnorm

if _is_npu:
    from sgl_kernel import fused_add_rmsnorm

if _is_musa:
    from sgl_kernel import fused_add_rmsnorm
if not _is_cpu:
    from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn

After:

_is_cuda = current_platform.is_cuda()
_is_hip = current_platform.is_hip()
_is_npu = current_platform.is_npu()
_is_musa = current_platform.is_musa()
_is_cpu = current_platform.is_cpu()
_is_xpu = current_platform.is_xpu()
_use_aiter = get_bool_env_var(

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