diff --git a/docs/docs/hardware-platforms/plugin.mdx b/docs/docs/hardware-platforms/plugin.mdx index 64de53ef9..dda26d0ef 100644 --- a/docs/docs/hardware-platforms/plugin.mdx +++ b/docs/docs/hardware-platforms/plugin.mdx @@ -544,7 +544,7 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat get_dispatch_key_name() "native" - MultiPlatformOp dispatch key name + BaseFusedOp (fused-op) dispatch key name diff --git a/python/sglang/kernels/README.md b/python/sglang/kernels/README.md index 598d1ddaa..7bbcc1797 100644 --- a/python/sglang/kernels/README.md +++ b/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) -are implemented as `BaseFusedOp` subclasses: one logical operator with one -`forward_` method per backend, all sharing one signature behind a -single `forward()`: +`BaseFusedOp` is a standard `torch.nn.Module` (it replaced the former +`sglang.srt.layers.utils.MultiPlatformOp`) that carries one logical operator +with interchangeable implementations along **two independent dimensions**: -- `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_jit` / `forward_aot` / - `forward_cute_dsl` / `forward_flashinfer` / `forward_deepgemm` — opt-in - overrides. A backend is *available* iff its method is overridden. +- **Kernel backends (provenance)** — one `forward_` method per + implementation source, all sharing one signature behind a single + `forward()`: + - `forward_native` — **required**; the pure-`torch` correctness reference + every other implementation is checked against. + - `forward_torch_compile` — inherited for free as + `torch.compile(forward_native)`. + - `forward_triton` / `forward_jit` / `forward_aot` / `forward_cute_dsl` / + `forward_flashinfer` / `forward_deepgemm` / `forward_aiter` / + `forward_torch_npu` — opt-in overrides. A backend is *available* iff its + method is overridden; it joins **auto-selection** only when also declared + in `capabilities` (device support is metadata, not guesswork). +- **Platforms / devices** — optional composite per-device paths: + `forward_cuda`, `forward_hip` (falls back to `forward_cuda`), + `forward_npu`, `forward_xpu`, `forward_musa` (no implicit CUDA fallback — + MUSA ops opt into the CUDA path with an explicit `forward_musa`), + `forward_cpu` (AMX CPUs), plus `forward_` / + `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`, -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.` functions stay thin wrappers over +Dispatch priority, highest first: explicit `forward(..., backend=...)` → +global forced backend (`SGLANG_FORCE_FUSED_OP_BACKEND`) → OOT platform +override → declared optimized kernel backends by `priority` (filtered by +`backend_eligible()`, a `CapabilityRequirement`-vs-`PlatformInfo` check +extensible with per-call shape/dtype gates) → platform-specific forward → +`forward_native`. The static part of the decision is resolved once and cached +on the instance, so the hot path stays a single indirect call. + +`BaseFusedOp` also owns the torch.compile mode protocol +(`enter_torch_compile(num_tokens)` / `leave_torch_compile()`, both +idempotent): while an outer model is compiled, ops switch to their +compile-safe native path so device kernels are never traced (TopK and Fused +MoE override `_torch_compile_forward()` to keep their bs>1 behavior). + +The public `ops.` 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. diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py index d39ed4d19..441bd3107 100644 --- a/python/sglang/kernels/fused_op.py +++ b/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 -``sglang.kernels.ops.*`` wrappers (RFC #29630): one logical operator, -implemented once, with multiple interchangeable backends behind a single -``forward()``. +:class:`BaseFusedOp` is the single operator abstraction of the unified +``sglang.kernels`` namespace: one logical operator, implemented once, with +multiple interchangeable implementations behind a single ``forward()``. It +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_`` method per backend it -supports: +- **Kernel backend (provenance)** — where an implementation comes from + (: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_`` + for out-of-tree (OOT) platform plugins. CUDA / HIP are **not** kernel + backends. -- ``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_jit`` / ``forward_aot`` / - ``forward_cute_dsl`` / ``forward_flashinfer`` / ``forward_deepgemm`` — - opt-in overrides. +Dispatch priority (highest first), resolved by :meth:`BaseFusedOp.forward`: -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. +1. **Explicit backend** — ``forward(..., backend=KernelBackend.X)``. +2. **Global forced backend** — ``SGLANG_FORCE_FUSED_OP_BACKEND`` / + :func:`set_fused_op_backend` (e.g. ``native`` to bisect numerical bugs). + Best-effort: an op that does not implement the forced backend falls back + to normal dispatch with a one-time warning, so the debug switch works on + whole models that contain device-only ops. +3. **OOT platform override** — on an out-of-tree platform, a forward + registered via :meth:`BaseFusedOp.register_oot_forward`, then a + ``forward_`` 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 -lazily inside the ``forward_`` methods. +``sglang.kernels.jit``), performs platform detection, or triggers JIT +compilation; backends are imported lazily inside the ``forward_`` +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_ method name. +logger = logging.getLogger(__name__) + +# backend (provenance) -> forward_ 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): - """One logical operator with interchangeable backends behind ``forward()``. +class BaseFusedOp(nn.Module, ABC): + """One logical operator with interchangeable implementations behind + ``forward()``. - Subclasses set :attr:`op` and implement :meth:`forward_native` plus any - optimized ``forward_`` 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. + Subclasses implement :meth:`forward_native` plus any optimized + ``forward_`` methods and/or platform-specific + ``forward_`` methods. All implementations 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 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, ``"."`` (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 - :data:`DEFAULT_PRIORITY`. + Kernel-backend preference for auto-selection, best first. Defaults to + :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 - :meth:`backend_eligible` (and exported into the registry specs). Use the - ``CapabilityRequirement.CUDA`` / ``.HIP`` / ``.NPU`` shortcuts, e.g. - ``{KernelBackend.AOT: {CapabilityRequirement.CUDA, CapabilityRequirement.HIP}}``. + an empty set value = runs on any device), consulted by + :meth:`backend_eligible` and exported into the registry specs. A + kernel backend joins **auto-selection** only when it is declared here + (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_`` + 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 - # 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) - ) + super().__init__() + # Statically resolved dispatch target (priority steps 3-6). ``None`` + # means "not resolved yet": resolution is deferred to the first call + # so module-level op instances never trigger platform detection at + # import time. Subclass __init__ may overwrite it to pin a path. + self._forward_method: Optional[Callable] = None + # torch.compile mode bookkeeping (see enter/leave_torch_compile). + self._original_forward_method: Optional[Callable] = None + self.is_torch_compile = False 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 False + return getattr(self, method_name) + 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): - raise NotImplementedError(f"{self.op}: no npu backend") + def forward_torch_npu(self, *args, **kwargs): + 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).""" - return list(self._available) + """Kernel backends this op implements (structural, platform-agnostic).""" + 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,40 +470,191 @@ 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: - forced = get_fused_op_backend() - if forced is not None: - return forced - for backend in self._ordered: - if self.backend_eligible(backend, *args, **kwargs): + def _auto_backend_candidates(self) -> Tuple[KernelBackend, ...]: + """Kernel backends participating in auto-selection, best first. + + A backend qualifies when its method is overridden *and* the op + declares it in :attr:`capabilities` (``TORCH_COMPILE`` needs no + 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_`` 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_ / + # forward_ 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.""" - if backend is None: - backend = self._resolve_backend(*args, **kwargs) - result = getattr(self, BACKEND_METHODS[backend])(*args, **kwargs) + """Run the op on ``backend``, or on the best eligible path when omitted.""" + if backend is not None: + # 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: + _record_trace(self, backend.value, args, kwargs) + return result + forced = _forced_backend + if forced is _UNRESOLVED: + 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: - _trace_records.append( - FusedOpTraceRecord( - op=self.op, - backend=backend.value, - tensor_args=_describe_tensors(args, kwargs), - ) - ) + _record_trace(self, _dispatch_label(method), 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. diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 2e04ccdbe..fa17572fa 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/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, diff --git a/python/sglang/srt/compilation/torch_compile_decoration.py b/python/sglang/srt/compilation/torch_compile_decoration.py index 9c396b27c..c7c305ce5 100644 --- a/python/sglang/srt/compilation/torch_compile_decoration.py +++ b/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: diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 837b766e7..61ef0724d 100644 --- a/python/sglang/srt/layers/activation.py +++ b/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" diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 79136182b..1cf7a4c5f 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 26c932a67..f27eb00d1 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index 7ceb3e161..045e4fec3 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/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, diff --git a/python/sglang/srt/layers/conv.py b/python/sglang/srt/layers/conv.py index 6c72ae107..704e41c76 100644 --- a/python/sglang/srt/layers/conv.py +++ b/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__( diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 67661147f..92b26a592 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 02b426466..c7f8990a4 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 5a66d91b2..edc72f0a7 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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, diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index f009f4d1f..7f33f3159 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/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__( diff --git a/python/sglang/srt/layers/rotary_embedding/rope_variant.py b/python/sglang/srt/layers/rotary_embedding/rope_variant.py index 7b0cb87e3..d49c76128 100644 --- a/python/sglang/srt/layers/rotary_embedding/rope_variant.py +++ b/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, diff --git a/python/sglang/srt/layers/utils/__init__.py b/python/sglang/srt/layers/utils/__init__.py index e3101d534..044a19fc7 100644 --- a/python/sglang/srt/layers/utils/__init__.py +++ b/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 diff --git a/python/sglang/srt/layers/utils/multi_platform.py b/python/sglang/srt/layers/utils/multi_platform.py index 893248a21..9856e3280 100644 --- a/python/sglang/srt/layers/utils/multi_platform.py +++ b/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 -from sglang.srt.platforms import current_platform -from sglang.srt.utils import ( - cpu_has_amx_support, - is_cpu, - is_cuda, - is_hip, - is_musa, - is_npu, - is_xpu, -) +In-repo code must subclass ``BaseFusedOp`` directly. This alias exists only so +out-of-tree platform plugins and external users keep importing from the old +path while they migrate; it will be removed in a future release. +""" -_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() +import warnings + +from sglang.kernels.fused_op import BaseFusedOp -class MultiPlatformOp(nn.Module): +class MultiPlatformOp(BaseFusedOp): + """Deprecated alias of :class:`sglang.kernels.fused_op.BaseFusedOp`. - # OOT forward registry: maps dispatch_key -> {op_cls -> forward_fn} - _oot_forward_registry: ClassVar[dict[str, dict[type, Callable]]] = {} + 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. + """ - @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 __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) 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 diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index a416d0c74..7aa110f37 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/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( - inner_model, reverse=False, num_tokens=16 - ) + _toggle_fused_ops(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): diff --git a/python/sglang/srt/platforms/interface.py b/python/sglang/srt/platforms/interface.py index e95aa4c30..bcf034e69 100644 --- a/python/sglang/srt/platforms/interface.py +++ b/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_()`` method is selected. - E.g. "cuda", "npu", "hip", "xpu", "cpu". + Determines which ``forward_()`` method is selected on an + 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" diff --git a/test/manual/kernels/bench_fused_op_dispatch.py b/test/manual/kernels/bench_fused_op_dispatch.py new file mode 100644 index 000000000..52f2c2ceb --- /dev/null +++ b/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() diff --git a/test/registered/kernels/ops/layernorm/test_fused_op.py b/test/registered/kernels/ops/layernorm/test_fused_op.py index c98469068..b23485d06 100644 --- a/test/registered/kernels/ops/layernorm/test_fused_op.py +++ b/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 diff --git a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py index a61f79124..3c71dde2b 100644 --- a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py +++ 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(): diff --git a/test/registered/kernels/test_fused_op_dispatch.py b/test/registered/kernels/test_fused_op_dispatch.py new file mode 100644 index 000000000..5bf549e41 --- /dev/null +++ b/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_ 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__]))