Add Jit Kernel

Tutorial: Adding a New JIT Kernel to SGLang

Apply kernel-organization for the public operator namespace, logical grouping, lazy registry metadata, and test placement. The implementation tutorial below does not replace that API contract.

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.
• Tuning policy belongs in Python; C++ gets the answer, not the decision. For a hyperparameter — split factor, threads per item, block size, vector width, an algorithm variant — prefer (not required) making it a template parameter and letting the @cache_once module factory choose the value, over a runtime if/switch in the launcher that picks among pre-instantiated kernels. The kernel then compiles exactly one specialisation and asserts its own preconditions with static_assert, while the heuristic that produced the value sits in Python where it is readable, adjustable, and inspectable without recompiling CUDA. store_cache's num_threads is the worked example: a get_kernel(num_split) ladder guarded by byte-alignment if constexprs became one template argument plus a heuristic in _jit_kvcache_module. The exception is a knob that genuinely varies per call with the runtime shape — that has to stay a kernel argument.
• 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.
• namespace details is private. Anything inside one is an implementation detail of its own header, free to change without notice. Do not name details:: from another module, and never using namespace details. If you find yourself wanting something in there, that is the signal to promote it to a documented name instead.
• 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; prefer recomputing a value over letting it spill. Run SGLANG_JIT_LOG_RESOURCE_USAGE=1 SGLANG_JIT_FORCE_RECOMPILE=1 to log per-kernel registers, spills and shared memory at INFO. Both halves are load-bearing: a cache hit has no compiler output to report, and passing -Xptxas -v by hand does nothing on its own because the build captures compiler output and replays it only on failure.
• Alignment is a contract you enforce, not a property you hope for. A vectorized copy derives its access width from the row size (load_bytes picks 16B for a 1024B row), but the address it lands on is base + index * row_stride — and a stride is not constrained by the size. TensorMatcher admits any stride unless you say otherwise, so a padded cache row silently produces misaligned address. Every kernel that vectorizes is highly recommended to end its validation with .ensure_alignment(w), where w is the width the kernel actually uses.

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.

• SGL_DEVICE_HOST — __forceinline__ __device__ __host__. Use it when a constexpr helper must be callable from both the kernel and its launcher, so the two cannot drift.

• device::kWarpThreads — Constant 32 (the codebase's logical warp width). kWarpSize is an alias. Note this is a group size, not the hardware wave: on gfx950 warpSize is 64, and warp::kFullWidth reflects that.

• device::get_lane_id<kNumThreads = kWarpThreads>() — this thread's index within its logical group. Prefer it over threadIdx.x % kNumThreads whenever the value feeds an address. On CUDA the lane register is one read that folds into the address computation, whereas the modulo is a derived value the compiler re-materialises at every address scale — 8 instructions and 2 registers on a two-tile warp copy, measured on sm_100a. It picks the cheaper form per platform, so callers never need to care. Only equals the true in-warp lane when blockDim.x is a multiple of kNumThreads.

• 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)
  • .ensure_alignment(bytes) — require the data pointer and every stride except the innermost to be a multiple of bytes (a power of two). Size-1 dimensions are skipped, since their strides are arbitrary. This is how a vectorized kernel states its precondition; see the alignment convention above.
  • .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)
    .ensure_alignment(16)  // e.g. for 128-bit vectorized load/store
    .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
• device::LoadStoreBytes — named widths to pass around instead of bare integers: MAX_GMEM (arch-dependent, 32 on Blackwell), MAX_SMEM (16), MAX_PORTABLE (16, safe on CUDA and HIP), MIN_COALALESCED (4), and RAW_1B .. RAW_32B.

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.
• device::tile::Memory<T>::warp() — the 32-lane flavour, and what warp::load_bytes is built on. It takes its lane index from get_lane_id() rather than threadIdx.x % 32, for the addressing reason described under utils.cuh. Use warp(int n) for a narrower sub-group; that overload keeps the modulo, since a hardware lane index is the wrong grouping below 32.
• device::tile::Memory<T>::thread() — single-thread accessor, no cooperation.
• .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.

bits.h — Compile-time bit helpers (host::)

#include <sgl_kernel/bits.h>

constexpr wrappers over <bit>, usable in static_assert and in template arguments: host::is_pow2(x), log2_floor(x), log2_ceil(x) (both -1 for x == 0), round_up_pow2(x), round_down_pow2(x). They live in host:: but are constant-expression-only, so device-side static_assert can call them.

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>

Vectorized row copy — reach for this before hand-rolling. A warp cooperatively moving one contiguous row is the single most common shape in this tree, and it used to be re-implemented per kernel.

• warp::load_bytes<kBytes, kPattern>(src) / warp::store_bytes<kBytes, kPattern>(dst, val) — the whole warp moves kBytes from/to one contiguous row. kBytes is arbitrary, down to 1; the helper picks the widest vector that divides it and handles the ragged tail. The returned value is opaque and encodes the width, so a load and its matching store must use the same <kBytes, kPattern>.
• warp::LoadStorePattern — the kPattern values, and the sign carries the meaning. A negative one (WARP_UNIFORM_16B and friends: _GMEM, _SMEM, _4B/_8B/_32B) says the whole warp splits this row, so the width is derived from the per-lane share (gcd(kBytes / 32, |kPattern|)), falling back to gcd(kBytes, 4) when the row is too narrow to give every lane 4 bytes. A positive one (a bare LoadStoreBytes value) makes no warp-splitting assumption and just caps the per-thread vector at gcd(kBytes, kPattern). WARP_UNIFORM_16B is the usual choice for a row copy.
• LoadStorePattern::get_vec_bytes<kBytes, kPattern>() — the width actually chosen. SGL_DEVICE_HOST, so the launcher must call this to feed .ensure_alignment(...). Pass it the same value the kernel passes load_bytes — for a split row that is the per-warp share, not the whole row; the full row can resolve to a different width and would under-constrain the strides.

Reductions.

• warp::reduce<Op, kStart, kFinish>(value, active_mask) — Op is a device::ReductionOp (SUM/MAX/MIN). kStart and kFinish bound a range of lane-index bits: <N, 1> reduces within contiguous groups of N lanes, <kFullWidth, N> reduces across groups at the same offset. The operation is symmetric — <A, B> and <B, A> reduce the same lane set.
• warp::reduce_sum/reduce_max/reduce_min<kStart, kFinish>(value) — wrappers. Work for any type with a ReductionTrait: floats, integers, packed x2. This is more widely used in practice than generic reduce.
• warp::inclusive_reduce<Op, kWidth, kStart, kFinish>(val, lane_id, mask) and inclusive_sum / inclusive_max / inclusive_min — segmented inclusive scan: every lane keeps its own running total, unlike reduce. Forward when kStart < kFinish, backward when kStart > kFinish. lane_id defaults to get_lane_id<kWidth>(), which is correct by construction — only pass it if you already have the segment-relative index (threadIdx.x % kWidth); a warp-relative one silently corrupts every segment past the first when kWidth < 32.
• warp::broadcast<kWidth>(value, src_lane) — src_lane is segment-relative: each kWidth segment reads its own lane, not one warp-wide source.
• warp::get_lane_id, warp::elect_one_lane() (one elected lane, for gating a single-thread TMA issue; CUDA-only), warp::kFullMask and warp::kFullWidth (32 on CUDA, 64 on HIP — note this differs from kWarpThreads).

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).
• device::atomic::Event — a cross-CTA arrive/wait counter in one 32-bit word. Producers call arrive(); consumers call wait(num_producers) (exactly one consumer) or wait_multi<kConsumerBits>(num_producers, num_consumers, n) (several). The last released consumer resets the word, so one Event is reusable across launches with no host re-zero. Preconditions, all of them load-bearing:
  • the word must be zeroed by the host before first use, and must live in global memory — a shared or local Event is an illegal address;
  • arrive() is a release and the waits are acquires, so a producer's writes before arrive() are visible after the wait;
  • generations must be explicitly ordered — the word has no phase bit, so overlapping generation N+1 with a consumer still in generation N is undefined behaviour (a lagging consumer never exits);
  • wait() and wait_multi() lay the word out incompatibly; mixing them on one Event is undefined behaviour.