sglang/python/sglang/kernels

sglang.kernels — unified kernel namespace

This package is the public in-tree import surface for callable kernels, per RFC #29630.

from sglang.kernels.ops.layernorm import rmsnorm
from sglang.kernels.ops.activation import silu_and_mul
from sglang.kernels.ops.kvcache import reshape_and_cache_flash

Layout

sglang/kernels/
  spec.py        # KernelSpec, KernelBackend, FormatSignature,
                 # CapabilityRequirement, PlatformInfo
  registry.py    # process-wide KernelRegistry + register_kernel()
  selector.py    # heuristic select_kernel() and cached get_kernel()
  fused_op.py    # BaseFusedOp: per-operator multi-backend contract
  ops/
    <group>/     # one subpackage per operator group (see list below)
  jit/           # shared JIT CUDA build/runtime infra: utils/, csrc/,
                 # include/, __main__ (KERNEL_PATH resolves here)

Operator groups (all populated): activation, attention, communication, diffusion, elementwise, embeddings, gemm, grammar, kv_canary, kvcache, layernorm, lplb, mamba, memory, moe, quantization, sampling, speculative.

As of the RFC #29630 finale (#32072) the legacy sglang.jit_kernel package has been removed: its shared build/runtime infra moved to sglang.kernels.jit and each JIT-backed operator into its group as sglang.kernels.ops.<group>._jit_<op>. Tests and benchmarks live under test/registered/kernels/ (ops/<group>/ for tests, benchmark/<group>/ for benchmarks); shared test helpers are in sglang.test.kernels.

How it works

Each ops.<group> function is a thin wrapper that forwards to a chosen backend, and every backend is described by a KernelSpec in the registry so alternatives can be inventoried and compared:

• register_kernel(KernelSpec(...)) records metadata only — an operator id ("<group>.<name>"), a backend, and an import path ("module:attr"). No torch or kernel backend is imported, and no JIT compilation is triggered, until a kernel is actually called.
• select_kernel(op, backend=None) resolves an op to its fixed call path. There is no priority ranking or heuristic auto-selection: an op with a single backend resolves to it; an op with several backends must be resolved by naming one (backend=...). The extra backends are inventory only.
• get_kernel(op, backend) resolves and caches the callable; the public wrappers use it, pinned to the backend whose signature they document.

The public wrappers currently default to the AOT sgl_kernel implementation (the stable wheel boundary, broadest shape support). The JIT CUDA backend is registered alongside for inventory; where its signature differs, select it explicitly, e.g.:

from sglang.kernels import select_kernel, KernelBackend
jit_rmsnorm = select_kernel("layernorm.rmsnorm", backend=KernelBackend.JIT).load()

BaseFusedOp — the per-operator implementation contract

Multi-backend operators (currently the layernorm and activation groups) are implemented as BaseFusedOp subclasses: one logical operator with one forward_<backend> method per backend, all sharing one signature behind a single forward():

• forward_native — required; the pure-torch correctness reference every other backend is checked against.
• forward_torch_compile — inherited for free as torch.compile(forward_native).
• forward_triton / forward_jit / forward_aot / forward_cute_dsl / forward_flashinfer / forward_deepgemm — opt-in overrides. A backend is available iff its method is overridden.

forward() auto-selects the best available backend by the class's priority, filtered per call through backend_eligible() (a CapabilityRequirement-vs-PlatformInfo check, extensible with per-call shape/dtype gates), and degrades to the native reference when no optimized backend fits. The public ops.<group> functions stay thin wrappers over module-level instances, so the import surface is unchanged; each instance also registers all of its backends as KernelSpecs so the registry inventory and select_kernel(..., backend=...) keep working.

What this buys (see the RFC discussion):

• Unified correctness testing — a generic harness enumerates available_backends() and asserts each one matches forward_native (test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py); new backends are picked up automatically.
• One-switch debugging — SGLANG_FORCE_FUSED_OP_BACKEND=torch (or set_fused_op_backend(KernelBackend.TORCH)) flips every fused op to its reference implementation for numerical-bug bisection.
• Safe fallbacks — a missing / ineligible optimized kernel degrades to native instead of scattering if/else at call sites.
• Incremental optimization — land forward_native first, add triton / cuda_jit / cuda_aot later without touching call sites; alternative implementations of the same op live side by side for A/B.
• Tracing — enable_fused_op_trace() records every call's op, backend, and tensor shapes/dtypes, giving an accurate inventory of what a model actually exercises.

Review rule (RFC #29630)

SGLang runtime code and tests should import callable kernels from sglang.kernels.ops.*.

Implementation work can still happen in sglang.kernels.jit or sgl_kernel. When a PR adds a new callable kernel, add a sglang.kernels.ops.* entry point for it, and avoid growing sglang.kernels.jit as a long-term public operator namespace.