minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-06 08:52:09 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent ba12a16a62
commit 4c0a8940fa
23 changed files with 1275 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/README.md
+38 -17
View File
@@ -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_<backend>` 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_<backend>` method per
implementation source, all sharing one signature behind a single
`forward()`:
- `forward_native` — **required**; the pure-`torch` correctness reference
every other implementation is checked against.
- `forward_torch_compile` — inherited for free as
`torch.compile(forward_native)`.
- `forward_triton` / `forward_jit` / `forward_aot` / `forward_cute_dsl` /
`forward_flashinfer` / `forward_deepgemm` / `forward_aiter` /
`forward_torch_npu` — opt-in overrides. A backend is *available* iff its
method is overridden; it joins **auto-selection** only when also declared
in `capabilities` (device support is metadata, not guesswork).
- **Platforms / devices** — optional composite per-device paths:
`forward_cuda`, `forward_hip` (falls back to `forward_cuda`),
`forward_npu`, `forward_xpu`, `forward_musa` (no implicit CUDA fallback —
MUSA ops opt into the CUDA path with an explicit `forward_musa`),
`forward_cpu` (AMX CPUs), plus `forward_<key>` /
`register_oot_forward()` for out-of-tree platform plugins. CUDA / HIP are
**not** kernel backends.
`forward()` auto-selects the best available backend by the class's `priority`,
filtered per call through `backend_eligible()` (a
`CapabilityRequirement`-vs-`PlatformInfo` check, extensible with per-call
shape/dtype gates), and degrades to the native reference when no optimized
backend fits. The public `ops.<group>` functions stay thin wrappers over
Dispatch priority, highest first: explicit `forward(..., backend=...)` →
global forced backend (`SGLANG_FORCE_FUSED_OP_BACKEND`) → OOT platform
override → declared optimized kernel backends by `priority` (filtered by
`backend_eligible()`, a `CapabilityRequirement`-vs-`PlatformInfo` check
extensible with per-call shape/dtype gates) → platform-specific forward →
`forward_native`. The static part of the decision is resolved once and cached
on the instance, so the hot path stays a single indirect call.
`BaseFusedOp` also owns the torch.compile mode protocol
(`enter_torch_compile(num_tokens)` / `leave_torch_compile()`, both
idempotent): while an outer model is compiled, ops switch to their
compile-safe native path so device kernels are never traced (TopK and Fused
MoE override `_torch_compile_forward()` to keep their bs>1 behavior).
The public `ops.<group>` functions stay thin wrappers over
module-level instances, so the import surface is unchanged; each instance also
registers all of its backends as `KernelSpec`s so the registry inventory and
`select_kernel(..., backend=...)` keep working.
python/sglang/kernels/fused_op.py
+429 -94
View File
@@ -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_<backend>`` 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_<key>``
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_<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
lazily inside the ``forward_<backend>`` methods.
``sglang.kernels.jit``), performs platform detection, or triggers JIT
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):
"""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_<backend>`` methods. All backend methods of one op
must share the same signature and semantics — each override adapts its
underlying kernel's calling convention so call sites never care which
backend ran.
Subclasses implement :meth:`forward_native` plus any optimized
``forward_<backend>`` methods and/or platform-specific
``forward_<device>`` 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, ``"<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
: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_<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
# 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_<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."""
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.
python/sglang/kernels/ops/layernorm/__init__.py
+3 -3
View File
@@ -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
+3 -3
View File
@@ -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
+11 -8
View File
@@ -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
+20 -2
View File
@@ -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
+2 -2
View File
@@ -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 -2
View File
@@ -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
+3 -3
View File
@@ -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
+25 -7
View File
@@ -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
+14 -2
View File
@@ -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
+18 -3
View File
@@ -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
+2 -2
View File
@@ -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
+7 -2
View File
@@ -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
+3
View File
@@ -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
+37 -109
View File
@@ -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
python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py
+7 -9
View File
@@ -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):
python/sglang/srt/platforms/interface.py
+7 -4
View File
@@ -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.
E.g. "cuda", "npu", "hip", "xpu", "cpu".
Determines which ``forward_<key>()`` 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"