본문으로 건너뛰기

[sglang] DeepSeek-V4를 위한 MXFP4 Marlin MoE 커널 최적화 및 JIT 통합 분석

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

들어가며

최근 DeepSeek-V4와 같은 대규모 언어 모델(LLM)이 등장하면서, 모델의 크기를 줄이면서도 성능을 유지하기 위한 양자화 기술이 더욱 중요해졌습니다. 특히 MXFP4 (E8M0) 양자화는 높은 압축률과 정밀도를 동시에 제공하여 Hopper 아키텍처(H100, H200)에서 강력한 성능을 발휘합니다.

이번 PR은 기존에 AOT(Ahead-Of-Time) 방식으로 구현되었던 MXFP4 Marlin MoE(Mixture of Experts) 지원 기능을 SGLang의 JIT(Just-In-Time) 커널 경로로 포팅한 작업입니다. 이를 통해 DeepSeek-V4 MXFP4 모델을 --moe-runner-backend marlin 옵션으로 실행할 수 있게 되었으며, 커널 수준의 최적화를 통해 추론 효율을 극대화했습니다.

코드 분석: 핵심 변경 사항

1. CUDA 커널: marlin_template.h의 일반화 및 최적화

가장 큰 변화는 커널 내에서 하드코딩된 타입 체크를 제거하고, 다양한 스케일 타입을 처리할 수 있도록 일반화(Generalization)한 것입니다.

스케일 비트 크기 일반화

기존에는 특정 FP8 타입(kFE2M1f)에 대해서만 예외 처리가 되어 있었으나, 이를 is_8bit_scale이라는 컴파일 타임 상수로 추상화했습니다.

// Before
const int scales_expert_stride = prob_n * prob_k / group_size / (w_type == host::kFE2M1f ? 16 : 8);

// After
constexpr bool is_8bit_scale = s_type.size_bits() == 8;
const int scales_expert_stride = prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8);

이 변경을 통해 E4M3, E8M0 등 다양한 8비트 스케일 타입을 유연하게 지원할 수 있게 되었습니다.

MoE 블록 데이터 읽기 최적화 (cp_async 및 Warp Reduction)

MoE 커널에서 토큰 ID를 읽어오는 로직이 대폭 개선되었습니다. 기존의 순차적인 루프 대신 NVIDIA Hopper의 비동기 복사 명령어인 cp_async와 Warp 수준의 리덕션(__reduce_add_sync)을 도입했습니다.

// After: cp_async를 이용한 비동기 데이터 로드
cp_async4_pred(
    sh_block_sorted_ids_int4 + threadIdx.x,
    reinterpret_cast<const int4*>(sorted_token_ids_ptr) + (block_id * moe_block_size / 4 + threadIdx.x),
    threadIdx.x < moe_block_size / 4
);

// After: Warp-level reduction을 통한 유효 토큰 수 계산
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count);
#else
    // 하위 아키텍처를 위한 fallback (shfl_down_sync)
#endif

cp_async를 사용하면 데이터를 공유 메모리(Shared Memory)로 옮기는 동안 연산을 병렬로 수행할 수 있어 파이프라인 효율이 높아집니다. 또한 __reduce_add_sync는 하드웨어 가속을 통해 Warp 내의 값을 빠르게 합산하여 유효 토큰 수를 계산합니다.

2. Python 레이어: MXFP4 지원 및 가중치 리팩킹

MXFP4는 일반적인 FP16/BF16과 데이터 포맷이 다르기 때문에, 가중치를 Marlin 커널이 이해할 수 있는 형태로 재배치(Repack)하는 과정이 필요합니다.

  • marlin_utils_fp4.py: MXFP4 전용 가중치 리팩킹 및 스케일 순열(Permutation) 로직이 추가되었습니다.
  • fused_marlin_moe.py: float8_e8m0fnu 타입을 감지하고, SwiGLU 활성화 함수의 제한값(swiglu_limit_func) 등을 설정하는 로직이 포함되었습니다.

왜 이게 좋은 최적화인가?

1. 하드웨어 가속 활용 (Hopper Optimized)

cp_async와 __reduce_add_sync는 최신 NVIDIA GPU의 기능을 직접적으로 활용합니다. 특히 MoE 구조에서는 각 전문가(Expert)에게 할당된 토큰의 개수가 매번 달라지는데, 이를 Warp 수준에서 빠르게 집계함으로써 커널 오버헤드를 줄였습니다.

2. JIT 통합을 통한 유연성 확보

AOT 방식은 미리 컴파일된 바이너리에 의존하지만, JIT 경로로 포팅함으로써 다양한 런타임 설정(예: 서로 다른 양자화 파라미터)에 맞춰 커널을 동적으로 최적화할 수 있습니다. 이는 SGLang이 다양한 모델 아키텍처를 지원하는 데 큰 이점을 줍니다.

3. 검증된 성능 및 정확도

PR 설명에 따르면, DeepSeek-V4-Pro 모델을 사용한 AIME25 벤치마크에서 96.25%의 pass@1을 기록했습니다. 이는 최적화된 커널이 단순히 빠른 것뿐만 아니라, 수학적으로 정확하게 MXFP4 양자화 로직을 수행하고 있음을 증명합니다.

결론

이번 변경사항은 SGLang이 최신 양자화 트렌드인 MXFP4를 수용하고, 이를 MoE 아키텍처에서 효율적으로 실행할 수 있는 기반을 마련했습니다. 시니어 엔지니어 입장에서 볼 때, 하드웨어 특화 명령어(cp_async)를 적재적소에 활용하고 복잡한 인덱싱 로직을 div_ceil 패턴 등으로 깔끔하게 정리한 점이 매우 인상적인 PR입니다.

앞으로 Hopper GPU 환경에서 DeepSeek-V4를 서빙하려는 사용자들에게 이 최적화는 필수적인 업데이트가 될 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글