[sglang] SGLang: ROCm 환경에서 Qwen3-VL 디코딩 성능 극대화를 위한 커널 퓨전 최적화

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

들어가며

LLM 추론, 특히 디코딩(Decoding) 단계에서는 메모리 대역폭이 병목 현상의 주범입니다. SGLang의 이번 PR은 AMD ROCm 환경에서 Qwen3-VL 모델의 디코딩 성능을 높이기 위해, 기존에 분리되어 실행되던 4개의 연산(QKV split, QK RMSNorm, 3D mRoPE, KV cache write)을 단일 HIP 커널로 통합(Fusion)하는 최적화를 수행했습니다. 이 글에서는 해당 변경 사항이 어떻게 성능을 개선했는지 코드 수준에서 분석합니다.

코드 분석

1. 커널 통합 및 조건부 실행 (python/sglang/srt/models/qwen3.py)

핵심 변경 사항은 forward_prepare_aiter_fused_mrope 함수를 도입하여 여러 연산을 하나의 커널로 묶은 것입니다. 기존에는 각 연산이 개별적으로 호출되었으나, 이제는 aiter 라이브러리의 fused_qk_norm_mrope_3d_cache_pts_quant_shuffle을 사용하여 한 번의 커널 호출로 처리합니다.

Before (기존 방식):

# 개별 연산이 순차적으로 실행됨
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
# 이후 RMSNorm, RoPE 적용 및 KV 캐시 쓰기 연산이 각각 호출됨

After (최적화 방식):

# 퓨전 커널을 통해 한 번에 처리
fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(
    qkv_3d, self.q_norm.weight, self.k_norm.weight, cos_sin, positions, 
    # ... (생략) ...
    q_out, k_cache, v_cache, slot_mapping, 
    self._fused_k_scale, self._fused_v_scale, 
    # ...
)

2. 아키텍처 유지보수성 고려

리뷰 과정에서 제기된 "전체 어텐션 블록 로직을 한 함수에 복사하지 말라"는 피드백을 반영하여, forward_prepare -> forward_core 패턴을 유지했습니다. 퓨전 커널은 forward_prepare 단계에서 실행되어 q 텐서만 반환하고, save_kv_cache=False 플래그를 통해 이후 forward_core에서 불필요한 KV 캐시 쓰기를 방지합니다.

왜 이게 좋은가

1. 커널 호출 오버헤드 감소: GPU 커널을 실행할 때마다 발생하는 CPU-GPU 간의 통신 및 스케줄링 오버헤드를 4번에서 1번으로 줄였습니다.
2. 메모리 트래픽 최적화: 중간 결과물(Q, K, V)을 전역 메모리에 쓰고 다시 읽는 과정을 생략하고, 레지스터나 공유 메모리 수준에서 데이터를 처리함으로써 메모리 대역폭 사용량을 획기적으로 줄였습니다.
3. 일반적 교훈: LLM 추론 최적화의 핵심은 '연산의 퓨전'입니다. 특히 디코딩 단계처럼 연산량 대비 메모리 접근이 많은 경우, 커널 퓨전은 성능 향상의 가장 효과적인 전략입니다.

이번 최적화는 SGLANG_USE_AITER 환경 변수와 is_hip() 체크를 통해 ROCm 환경에서만 안전하게 동작하도록 설계되어, NVIDIA 환경과의 호환성을 완벽하게 유지했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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