[sglang] AMD ROCm 환경에서의 DeepSeek-V4 성능 최적화: Aiter MHC 커널 통합 분석

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

들어가며

최근 대규모 언어 모델(LLM) 시장에서 DeepSeek-V4와 같은 고성능 모델의 등장은 하드웨어 가속기 지원에 대한 새로운 요구사항을 만들어내고 있습니다. 특히 AMD GPU(ROCm) 환경에서 이러한 최신 모델을 효율적으로 구동하기 위해서는 하드웨어 특성에 최적화된 커널(Kernel) 수준의 개선이 필수적입니다.

이번 PR(amd/deepseek_v4 integration 10/N)은 SGLang 프로젝트에서 AMD 환경의 DeepSeek-V4 성능을 극대화하기 위한 노력의 일환입니다. 핵심은 Multi-Head Compression(MHC) 연산의 전처리(Pre) 및 후처리(Post) 과정을 기존의 범용적인 PyTorch 구현이나 TileLang 기반에서 AMD 전용 가속 라이브러리인 Aiter 커널로 전환한 것입니다.

코드 분석: 무엇이 바뀌었는가?

1. 환경 변수를 통한 Aiter 커널 활성화

먼저 docker/rocm.Dockerfile과 python/sglang/srt/environ.py에서 Aiter MHC 커널 사용 여부를 제어하는 플래그가 추가되었습니다. 기본적으로 AMD 환경에서 성능 이점을 가져가기 위해 True로 설정되었습니다.

# docker/rocm.Dockerfile
+ENV SGLANG_OPT_USE_AITER_MHC_PRE=true
+ENV SGLANG_OPT_USE_AITER_MHC_POST=true
+ENV SGLANG_OPT_USE_TILELANG_MHC_PRE=false
+ENV SGLANG_OPT_USE_TILELANG_MHC_POST=false

2. DeepSeek-V4 모델 레이어의 커널 교체

가장 중요한 변경 사항은 python/sglang/srt/models/deepseek_v4.py에 있습니다. 기존에는 envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM 등을 통해 DeepGEMM을 사용하거나 일반 PyTorch 코드로 처리하던 로직에 AMD 전용 aiter.ops.mhc가 삽입되었습니다.

Before (Generic/DeepGEMM path): 기존에는 AMD 환경(_is_hip)에서도 전용 최적화 커널이 부족하여 DeepGEMM이나 일반 구현에 의존해야 했습니다.

After (Aiter Integration):

# python/sglang/srt/models/deepseek_v4.py
if _is_hip and envs.SGLANG_OPT_USE_AITER_MHC_PRE.get():
    from aiter.ops.mhc import mhc_pre

    post, comb, y = mhc_pre(
        residual=x,
        fn=hc_fn,
        hc_scale=hc_scale,
        # ... 생략 ...
        sinkhorn_repeat=self.hc_sinkhorn_iters,
    )
    return y, post.squeeze(-1), comb

후처리 과정(hc_post)에서도 동일하게 Aiter 커널이 적용되었습니다.

elif _is_hip and envs.SGLANG_OPT_USE_AITER_MHC_POST.get():
    from aiter.ops.mhc import mhc_post

    result = torch.empty_like(residual)
    mhc_post(result, x, residual, post, comb)
    return result

3. TileLang 커널의 세부 튜닝

python/sglang/srt/layers/mhc.py에서는 hc_hidden_size에 따른 hidden_block 크기를 명시적으로 지정하여 메모리 접근 효율을 높였습니다. 이는 하드웨어의 L2 캐시나 공유 메모리 활용도를 최적화하기 위한 전형적인 기법입니다.

# python/sglang/srt/layers/mhc.py
if num_tokens <= 2048:
    assert n_splits == 1
    if hc_hidden_size == 16384:
        hidden_block = 256
    elif hc_hidden_size == 28672:
        hidden_block = 128
    # ...
    kernel_0, kernel_1 = mhc_pre_gemm_sqrsum_splitk_kernel(
        # ...
        hidden_block=hidden_block,
    )

왜 이게 좋은 최적화인가?

1. 하드웨어 특화 가속 (Aiter): Aiter는 AMD GPU의 아키텍처(CDNA 등)에 최적화된 연산들을 제공합니다. DeepSeek-V4의 MHC 연산은 복잡한 Sinkhorn 반복 연산과 정규화 과정을 포함하는데, 이를 단일 커널(Fused Kernel)로 처리함으로써 GPU 메모리 대역폭 소모를 줄이고 커널 런칭 오버헤드를 최소화합니다.
2. 메모리 효율성: mhc_post에서 torch.empty_like(residual)를 미리 할당하고 커널에 전달하는 방식은 불필요한 메모리 재할당을 방지합니다.
3. 유연한 백엔드 선택: run_dsv4.sh 수정 내역을 보면 SGLANG_HACK_FLASHMLA_BACKEND=tilelang과 SGLANG_OPT_USE_AITER_MHC_PRE=true를 조합하여 사용합니다. 이는 특정 연산은 TileLang으로, 특정 연산은 Aiter로 처리하는 'Best-of-breed' 전략을 취할 수 있게 합니다.

결론

이번 업데이트를 통해 SGLang은 AMD GPU 사용자들에게 DeepSeek-V4 모델에 대한 더 강력한 추론 성능을 제공할 수 있게 되었습니다. 특히 SGLANG_OPT_USE_FUSED_COMPRESS=true 설정과 함께 Aiter 커널이 통합되면서, 대규모 배치 처리 시의 병목 현상이 크게 개선될 것으로 기대됩니다.

AMD 환경에서 LLM 서비스를 운영 중이라면, 최신 SGLang Docker 이미지와 함께 이 Aiter 최적화 옵션을 적극적으로 활용해 보시기 바랍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.empty_like.html
• https://github.com/sgl-project/sglang

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