# `sglang.kernels` — unified kernel namespace

This package is the public in-tree import surface for callable kernels, per
[RFC #29630](https://github.com/sgl-project/sglang/issues/29630).

```python
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.:

```python
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 `KernelSpec`s so the registry inventory and
`select_kernel(..., backend=...)` keep working.

What this buys (see the
[RFC discussion](https://github.com/sgl-project/sglang/issues/29630#issuecomment-4920387930)):

- **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.
