[sglang] SGLang의 MoE 성능 최적화: 512 전문가 모델을 위한 커널 최적화

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

들어가며

최근 대규모 언어 모델(LLM) 분야에서는 Mixture of Experts (MoE) 아키텍처가 주류로 자리 잡았습니다. 특히 Qwen3.5-397B와 같이 수많은 전문가(experts)를 활용하는 모델들은 추론 성능을 극대화하기 위해 효율적인 커널 구현이 필수적입니다. 이번 SGLang의 업데이트는 512개의 전문가를 가진 모델을 지원하기 위해 topkGatingSoftmaxKernelLauncher를 최적화하여, 특정 상황에서 최대 4배 이상의 성능 향상을 이끌어냈습니다.

코드 분석

1. sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu

핵심 변경 사항은 CUDA 커널 런처에 512 전문가 케이스를 명시적으로 추가한 것입니다. 기존에는 256까지만 최적화된 경로가 존재했으나, 이를 512까지 확장했습니다.

Before:

case 256:
  LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
  break;
default: {
  // ...
}

After:

case 256:
  LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
  break;
case 512:
  LAUNCH_SOFTMAX(T, 512, WARPS_PER_TB);
  break;
default: {
  // ...
}

또한, needs_workspace 로직을 수정하여 512 전문가 모델이 더 효율적인 경로를 타도록 변경했습니다.

Before:

const bool needs_workspace = !is_pow_2 || num_experts > 256;

After:

const bool needs_workspace = !is_pow_2 || num_experts > 512;

2. 벤치마크 및 테스트 코드 업데이트

새로운 케이스를 검증하기 위해 벤치마크 범위와 테스트 케이스에 topk=10 및 num_experts=512 설정을 추가했습니다.

# sgl-kernel/benchmark/bench_moe_topk_softmax.py
topk_range = [1, 2, 4, 8, 10]

왜 이게 좋은가

이번 최적화는 단순히 케이스를 추가한 것이 아니라, Fused Single-Kernel Path를 활성화함으로써 메모리 접근과 커널 실행 오버헤드를 획기적으로 줄였습니다.

성능 프로파일링 결과, num_tokens가 1024일 때 약 3.45배, 4096일 때 무려 4.25배의 속도 향상을 보였습니다. 이는 대규모 배치 추론 시 병목 현상을 크게 완화할 수 있음을 의미합니다.

일반적 교훈

1. 특수 케이스 최적화: MoE 모델처럼 전문가 수가 고정된 경우, 해당 숫자에 맞춘 커널 런처 최적화는 성능에 직결됩니다.
2. Workspace 관리: needs_workspace와 같은 메모리 할당 로직을 최신 하드웨어 및 모델 사양에 맞춰 튜닝하는 것이 중요합니다.
3. 벤치마크의 중요성: 새로운 기능을 추가할 때 벤치마크 범위를 함께 확장하여 성능 회귀(regression)를 방지하고 최적화의 효과를 수치로 증명해야 합니다.

참고 자료

• https://docs.sglang.ai/developer_guide/contribution_guide.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.size.html

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