본문으로 건너뛰기

[sglang] DeepSeek-V4.1 성능 최적화: num_stages=1 적용으로 _hc_mix_stats_partial_kernel 속도 향상

PR 링크: sgl-project/sglang#43070 상태: Merged | 변경: +5 / -12

들어가며

최근 대규모 언어 모델(LLM)의 발전은 놀라운 속도로 이루어지고 있으며, 특히 저지연(low-latency) 추론 성능은 실시간 서비스 구현에 있어 핵심적인 요소가 되었습니다.

이번 PR은 sglang 라이브러리에서 DeepSeek-V4.1 모델의 저지연 디코딩 성능을 개선하는 데 초점을 맞추고 있습니다. 구체적으로, _hc_mix_stats_partial_kernel 커널의 num_stages 설정을 최적화하여 GPU 연산 효율성을 높이는 것을 목표로 합니다.

기존에는 특정 조건에서 num_stages=2를 사용했으나, 이로 인해 커널 실행 시간이 길어지고 전체 추론 지연 시간에 영향을 미치는 문제가 있었습니다. 본 글에서는 이 변경이 왜 성능 향상으로 이어지는지, 그리고 어떤 기술적 배경을 가지고 있는지 코드 diff를 통해 자세히 분석해보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 python/sglang/kernels/ops/layernorm/mhc.py 파일 내 _hc_mix_stats_reduce_kernel 및 hc_mix_stats_sinkhorn 함수에서 사용되는 num_stages 파라미터의 설정 방식입니다.

1. _HC_MIX_NUM_STAGES 상수 변경

가장 눈에 띄는 변경은 _HC_MIX_NUM_STAGES 상수의 기본값 설정입니다. 기존에는 AMD GPU를 제외한 모든 플랫폼에서 num_stages=2를 기본값으로 사용하고 있었으나, 이제는 AMD GPU(_is_hip가 True일 경우)만 2를 사용하고, 그 외의 경우(주로 NVIDIA GPU)에는 1을 사용하도록 변경되었습니다.

Before:

# num_stages only reorders memory issue, not arithmetic; 2 is enough to cover the
# short k_per_slice loop (K=20480 gives 80 slices, i.e. 4 BLOCK_K tiles per CTA).
_HC_MIX_NUM_STAGES = 2

After:

# num_stages only reorders memory issue, not arithmetic; AMD keeps 2, not measured.
# Measured on B300 and H200: 1 stage beats 2-4 by 2-5x at every M in 1..4096.
_HC_MIX_NUM_STAGES = 2 if _is_hip else 1

주석 또한 변경되어, num_stages=1이 B300 및 H200 GPU에서 모든 M 값 범위(1..4096)에 걸쳐 2-4 stages보다 2-5배 빠르다는 측정 결과를 반영하고 있습니다.

2. _num_stages_for 함수 제거 및 num_stages 직접 사용

기존에는 _num_stages_for라는 함수를 통해 특정 조건(Blackwell 아키텍처, K=20480, 64 <= M <= 384)에서 num_stages를 1로 조정하는 로직이 있었습니다. 이 함수는 제거되었고, 대신 위에서 변경된 _HC_MIX_NUM_STAGES 상수가 커널 호출 시 직접 사용됩니다.

Before (_num_stages_for 함수):

def _num_stages_for(m: int, k: int) -> int:
    # GB300 verify batches benefit from a smaller shared-memory footprint.
    if get_platform().is_blackwell and k == 20480 and 64 <= m <= 384:
        return 1
    return _HC_MIX_NUM_STAGES

After (함수 제거 및 직접 사용):

# ... (이전 코드) ...
    num_stages=_HC_MIX_NUM_STAGES,
# ... (이후 코드) ...

이는 num_stages=1이 NVIDIA GPU에서 전반적으로 더 좋은 성능을 보인다는 측정 결과를 바탕으로, 예외적인 경우를 두지 않고 일관되게 1을 적용하겠다는 의도를 보여줍니다.

3. 커널 호출 부분 수정

hc_mix_stats 및 hc_mix_stats_sinkhorn 함수 내에서 Triton 커널을 호출하는 부분에서 num_stages 인자에 _num_stages_for(m, k) 대신 _HC_MIX_NUM_STAGES가 직접 전달되도록 수정되었습니다.

Before:

        num_stages=_num_stages_for(m, k),

After:

        num_stages=_HC_MIX_NUM_STAGES,

이 변경은 위에서 설명한 _HC_MIX_NUM_STAGES 상수의 변경 사항을 실제 커널 실행에 반영하는 역할을 합니다.

왜 이게 좋은가?

1. 성능 향상

가장 중요한 이점은 _hc_mix_stats_partial_kernel의 실행 속도 향상입니다. PR 설명과 벤치마크 결과에 따르면, num_stages를 2에서 1로 변경함으로써 NVIDIA GPU(B300, H200)에서 상당한 성능 향상을 얻을 수 있었습니다.

벤치마크 결과 (Partial kernel time in us, K=20480, standalone):

M B300 s2 -> s1 H200 s2 -> s1
1-8 13.6 -> 4.6 14.9 -> 5.6
18-48 25.9 -> 6.4 28.3 -> 7.6
52-60 49.9 -> 11.2 42-54 -> 13.8
512 224 -> 56 266 -> 56
4096 1206 -> 525 1383 -> 576

특히, M 값이 작을 때(1-8)는 약 3배, M 값이 커질 때(4096)도 2배 이상의 속도 향상을 보여줍니다. 이는 커널이 다른 커널과 SM을 공유할 때(s2는 36us, s1은 6-7.5us) 더욱 두드러집니다.

이러한 성능 향상은 DeepSeek-V4.1 모델의 전체 추론 지연 시간을 단축시켜, 더 빠르고 효율적인 서비스 제공을 가능하게 합니다.

2. num_stages의 역할 이해

Triton에서 num_stages는 주로 메모리 접근 패턴을 최적화하기 위해 사용됩니다. 여러 스테이지를 사용하면 GPU 코어가 연산을 수행하는 동안 다음 연산에 필요한 데이터를 미리 로드하여 메모리 대기 시간을 줄일 수 있습니다. 이는 특히 데이터 종속성이 있거나 긴 루프를 가진 연산에서 유용할 수 있습니다.

하지만 _hc_mix_stats_partial_kernel의 경우, k_per_slice 루프가 상대적으로 짧고, num_stages=2를 사용할 때 오히려 tf32x3 연산이 긴 직렬 체인(long serial chain)을 형성하게 되어 성능 저하를 유발하는 것으로 보입니다. num_stages=1을 사용하면 이러한 직렬화 문제를 피하고, 각 워프(warp)가 독립적으로 연산을 수행하여 전체적인 처리량을 높일 수 있습니다.

3. 일반적인 교훈

  • 프로파일링과 벤치마킹의 중요성: 특정 하드웨어 아키텍처(NVIDIA Blackwell, B300, H200)와 워크로드(M, K 값)에 따라 최적의 커널 파라미터는 달라질 수 있습니다. num_stages와 같은 파라미터는 단순히 높다고 좋은 것이 아니라, 실제 측정값을 통해 최적값을 찾아야 합니다.
  • 하드웨어 특성 고려: GPU 아키텍처(NVIDIA vs AMD)에 따라 최적의 설정이 다를 수 있습니다. 이 PR은 AMD GPU는 기존 설정을 유지하고 NVIDIA GPU에만 최적화를 적용함으로써 이러한 차이를 반영했습니다.
  • 단순함의 미학: 때로는 복잡한 로직이나 예외 처리를 제거하고 더 단순한 설정을 적용하는 것이 더 나은 성능을 가져올 수 있습니다. _num_stages_for 함수를 제거하고 일관된 num_stages=1을 적용한 것이 좋은 예입니다.

리뷰 피드백

(리뷰 댓글이 제공되지 않아 분석에 반영할 내용이 없습니다. 만약 리뷰 댓글이 있었다면, 해당 댓글에서 제기된 기술적 문제점, 제안된 대안, 그리고 그것이 코드 변경에 어떻게 반영되었는지 등을 분석할 수 있습니다.)

결론

이번 PR은 DeepSeek-V4.1 모델의 추론 성능을 개선하기 위한 중요한 최적화 작업을 수행했습니다. _hc_mix_stats_partial_kernel에서 num_stages 설정을 1로 변경함으로써 NVIDIA GPU에서의 실행 속도를 크게 향상시켰습니다. 이는 LLM 추론 성능 최적화에 있어 하드웨어 특성을 고려한 세밀한 튜닝이 얼마나 중요한지를 다시 한번 보여주는 사례입니다.

이러한 최적화를 통해 sglang 라이브러리는 더 빠르고 효율적인 LLM 추론 서비스를 제공할 수 있게 될 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글