[sglang] ROCm 환경에서의 성능 최적화: Triton을 활용한 Bit-exact QK RMSNorm 및 RoPE 융합

PR 링크: sgl-project/sglang#47569 상태: Merged | 변경: +330 / -2

들어가며

최근 sglang 프로젝트의 ROCm 환경에서 MiniMax-H3 모델의 추론 성능을 분석한 결과, QK(Query-Key) RMSNorm 연산이 전체 denoise 단계의 약 5.1%를 차지하는 병목 구간임을 확인했습니다. 기존에는 aten.rms_norm을 사용하면서 불필요한 .contiguous() 메모리 복사가 발생했고, 이는 MI355X GPU 환경에서 성능 저하의 주원인이었습니다. 본 PR은 이 문제를 해결하기 위해 Triton 커널을 도입하여 연산을 최적화하고, RoPE(Rotary Positional Embedding) 연산까지 하나의 커널로 융합하여 성능을 극대화했습니다.

코드 분석

1. Triton 커널 구현 (aten_head_rmsnorm_triton.py)

기존 aten의 vectorized_layer_norm_kernel은 256 스레드 블록을 사용하여 128-wide 행을 처리합니다. 이를 Triton으로 재구현하여 메모리 복사 없이 strided view를 직접 읽도록 개선했습니다.

# Before: aten.rms_norm 사용 (메모리 복사 발생)
# After: Triton 커널을 통한 직접 읽기
lanes = tl.reshape(x, (rows_per_prog, head_dim // 4, 2, 2), can_reorder=False)
even, odd = tl.split(lanes)
# ... (shfl-down을 이용한 병렬 합산)
acc = tl.sum(tl.reshape(acc, (rows_per_prog, 2, 16), can_reorder=False), axis=1)

이 커널은 rsqrtf를 사용하여 분산을 계산하고, w * (rstd * x) 연산을 수행하여 aten과 비트 단위로 동일한(Bit-exact) 결과를 보장합니다.

2. RoPE 융합 및 게이트 로직 (minimax_h3.py)

단순히 RMSNorm만 최적화한 것이 아니라, 뒤따르는 RoPE 연산까지 동일한 커널 내에서 처리하도록 융합했습니다. 또한 BitExactFusionGate를 사용하여 eager path와 결과값이 일치하는지 검증하는 안전장치를 마련했습니다.

# RoPE 융합 로직
if rope_half > 0:
    # ... (cos/sin 캐시를 활용한 회전 연산)
    a = (y.to(tl.float32) * cos.to(tl.float32)).to(tl.bfloat16).to(tl.float32)
    b = (yp.to(tl.float32) * sin.to(tl.float32)).to(tl.bfloat16).to(tl.float32)
    rope = tl.where(first[None, :], a - b, a + b).to(tl.bfloat16)

왜 이게 좋은가

이번 최적화는 단순히 코드 스타일을 바꾼 것이 아니라, 메모리 대역폭과 커널 실행 횟수를 획기적으로 줄였습니다.

• 성능 향상: 4x MI355X 환경에서 1280x704 해상도 기준, 전체 추론 시간이 44.78초에서 42.51초로 약 2.27초 단축되었습니다.
• 커널 효율성: QK norm 연산 자체는 43.7ms에서 7.1ms로 약 6배 이상의 속도 향상을 보였습니다.
• 교훈:
  1. 메모리 복사 제거: strided 텐서를 직접 처리하는 Triton 커널은 불필요한 .contiguous() 호출을 제거하여 대역폭 병목을 해결합니다.
  2. 연산 융합(Kernel Fusion): 여러 개의 작은 커널을 하나로 합치면 커널 실행 오버헤드를 줄이고, 중간 결과를 레지스터 수준에서 처리하여 메모리 접근을 최소화할 수 있습니다.
  3. 검증 전략: BitExactFusionGate와 같은 메커니즘을 통해 최적화된 커널이 기존 eager 구현과 동일한 결과를 내는지 실시간으로 검증하는 것은 복잡한 커널 최적화에서 필수적입니다.

결론

이번 PR은 ROCm 환경에서의 하드웨어 특성을 깊이 이해하고, Triton을 통해 연산과 메모리 접근 패턴을 최적화한 훌륭한 사례입니다. 특히 비트 단위의 정확성을 유지하면서도 성능을 끌어올린 점은 프로덕션 환경에서 매우 중요한 가치를 가집니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://pytorch.org/docs/stable/generated/torch.nn.functional.rms_norm.html

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