[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_pre와 mhc_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_PRE가 False이면 무조건 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)를 정확히 이해하는 것이 더 큰 임팩트를 낼 수 있다는 점입니다.
참고 자료
- https://github.com/sgl-project/sglang
- https://github.com/tile-ai/tilelang
- https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
- [sglang] SM120 Blackwell에서 DeepSeek-V4 모델 서빙 최적화: FlashInfer MXFP4 MoE 도입 및 메모리 절감
- [sglang] SGLang HiCache: Mamba 브랜칭을 위한 증분 백업 최적화
- [sglang] SGLang 캐시 시뮬레이터의 병목 해결: 직렬 전송에서 동시성 전송으로의 최적화
- [sglang] Diffusion DiT 모델의 FFN 성능 최적화: cublasLt GELU Epilogue 융합
- [sglang] DeepSeek-V4 Flash 공식 모델 최적화: H200 및 B200을 위한 SGLang 서빙 전략
PR Analysis 의 다른글
- 이전글 [ultralytics] Ultralytics FLOPs 프로파일링 최적화: deepcopy 제거를 통한 성능 향상
- 현재글 : [sglang] [DeepSeek-V4 최적화] SM120(Blackwell)에서 Fused MHC 커널 활성화로 성능 8.9% 향상시키기
- 다음글 [sglang] SGLang: HPC-Ops 백엔드에서 BF16 디코딩을 위한 동적 스케줄링 도입
댓글