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 python/sglang/kernels/aot/ 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.

Conventions

These hold for every step below.

• namespace sglang is where JIT code lives. Open it after the include block and close it at the end of the file, with the device kernels, traits and host wrapper inside. The shared host:: / device:: helpers are nested in it too, so they resolve unqualified. load_jit emits the TVM_FFI_DLL_EXPORT_TYPED_FUNC wrapper inside namespace sglang as well, so the kernel_name you pass from Python needs no sglang:: prefix.
• Check where the check is cheapest: static_assert > C++ host check > cached Python > per-call Python. Anything fixed at compile time is a static_assert. Anything about the tensors is a TensorMatcher / CHECK_HOST in the C++ launcher, free next to a kernel launch. A check Python cannot delegate goes inside the @cache_once module factory, where it runs once per specialisation. What remains in the per-call entry point costs interpreter time on every forward, so it should be nothing but picking the module and allocating out.
• Fixed-width integer types. Prefer int32_t / int64_t / uint32_t / size_t over int, long, or long long, so an index has the same width on both sides of the FFI boundary. Bare int is fine only where the width plainly cannot matter — an unrolled loop counter over a constexpr bound, a template int parameter. Shapes arrive as int64_t (SymbolicSize::unwrap()); narrowing to uint32_t for in-kernel indexing is a deliberate act, so write the static_cast explicitly and only where the range is known.
• Doxygen comments in C++. Document exported entities with /// or /** ... */ blocks using \brief, \param, \tparam, \return, the way include/sgl_kernel/ does. python -m sglang.kernels.jit writes CommentFormat: Doxygen into .clangd when clangd is 21 or newer, so these render on hover in the editor. Plain // remains fine for implementation notes inside a function body.
• ASCII only in C++ and CUDA sources. Write --, ->, <= instead of —, →, ≤, including in comments. grep -nP '[^\x00-\x7F]' <file> before committing.
• const T* __restrict__ for read-only pointers. This is what csrc/ does throughout, and it lets the compiler emit non-coherent (LDG) loads.
• Watch the register budget. For memory-bound kernels, keep to roughly 64 registers per thread so occupancy does not become the limit. Build once with extra_cuda_cflags=["-Xptxas", "-v"] to see the actual count, and prefer recomputing a value over letting it spill.

Common Abstractions in python/sglang/kernels/jit/include/sgl_kernel/

Always prefer these abstractions over raw CUDA primitives. They provide safety, readability, and consistency with the rest of the codebase. The only reason to drop to raw primitives is performance the abstraction cannot reach — a trade you make deliberately, and justify in a comment.

utils.h — Host-side utilities

#include <sgl_kernel/utils.h>

• CHECK_HOST(cond) << "msg " << value — Preferred runtime check: stream-style, throws PanicError with file/line info on failure. Zero overhead on the true path — the message expressions are only evaluated when the check fails.
• host::RuntimeCheck(cond, args...) — Function-style alternative to CHECK_HOST. Note its message args are always evaluated (even when the check passes), so prefer CHECK_HOST — especially on hot paths.
• 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+).

• device::PDLWaitPrimary<kUsePDL>() / device::PDLTriggerSecondary<kUsePDL>() — The two halves of PDL, on sm_90+ (no-ops on older archs and ROCm). Their guarantees are not symmetric:

  • PDLTriggerSecondary (griddepcontrol.launch_dependents) only lets the next kernel in the stream start early. It carries no memory ordering and publishes nothing — matching that, the header's asm has no "memory" clobber.
  • PDLWaitPrimary (griddepcontrol.wait) is the ordering point: it waits until the preceding kernel has fully finished and its writes are visible.

  So every read of data the preceding kernel produced must come after PDLWaitPrimary(). What overlaps with the primary's tail is whatever you put before the wait — loading parameters, computing indices, touching buffers the primary never wrote — so a kernel that waits on its first line gains nothing. Neither call is a barrier: threads may reach or skip them independently. See "Programmatic Dependent Launch and Synchronization" in the CUDA C++ Programming Guide.

• CHECK_CUDA(expr) << "context" — Stream-style CUDA error check; evaluates expr once and throws PanicError with cudaGetErrorString + file/line info if it is not cudaSuccess. Extra streamed context is optional.

• host::RuntimeDeviceCheck(cudaError_t) — Function-style alternative to CHECK_CUDA. It takes no context message, so prefer CHECK_CUDA, which builds its error object only on the failure path.

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)
• host::is_type<T>(dtype) — whether a DLDataType denotes the C++ type T (e.g. fp16_t).

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 int64_t n = N.unwrap();
const DLDevice dev = device.unwrap();
const int64_t last_dim = 128;
TensorMatcher({N, last_dim})  // a fixed dimension can be a plain integer
    .with_dtype<fp16_t>()
    .with_device<kDLCUDA>(device)
    .verify(tensor_2d);

ffi.h — Tensor allocation and blob wrapping (host::ffi::)

#include <sgl_kernel/ffi.h>

The counterpart to tensor.h: that one validates what came in, this one produces new tvm::ffi::Tensor values. Allocation goes through the environment allocator (TVMFFIEnvTensorAlloc), so buffers come from PyTorch's caching allocator rather than a raw cudaMalloc.

• host::alloc_workspace_tensor(nbytes, device) (declared in utils.cuh) — the way to get scratch memory: a 1-D uint8 tensor of nbytes, or an empty tensor when nbytes == 0. Hold the returned Tensor in a local across every launch that touches it — it frees on destruction.
• host::ffi::empty(shape, dtype, device) — Uninitialized tensor; shape accepts a braced list, so ffi::empty({rows, sizeof(Plan)}, dtype, device) works for a typed scratch array.
• host::ffi::empty_like(tensor_view) — Same shape, dtype, and device as an existing tensor.
• host::ffi::from_blob(data, shape, dtype, device[, deleter, stride, byte_offset]) / from_blob_like(data, tensor_view, ...) — View memory you already own as a Tensor, no copy. The default deleter does nothing, so ownership stays with the caller; pass one only when the Tensor should own the block. Strides default to contiguous.

type.cuh — DTypeTrait<T>, packed_t<T>, and reduction traits

#include <sgl_kernel/type.cuh>

• DTypeTrait<T> — Static trait struct, specialized for integral types, fp32_t, fp16_t, bf16_t, fp8_e4m3_t, and their packed x2/x4 variants. Provides:
  • DTypeTrait<T>::from(value) — convert from another type via the right CUDA intrinsic (e.g. fp32_t → fp16_t)
  • DTypeTrait<T>::abs/max/min — type-dispatched math (fp32, fp16/bf16 scalar and x2, integrals)
  • DTypeTrait<T>::sqrt/rsqrt/exp/sin/cos(x) — fp32_t only
  • Metadata: packed_t / unpacked_t / kVecSize (packed layout), kFloatMax (dtype max as float, e.g. 448.0f for fp8-e4m3), kZeroBits
• 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 DTypeTrait, e.g. cast<fp32x2_t, fp16x2_t>(v).
• device::unpack(value) — View a packed value as an unpacked_t[kVecSize] array reference (e.g. fp32x2_t → fp32_t[2]); element writes propagate back to the packed value.
• device::ReductionOp (SUM/MAX/MIN) and device::ReductionTrait<Op, T>::reduce(x, y) — One binary reduction step, dispatched through DTypeTrait (packed types reduce elementwise). This is the engine behind warp::reduce; use it directly when writing custom reductions.

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 N elements with value
  • 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 DTypeTrait
• device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x) — type-dispatched unary math via DTypeTrait

warp.cuh — Warp-level primitives

#include <sgl_kernel/warp.cuh>

• device::warp::reduce<Op, kNumThreads, kInner>(value, active_mask) — generic warp reduction via __shfl_xor_sync. Op is a device::ReductionOp (SUM/MAX/MIN); kNumThreads is a power-of-two group size (default 32 = full warp); kInner=true (default) reduces within each kNumThreads-sized group, kInner=false reduces across groups (lanes at the same offset in different groups).
• device::warp::reduce_sum/reduce_max/reduce_min<kNumThreads, kInner>(value) — convenience wrappers over reduce. Work for any type with a ReductionTrait: floats, integers, and packed x2 types.

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.kernels.jit -h  # for verbose help info about clangd configuration
python -m sglang.kernels.jit
python -m sglang.kernels.jit --dep cutlass flashinfer  # with cutlass/flashinfer dependency

Step 1: Implement the CUDA kernel in kernels/jit/csrc/

Create python/sglang/kernels/jit/csrc/elementwise/scale.cuh.

The implementation fully uses the project abstractions described above:

// NOTE: Comments for headers are not common in practice.
// It is only shown here for tutorial purposes to highlight the key abstractions.
#include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/type.cuh>   // For DTypeTrait, fp16_t, bf16_t, fp32_t
#include <sgl_kernel/utils.h>    // For CHECK_HOST, 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 sglang {

/**
 * \brief Element-wise scale using vectorized 128-bit loads/stores.
 *
 * \tparam T       Element type: fp16_t | bf16_t | fp32_t
 * \tparam kVecN   Elements per vector load (e.g. 8 for fp16)
 * \tparam kUsePDL Whether to emit the PDL wait/trigger pair
 * \param dst      Output buffer, `n_total` elements
 * \param src      Input buffer, `n_total` elements
 * \param factor   Runtime scale factor
 * \param n_total  Number of elements to scale
 */
template <typename T, int kVecN, bool kUsePDL>
__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;

  // If using PDL, wait for primary kernel before any global memory load.
  // This is NOT a synchronization point, which means some threads can early exit before this.
  device::PDLWaitPrimary<kUsePDL>();

  // --- 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);
  }

  // If using PDL, signal for the secondary kernel to start after all threads have finished
  // This is NOT a synchronization point, which means some threads can early exit before this.
  device::PDLTriggerSecondary<kUsePDL>();
}

/**
 * \brief Validate the tensors, select the vector width, launch `scale_kernel`.
 *
 * \tparam T       Element type: fp16_t | bf16_t | fp32_t
 * \tparam kUsePDL Whether to launch with PDL enabled
 * \param dst      Output tensor; same shape / dtype / device as `src`
 * \param src      Input tensor on CUDA
 * \param factor   Runtime scale factor
 */
template <typename T, bool kUsePDL>
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();

  CHECK_HOST(n > 0) << "scale: num_elements must be > 0, got " << n;

  // 2. Choose vector width for 128-bit loads (16 bytes)
  //    fp16/bf16: 8 elements x 2 bytes = 16 bytes
  //    fp32:      4 elements x 4 bytes = 16 bytes
  // We encourage using `device::kMaxVecBytes`, which will change according to
  // the target architecture and can enable 256-bit vectorization on SM100+ if desired.
  // But 128-bit is more commonly adapted for better compatibility,
  // so it's still ok to hardcode 16 here just for simplicity.
  constexpr int kVecN = 16 / sizeof(T);
  const uint32_t n_work_items = div_ceil(n, static_cast<uint32_t>(kVecN));

  // 3. Launch
  constexpr uint32_t kBlockSize = 256;
  const uint32_t grid = div_ceil(n_work_items, kBlockSize);

  // PDL feature is 100% optional. Without `enable_pdl`, the code should still be correct.
  // Try to enable it if profiling shows that it can benefit the performance of this kernel.
  LaunchKernel(grid, kBlockSize, device).enable_pdl(kUsePDL)(
      scale_kernel<T, kVecN, kUsePDL>,
      static_cast<T*>(dst.data_ptr()),
      static_cast<const T*>(src.data_ptr()),
      factor,
      n);
}

}  // namespace sglang

Key points:

• Include headers from sgl_kernel/ — not raw CUDA headers for anything already covered
• Use TensorMatcher for all tensor validation; never manually check shape/dtype/device
• Use AlignedVector for vectorised 128-bit loads/stores — significant bandwidth win
• Use LaunchKernel — it resolves the stream and checks errors automatically
• Use CHECK_HOST(cond) << ... for runtime assertions with useful error messages (zero overhead when the check passes)
• Prefer passing runtime scalars like factor directly unless compile-time specialisation is genuinely required
• fp16_t / bf16_t / fp32_t are the project's type aliases (from utils.cuh)
• device::cast<To, From> or DTypeTrait<T>::from(val) for cross-type conversions
• device::math:: functions for device math instead of bare __ intrinsics if possible.
• Consider PDL — it can help when the kernel has prologue work to overlap. Place PDLWaitPrimary() right before the first read of upstream data, not at the top of the kernel

Step 2: Add the Python wrapper in kernels/ops/

The wrapper lives next to its functional group under python/sglang/kernels/ops/, not beside the CUDA source — kernels/jit/ holds only the JIT infrastructure (csrc/, include/, utils/, benchmark/). Create python/sglang/kernels/ops/elementwise/scale.py:

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from sglang.kernels.jit.utils import (
    cache_once,
    is_arch_support_pdl,
    load_jit,
    make_cpp_args,
)

if TYPE_CHECKING:
    from tvm_ffi.module import Module


@cache_once
def _jit_scale_module(dtype: torch.dtype) -> Module:
    """Compile and cache the JIT scale module for a given dtype."""
    # Checks on the compile key live here, not in `scale`: `cache_once` keys on
    # `dtype`, so this runs once per specialisation instead of once per call.
    if dtype not in (torch.float16, torch.bfloat16, torch.float32):
        raise RuntimeError(
            f"Unsupported dtype {dtype}. Supported: float16, bfloat16, float32"
        )
    args = make_cpp_args(dtype, is_arch_support_pdl())
    return load_jit(
        "scale",
        *args,
        cuda_files=["elementwise/scale.cuh"],
        cuda_wrappers=[("scale", f"scale<{args}>")],
    )


def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> torch.Tensor:
    """
    Element-wise scale: dst = src * factor.

    Supported dtypes: torch.float16, torch.bfloat16, torch.float32.

    Parameters
    ----------
    src    : CUDA tensor (FP16 / BF16 / FP32)
    factor : scale factor
    out    : optional pre-allocated output tensor (same shape/dtype as src)

    Returns
    -------
    Scaled tensor (dst = src * factor).
    """
    # DO NOT add proactive validation here: every check costs interpreter time
    # on a per-forward path. Tensor invariants belong in the C++ launcher, and
    # anything about the compile key belongs in `_jit_scale_module`.
    if out is None:
        out = torch.empty_like(src)

    module = _jit_scale_module(src.dtype)
    module.scale(out, src, factor)
    return out

Key points:

• Use cache_once — not functools.lru_cache (incompatible with torch.compile)
• load_jit first arg(s) form the unique build marker; same marker = same cached binary
• Only include compile-time specialisation knobs in the build marker; runtime values like factor should stay runtime unless the kernel truly needs templating
• cuda_wrappers: (export_name, kernel_symbol) — export_name is called from Python
• make_cpp_args(dtype, ...) converts torch.dtype to C++ type alias:
• is_arch_support_pdl() checks if the current architecture supports PDL, which is typically passed as a template argument to the kernel.
• Keep the entry point thin (see Conventions). What Python must still check goes in the @cache_once module factory, not in the entry point: cache_once keys on its arguments, so a check there costs one evaluation per specialisation instead of one per call — that is where the supported-dtype guard lives. Tensor invariants belong in the C++ launcher; if something here never reaches a .verify(...), close that gap on the C++ side rather than in Python