futuremls-lab/oscar/add-jit-kernel

Tutorial: Adding a New JIT Kernel to SGLang

This tutorial walks through adding a simple element-wise scale operation as a JIT kernel. We'll implement scale(x, factor) = x * factor to demonstrate the complete workflow.

Goal

Add a new operation that scales each element of a tensor by a scalar factor:

• Input: tensor x (CUDA) and scalar factor (float, passed at runtime)
• Output: x * factor (element-wise), allocated internally
• Supported dtypes: FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)

When to use JIT vs AOT (sgl-kernel)

• JIT (`jit_kernel`): prefer this first for kernels that do not depend on CUTLASS or another large C++ project. It is the default choice for lightweight kernels that benefit from rapid iteration and first-use compilation.
• AOT (`sgl-kernel`): prefer this when the kernel does depend on CUTLASS or another large C++ project, or when it should live in sgl-kernel/ and participate in the wheel build / torch op registration flow.
• Exception: kernels that depend on flashinfer, or on CUTLASS that is already provided through flashinfer, can still be implemented as jit_kernel.

Common Abstractions in python/sglang/jit_kernel/include/sgl_kernel/

Always prefer these abstractions over raw CUDA primitives. They provide safety, readability, and consistency with the rest of the codebase.

Important include rule: for every #include <sgl_kernel/...> line, add a short trailing comment explaining why that header is included (for example // For TensorMatcher, SymbolicSize, SymbolicDevice). This matches the current JIT kernel style and keeps include usage self-documenting.

utils.h — Host-side utilities

#include <sgl_kernel/utils.h>

• `host::RuntimeCheck(cond, args...)` — Assert a condition at runtime; throws PanicError with file/line info on failure. Prefer this over bare assert.
• `host::Panic(args...)` — Unconditionally throw a PanicError with a descriptive message.
• `host::div_ceil(a, b)` — Integer ceiling division (a + b - 1) / b.
• `host::irange(n)` / `host::irange(start, end)` — Range views for cleaner loops.
• `host::pointer::offset(ptr, offsets...)` — Byte-safe pointer arithmetic on void*. Use this instead of raw casts.

utils.cuh — Device-side utilities + LaunchKernel

#include <sgl_kernel/utils.cuh>

• Type aliases: fp16_t, bf16_t, fp32_t, fp8_e4m3_t, fp8_e5m2_t and their packed variants fp16x2_t, bf16x2_t, fp32x2_t, etc.
• `SGL_DEVICE` — Expands to __forceinline__ __device__. Use on all device functions.
• `device::kWarpThreads` — Constant 32.
• `device::load_as<T>(ptr, offset)` / `device::store_as<T>(ptr, val, offset)` — Type-safe loads/stores from void*.
• `device::pointer::offset(ptr, offsets...)` — Pointer arithmetic on device.
• `host::LaunchKernel(grid, block, device_or_stream [, smem])` — RAII kernel launcher that:
• Resolves the CUDA stream from a DLDevice via TVM-FFI automatically.
• Checks the CUDA error with file/line info after launch via operator()(kernel, args...).
• Supports .enable_pdl(bool) for PDL (Programmatic Dependent Launch, SM90+).
• `host::RuntimeDeviceCheck(cudaError_t)` — Check a CUDA error; throw on failure.

tensor.h — Tensor validation (TensorMatcher, Symbolic types)

#include <sgl_kernel/tensor.h>

This is the primary validation API for all kernel launchers. Use it to validate every tvm::ffi::TensorView argument.

• `host::SymbolicSize{"name"}` — A named symbolic dimension. Call .set_value(n) to pin it, .unwrap() to extract after verification.
• `host::SymbolicDType` — Symbolic dtype. Use .set_options<Ts...>() to restrict allowed types.
• `host::SymbolicDevice` — Symbolic device. Use .set_options<kDLCUDA>() to restrict to CUDA.
• `host::TensorMatcher({dims...})` — Fluent builder for tensor validation:
• .with_dtype<T>() — require a specific C++ type (e.g. fp16_t)
• .with_dtype<T1, T2, ...>() — allow a set of types
• .with_device<kDLCUDA>(device_sym) — require CUDA and bind the checked device to a SymbolicDevice
• .with_strides({strides...}) — validate strides (omit to require contiguous)
• .verify(tensor_view) — execute the check; throws PanicError with full context on failure; chainable (verify(a).verify(b) to check multiple tensors with the same shape)

Typical pattern:

auto N = SymbolicSize{"num_elements"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
TensorMatcher({N})  //
    .with_dtype<fp16_t>()
    .with_device<kDLCUDA>(device)
    .verify(dst)
    .verify(src);  // same shape, dtype, device as dst
const size_t n = N.unwrap();
const DLDevice dev = device.unwrap();

type.cuh — dtype_trait<T> and packed_t<T>

#include <sgl_kernel/type.cuh>

• `dtype_trait<T>` — Static trait struct for each scalar type. Provides:
• dtype_trait<T>::from(value) — convert from another type (e.g. fp32_t → fp16_t)
• dtype_trait<T>::abs/sqrt/rsqrt/exp/sin/cos(x) — type-dispatched unary math (primarily for fp32_t)
• dtype_trait<T>::max/min(x, y) — type-dispatched binary math (primarily for fp32_t)
• `packed_t<T>` — Two-element packed alias: packed_t<fp16_t> = fp16x2_t, packed_t<bf16_t> = bf16x2_t, packed_t<fp32_t> = fp32x2_t. Use for vectorized loads/stores.
• `device::cast<To, From>(value)` — Type-safe cast using dtype_trait, e.g. cast<fp32x2_t, fp16x2_t>(v).

vec.cuh — Vectorized memory access (AlignedVector)

#include <sgl_kernel/vec.cuh>

• `device::AlignedVector<T, N>` — Aligned storage for N elements of type T. N must be a power of two, sizeof(T)*N <= 32. Enables vectorized loads/stores for bandwidth efficiency. In terms of API/codegen constraints, the upper bound is 256-bit; in practice, 128-bit is the portable default, while 256-bit vectorization is typically only viable on SM100+ and should be gated by an architecture check when needed.
• .load(ptr, offset) — vectorized load from ptr[offset]
• .store(ptr, offset) — vectorized store to ptr[offset]
• .fill(value) — fill all lanes
• operator[](i) — element access

tile.cuh — tile::Memory (strided memory access pattern)

#include <sgl_kernel/tile.cuh>

• tile::Memory<T> is fundamentally a 1D cooperative accessor over a contiguous region.
• `device::tile::Memory<T>::cta(blockDim.x)` — Creates a tile accessor where each thread handles tid = threadIdx.x with stride tsize (for cta(blockDim.x), this is blockDim.x). Common for loops over a 1D array.
• `.load(ptr, offset)` — loads ptr[tid + offset * tsize]
• `.store(ptr, val, offset)` — stores to ptr[tid + offset * tsize]
• `.in_bound(n, offset)` — boundary check

For a 2D tile, either flatten (row, col) into a linear tile index first, or compute the address manually with ptr[row * stride + col] using your thread/block coordinates.

math.cuh — Device math (device::math::)

#include <sgl_kernel/math.cuh>

• device::math::max/min<T>(a, b) — type-dispatched binary math via dtype_trait
• device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x) — type-dispatched unary math via dtype_trait

warp.cuh — Warp-level primitives

#include <sgl_kernel/warp.cuh>

• device::warp::reduce_sum<T>(value) — warp-level sum reduction via __shfl_xor_sync
• device::warp::reduce_max<T>(value) — warp-level max reduction

cta.cuh — CTA-level primitives

#include <sgl_kernel/cta.cuh>

• device::cta::reduce_max<T>(value, smem, min_value) — CTA-wide max using shared memory + warp reduction. Caller is responsible for a __syncthreads() after if the result in smem[0] is needed.

atomic.cuh — Atomic operations

#include <sgl_kernel/atomic.cuh>

• device::atomic::max(float* addr, float value) — float atomic max (handles negative values correctly via bit tricks).

runtime.cuh — Occupancy and device info

#include <sgl_kernel/runtime.cuh>

• host::runtime::get_blocks_per_sm(kernel, block_dim) — max active blocks per SM (occupancy)
• host::runtime::get_sm_count(device_id) — number of SMs on the device
• host::runtime::get_cc_major(device_id) — compute capability major version

Persistent kernel pattern (cap blocks to SM count × occupancy):

static const uint32_t max_occ = runtime::get_blocks_per_sm(kernel, kBlockSize);
static const uint32_t num_sm  = runtime::get_sm_count(device.unwrap().device_id);
const auto num_blocks = std::min(num_sm * max_occ, div_ceil(n, kBlockSize));
LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params);

Step 0 (optional): Generate a .clangd config for better IDE support

python -m sglang.jit_kernel

Step 1: Implement the CUDA kernel in jit_kernel/csrc/

Create python/sglang/jit_kernel/csrc/elementwise/scale.cuh.

The implementation fully uses the project abstractions described above:

#include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/type.cuh>   // For dtype_trait, fp16_t, bf16_t, fp32_t
#include <sgl_kernel/utils.h>    // For RuntimeCheck, div_ceil
#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
#include <sgl_kernel/vec.cuh>    // For AlignedVector
 
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
 
namespace {
 
// ----------------------------------------------------------------
// Kernel: element-wise scale using vectorized 128-bit loads/stores
// T       = fp16_t | bf16_t | fp32_t
// kVecN   = number of elements per vector load (e.g. 8 for fp16)
// factor  = runtime scale factor
// ----------------------------------------------------------------
template <typename T, int kVecN>
__global__ void scale_kernel(T* __restrict__ dst,
                              const T* __restrict__ src,
                              float factor,
                              uint32_t n_total) {
  using vec_t = device::AlignedVector<T, kVecN>;
  const uint32_t n_vecs = n_total / kVecN;
 
  // --- vectorised body ---
  const uint32_t vec_stride = blockDim.x * gridDim.x;
  for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x;
       vi < n_vecs;
       vi += vec_stride) {
    vec_t v;
    v.load(src, vi);
#pragma unroll
    for (int i = 0; i < kVecN; ++i) {
      v[i] = static_cast<T>(static_cast<float>(v[i]) * factor);
    }
    v.store(dst, vi);
  }
 
  // --- scalar tail ---
  const uint32_t base = n_vecs * kVecN;
  const uint32_t scalar_stride = blockDim.x * gridDim.x;
  for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
       base + i < n_total;
       i += scalar_stride) {
    dst[base + i] = static_cast<T>(static_cast<float>(src[base + i]) * factor);
  }
}
 
// ----------------------------------------------------------------
// Launcher: validates tensors, selects vector width, launches kernel
// ----------------------------------------------------------------
template <typename T>
void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
  using namespace host;
 
  // 1. Validate input tensors with TensorMatcher
  SymbolicSize N = {"num_elements"};
  SymbolicDevice device_;
  device_.set_options<kDLCUDA>();
 
  TensorMatcher({N})  //
      .with_dtype<T>()
      .with_device<kDLCUDA>(device_)
      .verify(dst)
      .verify(src);  // same shape / dtype / device as dst
 
  const uint32_t n = static_cast<uint32_t>(N.unwrap());
  const DLDevice device = device_.unwrap();
 
  RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n);