[sglang] SGLang의 Marlin MoE 커널 최적화: JIT 컴파일 시간 특화와 점유율 기반 스케줄링

PR 링크: sgl-project/sglang#31552 상태: Merged | 변경: +175 / -60

들어가며

대규모 언어 모델(LLM)의 추론 성능, 특히 Mixture-of-Experts(MoE) 구조에서의 성능은 연산 커널의 효율성에 크게 좌우됩니다. SGLang 프로젝트의 이번 PR은 Marlin MoE 커널을 대상으로 두 가지 핵심 최적화를 적용하여, Kimi-K2.7 모델의 Prefill 성능을 약 4% 향상시켰습니다. 본 글에서는 이 커널 최적화가 어떻게 이루어졌는지, 그리고 왜 이것이 성능 향상으로 이어졌는지 분석합니다.

코드 분석

1. JIT 컴파일 시간 특화 (Compile-time Specialization)

기존 Marlin 커널은 런타임에 if 분기를 통해 expert-parallel(EP) 여부와 bias 존재 여부를 확인했습니다. 이는 커널의 'hot loop' 내에서 불필요한 분기 예측 실패를 유발합니다.

Before:

if (has_bias) {
  b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
}

After:

if constexpr (kHasBias) {
  b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
}

constexpr를 사용하여 템플릿 인자로 전달된 kIsEP와 kHasBias를 기반으로 컴파일 시점에 코드를 생성합니다. 이를 통해 런타임 분기를 완전히 제거하여 명령어 파이프라인 효율을 극대화했습니다.

2. 점유율 기반 런타임 스케줄링 (Occupancy-aware Launch)

기존에는 유효한 설정 중 첫 번째 것을 무조건 선택했으나, 이제는 레지스터 사용량과 공유 메모리 점유율을 계산하여 최적의 allow_count를 동적으로 결정합니다.

After (moe_wna16_marlin.cuh):

cudaFuncAttributes attr;
cudaFuncGetAttributes(&attr, kernel);
int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
int allow_count = min(device_max_reg_size / reg_size, max_shared_mem / (cache_size + kSharedMemoryValidityMargin + kSharedMemoryLaunchReserve));

이 변경은 GPU의 하드웨어 리소스(레지스터, 공유 메모리)를 고려하여, 단순히 설정 가능한 커널을 선택하는 것을 넘어, 하드웨어 점유율을 극대화할 수 있는 launch configuration을 선택하게 합니다.

왜 이게 좋은가

이번 최적화의 핵심 교훈은 '런타임 오버헤드의 정적 제거'와 '하드웨어 리소스 최적화'입니다.

1. 분기 제거의 효과: kIsEP와 kHasBias를 템플릿화함으로써, 커널 내부의 루프가 단순화되었습니다. 이는 특히 MoE와 같이 수많은 작은 행렬 연산이 반복되는 구조에서 명령어 캐시 적중률과 분기 예측 정확도를 높입니다.
2. 점유율 최적화: GPU 커널 성능은 종종 레지스터 압박(register pressure)에 의해 제한됩니다. cudaFuncGetAttributes를 통해 실제 레지스터 사용량을 계산하고, 이를 바탕으로 allow_count를 조절함으로써 GPU의 병렬성을 최대한 활용할 수 있게 되었습니다.

성능 지표를 보면, Prefill 단계에서 약 4.01%의 속도 향상을 보였으며, 특히 Marlin 커널 자체의 실행 시간은 약 7.77% 단축되었습니다. 이는 복잡한 연산 커널일수록 런타임 판단을 줄이고 컴파일 타임에 최적화하는 전략이 유효함을 보여줍니다.

결론

이 PR은 단순히 코드를 정리하는 것을 넘어, GPU 아키텍처의 특성을 깊이 이해하고 이를 커널 스케줄링에 반영한 훌륭한 사례입니다. 템플릿 메타프로그래밍을 활용한 분기 제거와 하드웨어 리소스 기반의 동적 스케줄링은 고성능 추론 엔진을 개발하는 엔지니어들에게 중요한 시사점을 줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EXECUTION.html#group__CUDART__EXECUTION_1g97316681023774849202535793132616

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