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 unified operator contract

BaseFusedOp is a standard torch.nn.Module (it replaced the former sglang.srt.layers.utils.MultiPlatformOp) that carries one logical operator with interchangeable implementations along two independent dimensions:

• Kernel backends (provenance) — one forward_<backend> method per implementation source, all sharing one signature behind a single forward():
  • forward_native — required; the pure-torch correctness reference every other implementation 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 / forward_aiter / forward_torch_npu — opt-in overrides. A backend is available iff its method is overridden; it joins auto-selection only when also declared in capabilities (device support is metadata, not guesswork).
• Platforms / devices — optional composite per-device paths: forward_cuda, forward_hip (falls back to forward_cuda), forward_npu, forward_xpu, forward_musa (no implicit CUDA fallback — MUSA ops opt into the CUDA path with an explicit forward_musa), forward_cpu (AMX CPUs), plus forward_<key> / register_oot_forward() for out-of-tree platform plugins. CUDA / HIP are not kernel backends.

Dispatch priority, highest first: explicit forward(..., backend=...) → global forced backend (SGLANG_FORCE_FUSED_OP_BACKEND) → OOT platform override → declared optimized kernel backends by priority (filtered by backend_eligible(), a CapabilityRequirement-vs-PlatformInfo check extensible with per-call shape/dtype gates) → platform-specific forward → forward_native. The static part of the decision is resolved once and cached on the instance, so the hot path stays a single indirect call.

BaseFusedOp also owns the torch.compile mode protocol (enter_torch_compile(num_tokens) / leave_torch_compile(), both idempotent): while an outer model is compiled, ops switch to their compile-safe native path so device kernels are never traced (TopK and Fused MoE override _torch_compile_forward() to keep their bs>1 behavior).

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.