sgl-kernel

sgl-kernel

Kernel Library for LLM inference engines

sgl-kernel provides optimized compute primitives for LLM inference engines, enabling efficient inference for large language models and vision-language models through custom kernel operations. It has been used by LightLLM, SGLang and so on.

Installation

Requires torch == 2.9.1

# Latest version
pip3 install sgl-kernel --upgrade

Building from Source

Requires

• CMake ≥3.31,
• Python ≥3.10
• scikit-build-core
• ninja(optional)

Use Makefile to build sgl-kernel

make build

Limit build resource usage (CPU / parallelism)

By default, make build uses all available CPU cores. You can override build parallelism and NVCC compile threads:

# Limit parallel jobs (controls both make and cmake parallelism)
make build MAX_JOBS=2

# Additionally limit NVCC internal threads (reduces CPU and peak memory)
make build MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"

Contribution

Steps to add a new kernel:

• Implement the kernel in csrc
• Expose the interface in include/sgl_kernel_ops.h
• Create torch extension in csrc/common_extension.cc
• Update CMakeLists.txt to include new CUDA source
• Expose Python interface in python
• Add test and benchmark

Development Tips

• When creating torch extensions, add the function definition with m.def, and device binding with m.impl:

• How to write schema: Schema reference

  // We need def with schema here for 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);
  

Adapting C++ Native Types for Torch Compatibility

Third-party C++ libraries often use int and float, but PyTorch bindings require int64_t and double due to Python's type mapping.

Use make_pytorch_shim from sgl_kernel_torch_shim.h to handle conversions automatically:

// Add type conversion for int -> int64_t
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(), "value too large");
    TORCH_CHECK(arg >= std::numeric_limits<int>::min(), "value too small");
    return arg;
  }
};

// Wrap your function
m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));

Testing & Benchmarking

• Add pytest tests in tests/, if you need to skip some test, please use @pytest.mark.skipif

@pytest.mark.skipif(
    skip_condition, reason="Nvfp4 Requires compute capability of 10 or above."
)

• Add benchmarks using triton benchmark in benchmark/

  We recommend using triton.testing.do_bench_cudagraph for kernel benchmarking:

  Compared to triton.testing.do_bench, do_bench_cudagraph provides:

  • Reduced CPU overhead impact for more accurate kernel performance measurements
  • Incorporation of PDL (Programmatic Dependent Launch) effects into individual kernel results
  • More realistic performance data on PDL-supported architectures (SM >= 90)

• Run test suite

Kernel Size Analysis

Analyze CUDA kernel sizes in compiled wheel files to identify oversized kernels and template-instantiation bloat:

This tool requires cubloaty (install with pip install cubloaty) to work.

# Install cubloaty
pip install cubloaty

# Analyze a wheel file
python analyze_whl_kernel_sizes.py path/to/sgl_kernel-*.whl

# Custom output file
python analyze_whl_kernel_sizes.py path/to/sgl_kernel-*.whl --output my_analysis.txt

The tool generates:

• A text report with:
  • Kernel groups (by name prefix)
  • Individual kernel sizes (sorted by size)

Use this to identify large kernels and potential template instantiation bloat.

FAQ

• Q: Segmentation fault with CUDA 12.6
• A: Update ptxas to 12.8, reference: segment fault error