[sglang] SGLang: AMD GPU 환경에서의 DeepSeek-V4 성능 최적화 분석

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

들어가며

최근 SGLang 레포지토리에 AMD GPU 환경(특히 MI355X 등 gfx950 아키텍처)에서의 DeepSeek-V4 모델 성능을 극대화하기 위한 최적화 PR이 병합되었습니다. 이 PR은 크게 두 가지 병목 구간을 해결합니다. 첫째, MLA(Multi-Head Latent Attention) 투영 연산에서 사용되는 FP8 GEMM 커널의 효율성 문제, 둘째, HIP 환경에서 비효율적으로 동작하던 RoPE(Rotary Positional Embedding) 연산의 커널 실행 방식입니다. 본 글에서는 이 최적화가 어떻게 구현되었고, 왜 성능 향상을 가져왔는지 분석합니다.

코드 분석

1. FP8 GEMM 최적화 (fp8_utils.py)

기존에는 MLA 투영 연산 시 Triton 기반의 triton_gemm_a8w8_blockscale을 사용했으나, 특정 하드웨어(gfx950)에서는 AMD의 Compute Kernel(CK) 라이브러리가 제공하는 bpreshuffle 방식이 훨씬 빠릅니다. 이를 위해 모듈 수준의 토글을 도입했습니다.

Before:

def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool:
    return (n, k) in [(1024, 8192), ...]

After:

_FORCE_CK_W8A8: bool = False

def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool:
    if _FORCE_CK_W8A8:
        return False
    return (n, k) in [...]

_FORCE_CK_W8A8 플래그를 통해 Triton 대신 CK 커널을 강제로 선택하도록 경로를 변경했습니다.

2. RoPE 연산 최적화 (deepseek_v4_rope.py)

기존 RoPE 구현은 토큰당 하나의 Triton 프로그램을 실행하여 오버헤드가 컸습니다. 이를 배치 처리 및 연속 메모리 접근(Contiguous load) 방식으로 개선했습니다.

Before:

# 기존 방식: 토큰별로 커널 실행 (Launch overhead 발생)
apply_rotary_emb_triton(x, ...)

After:

# 개선 방식: BLOCK_M 단위로 토큰을 배치 처리
apply_rotary_emb_contig_kernel[grid](x, ..., BLOCK_M=BLOCK_M, ...)

apply_rotary_emb_contig_kernel은 메모리 접근을 Coalesced하게 만들어 대역폭 효율을 극대화했습니다.

왜 이게 좋은가

이번 최적화의 핵심 교훈은 '커널 실행 오버헤드 최소화'와 '하드웨어 특화 라이브러리 활용'입니다.

1. 성능 수치: Pure-prefill 단계에서 기존 대비 약 8.6%~8.8%의 토큰 처리량(tok/s) 향상을 보였습니다. 이는 GPU 연산 유닛의 가동률을 높이고, 메모리 접근 패턴을 하드웨어 친화적으로 변경한 결과입니다.
2. 일반적 교훈:
   • 고성능 추론 엔진에서는 범용적인 Triton 커널도 좋지만, 특정 아키텍처(gfx950)에서는 벤더 제공 라이브러리(CK)가 더 최적화되어 있을 수 있습니다.
   • RoPE와 같은 요소 연산(Element-wise)은 토큰 단위 실행보다 배치 단위 실행이 커널 런치 오버헤드를 줄이는 데 필수적입니다.

리뷰 과정에서 BLOCK_M=32의 최적성에 대한 논의가 있었으나, 테스트 결과 gfx950과 gfx942 모두에서 CK 커널이 일관되게 우수한 성능을 보임이 확인되었습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://rocm.docs.amd.com/en/latest/conceptual/ai-libraries.html

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