[Kernel] Unify BaseFusedOp and MultiPlatformOp dispatch (#33205)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

docs/docs/hardware-platforms/plugin.mdx

@@ -544,7 +544,7 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
     <tr>
       <td><code>get_dispatch_key_name()</code></td>
       <td><code>"native"</code></td>
-      <td>MultiPlatformOp dispatch key name</td>
+      <td>BaseFusedOp (fused-op) dispatch key name</td>
     </tr>
   </tbody>
 </table>

python/sglang/kernels/README.md

@@ -63,26 +63,47 @@ from sglang.kernels import select_kernel, KernelBackend
 jit_rmsnorm = select_kernel("layernorm.rmsnorm", backend=KernelBackend.JIT).load()
 ```
 
-## `BaseFusedOp` — the per-operator implementation contract
+## `BaseFusedOp` — the unified operator contract
 
-Multi-backend operators (currently the `layernorm` and `activation` groups)
+`BaseFusedOp` is a standard `torch.nn.Module` (it replaced the former
-are implemented as `BaseFusedOp` subclasses: one logical operator with one
+`sglang.srt.layers.utils.MultiPlatformOp`) that carries one logical operator
-`forward_<backend>` method per backend, all sharing one signature behind a
+with interchangeable implementations along **two independent dimensions**:
-single `forward()`:
 
+- **Kernel backends (provenance)** — one `forward_<backend>` method per
+  implementation source, all sharing one signature behind a single
+  `forward()`:
   - `forward_native` — **required**; the pure-`torch` correctness reference
-  every other backend is checked against.
+    every other implementation is checked against.
   - `forward_torch_compile` — inherited for free as
     `torch.compile(forward_native)`.
-- `forward_triton` / `forward_jit` / `forward_aot` /
+  - `forward_triton` / `forward_jit` / `forward_aot` / `forward_cute_dsl` /
-  `forward_cute_dsl` / `forward_flashinfer` / `forward_deepgemm` — opt-in
+    `forward_flashinfer` / `forward_deepgemm` / `forward_aiter` /
-  overrides. A backend is *available* iff its method is overridden.
+    `forward_torch_npu` — opt-in overrides. A backend is *available* iff its
+    method is overridden; it joins **auto-selection** only when also declared
+    in `capabilities` (device support is metadata, not guesswork).
+- **Platforms / devices** — optional composite per-device paths:
+  `forward_cuda`, `forward_hip` (falls back to `forward_cuda`),
+  `forward_npu`, `forward_xpu`, `forward_musa` (no implicit CUDA fallback —
+  MUSA ops opt into the CUDA path with an explicit `forward_musa`),
+  `forward_cpu` (AMX CPUs), plus `forward_<key>` /
+  `register_oot_forward()` for out-of-tree platform plugins. CUDA / HIP are
+  **not** kernel backends.
 
-`forward()` auto-selects the best available backend by the class's `priority`,
+Dispatch priority, highest first: explicit `forward(..., backend=...)` →
-filtered per call through `backend_eligible()` (a
+global forced backend (`SGLANG_FORCE_FUSED_OP_BACKEND`) → OOT platform
-`CapabilityRequirement`-vs-`PlatformInfo` check, extensible with per-call
+override → declared optimized kernel backends by `priority` (filtered by
-shape/dtype gates), and degrades to the native reference when no optimized
+`backend_eligible()`, a `CapabilityRequirement`-vs-`PlatformInfo` check
-backend fits. The public `ops.<group>` functions stay thin wrappers over
+extensible with per-call shape/dtype gates) → platform-specific forward →
+`forward_native`. The static part of the decision is resolved once and cached
+on the instance, so the hot path stays a single indirect call.
+
+`BaseFusedOp` also owns the torch.compile mode protocol
+(`enter_torch_compile(num_tokens)` / `leave_torch_compile()`, both
+idempotent): while an outer model is compiled, ops switch to their
+compile-safe native path so device kernels are never traced (TopK and Fused
+MoE override `_torch_compile_forward()` to keep their bs>1 behavior).
+
+The public `ops.<group>` functions stay thin wrappers over
 module-level instances, so the import surface is unchanged; each instance also
 registers all of its backends as `KernelSpec`s so the registry inventory and
 `select_kernel(..., backend=...)` keep working.

python/sglang/kernels/fused_op.py

@@ -1,44 +1,77 @@
-"""Multi-backend operator contract for the unified kernels namespace.
+"""Unified multi-backend / multi-platform operator contract (RFC #29630, #26426).
 
-:class:`BaseFusedOp` is the per-operator implementation object behind the
+:class:`BaseFusedOp` is the single operator abstraction of the unified
-``sglang.kernels.ops.*`` wrappers (RFC #29630): one logical operator,
+``sglang.kernels`` namespace: one logical operator, implemented once, with
-implemented once, with multiple interchangeable backends behind a single
+multiple interchangeable implementations behind a single ``forward()``. It
-``forward()``.
+subsumes the former ``MultiPlatformOp`` (``sglang.srt.layers.utils``), so it is
+a proper :class:`torch.nn.Module` and covers **two independent dimensions**:
 
-Each subclass implements one ``forward_<backend>`` method per backend it
+- **Kernel backend (provenance)** — where an implementation comes from
-supports:
+  (:class:`~sglang.kernels.spec.KernelBackend`): ``forward_native`` (pure
+  torch, required), ``forward_torch_compile``, ``forward_triton``,
+  ``forward_jit``, ``forward_aot``, ``forward_cute_dsl``,
+  ``forward_flashinfer``, ``forward_deepgemm``, ``forward_aiter``,
+  ``forward_torch_npu``. Which *devices* a backend supports is per-``(op,
+  backend)`` metadata (:attr:`BaseFusedOp.capabilities`), never implied by the
+  backend name.
+- **Platform / device** — device-specific composite paths inherited from
+  ``MultiPlatformOp``: ``forward_cuda``, ``forward_hip``, ``forward_npu``,
+  ``forward_xpu``, ``forward_musa``, ``forward_cpu``, plus ``forward_<key>``
+  for out-of-tree (OOT) platform plugins. CUDA / HIP are **not** kernel
+  backends.
 
-- ``forward_native`` — **required**; the pure-``torch`` correctness reference
+Dispatch priority (highest first), resolved by :meth:`BaseFusedOp.forward`:
-  every other backend is checked against.
-- ``forward_torch_compile`` — provided by the base class as
-  ``torch.compile(forward_native)``.
-- ``forward_triton`` / ``forward_jit`` / ``forward_aot`` /
-  ``forward_cute_dsl`` / ``forward_flashinfer`` / ``forward_deepgemm`` —
-  opt-in overrides.
 
-A backend is *available* iff the subclass overrides its method (``native`` and
+1. **Explicit backend** — ``forward(..., backend=KernelBackend.X)``.
-``torch_compile`` are always available). ``forward()`` picks the best
+2. **Global forced backend** — ``SGLANG_FORCE_FUSED_OP_BACKEND`` /
-available backend by :attr:`BaseFusedOp.priority`, filtered per call through
+   :func:`set_fused_op_backend` (e.g. ``native`` to bisect numerical bugs).
-:meth:`BaseFusedOp.backend_eligible` (which checks
+   Best-effort: an op that does not implement the forced backend falls back
-:class:`~sglang.kernels.spec.CapabilityRequirement` against the detected
+   to normal dispatch with a one-time warning, so the debug switch works on
-:class:`~sglang.kernels.spec.PlatformInfo`). The ``SGLANG_FORCE_FUSED_OP_BACKEND``
+   whole models that contain device-only ops.
-env var (or :func:`set_fused_op_backend`) forces every fused op onto one
+3. **OOT platform override** — on an out-of-tree platform, a forward
-backend — e.g. ``native`` to bisect numerical bugs against the reference
+   registered via :meth:`BaseFusedOp.register_oot_forward`, then a
-implementations with a single switch.
+   ``forward_<dispatch_key>`` method, then ``forward_native``.
+4. **Optimized kernel backends** — the first backend in :attr:`priority`
+   whose method is overridden, that is *declared* in :attr:`capabilities`,
+   and whose :class:`~sglang.kernels.spec.CapabilityRequirement` set matches
+   the detected platform. Ops may extend :meth:`backend_eligible` with
+   per-call shape/dtype gates; overriding it switches this step from a
+   statically cached choice to per-call selection.
+5. **Platform-specific forward** — ``forward_cuda`` on CUDA, ``forward_hip``
+   (falling back to ``forward_cuda``) on ROCm, ``forward_musa`` on MUSA,
+   ``forward_npu`` / ``forward_xpu`` on Ascend / XPU, ``forward_cpu`` on
+   AMX-capable CPUs.
+6. **Native fallback** — ``forward_native``.
+
+Steps 3-6 are static per process, so their outcome is resolved once (lazily,
+on first call) and cached in ``self._forward_method``; subclass ``__init__``
+may pre-seed that attribute to pin an instance to a specific path (e.g. the
+env-gated aiter paths in ``srt/layers``). Steps 1-2 stay per-call so tests and
+tools can flip backends on live instances.
+
+torch.compile integration: :meth:`enter_torch_compile` switches the op to its
+compile-safe path (``forward_native`` by default; see
+:meth:`_torch_compile_forward` for the TopK / FusedMoE special cases) so an
+*outer* ``torch.compile`` never traces device-specific kernels, and
+:meth:`leave_torch_compile` restores the original dispatch. Both are
+idempotent because one module instance may be shared by many layers.
 
 Like the rest of ``sglang.kernels``, importing this module (and instantiating
 subclasses) never imports a kernel backend (``sgl_kernel`` /
-``sglang.kernels.jit``) or triggers JIT compilation; backends are imported
+``sglang.kernels.jit``), performs platform detection, or triggers JIT
-lazily inside the ``forward_<backend>`` methods.
+compilation; backends are imported lazily inside the ``forward_<backend>``
+methods and dispatch is resolved on first call.
 """
 
 from __future__ import annotations
 
 import functools
+import logging
 from abc import ABC, abstractmethod
 from typing import (
     AbstractSet,
     Any,
+    Callable,
     ClassVar,
     Dict,
     List,
@@ -48,7 +81,10 @@ from typing import (
 )
 
 import msgspec
+import torch
+from torch import nn
 
+from sglang.kernel_api_logging import debug_kernel_api
 from sglang.kernels.registry import register_kernel
 from sglang.kernels.spec import (
     CapabilityRequirement,
@@ -59,7 +95,11 @@ from sglang.kernels.spec import (
     capabilities_satisfied,
 )
 
-# backend (provenance) -> forward_<backend> method name.
+logger = logging.getLogger(__name__)
+
+# backend (provenance) -> forward_<backend> method name. ``forward_torch_npu``
+# (not ``forward_npu``) so the torch_npu *backend* method never collides with
+# the NPU *platform* method.
 BACKEND_METHODS: Dict[KernelBackend, str] = {
     KernelBackend.TORCH: "forward_native",
     KernelBackend.TORCH_COMPILE: "forward_torch_compile",
@@ -70,7 +110,11 @@ BACKEND_METHODS: Dict[KernelBackend, str] = {
     KernelBackend.FLASHINFER: "forward_flashinfer",
     KernelBackend.DEEPGEMM: "forward_deepgemm",
     KernelBackend.AITER: "forward_aiter",
-    KernelBackend.TORCH_NPU: "forward_npu",
+    KernelBackend.TORCH_NPU: "forward_torch_npu",
+}
+
+_METHOD_BACKEND_LABELS: Dict[str, str] = {
+    name: backend.value for backend, name in BACKEND_METHODS.items()
 }
 
 # best -> fallback. ``torch_compile`` is deliberately absent: auto-selection
@@ -92,12 +136,79 @@ DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = (
 # concrete subclass always has it) and forward_torch_compile derives from it.
 _ALWAYS_AVAILABLE = (KernelBackend.TORCH, KernelBackend.TORCH_COMPILE)
 
+# In-tree platform key -> platform forward candidates, best first. A candidate
+# counts only when the subclass actually overrides it; otherwise dispatch
+# falls through to forward_native. Only HIP keeps the implicit CUDA-path
+# fallback (ROCm kernels are hipified CUDA and sgl_kernel builds for both);
+# MUSA deliberately does not chain into forward_cuda — srt module-level
+# kernel imports are gated on is_cuda(), so a CUDA path reached implicitly on
+# a MUSA box can NameError instead of degrading. MUSA ops that want the CUDA
+# path opt in with an explicit forward_musa.
+_PLATFORM_METHODS: Dict[str, Tuple[str, ...]] = {
+    "cuda": ("forward_cuda",),
+    "hip": ("forward_hip", "forward_cuda"),
+    "musa": ("forward_musa",),
+    "npu": ("forward_npu",),
+    "xpu": ("forward_xpu",),
+    "cpu": ("forward_cpu",),
+}
+
 
 @functools.lru_cache(maxsize=1)
 def _platform() -> PlatformInfo:
     return PlatformInfo.detect()
 
 
+@functools.lru_cache(maxsize=1)
+def _platform_key() -> str:
+    """In-tree platform dispatch key, or ``""`` for plain-native platforms.
+
+    Checked in the same order as the former ``MultiPlatformOp``: CPU counts
+    only when AMX is available (otherwise the pure-torch reference is faster
+    than pretending there is an optimized CPU path).
+    """
+    from sglang.srt.utils import (
+        cpu_has_amx_support,
+        is_cpu,
+        is_cuda,
+        is_hip,
+        is_musa,
+        is_npu,
+        is_xpu,
+    )
+
+    if is_cuda():
+        return "cuda"
+    if is_hip():
+        return "hip"
+    if is_cpu() and cpu_has_amx_support():
+        return "cpu"
+    if is_npu():
+        return "npu"
+    if is_xpu():
+        return "xpu"
+    if is_musa():
+        return "musa"
+    return ""
+
+
+@functools.lru_cache(maxsize=1)
+def _oot_dispatch_key() -> Optional[str]:
+    """The active out-of-tree platform's dispatch key, or ``None`` in-tree."""
+    from sglang.srt.platforms import current_platform
+
+    if current_platform.is_out_of_tree():
+        return current_platform.get_dispatch_key_name()
+    return None
+
+
+def clear_platform_caches() -> None:
+    """Drop the cached platform detection (used by tests that mock platforms)."""
+    _platform.cache_clear()
+    _platform_key.cache_clear()
+    _oot_dispatch_key.cache_clear()
+
+
 # --- global backend override ------------------------------------------------
 
 # Sentinel distinguishing "not resolved yet" from "resolved to None (no force)".
@@ -166,8 +277,6 @@ def clear_fused_op_trace() -> None:
 
 
 def _describe_tensors(args: tuple, kwargs: dict) -> Tuple[str, ...]:
-    import torch
-
     described = []
     for value in (*args, *kwargs.values()):
         if isinstance(value, torch.Tensor):
@@ -175,38 +284,92 @@ def _describe_tensors(args: tuple, kwargs: dict) -> Tuple[str, ...]:
     return tuple(described)
 
 
+def _dispatch_label(method: Callable) -> str:
+    """Trace label for a resolved dispatch target.
+
+    Backend methods map to their :class:`KernelBackend` value (so
+    ``forward_native`` traces as ``"torch"``, matching explicit-backend
+    calls); platform forwards trace as their device key (``"cuda"``, ...).
+    """
+    name = getattr(method, "__name__", "")
+    label = _METHOD_BACKEND_LABELS.get(name)
+    if label is not None:
+        return label
+    if name.startswith("forward_"):
+        return name[len("forward_") :]
+    return name or "unknown"
+
+
+def _record_trace(op: BaseFusedOp, label: str, args: tuple, kwargs: dict) -> None:
+    _trace_records.append(
+        FusedOpTraceRecord(
+            op=op.op or type(op).__name__,
+            backend=label,
+            tensor_args=_describe_tensors(args, kwargs),
+        )
+    )
+
+
+_warned_forced_fallbacks: set = set()
+
+
+def _warn_forced_backend_unavailable(*, op: BaseFusedOp, backend: KernelBackend):
+    key = (type(op), backend)
+    if key not in _warned_forced_fallbacks:
+        _warned_forced_fallbacks.add(key)
+        logger.warning(
+            "Forced fused-op backend %r is not implemented by %s; "
+            "falling back to normal dispatch for this op.",
+            backend.value,
+            op._op_label(),
+        )
+
+
 # --- the per-operator contract ------------------------------------------------
 
 
-class BaseFusedOp(ABC):
+class BaseFusedOp(nn.Module, ABC):
-    """One logical operator with interchangeable backends behind ``forward()``.
+    """One logical operator with interchangeable implementations behind
+    ``forward()``.
 
-    Subclasses set :attr:`op` and implement :meth:`forward_native` plus any
+    Subclasses implement :meth:`forward_native` plus any optimized
-    optimized ``forward_<backend>`` methods. All backend methods of one op
+    ``forward_<backend>`` methods and/or platform-specific
-    must share the same signature and semantics — each override adapts its
+    ``forward_<device>`` methods. All implementations of one op must share the
-    underlying kernel's calling convention so call sites never care which
+    same signature and semantics — each override adapts its underlying
-    backend ran.
+    kernel's calling convention so call sites never care which one ran.
+
+    This is a standard :class:`torch.nn.Module`: it participates in module
+    traversal / state dicts, and calls go through ``nn.Module.__call__`` so
+    forward hooks keep working.
 
     Class attributes
     ----------------
     op:
         Operator id, ``"<group>.<name>"`` (e.g. ``"layernorm.rmsnorm"``).
+        Required for :func:`register_fused_op`; layer-style subclasses that
+        are not registered in the kernel registry may leave it empty.
     priority:
-        Backend preference for auto-selection, best first. Defaults to
+        Kernel-backend preference for auto-selection, best first. Defaults to
-        :data:`DEFAULT_PRIORITY`.
+        :data:`DEFAULT_PRIORITY`. ``KernelBackend.TORCH`` entries are ignored:
+        the native reference is always the final fallback, after
+        platform-specific forwards.
     capabilities:
         Per-backend set of :class:`CapabilityRequirement` (OR semantics;
-        omitted / empty = runs on any device), consulted by
+        an empty set value = runs on any device), consulted by
-        :meth:`backend_eligible` (and exported into the registry specs). Use the
+        :meth:`backend_eligible` and exported into the registry specs. A
-        ``CapabilityRequirement.CUDA`` / ``.HIP`` / ``.NPU`` shortcuts, e.g.
+        kernel backend joins **auto-selection** only when it is declared here
-        ``{KernelBackend.AOT: {CapabilityRequirement.CUDA, CapabilityRequirement.HIP}}``.
+        (explicit ``backend=`` / forced selection can still target any
+        overridden method) — device support is metadata, not guesswork. Use
+        the ``CapabilityRequirement.CUDA`` / ``.HIP`` / ``.NPU`` shortcuts,
+        e.g. ``{KernelBackend.AOT: {CapabilityRequirement.CUDA,
+        CapabilityRequirement.HIP}}``.
     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]
+    op: ClassVar[str] = ""
     priority: ClassVar[Tuple[KernelBackend, ...]] = DEFAULT_PRIORITY
     capabilities: ClassVar[
         Mapping[KernelBackend, AbstractSet[CapabilityRequirement]]
@@ -214,30 +377,40 @@ class BaseFusedOp(ABC):
     format_signature: ClassVar[FormatSignature] = FormatSignature()
     descriptions: ClassVar[Mapping[KernelBackend, str]] = {}
 
+    # OOT forward registry shared by all fused ops: platform dispatch key ->
+    # {op class -> forward fn}. Populated by out-of-tree platform plugins.
+    _oot_forward_registry: ClassVar[Dict[str, Dict[type, Callable]]] = {}
+
+    @classmethod
+    def register_oot_forward(cls, op_cls: type, fn: Callable, platform_key: str):
+        """Register an OOT forward implementation for a specific op class and
+        platform. ``fn`` is bound to the instance at dispatch time (it receives
+        the op as ``self``) and takes precedence over ``forward_<platform_key>``
+        methods on that exact class."""
+        cls._oot_forward_registry.setdefault(platform_key, {})[op_cls] = fn
+
     def __init__(self) -> None:
-        # Cache the structural backend set and the priority-ordered subset once
+        super().__init__()
-        # so forward() avoids repeated introspection on the hot path.
+        # Statically resolved dispatch target (priority steps 3-6). ``None``
-        available = []
+        # means "not resolved yet": resolution is deferred to the first call
-        for backend in KernelBackend:
+        # so module-level op instances never trigger platform detection at
-            if backend in _ALWAYS_AVAILABLE or self._overrides(
+        # import time. Subclass __init__ may overwrite it to pin a path.
-                BACKEND_METHODS[backend]
+        self._forward_method: Optional[Callable] = None
-            ):
+        # torch.compile mode bookkeeping (see enter/leave_torch_compile).
-                available.append(backend)
+        self._original_forward_method: Optional[Callable] = None
-        self._available: Tuple[KernelBackend, ...] = tuple(available)
+        self.is_torch_compile = False
-        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:
+    def _defined_method(self, method_name: str) -> Optional[Callable]:
+        """The bound method if any class below ``BaseFusedOp`` defines it."""
         for klass in type(self).__mro__:
             if klass is BaseFusedOp:
-                return False
+                return None
             if method_name in klass.__dict__:
-                return True
+                return getattr(self, method_name)
-        return False
+        return None
 
-    # --- backends: native is required; the rest are opt-in overrides ---
+    # --- kernel backends: native is required; the rest are opt-in overrides ---
 
     @abstractmethod
     def forward_native(self, *args, **kwargs):
@@ -245,40 +418,51 @@ class BaseFusedOp(ABC):
 
     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")
+        raise NotImplementedError(f"{self._op_label()}: no triton backend")
 
     def forward_jit(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no jit backend")
+        raise NotImplementedError(f"{self._op_label()}: no jit backend")
 
     def forward_aot(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no aot backend")
+        raise NotImplementedError(f"{self._op_label()}: no aot backend")
 
     def forward_cute_dsl(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no cute_dsl backend")
+        raise NotImplementedError(f"{self._op_label()}: no cute_dsl backend")
 
     def forward_flashinfer(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no flashinfer backend")
+        raise NotImplementedError(f"{self._op_label()}: no flashinfer backend")
 
     def forward_deepgemm(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no deepgemm backend")
+        raise NotImplementedError(f"{self._op_label()}: no deepgemm backend")
 
     def forward_aiter(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no aiter backend")
+        raise NotImplementedError(f"{self._op_label()}: no aiter backend")
 
-    def forward_npu(self, *args, **kwargs):
+    def forward_torch_npu(self, *args, **kwargs):
-        raise NotImplementedError(f"{self.op}: no npu backend")
+        raise NotImplementedError(f"{self._op_label()}: no torch_npu backend")
+
+    def _op_label(self) -> str:
+        return self.op or type(self).__name__
+
+    # --- platform forwards (forward_cuda / forward_hip / forward_npu /
+    # forward_xpu / forward_musa / forward_cpu) are *not* defined here: a
+    # platform path exists exactly when a subclass defines it, and dispatch
+    # falls back to forward_native otherwise. ---
 
     # --- selection ---
 
     def available_backends(self) -> List[KernelBackend]:
-        """Backends this op implements (structural check, platform-agnostic)."""
+        """Kernel backends this op implements (structural, platform-agnostic)."""
-        return list(self._available)
+        return [
+            backend
+            for backend in KernelBackend
+            if backend in _ALWAYS_AVAILABLE
+            or self._defined_method(BACKEND_METHODS[backend]) is not None
+        ]
 
     def backend_eligible(self, backend: KernelBackend, *args, **kwargs) -> bool:
         """Whether ``backend`` may run *this* call.
@@ -286,39 +470,190 @@ class BaseFusedOp(ABC):
         The base implementation checks the backend's
         :class:`CapabilityRequirement` set (OR semantics) 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.
+        auto-selection bounces to the next backend instead of raising;
+        overriding this method switches backend auto-selection from a cached
+        static choice to per-call resolution.
         """
         return capabilities_satisfied(
             self.capabilities.get(backend, frozenset()), _platform()
         )
 
-    def _resolve_backend(self, *args, **kwargs) -> KernelBackend:
+    def _auto_backend_candidates(self) -> Tuple[KernelBackend, ...]:
-        forced = get_fused_op_backend()
+        """Kernel backends participating in auto-selection, best first.
-        if forced is not None:
+
-            return forced
+        A backend qualifies when its method is overridden *and* the op
-        for backend in self._ordered:
+        declares it in :attr:`capabilities` (``TORCH_COMPILE`` needs no
-            if self.backend_eligible(backend, *args, **kwargs):
+        declaration — it always exists and is device-agnostic — but must be
+        listed in :attr:`priority` explicitly). ``TORCH`` entries are skipped:
+        the native reference is the final fallback after platform forwards.
+        """
+        candidates = []
+        for backend in self.priority:
+            if backend is KernelBackend.TORCH:
+                continue
+            if backend is KernelBackend.TORCH_COMPILE:
+                candidates.append(backend)
+                continue
+            if (
+                backend in self.capabilities
+                and self._defined_method(BACKEND_METHODS[backend]) is not None
+            ):
+                candidates.append(backend)
+        return tuple(candidates)
+
+    def auto_selected_backend(self) -> Optional[KernelBackend]:
+        """The kernel backend auto-selection picks on this platform, or
+        ``None`` when dispatch falls through to a platform forward / native
+        (introspection only — per-call :meth:`backend_eligible` gates are not
+        consulted)."""
+        platform = _platform()
+        for backend in self._auto_backend_candidates():
+            if capabilities_satisfied(
+                self.capabilities.get(backend, frozenset()), platform
+            ):
                 return backend
-        return KernelBackend.TORCH
+        return None
+
+    def _platform_method(self, platform_key: str) -> Optional[Callable]:
+        """The platform forward for ``platform_key``, or ``None``.
+
+        In-tree keys use :data:`_PLATFORM_METHODS` (with the HIP → CUDA
+        fallback chain); OOT keys look up ``forward_<key>`` directly.
+        """
+        if not platform_key:
+            return None
+        names = _PLATFORM_METHODS.get(platform_key, (f"forward_{platform_key}",))
+        for name in names:
+            method = self._defined_method(name)
+            if method is not None:
+                return method
+        return None
+
+    def _resolve_forward_method(self) -> Callable:
+        """Resolve dispatch steps 3-6 (see module docstring) to one callable."""
+        # 3) OOT platform override.
+        oot_key = _oot_dispatch_key()
+        if oot_key is not None:
+            registered = self._oot_forward_registry.get(oot_key, {}).get(type(self))
+            if registered is not None:
+                return registered.__get__(self)
+            method = self._platform_method(oot_key)
+            if method is not None:
+                return method
+            return self.forward_native
+
+        # 4) Optimized kernel backends by priority.
+        candidates = self._auto_backend_candidates()
+        if candidates:
+            if self._defined_method("backend_eligible") is not None:
+                # Per-call shape/dtype gates: keep selection dynamic.
+                self._dynamic_backend_candidates = candidates
+                return self._forward_backend_dynamic
+            platform = _platform()
+            for backend in candidates:
+                if capabilities_satisfied(
+                    self.capabilities.get(backend, frozenset()), platform
+                ):
+                    return getattr(self, BACKEND_METHODS[backend])
+
+        # 5) Platform-specific forward; 6) native fallback.
+        method = self._platform_method(_platform_key())
+        return method if method is not None else self.forward_native
+
+    def _forward_backend_dynamic(self, *args, **kwargs):
+        """Per-call backend selection for ops with input-dependent gates."""
+        for backend in self._dynamic_backend_candidates:
+            if self.backend_eligible(backend, *args, **kwargs):
+                return getattr(self, BACKEND_METHODS[backend])(*args, **kwargs)
+        method = self._platform_method(_platform_key())
+        return (method or self.forward_native)(*args, **kwargs)
+
+    def dispatch_forward(self) -> Callable:
+        """The static dispatch target for this op on the current platform."""
+        return self._resolve_forward_method()
+
+    # --- torch.compile mode ---
+
+    def _torch_compile_forward(self, num_tokens: int) -> Optional[Callable]:
+        """The forward to use while the outer model is under ``torch.compile``.
+
+        Returning ``None`` keeps the current dispatch (used by the TopK /
+        FusedMoE overrides, whose optimized paths stay active for
+        ``num_tokens > 1`` where compiling the native path is a regression).
+        The default routes to the compile-safe pure-torch reference — the op
+        itself is *not* wrapped in a nested per-op ``torch.compile``.
+        """
+        return self.forward_native
+
+    def enter_torch_compile(self, num_tokens: int) -> None:
+        """Switch to the compile-safe forward. Idempotent.
+
+        Some ops (e.g. RotaryEmbedding) are reused among layers, so this may
+        be called many times; only the first call snapshots the original
+        dispatch, otherwise :meth:`leave_torch_compile` could not restore it.
+        """
+        if self.is_torch_compile:
+            return
+        # Warm the lazy globals now (still eager) so a compiled trace of
+        # forward() only reads plain module globals / instance attributes and
+        # never graph-breaks on an import inside get_fused_op_backend() or
+        # _resolve_forward_method().
+        get_fused_op_backend()
+        self._original_forward_method = self._forward_method
+        compile_forward = self._torch_compile_forward(num_tokens=num_tokens)
+        if compile_forward is not None:
+            self._forward_method = compile_forward
+        elif self._forward_method is None:
+            self._forward_method = self._resolve_forward_method()
+        self.is_torch_compile = True
+
+    def leave_torch_compile(self) -> None:
+        """Restore the pre-compile forward. Idempotent."""
+        if not self.is_torch_compile:
+            return
+        self._forward_method = self._original_forward_method
+        self._original_forward_method = None
+        self.is_torch_compile = False
 
     # --- dispatch ---
 
+    # Do not override forward(): implement forward_native / forward_<backend> /
+    # forward_<platform> instead, so backend forcing, OOT overrides and
+    # torch.compile mode switching keep working.
+    @debug_kernel_api
     def forward(self, *args, backend: Optional[KernelBackend] = None, **kwargs):
-        """Run the op on ``backend``, or on the best eligible one when omitted."""
+        """Run the op on ``backend``, or on the best eligible path when omitted."""
-        if backend is None:
+        if backend is not None:
-            backend = self._resolve_backend(*args, **kwargs)
+            # Explicit per-call selection is strict: an unimplemented backend
+            # raises so the caller's intent never degrades silently.
             result = getattr(self, BACKEND_METHODS[backend])(*args, **kwargs)
             if _trace_enabled:
-            _trace_records.append(
+                _record_trace(self, backend.value, args, kwargs)
-                FusedOpTraceRecord(
+            return result
-                    op=self.op,
+        forced = _forced_backend
-                    backend=backend.value,
+        if forced is _UNRESOLVED:
-                    tensor_args=_describe_tensors(args, kwargs),
+            forced = get_fused_op_backend()
-                )
+        if forced is not None:
-            )
+            # The global debug switch is best-effort: ops that do not
+            # implement the forced backend (e.g. forcing "torch" on an op
+            # whose only paths are device-specific) fall back to normal
+            # dispatch with a one-time warning instead of taking the whole
+            # model down.
+            try:
+                result = getattr(self, BACKEND_METHODS[forced])(*args, **kwargs)
+            except NotImplementedError:
+                _warn_forced_backend_unavailable(op=self, backend=forced)
+            else:
+                if _trace_enabled:
+                    _record_trace(self, forced.value, args, kwargs)
+                return result
+        method = self._forward_method
+        if method is None:
+            method = self._forward_method = self._resolve_forward_method()
+        result = method(*args, **kwargs)
+        if _trace_enabled:
+            _record_trace(self, _dispatch_label(method), args, kwargs)
         return result
-
-    __call__ = forward
 
 
 def register_fused_op(instance: BaseFusedOp, module: str, attr: str) -> BaseFusedOp:

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

@@ -138,7 +138,7 @@ class RMSNormOp(BaseFusedOp):
         rmsnorm2d_fwd(out, input, weight, eps)
         return out
 
-    def forward_npu(
+    def forward_torch_npu(
         self,
         input: torch.Tensor,
         weight: torch.Tensor,
@@ -253,7 +253,7 @@ class FusedAddRMSNormOp(BaseFusedOp):
         input.copy_(out)
         residual.copy_(residual_out)
 
-    def forward_npu(
+    def forward_torch_npu(
         self,
         input: torch.Tensor,
         residual: torch.Tensor,
@@ -344,7 +344,7 @@ class GemmaRMSNormOp(BaseFusedOp):
         out.copy_(result)
         return out
 
-    def forward_npu(
+    def forward_torch_npu(
         self,
         input: torch.Tensor,
         weight: torch.Tensor,

python/sglang/srt/compilation/torch_compile_decoration.py

@@ -9,7 +9,7 @@ flags expected by that path.
 Note: the prefill-tc_piecewise path (``TcPiecewiseCudaGraphBackend``) does NOT
 use ``patch_model`` — it goes through ``compilation/compile.py``'s
 ``install_torch_compiled``. ``_to_torch`` here is duplicated by
-tc_piecewise's local ``_toggle_multi_platform_ops``; the duplication is kept
+tc_piecewise's local ``_toggle_fused_ops``; the duplication is kept
 because the two paths have different lifecycle requirements.
 """
 
@@ -20,8 +20,8 @@ from contextlib import contextmanager
 
 import torch
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.distributed.parallel_state import GroupCoordinator
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.utils import get_bool_env_var, is_hip
 from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
 
@@ -30,7 +30,7 @@ _is_hip = is_hip()
 
 def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int) -> None:
     for sub in model._modules.values():
-        if isinstance(sub, MultiPlatformOp):
+        if isinstance(sub, BaseFusedOp):
             if reverse:
                 sub.leave_torch_compile()
             else:

python/sglang/srt/layers/activation.py

@@ -22,12 +22,12 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PretrainedConfig
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.distributed import (
     divide,
 )
 from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
     Phase,
@@ -127,7 +127,7 @@ if is_npu():
 logger = logging.getLogger(__name__)
 
 
-class SiluAndMul(MultiPlatformOp):
+class SiluAndMul(BaseFusedOp):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if get_exec().deterministic.rl_on_policy_target is not None:
@@ -181,7 +181,7 @@ class SiluAndMul(MultiPlatformOp):
         return self._musa_swish_glu(x)
 
 
-class SituAndMul(MultiPlatformOp):
+class SituAndMul(BaseFusedOp):
     """SituGLU activation used by Kimi K3.
 
     Computes beta * tanh(gate / beta) * sigmoid(gate) * up.
@@ -212,7 +212,7 @@ class SituAndMul(MultiPlatformOp):
         return self.forward_native(x)
 
 
-class GeluAndMul(MultiPlatformOp):
+class GeluAndMul(BaseFusedOp):
     def __init__(self, approximate="tanh"):
         super().__init__()
         self.approximate = approximate
@@ -259,7 +259,7 @@ class GeluAndMul(MultiPlatformOp):
         return y_npu
 
 
-class NewGELU(MultiPlatformOp):
+class NewGELU(BaseFusedOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         c = math.sqrt(2.0 / math.pi)
         return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0))))
@@ -269,7 +269,7 @@ class NewGELU(MultiPlatformOp):
         return self.forward_native(x)
 
 
-class ReLU2(MultiPlatformOp):
+class ReLU2(BaseFusedOp):
     """
     Applies the squared Rectified Linear Unit function.
     y = max(0, x)^2
@@ -283,7 +283,7 @@ class ReLU2(MultiPlatformOp):
         return relu2(x)
 
 
-class QuickGELU(MultiPlatformOp):
+class QuickGELU(BaseFusedOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         return x * torch.sigmoid(1.702 * x)
 
@@ -299,7 +299,7 @@ class QuickGELU(MultiPlatformOp):
         return torch_npu.npu_fast_gelu(x)
 
 
-class XIELU(MultiPlatformOp):
+class XIELU(BaseFusedOp):
     """
     Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010
     If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA
@@ -363,6 +363,9 @@ class XIELU(MultiPlatformOp):
             (torch.expm1(torch.min(x, self.eps)) - x) * alpha_n + self.beta * x,
         )
 
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        return self._xielu_python(x)
+
     def _xielu_cuda(self, x: torch.Tensor) -> torch.Tensor:
         """Firewall function to prevent torch.compile from seeing .item()"""
         assert self._xielu_cuda_obj is not None, "XIELU CUDA object must not be None"

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

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 import torch
 from einops import rearrange
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.kernels.ops.attention.fused_store_index_cache import (
     can_use_dsa_fused_store,
     fused_store_index_k_cache,
@@ -32,7 +33,6 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_graph_dsa_split_op_surface,
 )
 from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
     is_in_breakable_cuda_graph,
 )
@@ -197,7 +197,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     return hadamard_transform(x, scale=hidden_size**-0.5)
 
 
-class Indexer(DSANPUIndexerMixin, MultiPlatformOp):
+class Indexer(DSANPUIndexerMixin, BaseFusedOp):
     _MQA_LOGITS_BYTES_PER_ELEM = 4
     _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
     _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3
@@ -1474,6 +1474,24 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp):
             index_k_scale=k_scale,
         )
 
+    def forward_native(self, *args, **kwargs):
+        # The indexer has no pure-torch reference path; it only runs on
+        # platforms with a dedicated forward below.
+        raise NotImplementedError("Indexer has no native (pure-torch) path")
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        q_lora: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        return_indices: bool = True,
+    ) -> Optional[torch.Tensor]:
+        return self.forward_cuda(
+            x, q_lora, positions, forward_batch, layer_id, return_indices
+        )
+
     def forward_cuda(
         self,
         x: torch.Tensor,

python/sglang/srt/layers/attention/dsv4/compressor.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Union
 import torch
 import torch.nn as nn
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant
 from sglang.kernels.ops.attention.dsv4 import (
     linear_bf16_fp32,
@@ -25,7 +26,6 @@ from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
-from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
 from sglang.srt.mem_cache.deepseek_v4_compress_state import (
     CompressStatePool,
 )
@@ -344,7 +344,7 @@ def create_paged_compressor_data(
     return FusedCompressMetadata(write_loc=write_loc, extra_data=extra_data, plan=plan)
 
 
-class Compressor(MultiPlatformOp):
+class Compressor(BaseFusedOp):
     def __init__(
         self,
         config: DeepSeekV4Config,

python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py

@@ -2,6 +2,7 @@ from typing import Union
 
 import torch
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.kernels.ops.attention.fla.layernorm_gated import rms_norm_gated
 from sglang.srt.distributed.communication_op import (
     tensor_model_parallel_all_gather,
@@ -11,13 +12,12 @@ from sglang.srt.layers.dp_attention import (
     attn_tp_all_reduce,
     is_dp_attention_enabled,
 )
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_loader.weight_utils import sharded_weight_loader
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import set_weight_attrs
 
 
-class Mixer2RMSNormGated(MultiPlatformOp):
+class Mixer2RMSNormGated(BaseFusedOp):
     def __init__(
         self,
         full_hidden_size: int,

python/sglang/srt/layers/conv.py

@@ -14,8 +14,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.layers.amx_utils import PackWeightMethod
-from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
 from sglang.srt.utils import cpu_has_amx_support, is_cpu, use_intel_amx_backend
 
 _is_cpu = is_cpu()
@@ -98,7 +98,7 @@ def _validate_conv_args(
             raise ValueError("padding='same' is not supported for strided convolutions")
 
 
-class Conv2dLayer(MultiPlatformOp):
+class Conv2dLayer(BaseFusedOp):
     """Drop-in replacement for nn.Conv2d. Linear optimization disabled by default."""
 
     def __init__(
@@ -204,7 +204,7 @@ class Conv2dLayer(MultiPlatformOp):
         return self._forward_conv(x)
 
 
-class Conv3dLayer(MultiPlatformOp):
+class Conv3dLayer(BaseFusedOp):
     """Drop-in replacement for nn.Conv3d with automatic linear optimization."""
 
     def __init__(

python/sglang/srt/layers/layernorm.py

@@ -21,12 +21,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.batch_invariant_ops import (
     is_batch_invariant_mode_enabled,
     rms_norm_batch_invariant,
 )
 from sglang.srt.environ import envs
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
     Phase,
@@ -354,7 +354,7 @@ def _forward_with_allreduce_fusion_quant_per_group(
     return (bf16_out, fp8_out, scale_out), residual_out
 
 
-class RMSNorm(MultiPlatformOp):
+class RMSNorm(BaseFusedOp):
     def __init__(
         self,
         hidden_size: int,
@@ -770,7 +770,7 @@ class RMSNorm(MultiPlatformOp):
         )
 
 
-class LayerNorm(MultiPlatformOp):
+class LayerNorm(BaseFusedOp):
     def __init__(
         self,
         hidden_size: int,
@@ -854,7 +854,7 @@ class LayerNorm(MultiPlatformOp):
             return self.forward_native(x)
 
 
-class GemmaRMSNorm(MultiPlatformOp):
+class GemmaRMSNorm(BaseFusedOp):
     def __init__(
         self,
         hidden_size: int,
@@ -1014,6 +1014,16 @@ class GemmaRMSNorm(MultiPlatformOp):
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         return self._forward_impl(x, residual, post_residual_addition)
 
+    def forward_musa(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+        post_residual_addition: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # sgl_kernel's gemma norm ops are built for MUSA (see the import gate
+        # above); opt into the CUDA-path implementation explicitly.
+        return self._forward_impl(x, residual, post_residual_addition)
+
     def forward_with_allreduce_fusion(
         self,
         x: torch.Tensor,
@@ -1051,7 +1061,7 @@ class GemmaRMSNorm(MultiPlatformOp):
         )
 
 
-class Gemma3RMSNorm(MultiPlatformOp):
+class Gemma3RMSNorm(BaseFusedOp):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.eps = eps
@@ -1090,6 +1100,10 @@ class Gemma3RMSNorm(MultiPlatformOp):
             return gemma_rmsnorm(x, self.weight.data, self.eps)
         return self.forward_native(x)
 
+    def forward_musa(self, x, residual: Optional[torch.Tensor] = None):
+        # sgl_kernel's gemma norm ops are built for MUSA; follow the CUDA path.
+        return self.forward_cuda(x, residual)
+
     def forward_hip(self, x, residual: Optional[torch.Tensor] = None):
         # sgl_kernel's gemma_rmsnorm/gemma_fused_add_rmsnorm are not available on
         # ROCm; delegate to the pure-PyTorch implementation.
@@ -1105,7 +1119,7 @@ class Gemma3RMSNorm(MultiPlatformOp):
         return f"{tuple(self.weight.shape)}, eps={self.eps}"
 
 
-class Gemma4RMSNorm(MultiPlatformOp):
+class Gemma4RMSNorm(BaseFusedOp):
     def __init__(
         self,
         dim: int,
@@ -1177,13 +1191,17 @@ class Gemma4RMSNorm(MultiPlatformOp):
             out = rmsnorm(x, self.weight.data, self.eps)
         return out
 
+    def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
+        # sgl_kernel's gemma norm ops are built for MUSA; follow the CUDA path.
+        return self.forward_cuda(x)
+
     def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
         # sgl_kernel's gemma_rmsnorm is not available on ROCm;
         # delegate to the pure-PyTorch implementation.
         return self.forward_native(x)
 
 
-class RMSNormWithoutScale(MultiPlatformOp):
+class RMSNormWithoutScale(BaseFusedOp):
     def __init__(self, hidden_size: int, eps=1e-6):
         super().__init__()
         self.hidden_size = hidden_size

python/sglang/srt/layers/moe/topk.py

@@ -82,6 +82,7 @@ try:
 except ImportError:
     pass
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.kernels.ops.attention.dsv4 import mask_topk_ids
 from sglang.srt.distributed import (
     get_tp_group,
@@ -101,7 +102,6 @@ from sglang.srt.layers.moe import get_moe_runner_backend
 from sglang.srt.layers.moe.utils import (
     has_per_rank_fused_shared_slots,
 )
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -392,7 +392,7 @@ def _make_round_robin_expert_ids(
 # -------------------------------- TopK ---------------------------------------
 
 
-class TopK(MultiPlatformOp):
+class TopK(BaseFusedOp):
     """
     Parameters:
     --top_k: The all number of top experts selected per token, including the fused shared expert(s).
@@ -472,6 +472,18 @@ class TopK(MultiPlatformOp):
         assert TopKOutputChecker.format_is_standard(topk_output)
         return self.waterfill_balancer.expand_topk(topk_output, num_tokens)
 
+    def forward_musa(self, *args, **kwargs) -> TopKOutput:
+        # MUSA follows the CUDA path explicitly: select_experts branches on
+        # _is_musa internally (hardware_backend.musa topk kernels), so the
+        # native path would bypass them.
+        return self.forward_cuda(*args, **kwargs)
+
+    def _torch_compile_forward(self, num_tokens: int) -> Optional[Callable]:
+        # torch.compile of the native TopK only pays off at bs=1; for larger
+        # batches keep the current optimized dispatch (see MultiPlatformOp
+        # history: the compiled path regressed bs > 1).
+        return self.forward_native if num_tokens == 1 else None
+
     def forward_native(
         self,
         hidden_states: torch.Tensor,

python/sglang/srt/layers/quantization/unquant.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 logger = logging.getLogger(__name__)
 
@@ -10,6 +10,7 @@ import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.environ import envs
 from sglang.srt.layers.amx_utils import (
     CPUQuantMethod,
@@ -28,7 +29,7 @@ from sglang.srt.layers.quantization.base_config import (
     LinearMethodBase,
     QuantizeMethodBase,
 )
-from sglang.srt.layers.utils import MultiPlatformOp, copy_or_rebind_param
+from sglang.srt.layers.utils import copy_or_rebind_param
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -297,7 +298,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
         return output
 
 
-class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
+class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
     """MoE method without quantization."""
 
     def __init__(
@@ -610,6 +611,20 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
             dispatch_output=dispatch_output,
         )
 
+    # forward_native is aliased to forward_cpu at the end of the class body
+    # (pre-existing behavior); under torch.compile the dedicated
+    # fused_moe_forward_native is installed instead via this hook.
+    def _torch_compile_forward(self, num_tokens: int) -> Optional[Callable]:
+        # torch.compile on this layer only pays off at bs=1; keep the
+        # optimized dispatch otherwise.
+        if num_tokens == 1:
+            from sglang.srt.layers.moe.fused_moe_native import (
+                fused_moe_forward_native,
+            )
+
+            return fused_moe_forward_native
+        return None
+
     def forward_cuda(
         self,
         layer: torch.nn.Module,

python/sglang/srt/layers/rotary_embedding/base.py

@@ -7,9 +7,9 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 import torch
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.environ import envs
 from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_exec
 from sglang.srt.utils import (
@@ -75,7 +75,7 @@ if _is_xpu:
     from sgl_kernel import fused_qk_rope_with_cos_sin_cache_inplace
 
 
-class RotaryEmbedding(MultiPlatformOp):
+class RotaryEmbedding(BaseFusedOp):
     """Original rotary positional embedding."""
 
     def __init__(

python/sglang/srt/layers/rotary_embedding/rope_variant.py

@@ -10,6 +10,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding
 from sglang.srt.layers.rotary_embedding.utils import (
     apply_rotary_pos_emb_native,
@@ -21,7 +22,6 @@ from sglang.srt.layers.rotary_embedding.yarn import (
     yarn_get_mscale,
     yarn_linear_ramp_mask,
 )
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.utils import cpu_has_amx_support, get_device, is_cuda, is_hip, is_npu
 
 _is_cuda = is_cuda()
@@ -675,7 +675,7 @@ class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):
         return cache
 
 
-class DualChunkRotaryEmbedding(MultiPlatformOp):
+class DualChunkRotaryEmbedding(BaseFusedOp):
     """Rotary positional embedding for Dual Chunk Attention."""
 
     def __init__(
@@ -753,6 +753,11 @@ class DualChunkRotaryEmbedding(MultiPlatformOp):
         ).to(dtype=self.dtype, device=self.device)
         return q_cache, qc_cache, k_cache, qc_no_clamp_cache, q_inter_cache
 
+    def forward_native(self, *args, **kwargs):
+        # This op overrides forward() directly; there is no separate
+        # pure-torch reference path.
+        raise NotImplementedError("DualChunkRotaryEmbedding has no native path")
+
     def forward(
         self,
         positions: torch.Tensor,

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

@@ -1,3 +1,6 @@
 # Temp workaround, make layer utils more fine-grained later
 from sglang.srt.layers.utils.common import *
+
+# Deprecated re-export kept for external users; in-repo code subclasses
+# sglang.kernels.fused_op.BaseFusedOp directly (RFC #29630).
 from sglang.srt.layers.utils.multi_platform import MultiPlatformOp

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

@@ -1,86 +1,41 @@
-from typing import Callable, ClassVar
+"""Deprecated compatibility shim for the former ``MultiPlatformOp``.
 
-from torch import nn
+The multi-platform operator abstraction was unified into
+:class:`sglang.kernels.fused_op.BaseFusedOp` (RFC #29630): one class now
+covers kernel-backend selection (``forward_aot`` / ``forward_jit`` / ...),
+platform dispatch (``forward_cuda`` / ``forward_hip`` / ``forward_npu`` /
+...), out-of-tree platform overrides (:meth:`BaseFusedOp.register_oot_forward`),
+and the torch.compile enter/leave protocol.
 
-from sglang.kernel_api_logging import debug_kernel_api
+In-repo code must subclass ``BaseFusedOp`` directly. This alias exists only so
-from sglang.srt.platforms import current_platform
+out-of-tree platform plugins and external users keep importing from the old
-from sglang.srt.utils import (
+path while they migrate; it will be removed in a future release.
-    cpu_has_amx_support,
+"""
-    is_cpu,
+
-    is_cuda,
+import warnings
-    is_hip,
+
-    is_musa,
+from sglang.kernels.fused_op import BaseFusedOp
-    is_npu,
+
-    is_xpu,
+
+class MultiPlatformOp(BaseFusedOp):
+    """Deprecated alias of :class:`sglang.kernels.fused_op.BaseFusedOp`.
+
+    Kept attribute-compatible with the original class: ``forward_native`` is
+    concrete here (raising ``NotImplementedError``) so existing plugin
+    subclasses that only define platform forwards keep instantiating, and the
+    old per-platform default methods (``forward_hip`` -> ``forward_cuda``,
+    ``forward_cpu`` -> ``forward_native``, ...) remain callable for plugin
+    code that invokes them directly.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        warnings.warn(
+            "MultiPlatformOp is deprecated; subclass "
+            "sglang.kernels.fused_op.BaseFusedOp instead (RFC #29630).",
+            DeprecationWarning,
+            stacklevel=2,
         )
-
+        super().__init_subclass__(**kwargs)
-_is_cuda = is_cuda()
-_is_hip = is_hip()
-_is_cpu = is_cpu()
-_is_cpu_amx_available = cpu_has_amx_support()
-_is_npu = is_npu()
-_is_xpu = is_xpu()
-_is_musa = is_musa()
-
-
-class MultiPlatformOp(nn.Module):
-
-    # OOT forward registry: maps dispatch_key -> {op_cls -> forward_fn}
-    _oot_forward_registry: ClassVar[dict[str, dict[type, Callable]]] = {}
-
-    @classmethod
-    def register_oot_forward(cls, op_cls: type, fn: Callable, platform_key: str):
-        """Register an OOT forward implementation for a specific op class and platform."""
-        cls._oot_forward_registry.setdefault(platform_key, {})[op_cls] = fn
-
-    def __init__(self):
-        super().__init__()
-        self._forward_method: Callable = self.dispatch_forward()
-
-        # States for torch.compile
-        self._original_forward_method = None
-        self.is_torch_compile = False
-
-    def enter_torch_compile(self, num_tokens: int):
-        # Skip if Op is already entered compile mode.
-        # NOTE(alcanderian): Some Ops(for example RotaryEmbedding) will be reused
-        # among layers and `enter_torch_compile` will be called many times.
-        # We should prevent `self._original_forward_method` from being overridden when
-        # it is not the first time `enter_torch_compile` called.
-        if self.is_torch_compile:
-            return
-
-        self._original_forward_method = self._forward_method
-        # NOTE: Temporarily workaround MoE
-        # The performance of torch.compile on this layer is not always good when bs > 1,
-        # so we decide to only use torch.compile when bs=1
-        if "FusedMoE" in self.__class__.__name__:
-            if num_tokens == 1:
-                from sglang.srt.layers.moe.fused_moe_native import (
-                    fused_moe_forward_native,
-                )
-
-                self._forward_method = fused_moe_forward_native
-        elif "TopK" in self.__class__.__name__:
-            if num_tokens == 1:
-                self._forward_method = self.forward_native
-        else:
-            self._forward_method = self.forward_native
-        self.is_torch_compile = True
-
-    def leave_torch_compile(self):
-        # Skip if Op is already exited compile mode.
-        if not self.is_torch_compile:
-            return
-
-        self._forward_method = self._original_forward_method
-        self._original_forward_method = None
-        self.is_torch_compile = False
-
-    # Please do not override this method, because `self._forward_method` can change when in torch compile mode
-    @debug_kernel_api
-    def forward(self, *args, **kwargs):
-        return self._forward_method(*args, **kwargs)
 
     def forward_native(self, *args, **kwargs):
         raise NotImplementedError
@@ -88,47 +43,20 @@ class MultiPlatformOp(nn.Module):
     def forward_cuda(self, *args, **kwargs):
         raise NotImplementedError
 
-    def forward_npu(self, *args, **kwargs):
-        return self.forward_native(*args, **kwargs)
-
     def forward_hip(self, *args, **kwargs):
         return self.forward_cuda(*args, **kwargs)
 
-    def forward_xpu(self, *args, **kwargs):
-        return self.forward_native(*args, **kwargs)
-
     def forward_musa(self, *args, **kwargs):
         return self.forward_cuda(*args, **kwargs)
 
+    def forward_npu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
     def forward_hpu(self, *args, **kwargs):
         return self.forward_native(*args, **kwargs)
 
     def forward_cpu(self, *args, **kwargs):
         return self.forward_native(*args, **kwargs)
-
-    def dispatch_forward(self):
-        # OOT platform dispatch: check registry then method lookup
-        if current_platform.is_out_of_tree():
-            key = current_platform.get_dispatch_key_name()
-            oot = self._oot_forward_registry.get(key, {})
-            if type(self) in oot:
-                return oot[type(self)].__get__(self)
-            method = getattr(self, f"forward_{key}", None)
-            if method is not None:
-                return method
-            return self.forward_native
-
-        if _is_cuda:
-            return self.forward_cuda
-        elif _is_hip:
-            return self.forward_hip
-        elif _is_cpu and _is_cpu_amx_available:
-            return self.forward_cpu
-        elif _is_npu:
-            return self.forward_npu
-        elif _is_xpu:
-            return self.forward_xpu
-        elif _is_musa:
-            return self.forward_musa
-        else:
-            return self.forward_native

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
 import torch
 import tqdm
 
+from sglang.kernels.fused_op import BaseFusedOp
 from sglang.srt.compilation.compilation_config import CompilationConfig
 from sglang.srt.compilation.compile import install_torch_compiled
 from sglang.srt.compilation.compile_phase import (
@@ -39,7 +40,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
 )
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
-from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
@@ -68,19 +68,19 @@ def _suppress_lru_cache_dynamo_warning() -> None:
     warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
 
 
-def _toggle_multi_platform_ops(
+def _toggle_fused_ops(
     model: torch.nn.Module, *, reverse: bool, num_tokens: int
 ) -> None:
-    """Recursively flip MultiPlatformOp submodules into / out of
+    """Recursively flip BaseFusedOp submodules into / out of
     torch.compile mode."""
     for sub in model._modules.values():
-        if isinstance(sub, MultiPlatformOp):
+        if isinstance(sub, BaseFusedOp):
             if reverse:
                 sub.leave_torch_compile()
             else:
                 sub.enter_torch_compile(num_tokens=num_tokens)
         if isinstance(sub, torch.nn.Module):
-            _toggle_multi_platform_ops(sub, reverse=reverse, num_tokens=num_tokens)
+            _toggle_fused_ops(sub, reverse=reverse, num_tokens=num_tokens)
 
 
 class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
@@ -163,9 +163,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
         with enable_tc_piecewise_cuda_graph():
             try:
                 if compiler != "eager":
-                    _toggle_multi_platform_ops(
+                    _toggle_fused_ops(inner_model, reverse=False, num_tokens=16)
-                        inner_model, reverse=False, num_tokens=16
-                    )
 
                 cuda_graph_runner._run_dummy_forward(
                     num_tokens=cuda_graph_runner.capture_num_tokens[0]
@@ -215,7 +213,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                     inner_model, cuda_graph_runner.capture_num_tokens[-1]
                 )
             finally:
-                _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16)
+                _toggle_fused_ops(inner_model, reverse=True, num_tokens=16)
 
     @contextmanager
     def capture_session(self, stream: torch.cuda.Stream):

python/sglang/srt/platforms/interface.py

@@ -127,13 +127,16 @@ class SRTPlatform(DeviceMixin):
         pass
 
     # ------------------------------------------------------------------
-    # MultiPlatformOp integration
+    # BaseFusedOp integration
     # ------------------------------------------------------------------
 
     def get_dispatch_key_name(self) -> str:
-        """Return the dispatch key name for MultiPlatformOp.
+        """Return the dispatch key name for BaseFusedOp
+        (``sglang.kernels.fused_op``).
 
-        Determines which ``forward_<key>()`` method is selected.
+        Determines which ``forward_<key>()`` method is selected on an
-        E.g. "cuda", "npu", "hip", "xpu", "cpu".
+        out-of-tree platform. E.g. "cuda", "npu", "hip", "xpu", "cpu".
+        Forwards registered via ``BaseFusedOp.register_oot_forward`` with
+        this key take precedence over the method lookup.
         """
         return "native"

test/manual/kernels/bench_fused_op_dispatch.py

@@ -0,0 +1,81 @@
+"""Hot-path dispatch overhead microbenchmark for the unified ``BaseFusedOp``.
+
+Compares per-call overhead of:
+
+1. a plain bound-method call (theoretical floor),
+2. a minimal reproduction of the former ``MultiPlatformOp`` hot path
+   (``nn.Module.__call__`` -> ``self._forward_method(*args)``),
+3. the unified ``BaseFusedOp`` (adds the forced-backend check, the cached
+   dispatch lookup, and the trace flag check).
+
+The op body is a no-op so the numbers isolate pure dispatch overhead; real
+kernels are microseconds+, so the delta reported here is the worst case.
+
+Run locally (CPU is fine):
+    python test/manual/kernels/bench_fused_op_dispatch.py
+"""
+
+import time
+
+import torch
+from torch import nn
+
+from sglang.kernels.fused_op import BaseFusedOp
+
+N_WARMUP = 10_000
+N_ITERS = 200_000
+
+
+class _OldStyleOp(nn.Module):
+    """Minimal replica of the retired MultiPlatformOp hot path."""
+
+    def __init__(self):
+        super().__init__()
+        self._forward_method = self.forward_cuda
+
+    def forward(self, *args, **kwargs):
+        return self._forward_method(*args, **kwargs)
+
+    def forward_cuda(self, x):
+        return x
+
+
+class _NewOp(BaseFusedOp):
+    op = "bench.dispatch"
+
+    def forward_native(self, x):
+        return x
+
+    def forward_cuda(self, x):
+        return x
+
+
+def _bench(fn, x) -> float:
+    for _ in range(N_WARMUP):
+        fn(x)
+    start = time.perf_counter()
+    for _ in range(N_ITERS):
+        fn(x)
+    return (time.perf_counter() - start) / N_ITERS * 1e9  # ns/call
+
+
+def main():
+    x = torch.zeros(1)
+
+    old_op = _OldStyleOp()
+    new_op = _NewOp()
+    new_op(x)  # resolve + cache dispatch
+    bound = new_op.forward_native
+
+    floor_ns = _bench(bound, x)
+    old_ns = _bench(old_op, x)
+    new_ns = _bench(new_op, x)
+
+    print(f"plain bound method      : {floor_ns:8.1f} ns/call")
+    print(f"MultiPlatformOp replica : {old_ns:8.1f} ns/call")
+    print(f"unified BaseFusedOp     : {new_ns:8.1f} ns/call")
+    print(f"delta (new - old)       : {new_ns - old_ns:8.1f} ns/call")
+
+
+if __name__ == "__main__":
+    main()

test/registered/kernels/ops/layernorm/test_fused_op.py

@@ -22,6 +22,8 @@ register_cpu_ci(est_time=30, suite="base-a-test-cpu")
 class _ToyAdd(BaseFusedOp):
     op = "test.toy_add"
     priority = (KernelBackend.TRITON, KernelBackend.TORCH)
+    # Auto-selection requires backends to be declared (empty set = any device).
+    capabilities = {KernelBackend.TRITON: frozenset()}
 
     def forward_native(self, a, b):
         return a + b

test/registered/kernels/ops/layernorm/test_kernels_namespace.py

@@ -111,7 +111,7 @@ def test_activation_default_backend(monkeypatch, device, expect):
     from sglang.kernels.ops.activation import _SILU_AND_MUL
 
     monkeypatch.setattr(fo, "_platform", lambda: PlatformInfo(device_type=device))
-    assert _SILU_AND_MUL._resolve_backend().value == expect
+    assert _SILU_AND_MUL.auto_selected_backend().value == expect
 
 
 @pytest.mark.parametrize(
@@ -130,7 +130,7 @@ def test_layernorm_default_backend(monkeypatch, op_attr, device, expect):
     # CUDA-only, so HIP falls to aiter and NPU to torch_npu.
     ln = importlib.import_module("sglang.kernels.ops.layernorm")
     monkeypatch.setattr(fo, "_platform", lambda: PlatformInfo(device_type=device))
-    assert getattr(ln, op_attr)._resolve_backend().value == expect
+    assert getattr(ln, op_attr).auto_selected_backend().value == expect
 
 
 def test_per_op_backend_subset():

test/registered/kernels/test_fused_op_dispatch.py

@@ -0,0 +1,558 @@
+"""Dispatch-contract tests for the unified ``BaseFusedOp`` (RFC #29630, #26426).
+
+``BaseFusedOp`` replaced ``MultiPlatformOp`` as the single operator
+abstraction; these tests pin down the parts of that contract that a refactor
+could silently break:
+
+- the priority ladder: explicit ``backend=`` > global forced backend > OOT
+  platform override > declared optimized kernel backends > platform-specific
+  forward > native fallback;
+- the standard ``nn.Module`` behavior (hooks, traversal);
+- static-dispatch caching and per-call ``backend_eligible`` gating;
+- the torch.compile enter/leave protocol (idempotency, TopK / FusedMoE
+  special paths);
+- the deprecated ``MultiPlatformOp`` alias and its OOT plugin surface.
+
+Platform detection is mocked, so everything here runs on a CPU-only box.
+"""
+
+import warnings
+
+import pytest
+import torch
+from torch import nn
+
+import sglang.kernels.fused_op as fo
+from sglang.kernels.fused_op import BaseFusedOp
+from sglang.kernels.spec import CapabilityRequirement as Cap
+from sglang.kernels.spec import KernelBackend, PlatformInfo
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=60, suite="base-a-test-cpu")
+
+_CUDA = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
+_HIP = PlatformInfo(device_type="hip")
+_CPU = PlatformInfo()
+
+
+@pytest.fixture(autouse=True)
+def _reset_global_state():
+    saved_oot = {k: dict(v) for k, v in BaseFusedOp._oot_forward_registry.items()}
+    yield
+    fo.set_fused_op_backend(None)
+    fo.disable_fused_op_trace()
+    fo.clear_fused_op_trace()
+    BaseFusedOp._oot_forward_registry.clear()
+    BaseFusedOp._oot_forward_registry.update(saved_oot)
+
+
+def _mock_platform(monkeypatch, *, key="", info=_CPU, oot_key=None):
+    monkeypatch.setattr(fo, "_platform_key", lambda: key)
+    monkeypatch.setattr(fo, "_platform", lambda: info)
+    monkeypatch.setattr(fo, "_oot_dispatch_key", lambda: oot_key)
+
+
+class _AllPlatformsOp(BaseFusedOp):
+    """Marks which path ran by returning its name."""
+
+    op = "test.all_platforms"
+
+    def forward_native(self, x):
+        return "native"
+
+    def forward_cuda(self, x):
+        return "cuda"
+
+    def forward_hip(self, x):
+        return "hip"
+
+    def forward_npu(self, x):
+        return "npu"
+
+    def forward_xpu(self, x):
+        return "xpu"
+
+    def forward_musa(self, x):
+        return "musa"
+
+    def forward_cpu(self, x):
+        return "cpu"
+
+
+class _CudaOnlyPlatformOp(BaseFusedOp):
+    op = "test.cuda_only_platform"
+
+    def forward_native(self, x):
+        return "native"
+
+    def forward_cuda(self, x):
+        return "cuda"
+
+
+class _NativeOnlyOp(BaseFusedOp):
+    op = "test.native_only"
+
+    def forward_native(self, x):
+        return "native"
+
+
+class _BackendAndPlatformOp(BaseFusedOp):
+    """Declared JIT backend + a CUDA platform forward."""
+
+    op = "test.backend_and_platform"
+    priority = (KernelBackend.JIT, KernelBackend.TORCH)
+    capabilities = {KernelBackend.JIT: frozenset({Cap.CUDA})}
+
+    def forward_native(self, x):
+        return "native"
+
+    def forward_jit(self, x):
+        return "jit"
+
+    def forward_cuda(self, x):
+        return "cuda"
+
+
+class _UndeclaredBackendOp(BaseFusedOp):
+    """Overrides forward_aiter but does not declare it in ``capabilities``."""
+
+    op = "test.undeclared_backend"
+
+    def forward_native(self, x):
+        return "native"
+
+    def forward_aiter(self, x):
+        return "aiter"
+
+
+# --- nn.Module contract -------------------------------------------------------
+
+
+def test_is_standard_nn_module(monkeypatch):
+    _mock_platform(monkeypatch)
+    op = _NativeOnlyOp()
+    assert isinstance(op, nn.Module)
+
+    parent = nn.Module()
+    parent.act = op
+    assert dict(parent.named_modules())["act"] is op
+
+    seen = []
+    op.register_forward_hook(lambda module, args, output: seen.append(output))
+    assert op(torch.zeros(1)) == "native"
+    assert seen == ["native"]  # __call__ goes through nn.Module, hooks fire
+
+
+# --- platform dispatch + native fallback ---------------------------------------
+
+
+@pytest.mark.parametrize(
+    "key, expect",
+    [
+        ("cuda", "cuda"),
+        ("hip", "hip"),
+        ("npu", "npu"),
+        ("xpu", "xpu"),
+        ("musa", "musa"),
+        ("cpu", "cpu"),
+        ("", "native"),
+    ],
+)
+def test_platform_forward_dispatch(monkeypatch, key, expect):
+    _mock_platform(monkeypatch, key=key)
+    assert _AllPlatformsOp()(torch.zeros(1)) == expect
+
+
+@pytest.mark.parametrize(
+    "key, expect",
+    [
+        ("hip", "cuda"),  # HIP falls back to the CUDA path (hipified kernels)
+        # MUSA has no implicit CUDA fallback: srt kernel imports are gated on
+        # is_cuda(), so silently entering forward_cuda on a MUSA box can
+        # NameError; ops opt in with an explicit forward_musa instead.
+        ("musa", "native"),
+        ("npu", "native"),  # no NPU path -> native
+        ("cpu", "native"),
+        ("cuda", "cuda"),
+    ],
+)
+def test_platform_default_chains(monkeypatch, key, expect):
+    _mock_platform(monkeypatch, key=key)
+    assert _CudaOnlyPlatformOp()(torch.zeros(1)) == expect
+
+
+def test_native_fallback_without_any_override(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    assert _NativeOnlyOp()(torch.zeros(1)) == "native"
+
+
+# --- optimized-backend selection ------------------------------------------------
+
+
+def test_declared_backend_beats_platform_forward(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    assert _BackendAndPlatformOp()(torch.zeros(1)) == "jit"
+
+
+def test_capability_filters_backend_to_platform_forward(monkeypatch):
+    # JIT is declared CUDA-only; on HIP the platform chain (-> forward_cuda) runs.
+    _mock_platform(monkeypatch, key="hip", info=_HIP)
+    assert _BackendAndPlatformOp()(torch.zeros(1)) == "cuda"
+
+
+def test_undeclared_backend_not_auto_selected(monkeypatch):
+    _mock_platform(monkeypatch, key="", info=_CPU)
+    op = _UndeclaredBackendOp()
+    assert op(torch.zeros(1)) == "native"
+    # ... but stays reachable by explicit request.
+    assert op(torch.zeros(1), backend=KernelBackend.AITER) == "aiter"
+
+
+def test_priority_order_decides_between_backends(monkeypatch):
+    class _TwoBackends(BaseFusedOp):
+        op = "test.two_backends"
+        priority = (KernelBackend.TRITON, KernelBackend.JIT, KernelBackend.TORCH)
+        capabilities = {
+            KernelBackend.TRITON: frozenset(),
+            KernelBackend.JIT: frozenset(),
+        }
+
+        def forward_native(self, x):
+            return "native"
+
+        def forward_triton(self, x):
+            return "triton"
+
+        def forward_jit(self, x):
+            return "jit"
+
+    class _Flipped(_TwoBackends):
+        priority = (KernelBackend.JIT, KernelBackend.TRITON, KernelBackend.TORCH)
+
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    assert _TwoBackends()(torch.zeros(1)) == "triton"
+    assert _Flipped()(torch.zeros(1)) == "jit"
+
+
+def test_explicit_backend_beats_forced_global(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _BackendAndPlatformOp()
+    fo.set_fused_op_backend(KernelBackend.TORCH)
+    assert op(torch.zeros(1)) == "native"  # forced global
+    assert op(torch.zeros(1), backend=KernelBackend.JIT) == "jit"  # explicit wins
+
+
+def test_forced_global_falls_back_when_unimplemented(monkeypatch):
+    # The global debug switch must not take down ops that lack the forced
+    # backend (e.g. forcing "torch" on a device-only op like the DSA indexer,
+    # whose forward_native raises NotImplementedError).
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+
+    class _DeviceOnly(BaseFusedOp):
+        op = "test.device_only"
+
+        def forward_native(self, x):
+            raise NotImplementedError
+
+        def forward_cuda(self, x):
+            return "cuda"
+
+    op = _DeviceOnly()
+    fo.set_fused_op_backend(KernelBackend.TORCH)
+    assert op(torch.zeros(1)) == "cuda"  # fell back to normal dispatch
+    fo.set_fused_op_backend(KernelBackend.JIT)
+    assert op(torch.zeros(1)) == "cuda"  # no jit backend -> fall back too
+    # Explicit per-call selection stays strict.
+    fo.set_fused_op_backend(None)
+    with pytest.raises(NotImplementedError):
+        op(torch.zeros(1), backend=KernelBackend.JIT)
+
+
+def test_forced_global_beats_platform_and_oot(monkeypatch):
+    _mock_platform(monkeypatch, key="", oot_key="myplat")
+    BaseFusedOp.register_oot_forward(
+        _CudaOnlyPlatformOp, lambda self, x: "oot", "myplat"
+    )
+    op = _CudaOnlyPlatformOp()
+    fo.set_fused_op_backend(KernelBackend.TORCH)
+    assert op(torch.zeros(1)) == "native"
+    fo.set_fused_op_backend(None)
+    assert op(torch.zeros(1)) == "oot"
+
+
+# --- OOT platform overrides ----------------------------------------------------
+
+
+def test_oot_registered_forward_wins_over_method(monkeypatch):
+    class _OotOp(BaseFusedOp):
+        op = "test.oot"
+
+        def forward_native(self, x):
+            return "native"
+
+        def forward_myplat(self, x):
+            return "method"
+
+    _mock_platform(monkeypatch, oot_key="myplat")
+    assert _OotOp()(torch.zeros(1)) == "method"  # forward_<key> lookup
+
+    BaseFusedOp.register_oot_forward(_OotOp, lambda self, x: "registered", "myplat")
+    assert _OotOp()(torch.zeros(1)) == "registered"  # registry beats method
+
+
+def test_oot_registration_is_exact_type(monkeypatch):
+    _mock_platform(monkeypatch, oot_key="myplat")
+    BaseFusedOp.register_oot_forward(
+        _CudaOnlyPlatformOp, lambda self, x: "oot", "myplat"
+    )
+
+    class _Sub(_CudaOnlyPlatformOp):
+        pass
+
+    assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "oot"
+    # Subclasses do not inherit the registered forward (pre-existing
+    # MultiPlatformOp semantics: lookup is by exact type).
+    assert _Sub()(torch.zeros(1)) == "native"
+
+
+def test_oot_falls_back_to_native(monkeypatch):
+    _mock_platform(monkeypatch, oot_key="myplat")
+    assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "native"
+
+
+def test_oot_registered_fn_is_bound(monkeypatch):
+    _mock_platform(monkeypatch, oot_key="myplat")
+    BaseFusedOp.register_oot_forward(
+        _CudaOnlyPlatformOp, lambda self, x: type(self).__name__, "myplat"
+    )
+    assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "_CudaOnlyPlatformOp"
+
+
+# --- dispatch caching + per-call gates ------------------------------------------
+
+
+def test_static_dispatch_resolved_once(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _CudaOnlyPlatformOp()
+    calls = []
+    original = op._resolve_forward_method
+    monkeypatch.setattr(
+        op,
+        "_resolve_forward_method",
+        lambda: calls.append(1) or original(),
+    )
+    op(torch.zeros(1))
+    op(torch.zeros(1))
+    assert len(calls) == 1  # hot path must not re-resolve per call
+
+
+def test_init_preseeded_forward_method_is_kept(monkeypatch):
+    # srt layers pin instance paths in __init__ (e.g. env-gated aiter modes);
+    # lazy resolution must not clobber that.
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _AllPlatformsOp()
+    op._forward_method = op.forward_xpu
+    assert op(torch.zeros(1)) == "xpu"
+
+
+def test_backend_eligible_override_gates_per_call(monkeypatch):
+    class _Gated(BaseFusedOp):
+        op = "test.gated"
+        priority = (KernelBackend.JIT, KernelBackend.TORCH)
+        capabilities = {KernelBackend.JIT: frozenset()}
+
+        def forward_native(self, x):
+            return "native"
+
+        def forward_jit(self, x):
+            return "jit"
+
+        def backend_eligible(self, backend, *args, **kwargs):
+            if not super().backend_eligible(backend, *args, **kwargs):
+                return False
+            if backend is KernelBackend.JIT:
+                return args[0].shape[-1] % 2 == 0
+            return True
+
+    _mock_platform(monkeypatch, key="", info=_CPU)
+    op = _Gated()
+    assert op(torch.zeros(4)) == "jit"
+    assert op(torch.zeros(3)) == "native"  # same instance, per-call bounce
+    assert op(torch.zeros(8)) == "jit"
+
+
+# --- torch.compile protocol -----------------------------------------------------
+
+
+def test_enter_leave_torch_compile_roundtrip(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _CudaOnlyPlatformOp()
+    assert op(torch.zeros(1)) == "cuda"
+
+    op.enter_torch_compile(num_tokens=16)
+    assert op.is_torch_compile
+    assert op(torch.zeros(1)) == "native"
+
+    # Reused-module idempotency: a second enter must not overwrite the saved
+    # original forward, otherwise leave() cannot restore it.
+    op.enter_torch_compile(num_tokens=16)
+    op.leave_torch_compile()
+    assert not op.is_torch_compile
+    assert op(torch.zeros(1)) == "cuda"
+    op.leave_torch_compile()  # double leave is a no-op
+    assert op(torch.zeros(1)) == "cuda"
+
+
+def test_torch_compile_hook_none_keeps_dispatch(monkeypatch):
+    class _KeepOptimized(_CudaOnlyPlatformOp):
+        def _torch_compile_forward(self, num_tokens):
+            return None if num_tokens > 1 else self.forward_native
+
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _KeepOptimized()
+    op.enter_torch_compile(num_tokens=8)
+    assert op.is_torch_compile
+    assert op(torch.zeros(1)) == "cuda"  # dispatch unchanged for bs > 1
+    op.leave_torch_compile()
+
+    op.enter_torch_compile(num_tokens=1)
+    assert op(torch.zeros(1)) == "native"
+    op.leave_torch_compile()
+
+
+def test_topk_compile_hook_is_bs1_only():
+    from sglang.srt.layers.moe.topk import TopK
+
+    class _Probe:
+        forward_native = "native-sentinel"
+
+    assert TopK._torch_compile_forward(_Probe(), num_tokens=1) == "native-sentinel"
+    assert TopK._torch_compile_forward(_Probe(), num_tokens=2) is None
+
+
+def test_fused_moe_compile_hook_is_bs1_only():
+    from sglang.srt.layers.moe.fused_moe_native import fused_moe_forward_native
+    from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
+
+    probe = object.__new__(UnquantizedFusedMoEMethod)
+    assert (
+        UnquantizedFusedMoEMethod._torch_compile_forward(probe, num_tokens=1)
+        is fused_moe_forward_native
+    )
+    assert UnquantizedFusedMoEMethod._torch_compile_forward(probe, num_tokens=2) is None
+
+
+# --- tracing --------------------------------------------------------------------
+
+
+def test_trace_labels_platform_and_backend(monkeypatch):
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _CudaOnlyPlatformOp()
+    fo.enable_fused_op_trace()
+    op(torch.zeros(2, 3))
+    op(torch.zeros(2, 3), backend=KernelBackend.TORCH)
+    auto_rec, explicit_rec = fo.get_fused_op_trace()
+    assert auto_rec.op == "test.cuda_only_platform"
+    assert auto_rec.backend == "cuda"
+    assert auto_rec.tensor_args == ("torch.float32[2, 3]",)
+    assert explicit_rec.backend == "torch"
+
+
+# --- deprecated MultiPlatformOp alias --------------------------------------------
+
+
+def test_deprecated_alias_contract(monkeypatch):
+    from sglang.srt.layers.utils import MultiPlatformOp
+
+    assert issubclass(MultiPlatformOp, BaseFusedOp)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        class _LegacyOp(MultiPlatformOp):
+            # Old-style subclass: platform forwards only, no forward_native.
+            def forward_cuda(self, x):
+                return "cuda"
+
+    assert any(issubclass(w.category, DeprecationWarning) for w in caught)
+
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    op = _LegacyOp()  # instantiable without forward_native (lenient alias)
+    assert op(torch.zeros(1)) == "cuda"
+    with pytest.raises(NotImplementedError):
+        op.forward_native(torch.zeros(1))
+
+    # register_oot_forward via the alias lands in the shared registry.
+    MultiPlatformOp.register_oot_forward(_LegacyOp, lambda self, x: "oot", "aliasplat")
+    _mock_platform(monkeypatch, key="", oot_key="aliasplat")
+    assert _LegacyOp()(torch.zeros(1)) == "oot"
+
+
+def test_deprecated_alias_keeps_legacy_platform_defaults(monkeypatch):
+    """Old MultiPlatformOp defined per-platform default methods (hip/musa ->
+    cuda, npu/xpu/cpu -> native); plugin code may call them directly, and a
+    subclass without forward_cuda must still raise on CUDA like before."""
+    from sglang.srt.layers.utils import MultiPlatformOp
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+
+        class _NativeOnlyLegacy(MultiPlatformOp):
+            def forward_native(self, x):
+                return "native"
+
+    op = _NativeOnlyLegacy()
+    assert op.forward_cpu(torch.zeros(1)) == "native"
+    assert op.forward_npu(torch.zeros(1)) == "native"
+    with pytest.raises(NotImplementedError):
+        op.forward_hip(torch.zeros(1))  # chains to the raising forward_cuda
+
+    _mock_platform(monkeypatch, key="cuda", info=_CUDA)
+    with pytest.raises(NotImplementedError):
+        _NativeOnlyLegacy()(torch.zeros(1))  # old CUDA behavior preserved
+
+
+# --- migration completeness -------------------------------------------------------
+
+_MIGRATED_OPS = [
+    ("sglang.srt.layers.activation", "SiluAndMul"),
+    ("sglang.srt.layers.activation", "GeluAndMul"),
+    ("sglang.srt.layers.activation", "NewGELU"),
+    ("sglang.srt.layers.activation", "ReLU2"),
+    ("sglang.srt.layers.activation", "QuickGELU"),
+    ("sglang.srt.layers.activation", "XIELU"),
+    ("sglang.srt.layers.layernorm", "RMSNorm"),
+    ("sglang.srt.layers.layernorm", "LayerNorm"),
+    ("sglang.srt.layers.layernorm", "GemmaRMSNorm"),
+    ("sglang.srt.layers.layernorm", "Gemma3RMSNorm"),
+    ("sglang.srt.layers.layernorm", "Gemma4RMSNorm"),
+    ("sglang.srt.layers.layernorm", "RMSNormWithoutScale"),
+    ("sglang.srt.layers.conv", "Conv2dLayer"),
+    ("sglang.srt.layers.conv", "Conv3dLayer"),
+    ("sglang.srt.layers.moe.topk", "TopK"),
+    ("sglang.srt.layers.rotary_embedding.base", "RotaryEmbedding"),
+    ("sglang.srt.layers.rotary_embedding.rope_variant", "DualChunkRotaryEmbedding"),
+    ("sglang.srt.layers.attention.dsa.dsa_indexer", "Indexer"),
+    ("sglang.srt.layers.attention.dsv4.compressor", "Compressor"),
+    ("sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated", "Mixer2RMSNormGated"),
+    ("sglang.srt.layers.quantization.unquant", "UnquantizedFusedMoEMethod"),
+]
+
+
+@pytest.mark.parametrize("module_name, cls_name", _MIGRATED_OPS)
+def test_migrated_ops_subclass_base_fused_op(module_name, cls_name):
+    """Production ops must extend BaseFusedOp directly, never the deprecated
+    MultiPlatformOp alias (which exists only for out-of-tree users)."""
+    import importlib
+
+    from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
+
+    cls = getattr(importlib.import_module(module_name), cls_name)
+    assert issubclass(cls, BaseFusedOp)
+    assert MultiPlatformOp not in cls.__mro__
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))