[sglang] DeepSeek-V4.1 성능 최적화: mHC 및 메타데이터 오버헤드 개선 분석

PR 링크: sgl-project/sglang#39704 상태: Merged | 변경: +1039 / -59

들어가며

최근 SGLang 레포지토리에 병합된 DeepSeek-V4.1 관련 최적화 PR은 대규모 MoE(Mixture-of-Experts) 모델 추론 시 발생하는 병목 현상을 해결하는 데 초점을 맞추고 있습니다. 특히 DSpark 검증 배치(concurrency 32/64)에서 발생하는 mHC(Multi-Head Context) 경계 처리, 통신 오버헤드, 그리고 메타데이터 처리 비용을 줄여 실질적인 토큰 처리량(throughput)을 5% 이상 개선했습니다.

코드 분석

1. mHC Post/Combine/RMSNorm 융합 (Prefill 단계)

기존에는 여러 단계로 나뉘어 있던 연산을 하나의 커널로 융합하여 레지스터 사용량을 최적화하고 메모리 접근을 줄였습니다. mhc_post_combine_norm_prefill.cuh 파일에서 확인할 수 있듯이, BF16 중간값을 유지하면서도 Triton 스타일의 정규화 연산을 PTX로 직접 구현했습니다.

// Before: 개별 커널로 실행되던 연산들이 다중 메모리 접근을 유발함
// After: 공유 메모리를 활용한 융합 커널 구현
__shared__ __align__(16) bf16_t collapsed[5120];
// ... (중략) ...
// Vectorize loads but keep only one 8-element tile live at a time.
#pragma unroll 1
for (int vid = tid; vid < 640; vid += Threads) {
  V x, residual[4];
  x.load(p.x + row * 5120, vid);
  // ... 연산 로직 ...
}

2. All-Reduce 커널의 확장 및 융합

all_reduce_fusion.cuh에서는 moe_finalize_all_reduce_kernel을 확장하여 384개 행까지 처리가 가능하도록 수정했습니다. 특히 kCollapse 템플릿 인자를 추가하여 mHC 후처리와 다음 입력 결합(combine) 과정을 융합했습니다.

// Before: mHC와 결합 연산이 분리되어 통신 오버헤드 발생
// After: kCollapse 플래그를 통해 융합된 연산 수행
if constexpr (kMhc) {
  if constexpr (kCollapse) {
    const auto combined = mhc_post_vec<kHiddenDim, true>(params, red, row_idx, hvec);
    ptx::st_global_16B(combined, params.normalized, vid);
  } else {
    mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
  }
}

왜 이게 좋은가

이번 최적화의 핵심은 '커널 융합을 통한 메모리 트래픽 감소'와 '병렬 처리 단위의 확장'입니다.

1. 성능 향상: 벤치마크 결과, Concurrency 32에서 약 5.15%, 64에서 약 5.72%의 토큰 처리량 향상을 보였습니다. 이는 GPU 커널 호출 횟수를 1,835회에서 1,581회로 대폭 줄인 결과입니다.
2. 교훈: 대규모 모델 추론에서는 연산 자체보다 커널 간의 데이터 이동(Memory Bound)이 병목이 되는 경우가 많습니다. 공유 메모리(Shared Memory)를 적극 활용하여 중간 결과를 캐싱하고, 여러 연산을 하나의 커널로 융합하는 방식은 고성능 추론 엔진 설계의 필수적인 패턴입니다.

리뷰 피드백

리뷰어들은 성능 최적화와 함께 정확성 검증을 매우 중요하게 다루었습니다. 특히 test_fp4_indexer.py나 test_fused_moe.py 등 다양한 단위 테스트를 반복적으로 수행하여, 융합된 커널이 기존의 수치적 정확도(NaN/Inf 처리, tie-breaking 등)를 완벽히 보존하는지 검증했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html

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