sglang: Skill for Claude Code

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.