sgl-project--sglang/python/sglang/kernels/README.md

# `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
```

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

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