본문으로 건너뛰기

[sglang] [DeepSeek-V4 최적화] SM120(Blackwell)에서 Fused MHC 커널 활성화로 성능 8.9% 향상시키기

PR 링크: sgl-project/sglang#30954 상태: Merged | 변경: +84 / -3

들어가며

최신 대규모 언어 모델(LLM)인 DeepSeek-V4는 성능 최적화를 위해 Multi-Head Clamping (MHC)이라는 기법을 사용합니다. SGLang에서는 이를 효율적으로 처리하기 위해 TileLang 기반의 커널을 사용하는데, 여기에는 mhc_premhc_post라는 두 가지 주요 단계가 있습니다.

최근 NVIDIA의 차세대 아키텍처인 SM120(Blackwell) 지원이 추가되면서 한 가지 문제가 발견되었습니다. SM120에서는 안정성이나 호환성 이유로 독립형(standalone) TileLang MHC pre 경로가 기본적으로 비활성화(False)되어 있습니다. 그런데 기존 코드의 논리 구조상, 이 독립형 플래그가 꺼져 있으면 두 단계를 하나로 합친 Fused MHC 커널까지도 함께 비활성화되는 문제가 있었습니다.

이번 PR은 이 잘못된 게이팅(Gating) 로직을 수정하여, SM120에서도 Fused MHC 커널을 명시적으로 사용할 수 있도록 허용함으로써 상당한 성능 이득을 얻어낸 사례입니다.

코드 분석: 잘못된 의존성 연결 고리 끊기

1. python/sglang/srt/models/deepseek_v4.py: 게이팅 로직 수정

가장 핵심적인 변경 사항은 Fused MHC 활성화 여부를 결정하는 _is_fused_mhc_post_pre_enabled() 함수입니다.

Before:

def _is_fused_mhc_post_pre_enabled() -> bool:
    # Fused 경로는 TileLang mhc_post/mhc_pre 커널을 직접 재사용하므로
    # 두 의존성 중 하나라도 꺼져 있으면 비활성화 상태를 유지함
    return (
        envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
        and envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get()
        and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
    )

After:

def _is_fused_mhc_post_pre_enabled() -> bool:
    # SM120은 독립형 TileLang pre 경로를 비활성화함.
    # 하지만 mhc_fused_post_pre는 해당 플래그를 읽지 않고 독립적으로 디스패치하므로,
    # 독립형 pre 플래그가 Fused 옵트인을 거부(veto)해서는 안 됨.
    return (
        envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get()
        and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
        and (envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() or is_sm120_supported())
    )

분석: 기존 코드는 SGLANG_OPT_USE_TILELANG_MHC_PREFalse이면 무조건 Fused 커널을 못 쓰게 막았습니다. 하지만 PR 작성자인 ormandj에 따르면, Fused 커널(mhc_fused_post_pre)은 내부적으로 토큰 수에 따라 독립적인 디스패치 로직을 가지고 있어 독립형 플래그의 상태와 무관하게 동작할 수 있습니다. 따라서 SM120인 경우(is_sm120_supported())에는 독립형 플래그가 꺼져 있더라도 Fused 커널을 사용할 수 있도록 OR 조건을 추가했습니다.

2. test/registered/unit/models/test_deepseek_v4_fused_mhc_policy.py: 정책 검증 테스트 추가

이러한 로직 변경이 다른 아키텍처에 영향을 주지 않는지 확인하기 위해 단위 테스트가 추가되었습니다.

    def test_sm120_allows_fused_opt_in_with_standalone_pre_disabled(self):
        self.assertTrue(
            self._is_enabled(
                fuse=True,
                tilelang_pre=False, # 독립형은 꺼져 있어도
                tilelang_post=True,
                sm120=True,         # SM120이라면 True여야 함
            )
        )

    def test_other_platform_still_requires_tilelang_pre(self):
        self.assertFalse(
            self._is_enabled(
                fuse=True,
                tilelang_pre=False, # 독립형이 꺼져 있으면
                tilelang_post=True,
                sm120=False,        # 일반 플랫폼에서는 False여야 함
            )
        )

이 테스트는 unittest.mock.patch를 사용하여 아키텍처 환경을 시뮬레이션하며, 복잡해진 조건문이 의도대로 작동하는지 보장합니다.

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상 (8.9% throughput 증가)

SM120(RTX PRO 6000 Blackwell) 환경에서 DeepSeek-V4-Flash 모델로 테스트한 결과는 놀랍습니다.

설정 Median C1 tok/s 비고
Fused Off (기존 SM120 기본값) 172.78 hc_pre_torch_impl 사용
Fused On (본 PR 적용) 188.16 +8.9% 향상

단순히 게이팅 로직 하나를 바로잡았을 뿐인데 약 9%에 가까운 처리량 향상을 얻었습니다. 이는 Fused 커널이 활성화되지 않았을 때 시스템이 느린 hc_pre_torch_impl(PyTorch 레벨의 F.linear/rsqrt 구현체)로 폴백(Fallback)하고 있었기 때문입니다.

2. CUDA Graph 노드 감소 및 오버헤드 제거

프로파일링 결과에 따르면, Fused 커널을 사용함으로써 디코드 그래프 내의 커널 노드 수가 3,306개에서 2,796개로 510개나 감소했습니다.

  • Summed Kernel Time: 16.230 ms -> 14.587 ms (-1.643 ms)
  • GPU Span: 12.287 ms -> 11.549 ms

커널 수가 줄어들면 CPU-GPU 간의 런타임 오버헤드가 줄어들고, 특히 CUDA Graph를 사용하는 환경에서 그래프의 복잡도를 낮춰 전체적인 안정성과 속도를 높이는 데 기여합니다.

3. 아키텍처별 최적화의 유연성 확보

특정 하드웨어(SM120)에서 특정 기능(standalone pre)을 끄는 것은 안정성을 위한 선택일 수 있습니다. 하지만 이로 인해 연쇄적으로 다른 최적화(fusion)까지 막히는 것은 전형적인 '과잉 차단' 문제입니다. 이번 수정은 하드웨어 특성을 고려하면서도 최적화 경로를 최대한 확보하는 유연한 설계를 보여줍니다.

결론

이번 PR은 복잡한 커널 코드를 한 줄도 수정하지 않고, 단지 "어떤 상황에서 최적화된 커널을 사용할 것인가"에 대한 정책(Policy)을 정교화함으로써 엄청난 성능 이득을 얻었습니다. 시니어 엔지니어로서 배울 수 있는 교훈은, 때로는 복잡한 알고리즘 개선보다 시스템의 구성 요소들이 서로 어떻게 맞물려 있는지(Gating logic)를 정확히 이해하는 것이 더 큰 임팩트를 낼 수 있다는 점입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글