[SGLang] sgl-kernel: 커스텀 C++/CUDA 커널 라이브러리

들어가며

SGLang의 성능은 커스텀 CUDA 커널에 크게 의존한다. sgl-kernel(PyPI에서는 sglang-kernel으로 배포)은 LLM 추론에 최적화된 C++/CUDA 커널 라이브러리로, AllReduce, Attention, Quantization, Sampling 등 다양한 연산을 제공한다. Python에서는 sgl_kernel으로 임포트한다.

구조도

sgl-kernel/
├── csrc/                          ── C++/CUDA 소스 코드
│   ├── allreduce/                 ── 분산 통신 커널
│   ├── attention/                 ── 어텐션 커널
│   ├── elementwise/               ── 활성화, 노름 등
│   ├── gemm/                      ── GEMM 커널
│   ├── moe/                       ── MoE 커널
│   ├── quantization/              ── 양자화 커널
│   ├── memory/                    ── 메모리 관리
│   ├── mamba/                     ── Mamba SSM 커널
│   ├── grammar/                   ── Grammar 관련
│   ├── kvcacheio/                 ── KV 캐시 I/O
│   ├── expert_specialization/     ── 전문가 특화
│   ├── spatial/                   ── 공간 연산
│   ├── cpu/                       ── CPU 커널 (AMX)
│   ├── cutlass_extensions/        ── CUTLASS 확장
│   ├── common_extension.cc        ── CUDA 바인딩
│   ├── common_extension_rocm.cc   ── ROCm 바인딩
│   └── flash_extension.cc         ── Flash 관련 바인딩
├── include/
│   └── sgl_kernel_ops.h           ── C++ 인터페이스 헤더
├── python/
│   └── sgl_kernel/                ── Python 바인딩
├── benchmark/                     ── 벤치마크
└── CMakeLists.txt                 ── 빌드 설정

핵심 카테고리 분석

Elementwise: Fused 활성화 함수

elementwise/ 디렉토리에는 silu_and_mul, gelu_and_mul, gelu_tanh_and_mul 등의 Fused 활성화 커널이 있다. Python에서는 다음과 같이 사용한다.

from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul

# 사용 예시 (activation.py에서)
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
silu_and_mul(x, out)

AllReduce: 분산 통신

Tensor Parallel에서 사용하는 커스텀 AllReduce 커널이다. NCCL 대비 소규모 텐서에서 더 빠른 성능을 보인다.

Attention: 어텐션 커널

FlashAttention 변형, MLA(Multi-head Latent Attention) 커널, 페이지드 어텐션 등 다양한 어텐션 구현을 제공한다.

Quantization: 양자화 커널

FP8, INT8, INT4 등 다양한 양자화 포맷의 양자화/역양자화 커널이다. Block-wise, Per-tensor, Per-channel 스케일링을 지원한다.

MoE: Mixture of Experts

MoE 모델의 전문가 라우팅, 토큰 디스패치, Fused MoE GEMM 등의 커널이다.

커널 개발 규칙

Torch Extension 등록

새 커널을 추가할 때는 스키마 정의와 디바이스 바인딩을 분리한다.

// csrc/common_extension.cc
// 스키마 정의 (torch.compile 지원)
m.def(
    "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, "
    "Tensor B_scale, Tensor workspace_buffer, int cublas_handle) -> ()");
m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);

C++ 타입 변환

Python의 int는 C++에서 int64_t로 매핑된다. make_pytorch_shim을 사용하면 자동 변환이 가능하다.

template <>
struct pytorch_library_compatible_type<int> {
    using type = int64_t;
    static int convert_from_type(int64_t arg) {
        TORCH_CHECK(arg <= std::numeric_limits<int>::max());
        return arg;
    }
};

m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));

멀티 플랫폼 빌드

CMakeLists.txt를 통해 CUDA, ROCm, CPU, MUSA 등 다양한 플랫폼을 지원한다. 플랫폼별 extension 파일이 분리되어 있다.

common_extension.cc       ── CUDA
common_extension_rocm.cc  ── ROCm (AMD)
common_extension_musa.cc  ── MUSA (Moore Threads)

빌드 방법

# 기본 빌드 (모든 CPU 코어 사용)
make build

# 리소스 제한 빌드
make build MAX_JOBS=2

# NVCC 스레드 제한
make build MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"

CPU 커널 지원

Intel AMX(Advanced Matrix Extensions) 지원 CPU에서는 csrc/cpu/ 아래의 커널을 사용한다.

# activation.py에서의 CPU AMX 경로
if _is_cpu_amx_available:
    out = torch.ops.sgl_kernel.silu_and_mul_cpu(x)

관련 포스트

• Activation Functions: SiLU, GELU 커스텀 구현
• Deep GEMM Wrapper: 최적화 행렬 곱 라이브러리

참고

• 소스 코드: sgl-kernel/
• PyPI: sglang-kernel
• 요구 사항: torch == 2.9.1, CMake >= 3.31, Python >= 3.10