[sglang] SGLang의 MoE Top-K Softmax 커널: AOT에서 JIT로의 효율적인 전환

PR 링크: sgl-project/sglang#34509 상태: Merged | 변경: +1075 / -2

들어가며

최근 대규모 언어 모델(LLM)의 효율적인 추론을 위한 프레임워크인 SGLang에서 moe_topk_softmax 커널을 기존의 AOT(Ahead-of-Time) 방식에서 JIT(Just-in-Time) 방식으로 전환하는 의미 있는 개선이 이루어졌습니다. 기존 AOT 방식은 지원하는 모든 아키텍처에 대해 미리 컴파일된 바이너리를 포함해야 했기에 휠(wheel) 사이즈가 비대해지는 문제가 있었습니다. 이번 PR은 이를 JIT 방식으로 전환하여 런타임에 필요한 아키텍처에 맞춰 커널을 빌드함으로써, 배포 효율성을 높이고 유지보수 비용을 절감하는 것을 목표로 합니다.

코드 분석

1. JIT 커널 포팅 (python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh)

기존 AOT 커널의 CUDA 로직을 그대로 유지하면서, 호스트 측 런처를 TVM FFI(Foreign Function Interface) API를 사용하도록 수정했습니다. 헤더 전용(header-only) 방식으로 포팅하여 JIT 컴파일 환경에서 쉽게 통합될 수 있도록 설계되었습니다.

// Before: AOT 방식에서는 특정 아키텍처에 고정된 .so 파일로 빌드됨
// After: JIT 포팅을 통해 런타임에 컴파일 가능한 구조로 변경
#include <tvm/ffi/container/tensor.h>

template <typename T, int TPB>
__launch_bounds__(TPB) __global__ void moeSoftmax(...) {
  // ... 기존 CUDA 커널 로직 유지 ...
}

2. Python 래퍼 및 등록 (python/sglang/kernels/ops/moe/moe_topk_softmax.py)

load_jit 래퍼를 도입하여 커널이 처음 호출될 때 컴파일되도록 했습니다. 이를 통해 사용자는 별도의 복잡한 설정 없이 기존과 동일한 인터페이스를 유지하면서도 JIT의 이점을 누릴 수 있습니다.

# JIT 로더를 통한 커널 등록
def load_jit():
    return load_library("moe_topk_softmax.cuh")

# 커널 호출 시점에 컴파일 및 실행
register_custom_op("moe_topk_softmax", load_jit())

왜 이게 좋은가

이번 최적화의 핵심은 '성능 저하 없는 유연성 확보'입니다. 벤치마크 결과에 따르면, JIT 방식은 기존 AOT 방식 대비 성능 차이가 거의 없습니다(중앙값 1.0005배, 최대 0.8% 차이). 이는 디바이스 코드가 동일하고 호스트 측 디스패치 방식만 변경되었기 때문입니다.

• 휠 사이즈 최적화: 여러 아키텍처를 위한 바이너리를 하나로 합칠 필요가 없어, 결과적으로 .so 파일의 크기가 약 2배 이상 감소했습니다.
• 유지보수성: 아키텍처별로 별도의 빌드 타겟을 관리할 필요가 없어 CI/CD 파이프라인이 단순해졌습니다.
• 성능: PyTorch의 기본 구현(Softmax + TopK) 대비 약 4.9배 빠른 속도를 유지하여, 커스텀 커널의 이점을 그대로 가져갔습니다.

결론

이번 moe_topk_softmax의 JIT 전환은 라이브러리 배포 사이즈와 유연성 사이의 균형을 맞춘 훌륭한 사례입니다. 특히 성능 손실 없이 배포 효율을 극대화했다는 점에서, 향후 SGLang의 다른 커널들도 JIT 방식으로 전환될 가능성을 보여줍니다. 커널 개발 시 AOT와 JIT의 장단점을 고려하여, 성능이 중요한 핫패스는 유지하되 배포 유연성이 필요한 부분은 JIT로 전환하는 전략이 유효함을 다시 한번 확인했습니다.

참고 자료

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

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