[Kernel] Introduce sglang.kernels namespace and migrate scattered triton_ops kernels (RFC #29630, Phase 2) (#30044)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/jit_kernel/dsa/__init__.py

@@ -1,13 +1,3 @@
-from sglang.srt.utils import is_cuda
-
-_is_cuda = is_cuda()
-
-if _is_cuda:
-    from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
-else:
-    CuteDSLPagedMQALogitsRunner = None
-    pick_dsl_expand = None
-
 from .paged_mqa_logits import (
     aiter_paged_mqa_logits,
     cutedsl_paged_mqa_logits,
@@ -15,6 +5,21 @@ from .paged_mqa_logits import (
     deepgemm_paged_mqa_logits_split,
 )
 
+
+def pick_dsl_expand(*args, **kwargs):
+    from .cutedsl_paged_mqa_logits import pick_dsl_expand as _pick_dsl_expand
+
+    return _pick_dsl_expand(*args, **kwargs)
+
+
+def __getattr__(name: str):
+    if name == "CuteDSLPagedMQALogitsRunner":
+        from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner
+
+        return CuteDSLPagedMQALogitsRunner
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "CuteDSLPagedMQALogitsRunner",
     "pick_dsl_expand",

python/sglang/kernels/README.md

@@ -0,0 +1,110 @@
+# `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.

python/sglang/kernels/__init__.py

@@ -0,0 +1,72 @@
+"""Unified public kernel namespace for SGLang (RFC #29630).
+
+SGLang runtime code and tests should import callable kernels from
+``sglang.kernels.ops.<group>``, e.g.::
+
+    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
+
+Implementations still live in ``sglang.jit_kernel`` (JIT CUDA), the
+``sgl_kernel`` wheel (AOT CUDA/C++), Triton op modules, etc. The ``ops.*``
+functions are thin wrappers that forward to a chosen backend; the
+:data:`~sglang.kernels.registry.registry` provides an inventory of every
+backend so alternatives can be compared and selected. Multi-backend operators
+are structured as :class:`~sglang.kernels.fused_op.BaseFusedOp` subclasses —
+one ``forward_<backend>`` method per backend behind a single ``forward()``,
+with a required pure-``torch`` ``forward_native`` reference and a
+``SGLANG_FORCE_FUSED_OP_BACKEND`` global switch.
+
+Importing this package (and any ``ops.*`` group) does not import a kernel
+backend (``sgl_kernel`` / ``sglang.jit_kernel``) or trigger JIT compilation:
+registration is metadata-only and backends are imported lazily on first call.
+This keeps the namespace usable for inventory tooling on a CPU-only box.
+"""
+
+from sglang.kernels.fused_op import (
+    BaseFusedOp,
+    FusedOpTraceRecord,
+    clear_fused_op_trace,
+    disable_fused_op_trace,
+    enable_fused_op_trace,
+    get_fused_op_backend,
+    get_fused_op_trace,
+    register_fused_op,
+    set_fused_op_backend,
+)
+from sglang.kernels.registry import KernelRegistry, register_kernel, registry
+from sglang.kernels.selector import get_kernel, select_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+    PlatformInfo,
+)
+
+# Importing the operator groups populates the registry (metadata only). Kept
+# after the core imports above (and guarded from isort reordering) so those
+# modules are fully initialized first.
+from sglang.kernels import ops  # noqa: E402  # isort: skip
+
+__all__ = [
+    "ops",
+    "BaseFusedOp",
+    "CapabilityRequirement",
+    "FormatSignature",
+    "FusedOpTraceRecord",
+    "KernelBackend",
+    "KernelRegistry",
+    "KernelSpec",
+    "PlatformInfo",
+    "clear_fused_op_trace",
+    "disable_fused_op_trace",
+    "enable_fused_op_trace",
+    "get_fused_op_backend",
+    "get_fused_op_trace",
+    "get_kernel",
+    "register_fused_op",
+    "register_kernel",
+    "registry",
+    "select_kernel",
+]

python/sglang/kernels/fused_op.py

@@ -0,0 +1,318 @@
+"""Multi-backend operator contract for the unified kernels namespace.
+
+:class:`BaseFusedOp` is the per-operator implementation object behind the
+``sglang.kernels.ops.*`` wrappers (RFC #29630): one logical operator,
+implemented once, with multiple interchangeable backends behind a single
+``forward()``.
+
+Each subclass implements one ``forward_<backend>`` method per backend it
+supports:
+
+- ``forward_native`` — **required**; the pure-``torch`` correctness reference
+  every other backend is checked against.
+- ``forward_torch_compile`` — provided by the base class 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 the subclass overrides its method (``native`` and
+``torch_compile`` are always available). ``forward()`` picks the best
+available backend by :attr:`BaseFusedOp.priority`, filtered per call through
+:meth:`BaseFusedOp.backend_eligible` (which checks
+:class:`~sglang.kernels.spec.CapabilityRequirement` against the detected
+:class:`~sglang.kernels.spec.PlatformInfo`). The ``SGLANG_FORCE_FUSED_OP_BACKEND``
+env var (or :func:`set_fused_op_backend`) forces every fused op onto one
+backend — e.g. ``native`` to bisect numerical bugs against the reference
+implementations with a single switch.
+
+Like the rest of ``sglang.kernels``, importing this module (and instantiating
+subclasses) never imports a kernel backend (``sgl_kernel`` /
+``sglang.jit_kernel``) or triggers JIT compilation; backends are imported
+lazily inside the ``forward_<backend>`` methods.
+"""
+
+from __future__ import annotations
+
+import functools
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar, Dict, List, Mapping, Optional, Tuple
+
+import msgspec
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+    PlatformInfo,
+)
+
+# backend -> forward_<backend> method name.
+BACKEND_METHODS: Dict[KernelBackend, str] = {
+    KernelBackend.TORCH: "forward_native",
+    KernelBackend.TORCH_COMPILE: "forward_torch_compile",
+    KernelBackend.TRITON: "forward_triton",
+    KernelBackend.CUDA_JIT: "forward_cuda_jit",
+    KernelBackend.CUDA_AOT: "forward_cuda_aot",
+    KernelBackend.CUTE_DSL: "forward_cute_dsl",
+    KernelBackend.FLASHINFER: "forward_flashinfer",
+    KernelBackend.DEEPGEMM: "forward_deepgemm",
+}
+
+# best -> fallback. ``torch_compile`` is deliberately absent: auto-selection
+# must never trigger a surprise compilation in a serving process; force it
+# explicitly when wanted.
+DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = (
+    KernelBackend.CUDA_AOT,
+    KernelBackend.CUDA_JIT,
+    KernelBackend.FLASHINFER,
+    KernelBackend.DEEPGEMM,
+    KernelBackend.CUTE_DSL,
+    KernelBackend.TRITON,
+    KernelBackend.TORCH,
+)
+
+# Backends every op supports structurally: forward_native is abstract (so a
+# concrete subclass always has it) and forward_torch_compile derives from it.
+_ALWAYS_AVAILABLE = (KernelBackend.TORCH, KernelBackend.TORCH_COMPILE)
+
+
+@functools.lru_cache(maxsize=1)
+def _platform() -> PlatformInfo:
+    return PlatformInfo.detect()
+
+
+# --- global backend override ------------------------------------------------
+
+# Sentinel distinguishing "not resolved yet" from "resolved to None (no force)".
+_UNRESOLVED = object()
+_forced_backend: Any = _UNRESOLVED
+
+
+def get_fused_op_backend() -> Optional[KernelBackend]:
+    """The process-wide forced backend, or ``None`` for auto-selection.
+
+    Resolved once from ``SGLANG_FORCE_FUSED_OP_BACKEND`` on first use; tests
+    and tools flip it afterwards via :func:`set_fused_op_backend`.
+    """
+    global _forced_backend
+    if _forced_backend is _UNRESOLVED:
+        from sglang.srt.environ import envs
+
+        value = envs.SGLANG_FORCE_FUSED_OP_BACKEND.get()
+        _forced_backend = KernelBackend(value) if value is not None else None
+    return _forced_backend
+
+
+def set_fused_op_backend(backend: Optional[KernelBackend]) -> None:
+    """Force every :class:`BaseFusedOp` onto ``backend`` (``None`` = auto)."""
+    global _forced_backend
+    _forced_backend = backend
+
+
+# --- optional call tracing ----------------------------------------------------
+
+
+class FusedOpTraceRecord(msgspec.Struct, frozen=True):
+    """One traced ``forward()`` call: which op ran on which backend, and the
+    tensor shapes/dtypes it saw."""
+
+    op: str
+    backend: str
+    tensor_args: Tuple[str, ...]  # e.g. "torch.bfloat16[128, 4096]"
+
+
+_trace_enabled: bool = False
+_trace_records: List[FusedOpTraceRecord] = []
+
+
+def enable_fused_op_trace() -> None:
+    """Record every fused-op call (op, backend, tensor shapes/dtypes).
+
+    Gives an accurate inventory of which kernels a model actually exercises
+    and at what shapes — the raw material for representative-shape test sets.
+    """
+    global _trace_enabled
+    _trace_enabled = True
+
+
+def disable_fused_op_trace() -> None:
+    global _trace_enabled
+    _trace_enabled = False
+
+
+def get_fused_op_trace() -> List[FusedOpTraceRecord]:
+    return list(_trace_records)
+
+
+def clear_fused_op_trace() -> None:
+    _trace_records.clear()
+
+
+def _describe_tensors(args: tuple, kwargs: dict) -> Tuple[str, ...]:
+    import torch
+
+    described = []
+    for value in (*args, *kwargs.values()):
+        if isinstance(value, torch.Tensor):
+            described.append(f"{value.dtype}[{', '.join(map(str, value.shape))}]")
+    return tuple(described)
+
+
+# --- the per-operator contract ------------------------------------------------
+
+
+class BaseFusedOp(ABC):
+    """One logical operator with interchangeable backends behind ``forward()``.
+
+    Subclasses set :attr:`op` and implement :meth:`forward_native` plus any
+    optimized ``forward_<backend>`` methods. All backend methods of one op
+    must share the same signature and semantics — each override adapts its
+    underlying kernel's calling convention so call sites never care which
+    backend ran.
+
+    Class attributes
+    ----------------
+    op:
+        Operator id, ``"<group>.<name>"`` (e.g. ``"layernorm.rmsnorm"``).
+    priority:
+        Backend preference for auto-selection, best first. Defaults to
+        :data:`DEFAULT_PRIORITY`.
+    capabilities:
+        Per-backend :class:`CapabilityRequirement`, consulted by
+        :meth:`backend_eligible` (and exported into the registry specs).
+    format_signature:
+        Data-contract description shared by all backends of this op.
+    descriptions:
+        Optional per-backend one-liners for the registry inventory.
+    """
+
+    op: ClassVar[str]
+    priority: ClassVar[Tuple[KernelBackend, ...]] = DEFAULT_PRIORITY
+    capabilities: ClassVar[Mapping[KernelBackend, CapabilityRequirement]] = {}
+    format_signature: ClassVar[FormatSignature] = FormatSignature()
+    descriptions: ClassVar[Mapping[KernelBackend, str]] = {}
+
+    def __init__(self) -> None:
+        # Cache the structural backend set and the priority-ordered subset once
+        # so forward() avoids repeated introspection on the hot path.
+        available = []
+        for backend in KernelBackend:
+            if backend in _ALWAYS_AVAILABLE or self._overrides(
+                BACKEND_METHODS[backend]
+            ):
+                available.append(backend)
+        self._available: Tuple[KernelBackend, ...] = tuple(available)
+        self._ordered: Tuple[KernelBackend, ...] = tuple(
+            b for b in self.priority if b in set(available)
+        )
+        self._compiled_native = None
+
+    def _overrides(self, method_name: str) -> bool:
+        for klass in type(self).__mro__:
+            if klass is BaseFusedOp:
+                return False
+            if method_name in klass.__dict__:
+                return True
+        return False
+
+    # --- backends: native is required; the rest are opt-in overrides ---
+
+    @abstractmethod
+    def forward_native(self, *args, **kwargs):
+        """Pure-``torch`` reference implementation (correctness ground truth)."""
+
+    def forward_torch_compile(self, *args, **kwargs):
+        if self._compiled_native is None:
+            import torch
+
+            self._compiled_native = torch.compile(self.forward_native)
+        return self._compiled_native(*args, **kwargs)
+
+    def forward_triton(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no triton backend")
+
+    def forward_cuda_jit(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no cuda_jit backend")
+
+    def forward_cuda_aot(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no cuda_aot backend")
+
+    def forward_cute_dsl(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no cute_dsl backend")
+
+    def forward_flashinfer(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no flashinfer backend")
+
+    def forward_deepgemm(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.op}: no deepgemm backend")
+
+    # --- selection ---
+
+    def available_backends(self) -> List[KernelBackend]:
+        """Backends this op implements (structural check, platform-agnostic)."""
+        return list(self._available)
+
+    def backend_eligible(self, backend: KernelBackend, *args, **kwargs) -> bool:
+        """Whether ``backend`` may run *this* call.
+
+        The base implementation checks the backend's
+        :class:`CapabilityRequirement` against the detected platform.
+        Subclasses may extend it with per-call shape/dtype gates so
+        auto-selection bounces to the next backend instead of raising.
+        """
+        capability = self.capabilities.get(backend)
+        return capability is None or capability.is_satisfied_by(_platform())
+
+    def _resolve_backend(self, *args, **kwargs) -> KernelBackend:
+        forced = get_fused_op_backend()
+        if forced is not None:
+            return forced
+        for backend in self._ordered:
+            if self.backend_eligible(backend, *args, **kwargs):
+                return backend
+        return KernelBackend.TORCH
+
+    # --- dispatch ---
+
+    def forward(self, *args, backend: Optional[KernelBackend] = None, **kwargs):
+        """Run the op on ``backend``, or on the best eligible one when omitted."""
+        if backend is None:
+            backend = self._resolve_backend(*args, **kwargs)
+        result = getattr(self, BACKEND_METHODS[backend])(*args, **kwargs)
+        if _trace_enabled:
+            _trace_records.append(
+                FusedOpTraceRecord(
+                    op=self.op,
+                    backend=backend.value,
+                    tensor_args=_describe_tensors(args, kwargs),
+                )
+            )
+        return result
+
+    __call__ = forward
+
+
+def register_fused_op(instance: BaseFusedOp, module: str, attr: str) -> BaseFusedOp:
+    """Register every available backend of ``instance`` in the kernel registry.
+
+    ``module``/``attr`` locate the module-level instance so that
+    ``KernelSpec.load()`` can lazily resolve e.g.
+    ``"<module>:<attr>.forward_cuda_aot"`` to the bound backend method. Returns
+    ``instance`` so group packages can write
+    ``_RMSNORM = register_fused_op(_RMSNormOp(), __name__, "_RMSNORM")``.
+    """
+    for backend in instance.available_backends():
+        register_kernel(
+            KernelSpec(
+                op=instance.op,
+                backend=backend,
+                target=f"{module}:{attr}.{BACKEND_METHODS[backend]}",
+                capability=instance.capabilities.get(backend, CapabilityRequirement()),
+                format_signature=instance.format_signature,
+                description=instance.descriptions.get(backend, ""),
+            )
+        )
+    return instance

python/sglang/kernels/ops/__init__.py

@@ -0,0 +1,40 @@
+"""Public operator groups for the ``sglang.kernels`` namespace.
+
+Each submodule corresponds to one operator group from RFC #29630. Importing a
+group registers its :class:`~sglang.kernels.spec.KernelSpec` metadata and
+exposes thin, lazily-dispatched wrapper callables.
+
+Importing this package eagerly imports every group so the registry is fully
+populated for inventory tooling. Group imports are metadata-only and do not
+import ``torch`` or a kernel backend.
+"""
+
+from importlib import import_module
+
+# Operator groups from the RFC's proposed shape. Populated groups expose
+# callable wrappers today; the rest are reserved package placeholders that keep
+# the namespace shape stable for later phases.
+_GROUPS = (
+    "activation",
+    "attention",
+    "communication",
+    "diffusion",
+    "gemm",
+    "grammar",
+    "kvcache",
+    "layernorm",
+    "mamba",
+    "memory",
+    "moe",
+    "quantization",
+    "sampling",
+    "spatial",
+    "speculative",
+)
+
+for _group in _GROUPS:
+    import_module(f"{__name__}.{_group}")
+
+del import_module, _group
+
+__all__ = list(_GROUPS)

python/sglang/kernels/ops/activation/__init__.py

@@ -0,0 +1,184 @@
+"""Fused gated-activation kernels (``act(x[:h]) * x[h:]``).
+
+Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a
+pure-``torch`` reference (``forward_native``) plus AOT (``sgl_kernel``) and
+JIT CUDA backends behind one ``(input, out)`` signature. The JIT backend
+additionally accepts ``expert_ids`` / ``expert_step`` — call
+``forward_cuda_jit`` directly when those are needed.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.fused_op import BaseFusedOp, register_fused_op
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_ACT_DTYPES = ("float16", "bfloat16")
+_CUDA = CapabilityRequirement(requires_cuda=True)
+_ACT_PRIORITY = (
+    KernelBackend.CUDA_AOT,
+    KernelBackend.CUDA_JIT,
+    KernelBackend.TORCH,
+)
+
+
+class _GatedActivationOp(BaseFusedOp):
+    """Shared structure for ``act(x[..., :d]) * x[..., d:]`` operators."""
+
+    # Set by subclasses: sgl_kernel / jit_kernel attr name (same for both).
+    kernel_attr: str
+
+    priority = _ACT_PRIORITY
+    capabilities = {
+        KernelBackend.CUDA_AOT: _CUDA,
+        KernelBackend.CUDA_JIT: _CUDA,
+    }
+    format_signature = FormatSignature(
+        supported_dtypes=_ACT_DTYPES,
+        description="gated activation; returns tensor",
+    )
+
+    def _act(self, gate: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_native(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        d = input.shape[-1] // 2
+        result = self._act(input[..., :d]) * input[..., d:]
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_cuda_aot(
+        self, input: torch.Tensor, out: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        import sgl_kernel
+
+        return getattr(sgl_kernel, self.kernel_attr)(input, out)
+
+    def forward_cuda_jit(
+        self,
+        input: torch.Tensor,
+        out: Optional[torch.Tensor] = None,
+        expert_ids: Optional[torch.Tensor] = None,
+        expert_step: int = 1,
+    ) -> torch.Tensor:
+        import sglang.jit_kernel.activation as jit_activation
+
+        return getattr(jit_activation, self.kernel_attr)(
+            input, out, expert_ids, expert_step
+        )
+
+
+class SiluAndMulOp(_GatedActivationOp):
+    """``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``."""
+
+    op = "activation.silu_and_mul"
+    kernel_attr = "silu_and_mul"
+    descriptions = {
+        KernelBackend.CUDA_AOT: "silu_and_mul (sgl_kernel wheel).",
+        KernelBackend.CUDA_JIT: "silu_and_mul (sglang.jit_kernel).",
+        KernelBackend.TORCH: "silu_and_mul (pure-torch reference).",
+    }
+
+    def _act(self, gate: torch.Tensor) -> torch.Tensor:
+        import torch.nn.functional as F
+
+        return F.silu(gate)
+
+
+class GeluAndMulOp(_GatedActivationOp):
+    """``out = gelu(input[..., :d]) * input[..., d:]`` (erf-based GELU)."""
+
+    op = "activation.gelu_and_mul"
+    kernel_attr = "gelu_and_mul"
+    descriptions = {
+        KernelBackend.CUDA_AOT: "gelu_and_mul (sgl_kernel wheel).",
+        KernelBackend.CUDA_JIT: "gelu_and_mul (sglang.jit_kernel).",
+        KernelBackend.TORCH: "gelu_and_mul (pure-torch reference).",
+    }
+
+    def _act(self, gate: torch.Tensor) -> torch.Tensor:
+        import torch.nn.functional as F
+
+        return F.gelu(gate, approximate="none")
+
+
+class GeluTanhAndMulOp(_GatedActivationOp):
+    """``out = gelu_tanh(input[..., :d]) * input[..., d:]`` (tanh-approximated GELU)."""
+
+    op = "activation.gelu_tanh_and_mul"
+    kernel_attr = "gelu_tanh_and_mul"
+    descriptions = {
+        KernelBackend.CUDA_AOT: "gelu_tanh_and_mul (sgl_kernel wheel).",
+        KernelBackend.CUDA_JIT: "gelu_tanh_and_mul (sglang.jit_kernel).",
+        KernelBackend.TORCH: "gelu_tanh_and_mul (pure-torch reference).",
+    }
+
+    def _act(self, gate: torch.Tensor) -> torch.Tensor:
+        import torch.nn.functional as F
+
+        return F.gelu(gate, approximate="tanh")
+
+
+_SILU_AND_MUL = register_fused_op(SiluAndMulOp(), __name__, "_SILU_AND_MUL")
+_GELU_AND_MUL = register_fused_op(GeluAndMulOp(), __name__, "_GELU_AND_MUL")
+_GELU_TANH_AND_MUL = register_fused_op(
+    GeluTanhAndMulOp(), __name__, "_GELU_TANH_AND_MUL"
+)
+
+
+def silu_and_mul(
+    input: torch.Tensor, out: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``."""
+    return _SILU_AND_MUL(input, out)
+
+
+def gelu_and_mul(
+    input: torch.Tensor, out: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """``out = gelu(input[..., :d]) * input[..., d:]``."""
+    return _GELU_AND_MUL(input, out)
+
+
+def gelu_tanh_and_mul(
+    input: torch.Tensor, out: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """``out = gelu_tanh(input[..., :d]) * input[..., d:]``."""
+    return _GELU_TANH_AND_MUL(input, out)
+
+
+__all__ = [
+    "SiluAndMulOp",
+    "GeluAndMulOp",
+    "GeluTanhAndMulOp",
+    "silu_and_mul",
+    "gelu_and_mul",
+    "gelu_tanh_and_mul",
+]
+
+
+# Triton kernel migrated into this group (from layers/triton_ops/softcap);
+# registered for inventory. Import it from its module.
+for _fn in ("softcap_out", "softcap_inplace_logits"):
+    register_kernel(
+        KernelSpec(
+            op=f"activation.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.activation.softcap:{_fn}",
+        )
+    )
+del _fn

python/sglang/kernels/ops/attention/__init__.py

@@ -0,0 +1,43 @@
+"""Attention compute kernels (Triton): decode / extend / prefill / metadata.
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.attention.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+KV-cache index/write kernels went to the ``kvcache`` group instead.
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from layers/attention/triton_ops + model_executor.
+_TRITON_KERNELS = [
+    ("decode_attention", "decode_attention_fwd"),
+    ("extend_attention", "extend_attention_fwd"),
+    ("extend_attention", "build_unified_kv_indices"),
+    ("prefill_attention", "context_attention_fwd"),
+    ("merge_state", "merge_state_triton"),
+    ("metadata", "get_num_kv_splits_triton"),
+    ("metadata", "prepare_swa_spec_page_table_triton"),
+    ("metadata", "normal_decode_set_metadata"),
+    ("dsa_metadata", "fused_dsa_decode_metadata"),
+    ("dsa_metadata", "fused_dsa_target_verify_metadata"),
+    ("dsa_metadata", "fused_dsa_draft_extend_metadata"),
+    ("rocm_mla_decode_rope", "decode_attention_fwd_grouped_rope"),
+    ("verify_splitkv", "verify_splitkv_fwd"),
+    ("pad", "pad_sequence_with_mask"),
+    ("pad", "pad_draft_extend_query"),
+    ("pad", "unpad_draft_extend_output"),
+    ("pad", "seqlens_expand_triton"),
+    ("position", "compute_position_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"attention.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.attention.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/attention/extend_attention.py

@@ -20,8 +20,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.triton_ops.decode_attention import _extract_kv_strides
-from sglang.srt.layers.attention.triton_ops.prefill_attention import (
+from sglang.kernels.ops.attention.decode_attention import _extract_kv_strides
+from sglang.kernels.ops.attention.prefill_attention import (
     context_attention_fwd,
 )
 from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip

python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py

@@ -23,7 +23,7 @@ It supports page size = 1.
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.triton_ops.decode_attention import (
+from sglang.kernels.ops.attention.decode_attention import (
     _decode_softmax_reducev_fwd,
 )
 

python/sglang/kernels/ops/communication/__init__.py

@@ -0,0 +1,12 @@
+"""Collective-communication kernels (custom all-reduce, ...).
+
+Reserved group in the ``sglang.kernels`` namespace (RFC #29630). No thin
+wrappers are exposed here: the collective ops (custom all-reduce and friends)
+are stateful — they manage workspaces / IPC handles and are driven through a
+``CustomAllreduce``-style object and ``torch.ops.sgl_kernel.*`` bindings rather
+than standalone callable kernels, so a thin ``sglang.kernels.ops`` forwarder
+would be misleading. Import them from ``sgl_kernel`` / ``sglang.jit_kernel``
+directly until a proper stateful-op interface is designed.
+"""
+
+__all__ = []

python/sglang/kernels/ops/diffusion/__init__.py

@@ -0,0 +1,109 @@
+"""Diffusion-model kernels (group-norm+silu, residual-gate-add, qk-norm+rope).
+
+These are JIT CUDA kernels; the wrappers forward to ``sglang.jit_kernel.diffusion``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+    from torch import nn
+
+_CUDA = CapabilityRequirement(requires_cuda=True)
+
+register_kernel(
+    KernelSpec(
+        op="diffusion.apply_group_norm_silu",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.diffusion.group_norm_silu:apply_group_norm_silu",
+        capability=_CUDA,
+        format_signature=FormatSignature(description="fused GroupNorm + SiLU"),
+        description="Fused group-norm + SiLU (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="diffusion.residual_gate_add",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.diffusion.residual_gate_add:residual_gate_add_cuda",
+        capability=_CUDA,
+        format_signature=FormatSignature(description="residual + gate * update"),
+        description="Fused residual gate-add (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="diffusion.fused_inplace_qknorm_rope",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.diffusion.qknorm_rope:fused_inplace_qknorm_rope",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            in_place=True, description="fused in-place QK-norm + RoPE"
+        ),
+        description="Fused QK-norm + RoPE (sglang.jit_kernel).",
+    )
+)
+
+
+def apply_group_norm_silu(
+    x: torch.Tensor, norm: nn.Module, activation: nn.Module
+) -> torch.Tensor:
+    """Fused GroupNorm + SiLU (falls back to eager when unsupported)."""
+    return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.CUDA_JIT)(
+        x, norm, activation
+    )
+
+
+def residual_gate_add(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    """Fused ``residual + gate * update``."""
+    return get_kernel("diffusion.residual_gate_add", KernelBackend.CUDA_JIT)(
+        residual, update, gate
+    )
+
+
+def fused_inplace_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    *,
+    is_neox: bool,
+    eps: float = 1e-6,
+    head_dim: int = 0,
+    rope_dim: int = 0,
+) -> None:
+    """Fused in-place QK RMS-norm + RoPE."""
+    return get_kernel("diffusion.fused_inplace_qknorm_rope", KernelBackend.CUDA_JIT)(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        cos_sin_cache,
+        positions,
+        is_neox=is_neox,
+        eps=eps,
+        head_dim=head_dim,
+        rope_dim=rope_dim,
+    )
+
+
+__all__ = [
+    "apply_group_norm_silu",
+    "residual_gate_add",
+    "fused_inplace_qknorm_rope",
+]

python/sglang/kernels/ops/gemm/__init__.py

@@ -0,0 +1,138 @@
+"""GEMM and fused-GEMM kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_CUDA = CapabilityRequirement(requires_cuda=True)
+
+register_kernel(
+    KernelSpec(
+        op="gemm.fp8_scaled_mm",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:fp8_scaled_mm",
+        format_signature=FormatSignature(
+            supported_dtypes=("float8_e4m3fn",),
+            description="C = (A_fp8 @ B_fp8) * scales_a * scales_b (+ bias)",
+        ),
+        description="FP8 scaled matmul (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_fused_a_gemm",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:dsv3_fused_a_gemm",
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 fused QKV-A GEMM",
+        ),
+        description="DeepSeek-V3 fused-A GEMM (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_fused_a_gemm",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.dsv3_fused_a_gemm:dsv3_fused_a_gemm",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 fused QKV-A GEMM (drop-in with AOT signature)",
+        ),
+        description="DeepSeek-V3 fused-A GEMM (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="gemm.dsv3_router_gemm",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.dsv3_router_gemm:dsv3_router_gemm",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            supported_dtypes=("bfloat16",),
+            description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]",
+        ),
+        description="DeepSeek-V3 router GEMM (sglang.jit_kernel, JIT-only).",
+    )
+)
+
+
+def fp8_scaled_mm(
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    scales_a: torch.Tensor,
+    scales_b: torch.Tensor,
+    out_dtype: torch.dtype,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 scaled matmul: ``(mat_a @ mat_b) * scales_a * scales_b (+ bias)``."""
+    return get_kernel("gemm.fp8_scaled_mm", KernelBackend.CUDA_AOT)(
+        mat_a, mat_b, scales_a, scales_b, out_dtype, bias
+    )
+
+
+def dsv3_fused_a_gemm(
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V3 fused QKV-A GEMM."""
+    return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.CUDA_AOT)(
+        mat_a, mat_b, output
+    )
+
+
+def dsv3_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weights: torch.Tensor,
+    out_dtype: Optional[torch.dtype] = None,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V3 router GEMM (JIT-backed). ``out_dtype`` defaults to bfloat16."""
+    impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.CUDA_JIT)
+    if out_dtype is None:
+        return impl(hidden_states, router_weights, output=output)
+    return impl(hidden_states, router_weights, out_dtype, output)
+
+
+__all__ = ["fp8_scaled_mm", "dsv3_fused_a_gemm", "dsv3_router_gemm"]
+
+
+# LoRA SGMV Triton kernels migrated into this group (from lora/triton_ops);
+# registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("chunked_embedding_lora_a", "chunked_embedding_lora_a_forward"),
+    ("chunked_sgmv_expand", "chunked_sgmv_lora_expand_forward"),
+    ("chunked_sgmv_shrink", "chunked_sgmv_lora_shrink_forward"),
+    ("embedding_lora_a", "embedding_lora_a_fwd"),
+    ("gate_up_lora_b", "gate_up_lora_b_fwd"),
+    ("qkv_lora_b", "qkv_lora_b_fwd"),
+    ("sgemm_lora_a", "sgemm_lora_a_fwd"),
+    ("sgemm_lora_b", "sgemm_lora_b_fwd"),
+    ("kv_b_lora_absorbed", "step_a_q_fwd"),
+    ("kv_b_lora_absorbed", "step_b_q_fwd"),
+    ("kv_b_lora_absorbed", "step_a_v_fwd"),
+    ("kv_b_lora_absorbed", "step_b_v_fwd"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"gemm.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.gemm.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/gemm/chunked_sgmv_expand.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_expand_config
+from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_expand_config
 from sglang.srt.lora.utils import LoRABatchInfo
 from sglang.srt.utils import cached_triton_kernel
 

python/sglang/kernels/ops/gemm/chunked_sgmv_shrink.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.lora_tuning_config import get_lora_shrink_config
+from sglang.kernels.ops.gemm.lora_tuning_config import get_lora_shrink_config
 from sglang.srt.lora.utils import LoRABatchInfo
 from sglang.srt.utils import cached_triton_kernel
 

python/sglang/kernels/ops/gemm/gate_up_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/gemm/kv_b_lora_absorbed.py

@@ -48,7 +48,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
 from sglang.srt.lora.utils import LoRABatchInfo
 
 # ---------------------------------------------------------------------------

python/sglang/kernels/ops/gemm/lora_tuning_config.py

@@ -20,7 +20,7 @@ Usage:
     python3 benchmark/kernels/lora_csgmv/tune_lora_csgmv.py \
         --model Qwen/Qwen3-Embedding-0.6B --max-lora-rank 64
 
-    # Configs saved to python/sglang/srt/lora/triton_ops/configs/
+    # Configs saved to python/sglang/kernels/ops/gemm/configs/
 
     # Server automatically picks them up:
     python3 -m sglang.launch_server --model ... --enable-lora --lora-backend csgmv

python/sglang/kernels/ops/gemm/qkv_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/gemm/sgemm_lora_a.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/gemm/sgemm_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.triton_ops.kernel_utils import _resolve_token_positions
+from sglang.kernels.ops.gemm.kernel_utils import _resolve_token_positions
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/gemm/trtllm_lora_temp/__init__.py

@@ -0,0 +1,3 @@
+"""Experimental TRT-LLM LoRA kernel variants (gated by ``SGLANG_EXPERIMENTAL_LORA_OPTI`` / ``lora_envs``).
+
+Migrated from ``sglang.srt.lora.trtllm_lora_temp.triton_ops`` (RFC #29630)."""

python/sglang/kernels/ops/gemm/trtllm_lora_temp/gate_up_lora_b.py

@@ -2,11 +2,11 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     _resolve_token_positions,
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 from sglang.srt.lora.utils import LoRABatchInfo
 
 # Minimum total_tokens * rank for the single-adapter cuBLAS path; below this

python/sglang/kernels/ops/gemm/trtllm_lora_temp/kv_b_lora_absorbed.py

@@ -48,11 +48,11 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     _resolve_token_positions,
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 from sglang.srt.lora.utils import LoRABatchInfo
 
 # ---------------------------------------------------------------------------

python/sglang/kernels/ops/gemm/trtllm_lora_temp/qkv_lora_b.py

@@ -4,11 +4,11 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     _resolve_token_positions,
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 from sglang.srt.lora.utils import LoRABatchInfo
 
 # Minimum max_len (longest segment) for the single-adapter cuBLAS path; below

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_a.py

@@ -4,11 +4,11 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     _resolve_token_positions,
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/gemm/trtllm_lora_temp/sgemm_lora_b.py

@@ -2,14 +2,14 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.gate_up_lora_b import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.gate_up_lora_b import (
     _CUBLAS_MIN_S_RANK,
 )
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     _resolve_token_positions,
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 from sglang.srt.lora.utils import LoRABatchInfo
 
 

python/sglang/kernels/ops/grammar/__init__.py

@@ -0,0 +1,26 @@
+"""Constrained-decoding / grammar kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.grammar.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from constrained/triton_ops.
+_TRITON_KERNELS = [
+    ("bitmask_ops", "apply_token_bitmask_inplace_triton"),
+    ("token_filter_ops", "set_token_filter_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"grammar.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.grammar.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/kvcache/__init__.py

@@ -0,0 +1,89 @@
+"""KV-cache write/transfer kernels.
+
+This group wraps the Triton ``reshape_and_cache`` launcher, whose implementation
+now lives in this package (``sglang.kernels.ops.kvcache.cache_ops``) after being
+migrated out of ``sglang.srt.layers.attention.triton_ops`` (RFC #29630).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
+
+if TYPE_CHECKING:
+    import torch
+
+register_kernel(
+    KernelSpec(
+        op="kvcache.reshape_and_cache_flash",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.kvcache.cache_ops:launch_reshape_and_cache_flash",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="write token-major K/V into paged KV cache layout",
+        ),
+        description="Reshape-and-cache (Triton launcher).",
+    )
+)
+
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    swa_slot_mapping: Optional[torch.Tensor] = None,
+    k_scale: Optional[torch.Tensor] = None,
+    v_scale: Optional[torch.Tensor] = None,
+) -> None:
+    """Write token-major ``key``/``value`` into paged KV cache layout."""
+    return get_kernel("kvcache.reshape_and_cache_flash", KernelBackend.TRITON)(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        k_scale,
+        v_scale,
+    )
+
+
+__all__ = ["reshape_and_cache_flash"]
+
+
+# Other Triton kernels migrated into this group (from attention/mem_cache
+# triton_ops); registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("cache_ops", "concat_and_cast_mha_k_triton"),
+    ("cache_ops", "launch_reshape_and_cache_flash"),
+    ("kv_indices", "create_flashinfer_kv_indices_triton"),
+    ("kv_indices", "create_flashmla_kv_indices_triton"),
+    ("kv_indices", "create_chunked_prefix_cache_kv_indices"),
+    ("kv_indices", "get_num_kv_index_blocks_flashmla"),
+    ("kv_indices", "get_num_page_per_block_flashmla"),
+    ("rope_cache", "fused_qk_rope_reshape_and_cache"),
+    ("trtllm_fp8_kv_kernel", "fused_fp8_set_kv_buffer"),
+    ("trtllm_mha_page_table", "build_trtllm_mha_page_table"),
+    ("trtllm_mha_graph_metadata", "update_trtllm_mha_graph_metadata"),
+    ("aiter_unified_attention", "scatter_ragged_to_page_table_kernel"),
+    ("aiter_unified_attention", "scatter_req_to_token_to_page_table_kernel"),
+    ("cache_move", "store_cache_4d"),
+    ("cache_move", "set_kv_buffer_prefix_valid_tiled"),
+    ("cache_move", "copy_all_layer_kv_cache_tiled"),
+    ("mla_buffer", "set_mla_kv_buffer_triton"),
+    ("mla_buffer", "get_mla_kv_buffer_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"kvcache.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.kvcache.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/layernorm/__init__.py

@@ -0,0 +1,338 @@
+"""Layer-normalization kernels.
+
+Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a
+pure-``torch`` reference (``forward_native``) plus optimized CUDA backends,
+all behind one signature. The public module-level functions are thin wrappers
+over module-level instances; auto-selection prefers the AOT ``sgl_kernel``
+implementation on CUDA and falls back to the native reference elsewhere.
+Pick a specific backend with e.g.
+``_RMSNORM.forward(x, w, backend=KernelBackend.CUDA_JIT)`` or globally via
+``SGLANG_FORCE_FUSED_OP_BACKEND``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.fused_op import BaseFusedOp, register_fused_op
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_NORM_DTYPES = ("float16", "bfloat16")
+_CUDA = CapabilityRequirement(requires_cuda=True)
+_NORM_PRIORITY = (
+    KernelBackend.CUDA_AOT,
+    KernelBackend.CUDA_JIT,
+    KernelBackend.TORCH,
+)
+
+
+class RMSNormOp(BaseFusedOp):
+    """``out = (input / RMS(input)) * weight``; returns a tensor.
+
+    ``enable_pdl`` is honored by the AOT backend only.
+    """
+
+    op = "layernorm.rmsnorm"
+    priority = _NORM_PRIORITY
+    capabilities = {
+        KernelBackend.CUDA_AOT: _CUDA,
+        KernelBackend.CUDA_JIT: _CUDA,
+    }
+    format_signature = FormatSignature(
+        supported_dtypes=_NORM_DTYPES,
+        description="out = (x / RMS(x)) * weight; returns tensor",
+    )
+    descriptions = {
+        KernelBackend.CUDA_AOT: "RMS normalization (sgl_kernel wheel).",
+        KernelBackend.CUDA_JIT: "RMS normalization (sglang.jit_kernel).",
+        KernelBackend.TORCH: "RMS normalization (pure-torch reference).",
+    }
+
+    def forward_native(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch
+
+        x = input.to(torch.float32)
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + eps)
+        result = (x * weight).to(input.dtype)
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_cuda_aot(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import sgl_kernel
+
+        return sgl_kernel.rmsnorm(input, weight, eps, out, enable_pdl)
+
+    def forward_cuda_jit(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch
+
+        from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
+
+        if out is None:
+            out = torch.empty_like(input)
+        jit_rmsnorm(input, weight, out, eps)
+        return out
+
+
+class FusedAddRMSNormOp(BaseFusedOp):
+    """In-place ``residual += input; input = RMSNorm(residual) * weight``.
+
+    Writes the sum into ``residual`` and the normalized value into ``input``;
+    returns ``None``. ``enable_pdl`` is honored by the AOT backend only.
+    """
+
+    op = "layernorm.fused_add_rmsnorm"
+    priority = _NORM_PRIORITY
+    capabilities = {
+        KernelBackend.CUDA_AOT: _CUDA,
+        KernelBackend.CUDA_JIT: _CUDA,
+    }
+    format_signature = FormatSignature(
+        supported_dtypes=_NORM_DTYPES,
+        in_place=True,
+        description="residual += x; x = RMSNorm(residual) * weight",
+    )
+    descriptions = {
+        KernelBackend.CUDA_AOT: (
+            "Fused residual-add + RMS normalization (sgl_kernel wheel)."
+        ),
+        KernelBackend.CUDA_JIT: (
+            "Fused residual-add + RMS normalization (sglang.jit_kernel)."
+        ),
+        KernelBackend.TORCH: (
+            "Fused residual-add + RMS normalization (pure-torch reference)."
+        ),
+    }
+
+    def forward_native(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import torch
+
+        acc = input.to(torch.float32) + residual.to(torch.float32)
+        residual.copy_(acc.to(residual.dtype))
+        variance = acc.pow(2).mean(dim=-1, keepdim=True)
+        normed = acc * torch.rsqrt(variance + eps)
+        input.copy_((normed * weight).to(input.dtype))
+
+    def forward_cuda_aot(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import sgl_kernel
+
+        return sgl_kernel.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
+
+    def forward_cuda_jit(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
+
+        return jit_fused_add_rmsnorm(input, residual, weight, eps)
+
+
+class GemmaRMSNormOp(BaseFusedOp):
+    """``out = (input / RMS(input)) * (weight + 1)``; returns a tensor."""
+
+    op = "layernorm.gemma_rmsnorm"
+    priority = _NORM_PRIORITY
+    capabilities = {KernelBackend.CUDA_AOT: _CUDA}
+    format_signature = FormatSignature(
+        supported_dtypes=_NORM_DTYPES,
+        description="out = (x / RMS(x)) * (weight + 1); returns tensor",
+    )
+    descriptions = {
+        KernelBackend.CUDA_AOT: "Gemma-style RMS normalization (sgl_kernel wheel).",
+        KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).",
+    }
+
+    def forward_native(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import torch
+
+        x = input.to(torch.float32)
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + eps)
+        result = (x * (1.0 + weight.to(torch.float32))).to(input.dtype)
+        if out is None:
+            return result
+        out.copy_(result)
+        return out
+
+    def forward_cuda_aot(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        out: Optional[torch.Tensor] = None,
+        enable_pdl: Optional[bool] = None,
+    ) -> torch.Tensor:
+        import sgl_kernel
+
+        return sgl_kernel.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
+
+
+class GemmaFusedAddRMSNormOp(BaseFusedOp):
+    """In-place ``residual += input; input = GemmaRMSNorm(residual) * (weight + 1)``."""
+
+    op = "layernorm.gemma_fused_add_rmsnorm"
+    priority = _NORM_PRIORITY
+    capabilities = {KernelBackend.CUDA_AOT: _CUDA}
+    format_signature = FormatSignature(
+        supported_dtypes=_NORM_DTYPES,
+        in_place=True,
+        description="residual += x; x = GemmaRMSNorm(residual) * (weight + 1)",
+    )
+    descriptions = {
+        KernelBackend.CUDA_AOT: ("Gemma-style fused residual-add + RMS normalization."),
+        KernelBackend.TORCH: (
+            "Gemma-style fused residual-add + RMS normalization "
+            "(pure-torch reference)."
+        ),
+    }
+
+    def forward_native(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import torch
+
+        acc = input.to(torch.float32) + residual.to(torch.float32)
+        residual.copy_(acc.to(residual.dtype))
+        variance = acc.pow(2).mean(dim=-1, keepdim=True)
+        normed = acc * torch.rsqrt(variance + eps)
+        input.copy_((normed * (1.0 + weight.to(torch.float32))).to(input.dtype))
+
+    def forward_cuda_aot(
+        self,
+        input: torch.Tensor,
+        residual: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-6,
+        enable_pdl: Optional[bool] = None,
+    ) -> None:
+        import sgl_kernel
+
+        return sgl_kernel.gemma_fused_add_rmsnorm(
+            input, residual, weight, eps, enable_pdl
+        )
+
+
+_RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM")
+_FUSED_ADD_RMSNORM = register_fused_op(
+    FusedAddRMSNormOp(), __name__, "_FUSED_ADD_RMSNORM"
+)
+_GEMMA_RMSNORM = register_fused_op(GemmaRMSNormOp(), __name__, "_GEMMA_RMSNORM")
+_GEMMA_FUSED_ADD_RMSNORM = register_fused_op(
+    GemmaFusedAddRMSNormOp(), __name__, "_GEMMA_FUSED_ADD_RMSNORM"
+)
+
+
+def rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+    enable_pdl: Optional[bool] = None,
+) -> torch.Tensor:
+    """RMS normalization: ``out = (input / RMS(input)) * weight``."""
+    return _RMSNORM(input, weight, eps, out, enable_pdl)
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    """In-place fused residual add + RMS normalization."""
+    return _FUSED_ADD_RMSNORM(input, residual, weight, eps, enable_pdl)
+
+
+def gemma_rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+    enable_pdl: Optional[bool] = None,
+) -> torch.Tensor:
+    """Gemma-style RMS normalization: ``out = (input / RMS(input)) * (weight + 1)``."""
+    return _GEMMA_RMSNORM(input, weight, eps, out, enable_pdl)
+
+
+def gemma_fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    """In-place Gemma-style fused residual add + RMS normalization."""
+    return _GEMMA_FUSED_ADD_RMSNORM(input, residual, weight, eps, enable_pdl)
+
+
+__all__ = [
+    "RMSNormOp",
+    "FusedAddRMSNormOp",
+    "GemmaRMSNormOp",
+    "GemmaFusedAddRMSNormOp",
+    "rmsnorm",
+    "fused_add_rmsnorm",
+    "gemma_rmsnorm",
+    "gemma_fused_add_rmsnorm",
+]

python/sglang/kernels/ops/mamba/__init__.py

@@ -0,0 +1,86 @@
+"""State-space / Mamba kernels (causal conv1d)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
+
+if TYPE_CHECKING:
+    import torch
+
+register_kernel(
+    KernelSpec(
+        op="mamba.causal_conv1d_fwd",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.mamba:causal_conv1d_fwd",
+        format_signature=FormatSignature(
+            in_place=True, description="causal depthwise conv1d forward (prefill)"
+        ),
+        description="Causal conv1d forward (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="mamba.causal_conv1d_update",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.mamba:causal_conv1d_update",
+        format_signature=FormatSignature(
+            in_place=True, description="causal depthwise conv1d update (decode)"
+        ),
+        description="Causal conv1d update (sgl_kernel wheel).",
+    )
+)
+
+
+def causal_conv1d_fwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias_: Optional[torch.Tensor],
+    conv_states: Optional[torch.Tensor],
+    query_start_loc: Optional[torch.Tensor],
+    cache_indices: Optional[torch.Tensor],
+    has_initial_state: Optional[torch.Tensor],
+    silu_activation: bool,
+    pad_slot_id: int,
+):
+    """Causal depthwise conv1d forward (prefill)."""
+    return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.CUDA_AOT)(
+        x,
+        weight,
+        bias_,
+        conv_states,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        silu_activation,
+        pad_slot_id,
+    )
+
+
+def causal_conv1d_update(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias_: Optional[torch.Tensor],
+    silu_activation: bool,
+    cache_seqlens: Optional[torch.Tensor],
+    conv_state_indices: Optional[torch.Tensor],
+    pad_slot_id: int,
+):
+    """Causal depthwise conv1d update (decode)."""
+    return get_kernel("mamba.causal_conv1d_update", KernelBackend.CUDA_AOT)(
+        x,
+        conv_state,
+        weight,
+        bias_,
+        silu_activation,
+        cache_seqlens,
+        conv_state_indices,
+        pad_slot_id,
+    )
+
+
+__all__ = ["causal_conv1d_fwd", "causal_conv1d_update"]

python/sglang/kernels/ops/memory/__init__.py

@@ -0,0 +1,30 @@
+"""Memory / KV-slot allocation kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.memory.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from mem_cache/triton_ops.
+_TRITON_KERNELS = [
+    ("allocator", "alloc_extend_kernel"),
+    ("allocator", "alloc_decode_kernel"),
+    ("common", "write_req_to_token_pool_triton"),
+    ("common", "get_last_loc_triton"),
+    ("common", "get_last_loc_triton_safe"),
+    ("virtual_slot", "alloc_bind_inplace"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"memory.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.memory.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/moe/__init__.py

@@ -0,0 +1,119 @@
+"""Mixture-of-Experts routing / bookkeeping kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_CUDA = CapabilityRequirement(requires_cuda=True)
+
+register_kernel(
+    KernelSpec(
+        op="moe.moe_align_block_size",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:moe_align_block_size",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="align/sort expert token ids into block-padded buffers",
+        ),
+        description="MoE align-block-size (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="moe.moe_align_block_size",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.moe_align:moe_align_block_size",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            in_place=True,
+            description="MoE align-block-size (JIT variant, AOT signature)",
+        ),
+        description="MoE align-block-size (sglang.jit_kernel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="moe.topk_softmax",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:topk_softmax",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="top-k softmax routing weights/ids",
+        ),
+        description="MoE top-k softmax (sgl_kernel wheel).",
+    )
+)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+    cumsum_buffer: torch.Tensor,
+    pad_sorted_token_ids: bool = False,
+) -> None:
+    """Align and sort expert token ids into block-padded output buffers."""
+    return get_kernel("moe.moe_align_block_size", KernelBackend.CUDA_AOT)(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+        cumsum_buffer,
+        pad_sorted_token_ids,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    moe_softcapping: float = 0.0,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """Compute top-k softmax routing weights/ids for MoE."""
+    return get_kernel("moe.topk_softmax", KernelBackend.CUDA_AOT)(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        renormalize,
+        moe_softcapping,
+        correction_bias,
+    )
+
+
+__all__ = ["moe_align_block_size", "topk_softmax"]
+
+
+# Fused MoE-LoRA Triton kernels migrated into this group (from lora/triton_ops);
+# registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("fused_moe_lora_kernel", "fused_moe_lora"),
+    ("virtual_experts", "merged_experts_fused_moe_lora_add"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"moe.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.moe.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/moe/trtllm_lora_temp/__init__.py

@@ -0,0 +1,3 @@
+"""Experimental TRT-LLM LoRA kernel variants (gated by ``SGLANG_EXPERIMENTAL_LORA_OPTI`` / ``lora_envs``).
+
+Migrated from ``sglang.srt.lora.trtllm_lora_temp.triton_ops`` (RFC #29630)."""

python/sglang/kernels/ops/moe/trtllm_lora_temp/virtual_experts.py

@@ -10,10 +10,10 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.moe_align import moe_align_block_size as jit_moe_align_block_size
-from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
-from sglang.srt.lora.trtllm_lora_temp.triton_ops.kernel_utils import (
+from sglang.kernels.ops.gemm.trtllm_lora_temp.kernel_utils import (
     get_pdl_launch_metadata,
 )
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
 
 
 @triton.jit

python/sglang/kernels/ops/quantization/__init__.py

@@ -0,0 +1,122 @@
+"""Quantization kernels (per-token / per-token-group FP8 & INT8)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import (
+    CapabilityRequirement,
+    FormatSignature,
+    KernelBackend,
+    KernelSpec,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+_CUDA = CapabilityRequirement(requires_cuda=True)
+
+register_kernel(
+    KernelSpec(
+        op="quantization.sgl_per_token_quant_fp8",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel:sgl_per_token_quant_fp8",
+        format_signature=FormatSignature(
+            supported_dtypes=("float8_e4m3fn",),
+            in_place=True,
+            description="per-token FP8 quantization into output_q/output_s",
+        ),
+        description="Per-token FP8 quantization (sgl_kernel wheel).",
+    )
+)
+# fp8 / int8 are legacy aliases of the same 8bit kernel in the wheel; register
+# each public name so runtime imports resolve to a stable spec.
+for _name in (
+    "sgl_per_token_group_quant_8bit",
+    "sgl_per_token_group_quant_fp8",
+    "sgl_per_token_group_quant_int8",
+):
+    register_kernel(
+        KernelSpec(
+            op=f"quantization.{_name}",
+            backend=KernelBackend.CUDA_AOT,
+            target=f"sgl_kernel:{_name}",
+            format_signature=FormatSignature(
+                in_place=True,
+                description="per-token-group 8-bit quantization",
+            ),
+            description=f"{_name} (sgl_kernel wheel).",
+        )
+    )
+del _name
+
+register_kernel(
+    KernelSpec(
+        op="quantization.sgl_per_token_group_quant_8bit",
+        backend=KernelBackend.CUDA_JIT,
+        target="sglang.jit_kernel.per_token_group_quant_8bit:per_token_group_quant_8bit",
+        capability=_CUDA,
+        format_signature=FormatSignature(
+            in_place=True,
+            description="per-token-group 8-bit quantization (JIT variant)",
+        ),
+        description="Per-token-group 8-bit quantization (sglang.jit_kernel).",
+    )
+)
+
+
+def sgl_per_token_quant_fp8(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    """Per-token FP8 quantization, writing into ``output_q`` / ``output_s``."""
+    return get_kernel("quantization.sgl_per_token_quant_fp8", KernelBackend.CUDA_AOT)(
+        input, output_q, output_s
+    )
+
+
+def sgl_per_token_group_quant_8bit(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool = False,
+    fuse_silu_and_mul: bool = False,
+    masked_m: Optional[torch.Tensor] = None,
+    enable_v2: Optional[bool] = None,
+) -> None:
+    """Per-token-group 8-bit quantization, writing into ``output_q`` / ``output_s``."""
+    return get_kernel(
+        "quantization.sgl_per_token_group_quant_8bit", KernelBackend.CUDA_AOT
+    )(
+        input,
+        output_q,
+        output_s,
+        group_size,
+        eps,
+        fp8_min,
+        fp8_max,
+        scale_ue8m0,
+        fuse_silu_and_mul,
+        masked_m,
+        enable_v2,
+    )
+
+
+# Legacy aliases kept for source compatibility with existing call sites.
+sgl_per_token_group_quant_fp8 = sgl_per_token_group_quant_8bit
+sgl_per_token_group_quant_int8 = sgl_per_token_group_quant_8bit
+
+
+__all__ = [
+    "sgl_per_token_quant_fp8",
+    "sgl_per_token_group_quant_8bit",
+    "sgl_per_token_group_quant_fp8",
+    "sgl_per_token_group_quant_int8",
+]

python/sglang/kernels/ops/sampling/__init__.py

@@ -0,0 +1,56 @@
+"""Sampling kernels (top-k / top-p probability renormalization)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Union
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
+
+if TYPE_CHECKING:
+    import torch
+
+register_kernel(
+    KernelSpec(
+        op="sampling.top_k_renorm_probs",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.sampling:top_k_renorm_probs",
+        format_signature=FormatSignature(
+            description="renormalize probs by top-k thresholding; returns tensor"
+        ),
+        description="Top-k probability renormalization (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="sampling.top_p_renorm_probs",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.sampling:top_p_renorm_probs",
+        format_signature=FormatSignature(
+            description="renormalize probs by top-p thresholding; returns tensor"
+        ),
+        description="Top-p probability renormalization (sgl_kernel wheel).",
+    )
+)
+
+
+def top_k_renorm_probs(
+    probs: torch.Tensor, top_k: Union[torch.Tensor, int]
+) -> torch.Tensor:
+    """Renormalize ``probs`` by top-k thresholding."""
+    return get_kernel("sampling.top_k_renorm_probs", KernelBackend.CUDA_AOT)(
+        probs, top_k
+    )
+
+
+def top_p_renorm_probs(
+    probs: torch.Tensor, top_p: Union[torch.Tensor, float]
+) -> torch.Tensor:
+    """Renormalize ``probs`` by top-p thresholding."""
+    return get_kernel("sampling.top_p_renorm_probs", KernelBackend.CUDA_AOT)(
+        probs, top_p
+    )
+
+
+__all__ = ["top_k_renorm_probs", "top_p_renorm_probs"]

python/sglang/kernels/ops/spatial/__init__.py

@@ -0,0 +1,49 @@
+"""Spatial / green-context stream helpers."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
+
+register_kernel(
+    KernelSpec(
+        op="spatial.get_sm_available",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.spatial:get_sm_available",
+        format_signature=FormatSignature(
+            description="number of SMs available on device"
+        ),
+        description="Query available SM count (sgl_kernel wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="spatial.create_greenctx_stream_by_value",
+        backend=KernelBackend.CUDA_AOT,
+        target="sgl_kernel.spatial:create_greenctx_stream_by_value",
+        format_signature=FormatSignature(
+            description="create two green-context streams partitioned by SM count"
+        ),
+        description="Green-context stream creation (sgl_kernel wheel).",
+    )
+)
+
+
+def get_sm_available(device_id: Optional[int] = None) -> int:
+    """Return the number of SMs available on ``device_id``."""
+    return get_kernel("spatial.get_sm_available", KernelBackend.CUDA_AOT)(device_id)
+
+
+def create_greenctx_stream_by_value(
+    SM_a: int, SM_b: int, device_id: Optional[int] = None
+):
+    """Create two green-context streams partitioned by ``SM_a`` / ``SM_b``."""
+    return get_kernel(
+        "spatial.create_greenctx_stream_by_value", KernelBackend.CUDA_AOT
+    )(SM_a, SM_b, device_id)
+
+
+__all__ = ["get_sm_available", "create_greenctx_stream_by_value"]

python/sglang/kernels/ops/speculative/__init__.py

@@ -0,0 +1,33 @@
+"""Speculative-decoding kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.speculative.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from speculative/triton_ops.
+_TRITON_KERNELS = [
+    ("cache_locs", "assign_req_to_token_pool_func"),
+    ("cache_locs", "assign_extend_cache_locs_func"),
+    ("cache_locs", "generate_draft_decode_kv_indices"),
+    ("eagle", "fill_bonus_tokens"),
+    ("eagle", "fill_accept_out_cache_loc"),
+    ("gather_spec_extras", "gather_spec_extras"),
+    ("multi_layer_eagle", "rotate_input_ids_triton"),
+    ("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
+    ("spec_tree", "verify_tree_greedy_kernel_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"speculative.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.speculative.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/registry.py

@@ -0,0 +1,71 @@
+"""In-memory registry of :class:`KernelSpec` entries.
+
+The registry is the single inventory of "which operators have which backend
+implementations". It is populated at import time by the ``sglang.kernels.ops.*``
+group packages, using only metadata (import path strings) — registering a spec
+never imports ``torch`` or a kernel backend.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, List
+
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+
+class KernelRegistry:
+    """Maps ``"<group>.<name>"`` operator ids to their :class:`KernelSpec` list."""
+
+    def __init__(self) -> None:
+        self._by_op: Dict[str, List[KernelSpec]] = defaultdict(list)
+
+    def register(self, spec: KernelSpec) -> KernelSpec:
+        """Register ``spec``.
+
+        Re-registering the same ``(op, backend)`` pair replaces the previous
+        entry so that module reloads during tests stay idempotent.
+        """
+        existing = self._by_op[spec.op]
+        for i, other in enumerate(existing):
+            if other.backend == spec.backend:
+                existing[i] = spec
+                return spec
+        existing.append(spec)
+        return spec
+
+    def get(self, op: str) -> List[KernelSpec]:
+        """All registered specs for ``op`` (empty list if none)."""
+        return list(self._by_op.get(op, ()))
+
+    def get_backend(self, op: str, backend: KernelBackend) -> KernelSpec:
+        """The spec for ``op`` provided by ``backend``.
+
+        Raises ``KeyError`` if no such implementation is registered.
+        """
+        for spec in self._by_op.get(op, ()):
+            if spec.backend == backend:
+                return spec
+        raise KeyError(f"No '{backend.value}' backend registered for op {op!r}")
+
+    def has(self, op: str) -> bool:
+        return bool(self._by_op.get(op))
+
+    def ops(self) -> List[str]:
+        """Sorted list of all registered operator ids."""
+        return sorted(self._by_op.keys())
+
+    def all_specs(self) -> List[KernelSpec]:
+        specs: List[KernelSpec] = []
+        for op in self.ops():
+            specs.extend(self._by_op[op])
+        return specs
+
+
+# Process-wide registry. Group packages register into this instance on import.
+registry = KernelRegistry()
+
+
+def register_kernel(spec: KernelSpec) -> KernelSpec:
+    """Register ``spec`` in the process-wide :data:`registry`."""
+    return registry.register(spec)

python/sglang/kernels/selector.py

@@ -0,0 +1,78 @@
+"""Fixed-path kernel resolution over the :data:`registry`.
+
+There is no priority ranking or heuristic backend selection. Each operator has
+a fixed call path — its :attr:`KernelSpec.target`:
+
+- an op with a single registered backend resolves to it directly;
+- an op with several registered backends must be resolved by naming the backend
+  explicitly (``backend=...``). The extra backends exist only as inventory, and
+  are never silently auto-picked.
+
+:func:`get_kernel` is the fast path used by the public ``ops.*`` wrappers: it
+resolves the spec to its callable and caches the result so repeated calls do
+not re-run resolution or re-import.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Callable, Optional
+
+from sglang.kernels.registry import registry
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+
+def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpec:
+    """Return the :class:`KernelSpec` for ``op`` (its fixed call path).
+
+    Parameters
+    ----------
+    op:
+        Operator id, ``"<group>.<name>"``.
+    backend:
+        Required only when ``op`` has more than one registered backend; selects
+        which one. For single-backend ops it is optional.
+
+    Raises
+    ------
+    KeyError
+        If ``op`` is unknown, or if ``backend`` is requested but not registered.
+    ValueError
+        If ``op`` has multiple backends and ``backend`` is not given.
+    """
+    specs = registry.get(op)
+    if not specs:
+        raise KeyError(f"No kernels registered for op {op!r}")
+
+    if backend is not None:
+        for spec in specs:
+            if spec.backend == backend:
+                return spec
+        raise KeyError(f"No '{backend.value}' backend registered for op {op!r}")
+
+    if len(specs) == 1:
+        return specs[0]
+
+    raise ValueError(
+        f"op {op!r} has multiple registered backends "
+        f"({[s.backend.value for s in specs]}); pass backend=... to choose one"
+    )
+
+
+@lru_cache(maxsize=None)
+def _resolve(op: str, backend: Optional[KernelBackend]) -> Callable:
+    return select_kernel(op, backend=backend).load()
+
+
+def get_kernel(op: str, backend: Optional[KernelBackend] = None) -> Callable:
+    """Resolve ``op`` to a callable kernel and cache it.
+
+    This is what the public ``sglang.kernels.ops.*`` wrappers call. The first
+    call resolves and imports the backend; later calls hit the cache.
+    """
+    return _resolve(op, backend)
+
+
+def clear_cache() -> None:
+    """Drop the resolved-callable cache (used by tests)."""
+    _resolve.cache_clear()

python/sglang/kernels/spec.py

@@ -0,0 +1,185 @@
+"""Lightweight metadata for the unified ``sglang.kernels`` namespace.
+
+This module defines small, dependency-free descriptors used to *inventory*
+kernel implementations and drive a simple, heuristic dispatch. It intentionally
+does not import ``torch``, ``sgl_kernel`` or ``sglang.jit_kernel`` at module
+import time so that ``import sglang.kernels`` stays cheap and works on a CPU-only
+box (see RFC #29630, Phase 2).
+
+The concrete callable behind a :class:`KernelSpec` is resolved lazily through
+``KernelSpec.load()``; nothing is imported until a kernel is actually called.
+"""
+
+from __future__ import annotations
+
+import importlib
+from enum import Enum
+from typing import Callable, Optional, Tuple
+
+import msgspec
+
+
+class KernelBackend(str, Enum):
+    """Implementation backend for a kernel.
+
+    Values mirror the backends called out in RFC #29630: JIT CUDA, AOT
+    CUDA/C++ (the ``sgl_kernel`` wheel), Triton, CuTe DSL, FlashInfer, DeepGEMM,
+    and the pure-``torch`` fallback path.
+    """
+
+    TORCH = "torch"  # pure-torch reference (forward_native)
+    TORCH_COMPILE = "torch_compile"  # torch.compile(forward_native)
+    TRITON = "triton"
+    CUDA_JIT = "cuda_jit"  # sglang.jit_kernel
+    CUDA_AOT = "cuda_aot"  # sgl_kernel wheel
+    CUTE_DSL = "cute_dsl"
+    FLASHINFER = "flashinfer"
+    DEEPGEMM = "deepgemm"
+    # TODO(RFC #29630): backends for other hardware (hip_c / npu / cpu-avx, ...)
+
+
+class PlatformInfo(msgspec.Struct, frozen=True):
+    """A minimal snapshot of the runtime accelerator platform.
+
+    Kept torch-free at import time; use :meth:`detect` to build one from the
+    live process (which does import ``torch``).
+    """
+
+    device_type: str = "cpu"  # "cuda", "hip", "cpu", ...
+    cuda_arch_major: Optional[int] = None
+    cuda_arch_minor: Optional[int] = None
+
+    @property
+    def is_cuda(self) -> bool:
+        return self.device_type == "cuda"
+
+    @property
+    def is_hip(self) -> bool:
+        return self.device_type == "hip"
+
+    @classmethod
+    def detect(cls) -> PlatformInfo:
+        """Build a :class:`PlatformInfo` from the current process.
+
+        Never raises: if ``torch`` is missing or no accelerator is visible the
+        default CPU platform is returned.
+        """
+        try:
+            import torch
+        except Exception:
+            return cls()
+
+        try:
+            if torch.version.hip is not None and torch.cuda.is_available():
+                return cls(device_type="hip")
+            if torch.cuda.is_available():
+                major, minor = torch.cuda.get_device_capability()
+                return cls(
+                    device_type="cuda",
+                    cuda_arch_major=major,
+                    cuda_arch_minor=minor,
+                )
+        except Exception:
+            pass
+        return cls()
+
+
+class CapabilityRequirement(msgspec.Struct, frozen=True):
+    """Coarse hardware requirement used to filter out unusable backends.
+
+    ``min_cuda_arch`` / ``max_cuda_arch`` are ``(major, minor)`` tuples, e.g.
+    ``(9, 0)`` for SM90. They only apply when the kernel requires CUDA.
+    """
+
+    requires_cuda: bool = False
+    requires_hip: bool = False
+    min_cuda_arch: Optional[Tuple[int, int]] = None
+    max_cuda_arch: Optional[Tuple[int, int]] = None
+
+    def is_satisfied_by(self, platform: PlatformInfo) -> bool:
+        if self.requires_hip and not platform.is_hip:
+            return False
+        if self.requires_cuda and not platform.is_cuda:
+            return False
+        if platform.is_cuda and platform.cuda_arch_major is not None:
+            arch = (platform.cuda_arch_major, platform.cuda_arch_minor or 0)
+            if self.min_cuda_arch is not None and arch < self.min_cuda_arch:
+                return False
+            if self.max_cuda_arch is not None and arch > self.max_cuda_arch:
+                return False
+        return True
+
+
+class FormatSignature(msgspec.Struct, frozen=True):
+    """A light description of a kernel's data contract.
+
+    This is deliberately loose in the first version — enough to document intent
+    and support future inventory tooling, not a strict schema.
+    """
+
+    supported_dtypes: Tuple[str, ...] = ()
+    in_place: bool = False
+    description: str = ""
+
+
+class KernelSpec(msgspec.Struct, frozen=True):
+    """A single callable kernel implementation and its metadata.
+
+    Parameters
+    ----------
+    op:
+        Fully-qualified operator id, ``"<group>.<name>"`` (e.g.
+        ``"layernorm.rmsnorm"``). This is the public lookup key.
+    backend:
+        Which :class:`KernelBackend` provides this implementation.
+    target:
+        Import path of the callable in ``"module:attr"`` form, resolved lazily
+        by :meth:`load` (e.g. ``"sgl_kernel:rmsnorm"``). ``attr`` may be a
+        dotted path into a module-level object, e.g.
+        ``"sglang.kernels.ops.layernorm:_RMSNORM.forward_cuda_aot"`` for a
+        bound :class:`~sglang.kernels.fused_op.BaseFusedOp` backend method.
+    capability:
+        Hardware requirement used by the selector to skip unusable backends.
+    format_signature:
+        Optional data-contract description for inventory/documentation.
+    description:
+        Human-readable one-liner.
+    """
+
+    op: str
+    backend: KernelBackend
+    target: str
+    capability: CapabilityRequirement = msgspec.field(
+        default_factory=CapabilityRequirement
+    )
+    format_signature: FormatSignature = msgspec.field(default_factory=FormatSignature)
+    description: str = ""
+
+    @property
+    def group(self) -> str:
+        return self.op.split(".", 1)[0]
+
+    @property
+    def name(self) -> str:
+        return self.op.split(".", 1)[1] if "." in self.op else self.op
+
+    def is_available(self, platform: PlatformInfo) -> bool:
+        """Whether this backend can run on ``platform`` (metadata-only check)."""
+        return self.capability.is_satisfied_by(platform)
+
+    def load(self) -> Callable:
+        """Import and return the backing callable.
+
+        Raises the underlying ``ImportError`` / ``AttributeError`` if the
+        backend is not installed on this platform — call sites decide how to
+        handle that.
+        """
+        module_path, sep, attr = self.target.partition(":")
+        if not sep or not attr:
+            raise ValueError(
+                f"KernelSpec.target must be 'module:attr', got {self.target!r}"
+            )
+        obj = importlib.import_module(module_path)
+        for part in attr.split("."):
+            obj = getattr(obj, part)
+        return obj

python/sglang/srt/constrained/xgrammar_backend.py

@@ -43,14 +43,14 @@ _is_hip = is_hip()
 if _is_hip:
     from sgl_kernel import apply_token_bitmask_inplace_cuda
 else:
-    from sglang.srt.constrained.triton_ops.bitmask_ops import (
+    from sglang.kernels.ops.grammar.bitmask_ops import (
         apply_token_bitmask_inplace_triton,
     )
 
+from sglang.kernels.ops.grammar.token_filter_ops import set_token_filter_triton
 from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
     set_token_filter_torch,
 )
-from sglang.srt.constrained.triton_ops.token_filter_ops import set_token_filter_triton
 
 logger = logging.getLogger(__name__)
 MAX_ROLLBACK_TOKENS = 200

python/sglang/srt/environ.py

@@ -648,6 +648,8 @@ class Envs:
     SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)
     SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
     SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
+    SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
+    SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True)
 
     # sgl-kernel
     SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
@@ -656,6 +658,11 @@ class Envs:
     SGLANG_USE_SGL_FA3_KERNEL = EnvBool(True)
 
     # Kernels
+    # Force every sglang.kernels BaseFusedOp onto one backend (a KernelBackend
+    # value, e.g. "torch" / "torch_compile" / "triton" / "cuda_aot"); unset =
+    # auto-select by priority. "torch" flips all fused ops to their pure-torch
+    # reference implementations for numerical-bug bisection.
+    SGLANG_FORCE_FUSED_OP_BACKEND = EnvStr(None)
     USE_TRITON_W8A8_FP8_KERNEL = EnvBool(False)
     SGLANG_RETURN_ORIGINAL_LOGPROB = EnvBool(False)
     SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN = EnvBool(False)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -14,11 +14,11 @@ from typing import TYPE_CHECKING, Optional
 import torch
 import triton
 
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.triton_ops.aiter_unified_attention import (
+from sglang.kernels.ops.kvcache.aiter_unified_attention import (
     scatter_ragged_to_page_table_kernel,
     scatter_req_to_token_to_page_table_kernel,
 )
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.utils import (
     assert_buffer_fits,
     create_flashinfer_kv_indices_triton,
@@ -136,7 +136,7 @@ class AiterAttnBackend(AttentionBackend):
     ):
         super().__init__()
         # Lazy import to avoid the initialization of cuda context
-        from sglang.srt.layers.attention.triton_ops.extend_attention import (
+        from sglang.kernels.ops.attention.extend_attention import (
             extend_attention_fwd,
         )
 

python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py

@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens
 from sglang.srt.layers.attention.utils import seqlens_expand_triton
 from sglang.srt.utils import is_cuda, is_hip
@@ -20,6 +21,9 @@ if TYPE_CHECKING:
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_USE_FUSED_METADATA_GENERATION = (
+    envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip
+)
 
 
 @dataclass
@@ -121,8 +125,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         """Precompute metadata for normal decode mode."""
         max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
 
-        if _is_cuda and not _is_hip:
-            from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+        if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip:
+            from sglang.kernels.ops.attention.dsa_metadata import (
                 fused_dsa_decode_metadata,
             )
 
@@ -240,8 +244,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
         seqlens_expanded_size = bs * self.speculative_num_draft_tokens
 
-        if _is_cuda and not _is_hip:
-            from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+        if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip:
+            from sglang.kernels.ops.attention.dsa_metadata import (
                 fused_dsa_target_verify_metadata,
             )
 

python/sglang/srt/layers/attention/dsa_backend.py

@@ -138,6 +138,13 @@ def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tens
 
 # Reuse this workspace buffer across all DSA backend instances
 
+# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled)
+# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable
+_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip
+_USE_FUSED_METADATA_GENERATION = (
+    envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip
+)
+
 
 @dataclass(frozen=True)
 class DSAFlashMLAMetadata:
@@ -1374,8 +1381,8 @@ class DeepseekSparseAttnBackend(
             # Normal Decode
             max_len = self._graph_page_table_width(metadata)
 
-            if is_cuda() and not _is_hip:
-                from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+            if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip:
+                from sglang.kernels.ops.attention.dsa_metadata import (
                     fused_dsa_decode_metadata,
                 )
 
@@ -1416,8 +1423,8 @@ class DeepseekSparseAttnBackend(
         elif forward_mode.is_target_verify():
             max_seqlen_k = self._graph_page_table_width(metadata)
 
-            if is_cuda() and not _is_hip:
-                from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+            if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip:
+                from sglang.kernels.ops.attention.dsa_metadata import (
                     fused_dsa_target_verify_metadata,
                 )
 
@@ -1514,8 +1521,8 @@ class DeepseekSparseAttnBackend(
                 device=self.device,
             )
 
-            if is_cuda() and not _is_hip:
-                from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+            if _USE_FUSED_METADATA_GENERATION and is_cuda() and not _is_hip:
+                from sglang.kernels.ops.attention.dsa_metadata import (
                     fused_dsa_draft_extend_metadata,
                 )
 

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -8,15 +8,15 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.configs.model_config import AttentionArch
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.triton_ops.metadata import (
+from sglang.kernels.ops.attention.metadata import (
     normal_decode_set_metadata,
     prepare_swa_spec_page_table_triton,
 )
-from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import (
+from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
     build_trtllm_mha_page_table,
 )
+from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.utils import assert_buffer_fits
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_v2_active

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -86,7 +86,7 @@ if is_flashinfer_available():
     )
     from flashinfer.cascade import merge_state
 
-    from sglang.srt.layers.attention.triton_ops.merge_state import merge_state_triton
+    from sglang.kernels.ops.attention.merge_state import merge_state_triton
 
     # FlashInfer's MergeState CUDA kernel uses blockDim = (head_dim/vec_size, num_heads).
     # When num_heads is large (e.g. with DP attention where attention_tp_size=1), the

python/sglang/srt/layers/attention/merge_state.py

@@ -3,7 +3,7 @@ from typing import Optional, Tuple
 import torch
 from sgl_kernel import merge_state_v2
 
-from sglang.srt.layers.attention.triton_ops.merge_state import merge_state_triton
+from sglang.kernels.ops.attention.merge_state import merge_state_triton
 from sglang.srt.utils import is_cuda
 
 _is_cuda = is_cuda()

python/sglang/srt/layers/attention/triton_backend.py

@@ -6,6 +6,10 @@ from typing import TYPE_CHECKING, List, Optional
 import torch
 import triton
 
+from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton
+from sglang.kernels.ops.kvcache.kv_indices import (
+    create_flashinfer_kv_indices_triton,
+)
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
@@ -13,10 +17,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 from sglang.srt.distributed.parallel_state import get_dcp_group
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
-    create_flashinfer_kv_indices_triton,
-)
-from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton
 from sglang.srt.layers.dcp import (
     cp_lse_ag_out_rs_mha,
     create_triton_kv_indices_for_dcp_triton,
@@ -115,15 +115,15 @@ class TritonAttnBackend(AttentionBackend):
         kv_indptr_buf: Optional[torch.Tensor] = None,
     ):
         # Lazy import to avoid the initialization of cuda context
-        from sglang.srt.layers.attention.triton_ops.decode_attention import (
+        from sglang.kernels.ops.attention.decode_attention import (
             decode_attention_fwd,
         )
-        from sglang.srt.layers.attention.triton_ops.extend_attention import (
+        from sglang.kernels.ops.attention.extend_attention import (
             build_unified_kv_indices,
             extend_attention_fwd,
             extend_attention_fwd_unified,
         )
-        from sglang.srt.layers.attention.triton_ops.verify_splitkv import (
+        from sglang.kernels.ops.attention.verify_splitkv import (
             verify_splitkv_fwd,
         )
 

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -11,22 +11,22 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.srt.environ import envs
-from sglang.srt.layers.attention.flashinfer_backend import (
-    FlashInferAttnBackend,
-    FlashInferMultiStepDraftBackend,
-)
-from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
+from sglang.kernels.ops.kvcache.trtllm_fp8_kv_kernel import (
     fused_fp8_set_kv_buffer,
 )
-from sglang.srt.layers.attention.triton_ops.trtllm_mha_graph_metadata import (
+from sglang.kernels.ops.kvcache.trtllm_mha_graph_metadata import (
     Q_MODE_NONE,
     Q_MODE_STRIDED,
     update_trtllm_mha_graph_metadata,
 )
-from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import (
+from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
     build_trtllm_mha_page_table,
 )
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.flashinfer_backend import (
+    FlashInferAttnBackend,
+    FlashInferMultiStepDraftBackend,
+)
 from sglang.srt.layers.attention.utils import canonicalize_stride
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -13,21 +13,21 @@ import torch
 import triton
 
 from sglang.jit_kernel.fixup_zero_kv import fixup_zero_kv_rows
-from sglang.srt.environ import envs
-from sglang.srt.layers.attention.flashinfer_mla_backend import (
-    FlashInferMLAAttnBackend,
-    FlashInferMLAMultiStepDraftBackend,
+from sglang.kernels.ops.attention.pad import (
+    pad_draft_extend_query as pad_draft_extend_query_triton,
 )
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
+from sglang.kernels.ops.attention.pad import (
+    unpad_draft_extend_output as unpad_draft_extend_output_triton,
+)
+from sglang.kernels.ops.kvcache.kv_indices import (
     create_flashmla_kv_indices_triton,
     get_num_kv_index_blocks_flashmla,
     get_num_page_per_block_flashmla,
 )
-from sglang.srt.layers.attention.triton_ops.pad import (
-    pad_draft_extend_query as pad_draft_extend_query_triton,
-)
-from sglang.srt.layers.attention.triton_ops.pad import (
-    unpad_draft_extend_output as unpad_draft_extend_output_triton,
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.flashinfer_mla_backend import (
+    FlashInferMLAAttnBackend,
+    FlashInferMLAMultiStepDraftBackend,
 )
 from sglang.srt.layers.attention.utils import (
     concat_mla_absorb_q_general,

python/sglang/srt/layers/attention/utils.py

@@ -3,43 +3,43 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
-from sglang.srt.layers.attention.triton_ops.cache_ops import (
-    concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel,
-)
-from sglang.srt.layers.attention.triton_ops.cache_ops import (
-    concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton,
-)
-from sglang.srt.layers.attention.triton_ops.cache_ops import (
-    launch_reshape_and_cache_flash as launch_reshape_and_cache_flash,
-)
-from sglang.srt.layers.attention.triton_ops.cache_ops import (
-    reshape_and_cache_flash as reshape_and_cache_flash,
-)
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
-    create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton,
-)
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
-    create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton,
-)
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
-    get_num_kv_index_blocks_flashmla as get_num_kv_index_blocks_flashmla,
-)
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
-    get_num_page_per_block_flashmla as get_num_page_per_block_flashmla,
-)
-from sglang.srt.layers.attention.triton_ops.pad import (
+from sglang.kernels.ops.attention.pad import (
     pad_sequence_with_mask as pad_sequence_with_mask,
 )
-from sglang.srt.layers.attention.triton_ops.pad import (
+from sglang.kernels.ops.attention.pad import (
     pad_sequence_with_mask_kernel as pad_sequence_with_mask_kernel,
 )
-from sglang.srt.layers.attention.triton_ops.pad import (
+from sglang.kernels.ops.attention.pad import (
     seqlens_expand_kernel as seqlens_expand_kernel,
 )
-from sglang.srt.layers.attention.triton_ops.pad import (
+from sglang.kernels.ops.attention.pad import (
     seqlens_expand_triton as seqlens_expand_triton,
 )
-from sglang.srt.layers.attention.triton_ops.rope_cache import (
+from sglang.kernels.ops.kvcache.cache_ops import (
+    concat_and_cast_mha_k_kernel as concat_and_cast_mha_k_kernel,
+)
+from sglang.kernels.ops.kvcache.cache_ops import (
+    concat_and_cast_mha_k_triton as concat_and_cast_mha_k_triton,
+)
+from sglang.kernels.ops.kvcache.cache_ops import (
+    launch_reshape_and_cache_flash as launch_reshape_and_cache_flash,
+)
+from sglang.kernels.ops.kvcache.cache_ops import (
+    reshape_and_cache_flash as reshape_and_cache_flash,
+)
+from sglang.kernels.ops.kvcache.kv_indices import (
+    create_flashinfer_kv_indices_triton as create_flashinfer_kv_indices_triton,
+)
+from sglang.kernels.ops.kvcache.kv_indices import (
+    create_flashmla_kv_indices_triton as create_flashmla_kv_indices_triton,
+)
+from sglang.kernels.ops.kvcache.kv_indices import (
+    get_num_kv_index_blocks_flashmla as get_num_kv_index_blocks_flashmla,
+)
+from sglang.kernels.ops.kvcache.kv_indices import (
+    get_num_page_per_block_flashmla as get_num_page_per_block_flashmla,
+)
+from sglang.kernels.ops.kvcache.rope_cache import (
     fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache,
 )
 from sglang.srt.utils import is_cuda

python/sglang/srt/layers/attention/vision.py

@@ -61,14 +61,14 @@ if _is_npu:
 if _is_xpu:
     from sgl_kernel.flash_attn import flash_attn_varlen_func
 
+from sglang.kernels.ops.attention.prefill_attention import (
+    context_attention_fwd,
+)
 from sglang.srt.distributed import (
     split_tensor_along_last_dim,
     tensor_model_parallel_all_gather,
 )
 from sglang.srt.distributed import utils as dist_utils
-from sglang.srt.layers.attention.triton_ops.prefill_attention import (
-    context_attention_fwd,
-)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,

python/sglang/srt/layers/attention/wave_backend.py

@@ -7,11 +7,11 @@ from typing import TYPE_CHECKING, Optional
 import torch
 import triton
 
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.triton_ops.kv_indices import (
+from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton
+from sglang.kernels.ops.kvcache.kv_indices import (
     create_flashinfer_kv_indices_triton,
 )
-from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_bool_env_var, get_device_core_count

python/sglang/srt/layers/elementwise.py

@@ -5,7 +5,7 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
-from sglang.srt.layers.triton_ops.softcap import softcap_out as fused_softcap
+from sglang.kernels.ops.activation.softcap import softcap_out as fused_softcap
 from sglang.srt.utils import is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 

python/sglang/srt/layers/logits_processor.py

@@ -21,6 +21,9 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 from torch import nn
 
+from sglang.kernels.ops.activation.softcap import (
+    softcap_inplace_logits as fused_softcap,
+)
 from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
@@ -33,7 +36,6 @@ from sglang.srt.layers.dp_attention import (
     get_dp_dtype,
     get_dp_hidden_size,
 )
-from sglang.srt.layers.triton_ops.softcap import softcap_inplace_logits as fused_softcap
 from sglang.srt.layers.utils.logprob import (
     InputLogprobsResult,
     get_token_ids_logprobs_chunk,