sgl-project--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

Groups populated in this phase: activation, gemm, kvcache, layernorm, moe, quantization. The remaining groups (attention, communication, diffusion, grammar, mamba, memory, sampling, spatial, speculative) are reserved package placeholders whose implementations still live in sglang.jit_kernel / sgl_kernel / triton_ops and will migrate in later phases.

How it works

Implementations are not moved yet. 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.CUDA_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_cuda_jit / forward_cuda_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/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.jit_kernel or sgl_kernel. When a PR adds a new callable kernel, add a sglang.kernels.ops.* entry point for it, and avoid growing sglang.jit_kernel as a long-term public operator namespace.