| name | add-jit-kernel |
| description | Step-by-step tutorial for adding a new lightweight JIT CUDA kernel to sglang's jit_kernel module |
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.
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);
const int64_t n = N.unwrap();
const DLDevice dev = device.unwrap();
const int64_t last_dim = 128;
TensorMatcher({N, last_dim})
.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 valueoperator[](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 DTypeTraitdevice::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 devicehost::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
python -m sglang.kernels.jit
python -m sglang.kernels.jit --dep cutlass flashinfer
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:
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
namespace sglang {
template <typename T, int kVecN, bool kUsePDL>
__global__ void scale_kernel(T* __restrict__ dst,
const T* __restrict__ src,
float factor,
uint32_t n_total) {
= device::AlignedVector<T, kVecN>;
n_vecs = n_total / kVecN;
device::<kUsePDL>();
vec_stride = blockDim.x * gridDim.x;
( vi = blockIdx.x * blockDim.x + threadIdx.x;
vi < n_vecs;
vi += vec_stride) {
v;
v.(src, vi);
( i = ; i < kVecN; ++i) {
v[i] = <T>(<>(v[i]) * factor);
}
v.(dst, vi);
}
base = n_vecs * kVecN;
scalar_stride = blockDim.x * gridDim.x;
( i = blockIdx.x * blockDim.x + threadIdx.x;
base + i < n_total;
i += scalar_stride) {
dst[base + i] = <T>(<>(src[base + i]) * factor);
}
device::<kUsePDL>();
}
< T, kUsePDL>
{
host;
SymbolicSize N = {};
SymbolicDevice device_;
device_.<kDLCUDA>();
({N})
.<T>()
.<kDLCUDA>(device_)
.(dst)
.(src);
n = <>(N.());
DLDevice device = device_.();
(n > ) << << n;
kVecN = / (T);
n_work_items = (n, <>(kVecN));
kBlockSize = ;
grid = (n_work_items, kBlockSize);
(grid, kBlockSize, device).(kUsePDL)(
scale_kernel<T, kVecN, kUsePDL>,
<T*>(dst.()),
< T*>(src.()),
factor,
n);
}
}
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 conversionsdevice::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."""
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).
"""
out :
out = torch.empty_like(src)
module = _jit_scale_module(src.dtype)
module.scale(out, src, factor)
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 Pythonmake_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
torch.dtype | C++ type |
|---|
torch.float16 | fp16_t |
torch.bfloat16 | bf16_t |
torch.float32 | fp32_t |
Step 3 (optional): Tune JIT build flags
If your kernel uses some math functions like expf or sinf, consider enabling --use_fast_math for better performance (with a potential precision tradeoff):
return load_jit(
"scale",
*args,
cuda_files=["elementwise/scale.cuh"],
cuda_wrappers=[("scale", f"scale<{args}>")],
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
If your kernel requires SM90+, raise a clear Python error before calling load_jit. Arch gating is one of the checks that has to live in Python — it decides whether to compile at all, so the C++ launcher never gets to run:
if torch.cuda.get_device_capability()[0] < 9:
raise RuntimeError("This kernel requires SM90 (Hopper) or later")
Step 4: Write tests (required)
JIT kernel correctness tests and benchmarks live under test/registered/kernels/ops/<group>/ and test/registered/kernels/benchmark/<group>/, mirroring the wrapper's group under python/sglang/kernels/ops/ (NOT inside the sglang package -- a register_*_ci(...) call anywhere under python/sglang/ is rejected by the check-no-registered-tests-in-package pre-commit hook). Only their test-only helpers (e.g. benchmark/marker.py) stay alongside the kernel source under python/sglang/kernels/jit/ and are imported by absolute path. CI does not run pytest in those directories directly. The unified runner test/run_suite.py discovers every test_*.py and bench_*.py under test/registered/, collects register_*_ci(...) calls by statically parsing each file's AST, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check.
- PR / per-commit CUDA suites (see
test/run_suite.py → PER_COMMIT_SUITES): JIT unit tests use base-b-kernel-unit-test-1-gpu-large on H100 and base-b-kernel-unit-test-4-gpu-b200 on B200/SM100 paths (see .github/workflows/pr-test-jit-kernel.yml). Multi-GPU JIT tests use base-b-kernel-unit-test-8-gpu-h200.
- Nightly kernel suite: register with
stage="nightly" plus the runner_config of the machine it needs (e.g. 1-gpu-large), giving the nightly-test-1-gpu-large suite. .github/workflows/nightly-test-nvidia.yml sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 for the whole nightly run, so the expanded parameter grids apply automatically (see python/sglang/kernels/jit/utils/common.py → should_run_full_tests / get_ci_test_range). There is no separate kernel-only nightly job: every nightly test on one machine type shares that machine's suite.
Registration pattern (module level, literal est_time, stage, and runner_config values — required for AST parsing):
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
CI generates the suite name as {stage}-test-{runner_config}, so stage="base-b-kernel-unit", runner_config="1-gpu-large" becomes the base-b-kernel-unit-test-1-gpu-large suite you pass to run_suite.py below — don't put the -test- infix in register_cuda_ci. Nightly uses the same shape with stage="nightly"; the single-string suite= form is left only for stress and non-CUDA pools.
Keep est_time, stage, runner_config, and suite as literal values. run_suite.py collects them from the file AST, so computed values and helper wrappers can break CI discovery.
Use register_cuda_ci(..., disabled="reason") if the file must stay in-tree but should be skipped in CI (e.g. multi-GPU only).
Run like CI (from repo root):
(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-1-gpu-large)
(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-4-gpu-b200)
For fast iteration you can still run pytest on a single file locally; CI coverage is via run_suite.py.
Create test/registered/kernels/ops/elementwise/test_scale.py:
import pytest
import torch
from sglang.kernels.ops.elementwise.scale import scale
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("size", [1, 127, 128, 1024, 4097])
@pytest.mark.parametrize("factor", [0.5, 1.0, 2.0, 3.0])
def test_scale_correctness(dtype, size, factor):
src = torch.randn(size, dtype=dtype, device="cuda")
out = scale(src, factor)
expected = src * factor
rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-2)
torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_scale_out_param(dtype):
src = torch.randn(1024, dtype=dtype, device="cuda")
out = torch.empty_like(src)
result = scale(src, 2.0, out=out)
result out
torch.testing.assert_close(out, src * , rtol=, atol=)
():
src = torch.randn(, dtype=torch.float16)
pytest.raises(RuntimeError, =):
scale(src, )
():
src = torch.randint(, , (,), dtype=torch.int32, device=)
pytest.raises(RuntimeError, =):
scale(src, )
__name__ == :
sys
sys.exit(pytest.main([__file__, , ]))
Step 5: Add a benchmark (required)
Benchmarks are bench_*.py files under test/registered/kernels/benchmark/<group>/. They are picked up by the same run_suite.py machinery as unit tests. Register them for base-b-kernel-benchmark-test-1-gpu-large (PR JIT benchmark job: python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large).
Benchmarks use the project's own marker framework (in python/sglang/kernels/jit/benchmark/marker.py) — do not use triton.testing.perf_report / triton.testing.do_bench directly. The marker framework provides (public names: benchmark, parametrize, do_bench, skip, BenchResult, BenchSkip):
@marker.benchmark(line_arg, line_vals, *, unit="us") — the innermost decorator (bottom of the stack, directly above def benchmark). Declares the column axis: each value in line_vals becomes a result column, and line_arg is the parameter name passed into the benchmark function. unit is one of "us" | "ms" | "s".@marker.parametrize(names, vals, ci_vals=None) — stackable decorator that adds a row axis (pytest-style). Each @parametrize adds one (or more, correlated) parameter the benchmark is swept over (Cartesian product across all parametrize decorators). names may be a single name ("size") or a comma-separated correlated tuple axis ("h,d", with vals then a list of tuples like [(1, 64), (2, 128)]). Pass the optional third ci_vals for a smaller sweep that is auto-selected under is_in_ci() — this is the built-in CI-shrinking mechanism, so you usually don't need get_benchmark_range for swept axes.marker.do_bench(fn, *, input_args=(), input_kwargs={}, ...) — runs fn under CUDA graph (default) or a naive loop, returns a BenchResult. Key knobs:
memory_args: defaults to "all" (footprint derived from all input args/kwargs). Pass an explicit tuple of tensors (e.g. (k, v, indices)) to count only the inputs the kernel actually touches.memory_output: defaults to "out" — re-runs fn once to capture its returned tensor and counts it. For in-place kernels (which return None), pass the written tensors explicitly (e.g. memory_output=(k, v)); the re-run is then skipped. Set to None to count no output.- Together
memory_args + memory_output give the GB/s column; with both defaults a function out = f(src) already reports bytes(src) + bytes(out).
graph_clone_args / graph_clone_kwargs: which inputs to clone per CUDA-graph iteration to defeat L2 cache reuse. Defaults to — pass an iterable of indices/keys to limit to the args (writes don't need cloning).
Create test/registered/kernels/benchmark/elementwise/bench_scale.py:
import torch
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.jit.benchmark.utils import create_random
from sglang.kernels.ops.elementwise.scale import scale as jit_scale
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large")
@torch.compile()
def torch_impl_scale(src: torch.Tensor, factor: float) -> torch.Tensor:
return src * factor
FN_MAP = {
"jit": jit_scale,
"torch": torch_impl_scale,
}
@marker.parametrize("size", [2**n for n in range(10, 20)], [4096, 65536])
@marker.benchmark("impl", ["jit", "torch"])
def benchmark(size: int, impl: str):
src = create_random(size)
factor = 2.0
return marker.do_bench(
FN_MAP[impl],
input_args=(src, factor),
graph_clone_args=(,),
)
__name__ == :
benchmark.run()
Key points:
- The
line_arg name passed to benchmark ("impl" here) must match a parameter on benchmark(...); same for every parametrize name ("size").
- Stack
@parametrize once per swept axis. The required @marker.benchmark is the innermost decorator (bottom of the stack, directly above the function) — @parametrize rows go above it.
- Prefer
create_random / create_empty from utils.py over open-coding torch.randn(..., dtype=..., device=...).
- The GB/s column appears by default (
memory_args="all" + memory_output="out"). For memory-bound kernels it's the most informative number; scope memory_args / memory_output to the tensors actually touched if the defaults over- or under-count. For compute-bound kernels where bandwidth is misleading, set SGLANG_KERNEL_DISABLE_LOG_BANDWIDTH=1 (or disable_log_bandwidth=True).
- For in-place kernels (which return
None), pass the written tensors via memory_output=(...) since the "out" default would capture nothing.
- Tune
graph_clone_args / graph_clone_kwargs to all the arguments that might be read by the kernel. We can only skip cloning for write-only args. For in-place modified args, we still need to clone them to get accurate timing (reusing the same buffer keeps it L2-hot and skews results).
- Call
benchmark.run() (no print_data= kwarg — the marker framework prints directly).
Run locally:
python test/registered/kernels/benchmark/elementwise/bench_scale.py
Run the benchmark suite the way CI does:
cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large
Troubleshooting
No CI registry found in ... from run_suite.py: add a module-level register_cuda_ci(...) with literal est_time, stage, and runner_config; starred args and non-literal values break AST collection- JIT compilation fails: ensure the
.cuh file is under python/sglang/kernels/jit/csrc/; reduce template argument combinations
- CUDA crash / illegal memory access:
CUDA_LAUNCH_BLOCKING=1; compute-sanitizer --tool memcheck python ...
- Unstable benchmark results:
marker.do_bench uses CUDA-graph-based timing by default; set use_cuda_graph=False only if the kernel can't be captured. graph_clone_args defaults to "all"; if you narrow it, it must still cover every read tensor — reusing a single buffer keeps it L2-hot and skews results. Keep write tensors in it too: they are what sets the rotation count, and a shared output buffer stays L2-hot the same way.
- Missing GB/s column: the column is on by default; check that
SGLANG_KERNEL_DISABLE_LOG_BANDWIDTH is not 1 and disable_log_bandwidth is not True. For in-place kernels (return None) the memory_output="out" default counts nothing — pass the written tensors via memory_output=(...)
References
docs/docs/developer_guide/development_jit_kernel_guide.mdxtest/run_suite.py — suite names, discovery of test/registered/, execution entrypoint for CIpython/sglang/test/ci/ci_register.py — register_cuda_ci and AST registration rulespython/sglang/kernels/jit/utils/compile.py — load_jit, make_cpp_argspython/sglang/kernels/jit/utils/common.py — cache_once, should_run_full_tests, get_ci_test_rangepython/sglang/kernels/jit/include/sgl_kernel/tensor.h — TensorMatcher, SymbolicSize/DType/Device, is_typepython/sglang/kernels/jit/include/sgl_kernel/ffi.h — ffi::empty, ffi::empty_like, ffi::from_blobpython/sglang/kernels/jit/include/sgl_kernel/utils.cuh — type aliases, LaunchKernel, SGL_DEVICEpython/sglang/kernels/jit/include/sgl_kernel/vec.cuh — AlignedVectorpython/sglang/kernels/jit/include/sgl_kernel/tile.cuh — tile::Memorypython/sglang/kernels/jit/include/sgl_kernel/type.cuh — DTypeTrait, packed_t, device::cast, device::unpack, ReductionTraitpython/sglang/kernels/jit/include/sgl_kernel/math.cuh — device::math::python/sglang/kernels/jit/include/sgl_kernel/warp.cuh — warp::reduce<Op> and reduce_sum/max/min wrapperspython/sglang/kernels/jit/include/sgl_kernel/cta.cuh — cta::reduce_maxpython/sglang/kernels/jit/include/sgl_kernel/atomic.cuh — atomic::maxpython/sglang/kernels/jit/include/sgl_kernel/runtime.cuh — occupancy / SM count helperspython/sglang/kernels/jit/csrc/add_constant.cuh — minimal runnable reference
Summary of Files Created
python/sglang/kernels/jit/csrc/elementwise/scale.cuh # NEW: CUDA kernel
python/sglang/kernels/ops/elementwise/scale.py # NEW: Python wrapper
test/registered/kernels/ops/elementwise/test_scale.py # NEW: Tests
test/registered/kernels/benchmark/elementwise/bench_scale.py # NEW: Benchmark
