[Kernel] Decouple KernelBackend from device + device-based CapabilityRequirement (RFC #29630) (#31292)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-17 10:35:34 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 68d324d697
commit 8432eafd3d
17 changed files with 433 additions and 202 deletions
+2 -2
View File
@@ -53,7 +53,7 @@ explicitly, e.g.:
```python
from sglang.kernels import select_kernel, KernelBackend
jit_rmsnorm = select_kernel("layernorm.rmsnorm", backend=KernelBackend.CUDA_JIT).load()
jit_rmsnorm = select_kernel("layernorm.rmsnorm", backend=KernelBackend.JIT).load()
```
## `BaseFusedOp` — the per-operator implementation contract
@@ -67,7 +67,7 @@ single `forward()`:
every other backend is checked against.
- `forward_torch_compile` — inherited for free as
`torch.compile(forward_native)`.
- `forward_triton` / `forward_cuda_jit` / `forward_cuda_aot` /
- `forward_triton` / `forward_jit` / `forward_aot` /
`forward_cute_dsl` / `forward_flashinfer` / `forward_deepgemm` — opt-in
overrides. A backend is *available* iff its method is overridden.
+4
View File
@@ -38,10 +38,12 @@ from sglang.kernels.registry import KernelRegistry, register_kernel, registry
from sglang.kernels.selector import get_kernel, select_kernel
from sglang.kernels.spec import (
CapabilityRequirement,
DeviceType,
FormatSignature,
KernelBackend,
KernelSpec,
PlatformInfo,
capabilities_satisfied,
)
# Importing the operator groups populates the registry (metadata only). Kept
@@ -53,6 +55,8 @@ __all__ = [
"ops",
"BaseFusedOp",
"CapabilityRequirement",
"DeviceType",
"capabilities_satisfied",
"FormatSignature",
"FusedOpTraceRecord",
"KernelBackend",
+42 -21
View File
@@ -12,7 +12,7 @@ supports:
every other backend is checked against.
- ``forward_torch_compile`` — provided by the base class as
``torch.compile(forward_native)``.
- ``forward_triton`` / ``forward_cuda_jit`` / ``forward_cuda_aot`` /
- ``forward_triton`` / ``forward_jit`` / ``forward_aot`` /
``forward_cute_dsl`` / ``forward_flashinfer`` / ``forward_deepgemm`` —
opt-in overrides.
@@ -36,7 +36,16 @@ from __future__ import annotations
import functools
from abc import ABC, abstractmethod
from typing import Any, ClassVar, Dict, List, Mapping, Optional, Tuple
from typing import (
AbstractSet,
Any,
ClassVar,
Dict,
List,
Mapping,
Optional,
Tuple,
)
import msgspec
@@ -47,29 +56,32 @@ from sglang.kernels.spec import (
KernelBackend,
KernelSpec,
PlatformInfo,
capabilities_satisfied,
)
# backend -> forward_<backend> method name.
# backend (provenance) -> forward_<backend> method name.
BACKEND_METHODS: Dict[KernelBackend, str] = {
KernelBackend.TORCH: "forward_native",
KernelBackend.TORCH_COMPILE: "forward_torch_compile",
KernelBackend.TRITON: "forward_triton",
KernelBackend.CUDA_JIT: "forward_cuda_jit",
KernelBackend.CUDA_AOT: "forward_cuda_aot",
KernelBackend.JIT: "forward_jit",
KernelBackend.AOT: "forward_aot",
KernelBackend.CUTE_DSL: "forward_cute_dsl",
KernelBackend.FLASHINFER: "forward_flashinfer",
KernelBackend.DEEPGEMM: "forward_deepgemm",
KernelBackend.AITER: "forward_aiter",
}
# best -> fallback. ``torch_compile`` is deliberately absent: auto-selection
# must never trigger a surprise compilation in a serving process; force it
# explicitly when wanted.
# explicitly when wanted. Per-op priority overrides this (see BaseFusedOp).
DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = (
KernelBackend.CUDA_AOT,
KernelBackend.CUDA_JIT,
KernelBackend.AOT,
KernelBackend.JIT,
KernelBackend.FLASHINFER,
KernelBackend.DEEPGEMM,
KernelBackend.CUTE_DSL,
KernelBackend.AITER,
KernelBackend.TRITON,
KernelBackend.TORCH,
)
@@ -181,8 +193,11 @@ class BaseFusedOp(ABC):
Backend preference for auto-selection, best first. Defaults to
:data:`DEFAULT_PRIORITY`.
capabilities:
Per-backend :class:`CapabilityRequirement`, consulted by
:meth:`backend_eligible` (and exported into the registry specs).
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}}``.
format_signature:
Data-contract description shared by all backends of this op.
descriptions:
@@ -191,7 +206,9 @@ class BaseFusedOp(ABC):
op: ClassVar[str]
priority: ClassVar[Tuple[KernelBackend, ...]] = DEFAULT_PRIORITY
capabilities: ClassVar[Mapping[KernelBackend, CapabilityRequirement]] = {}
capabilities: ClassVar[
Mapping[KernelBackend, AbstractSet[CapabilityRequirement]]
] = {}
format_signature: ClassVar[FormatSignature] = FormatSignature()
descriptions: ClassVar[Mapping[KernelBackend, str]] = {}
@@ -234,11 +251,11 @@ class BaseFusedOp(ABC):
def forward_triton(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no triton backend")
def forward_cuda_jit(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no cuda_jit backend")
def forward_jit(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no jit backend")
def forward_cuda_aot(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no cuda_aot backend")
def forward_aot(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no aot backend")
def forward_cute_dsl(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no cute_dsl backend")
@@ -249,6 +266,9 @@ class BaseFusedOp(ABC):
def forward_deepgemm(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no deepgemm backend")
def forward_aiter(self, *args, **kwargs):
raise NotImplementedError(f"{self.op}: no aiter backend")
# --- selection ---
def available_backends(self) -> List[KernelBackend]:
@@ -259,12 +279,13 @@ class BaseFusedOp(ABC):
"""Whether ``backend`` may run *this* call.
The base implementation checks the backend's
:class:`CapabilityRequirement` against the detected platform.
Subclasses may extend it with per-call shape/dtype gates so
: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.
"""
capability = self.capabilities.get(backend)
return capability is None or capability.is_satisfied_by(_platform())
return capabilities_satisfied(
self.capabilities.get(backend, frozenset()), _platform()
)
def _resolve_backend(self, *args, **kwargs) -> KernelBackend:
forced = get_fused_op_backend()
@@ -300,7 +321,7 @@ def register_fused_op(instance: BaseFusedOp, module: str, attr: str) -> BaseFuse
``module``/``attr`` locate the module-level instance so that
``KernelSpec.load()`` can lazily resolve e.g.
``"<module>:<attr>.forward_cuda_aot"`` to the bound backend method. Returns
``"<module>:<attr>.forward_aot"`` to the bound backend method. Returns
``instance`` so group packages can write
``_RMSNORM = register_fused_op(_RMSNormOp(), __name__, "_RMSNORM")``.
"""
@@ -310,7 +331,7 @@ def register_fused_op(instance: BaseFusedOp, module: str, attr: str) -> BaseFuse
op=instance.op,
backend=backend,
target=f"{module}:{attr}.{BACKEND_METHODS[backend]}",
capability=instance.capabilities.get(backend, CapabilityRequirement()),
capabilities=frozenset(instance.capabilities.get(backend, ())),
format_signature=instance.format_signature,
description=instance.descriptions.get(backend, ""),
)
@@ -4,7 +4,7 @@ Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a
pure-``torch`` reference (``forward_native``) plus AOT (``sgl_kernel``) and
JIT CUDA backends behind one ``(input, out)`` signature. The JIT backend
additionally accepts ``expert_ids`` / ``expert_step`` — call
``forward_cuda_jit`` directly when those are needed.
``forward_jit`` directly when those are needed.
"""
from __future__ import annotations
@@ -24,10 +24,17 @@ if TYPE_CHECKING:
import torch
_ACT_DTYPES = ("float16", "bfloat16")
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
_HIP = frozenset({CapabilityRequirement.HIP})
# sgl_kernel's gated-activation ops build for CUDA *and* ROCm (production
# imports them from sgl_kernel on both), so the AOT backend spans both devices
# — the canonical OR-semantics case that a device-baked backend name couldn't.
_CUDA_HIP = frozenset({CapabilityRequirement.CUDA, CapabilityRequirement.HIP})
# JIT before AOT to match the production path (srt/layers/activation.py imports
# from sglang.jit_kernel.activation on CUDA); auto-selection must not invert it.
_ACT_PRIORITY = (
KernelBackend.CUDA_AOT,
KernelBackend.CUDA_JIT,
KernelBackend.JIT,
KernelBackend.AOT,
KernelBackend.TORCH,
)
@@ -40,8 +47,8 @@ class _GatedActivationOp(BaseFusedOp):
priority = _ACT_PRIORITY
capabilities = {
KernelBackend.CUDA_AOT: _CUDA,
KernelBackend.CUDA_JIT: _CUDA,
KernelBackend.AOT: _CUDA_HIP,
KernelBackend.JIT: _CUDA,
}
format_signature = FormatSignature(
supported_dtypes=_ACT_DTYPES,
@@ -61,14 +68,14 @@ class _GatedActivationOp(BaseFusedOp):
out.copy_(result)
return out
def forward_cuda_aot(
def forward_aot(
self, input: torch.Tensor, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
import sgl_kernel
return getattr(sgl_kernel, self.kernel_attr)(input, out)
def forward_cuda_jit(
def forward_jit(
self,
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
@@ -83,13 +90,37 @@ class _GatedActivationOp(BaseFusedOp):
class SiluAndMulOp(_GatedActivationOp):
"""``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``."""
"""``out = silu(input[..., :d]) * input[..., d:]`` with ``d = input.shape[-1] // 2``.
Adds an ``AITER`` backend on ``device=HIP``: on ROCm this op has a native
``aiter`` kernel (``srt/layers/activation.py`` uses it in production). Note
the sibling gelu ops below deliberately do *not* register AITER — ROCm/aiter
coverage is a per-``(op, backend)`` subset, which the decoupled backend/device
model expresses directly (a device-agnostic ``KernelBackend`` name plus a
per-backend ``CapabilityRequirement``).
"""
op = "activation.silu_and_mul"
kernel_attr = "silu_and_mul"
# AOT spans CUDA+HIP; JIT is CUDA; AITER is an opt-in HIP path. By priority,
# CUDA resolves to JIT and HIP resolves to AOT (matching production
# defaults); AITER is registered and HIP-eligible but sits below AOT, so it
# is available for explicit/forced selection without changing the default.
priority = (
KernelBackend.JIT,
KernelBackend.AOT,
KernelBackend.AITER,
KernelBackend.TORCH,
)
capabilities = {
KernelBackend.AOT: _CUDA_HIP,
KernelBackend.JIT: _CUDA,
KernelBackend.AITER: _HIP,
}
descriptions = {
KernelBackend.CUDA_AOT: "silu_and_mul (sgl_kernel wheel).",
KernelBackend.CUDA_JIT: "silu_and_mul (sglang.jit_kernel).",
KernelBackend.AOT: "silu_and_mul (sgl_kernel wheel).",
KernelBackend.JIT: "silu_and_mul (sglang.jit_kernel).",
KernelBackend.AITER: "silu_and_mul (aiter, ROCm).",
KernelBackend.TORCH: "silu_and_mul (pure-torch reference).",
}
@@ -98,6 +129,22 @@ class SiluAndMulOp(_GatedActivationOp):
return F.silu(gate)
def forward_aiter(
self, input: torch.Tensor, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
import torch
from aiter import silu_and_mul as _aiter_silu_and_mul
d = input.shape[-1] // 2
if out is None:
out = torch.empty(
(*input.shape[:-1], d), dtype=input.dtype, device=input.device
)
# aiter's ROCm silu_and_mul: (out, input, limit); limit=0.0 = no clamp,
# matching the standard (unclamped) gated-SiLU used elsewhere.
_aiter_silu_and_mul(out, input, 0.0)
return out
class GeluAndMulOp(_GatedActivationOp):
"""``out = gelu(input[..., :d]) * input[..., d:]`` (erf-based GELU)."""
@@ -105,8 +152,8 @@ class GeluAndMulOp(_GatedActivationOp):
op = "activation.gelu_and_mul"
kernel_attr = "gelu_and_mul"
descriptions = {
KernelBackend.CUDA_AOT: "gelu_and_mul (sgl_kernel wheel).",
KernelBackend.CUDA_JIT: "gelu_and_mul (sglang.jit_kernel).",
KernelBackend.AOT: "gelu_and_mul (sgl_kernel wheel).",
KernelBackend.JIT: "gelu_and_mul (sglang.jit_kernel).",
KernelBackend.TORCH: "gelu_and_mul (pure-torch reference).",
}
@@ -122,8 +169,8 @@ class GeluTanhAndMulOp(_GatedActivationOp):
op = "activation.gelu_tanh_and_mul"
kernel_attr = "gelu_tanh_and_mul"
descriptions = {
KernelBackend.CUDA_AOT: "gelu_tanh_and_mul (sgl_kernel wheel).",
KernelBackend.CUDA_JIT: "gelu_tanh_and_mul (sglang.jit_kernel).",
KernelBackend.AOT: "gelu_tanh_and_mul (sgl_kernel wheel).",
KernelBackend.JIT: "gelu_tanh_and_mul (sglang.jit_kernel).",
KernelBackend.TORCH: "gelu_tanh_and_mul (pure-torch reference).",
}
+10 -10
View File
@@ -20,14 +20,14 @@ if TYPE_CHECKING:
import torch
from torch import nn
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
register_kernel(
KernelSpec(
op="diffusion.apply_group_norm_silu",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.diffusion.group_norm_silu:apply_group_norm_silu",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(description="fused GroupNorm + SiLU"),
description="Fused group-norm + SiLU (sglang.jit_kernel).",
)
@@ -35,9 +35,9 @@ register_kernel(
register_kernel(
KernelSpec(
op="diffusion.residual_gate_add",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.diffusion.residual_gate_add:residual_gate_add_cuda",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(description="residual + gate * update"),
description="Fused residual gate-add (sglang.jit_kernel).",
)
@@ -45,9 +45,9 @@ register_kernel(
register_kernel(
KernelSpec(
op="diffusion.fused_inplace_qknorm_rope",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.diffusion.qknorm_rope:fused_inplace_qknorm_rope",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(
in_place=True, description="fused in-place QK-norm + RoPE"
),
@@ -60,7 +60,7 @@ def apply_group_norm_silu(
x: torch.Tensor, norm: nn.Module, activation: nn.Module
) -> torch.Tensor:
"""Fused GroupNorm + SiLU (falls back to eager when unsupported)."""
return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.CUDA_JIT)(
return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.JIT)(
x, norm, activation
)
@@ -69,7 +69,7 @@ def residual_gate_add(
residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
) -> torch.Tensor:
"""Fused ``residual + gate * update``."""
return get_kernel("diffusion.residual_gate_add", KernelBackend.CUDA_JIT)(
return get_kernel("diffusion.residual_gate_add", KernelBackend.JIT)(
residual, update, gate
)
@@ -88,7 +88,7 @@ def fused_inplace_qknorm_rope(
rope_dim: int = 0,
) -> None:
"""Fused in-place QK RMS-norm + RoPE."""
return get_kernel("diffusion.fused_inplace_qknorm_rope", KernelBackend.CUDA_JIT)(
return get_kernel("diffusion.fused_inplace_qknorm_rope", KernelBackend.JIT)(
q,
k,
q_weight,
+10 -12
View File
@@ -16,12 +16,12 @@ from sglang.kernels.spec import (
if TYPE_CHECKING:
import torch
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
register_kernel(
KernelSpec(
op="gemm.fp8_scaled_mm",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel:fp8_scaled_mm",
format_signature=FormatSignature(
supported_dtypes=("float8_e4m3fn",),
@@ -33,7 +33,7 @@ register_kernel(
register_kernel(
KernelSpec(
op="gemm.dsv3_fused_a_gemm",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel:dsv3_fused_a_gemm",
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
@@ -45,9 +45,9 @@ register_kernel(
register_kernel(
KernelSpec(
op="gemm.dsv3_fused_a_gemm",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.dsv3_fused_a_gemm:dsv3_fused_a_gemm",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
description="DeepSeek-V3 fused QKV-A GEMM (drop-in with AOT signature)",
@@ -58,9 +58,9 @@ register_kernel(
register_kernel(
KernelSpec(
op="gemm.dsv3_router_gemm",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.dsv3_router_gemm:dsv3_router_gemm",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(
supported_dtypes=("bfloat16",),
description="DeepSeek-V3 router GEMM; num_tokens in [1, 16]",
@@ -79,7 +79,7 @@ def fp8_scaled_mm(
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""FP8 scaled matmul: ``(mat_a @ mat_b) * scales_a * scales_b (+ bias)``."""
return get_kernel("gemm.fp8_scaled_mm", KernelBackend.CUDA_AOT)(
return get_kernel("gemm.fp8_scaled_mm", KernelBackend.AOT)(
mat_a, mat_b, scales_a, scales_b, out_dtype, bias
)
@@ -90,9 +90,7 @@ def dsv3_fused_a_gemm(
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""DeepSeek-V3 fused QKV-A GEMM."""
return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.CUDA_AOT)(
mat_a, mat_b, output
)
return get_kernel("gemm.dsv3_fused_a_gemm", KernelBackend.AOT)(mat_a, mat_b, output)
def dsv3_router_gemm(
@@ -102,7 +100,7 @@ def dsv3_router_gemm(
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""DeepSeek-V3 router GEMM (JIT-backed). ``out_dtype`` defaults to bfloat16."""
impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.CUDA_JIT)
impl = get_kernel("gemm.dsv3_router_gemm", KernelBackend.JIT)
if out_dtype is None:
return impl(hidden_states, router_weights, output=output)
return impl(hidden_states, router_weights, out_dtype, output)
+22 -22
View File
@@ -6,7 +6,7 @@ all behind one signature. The public module-level functions are thin wrappers
over module-level instances; auto-selection prefers the AOT ``sgl_kernel``
implementation on CUDA and falls back to the native reference elsewhere.
Pick a specific backend with e.g.
``_RMSNORM.forward(x, w, backend=KernelBackend.CUDA_JIT)`` or globally via
``_RMSNORM.forward(x, w, backend=KernelBackend.JIT)`` or globally via
``SGLANG_FORCE_FUSED_OP_BACKEND``.
"""
@@ -25,10 +25,10 @@ if TYPE_CHECKING:
import torch
_NORM_DTYPES = ("float16", "bfloat16")
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
_NORM_PRIORITY = (
KernelBackend.CUDA_AOT,
KernelBackend.CUDA_JIT,
KernelBackend.AOT,
KernelBackend.JIT,
KernelBackend.TORCH,
)
@@ -42,16 +42,16 @@ class RMSNormOp(BaseFusedOp):
op = "layernorm.rmsnorm"
priority = _NORM_PRIORITY
capabilities = {
KernelBackend.CUDA_AOT: _CUDA,
KernelBackend.CUDA_JIT: _CUDA,
KernelBackend.AOT: _CUDA,
KernelBackend.JIT: _CUDA,
}
format_signature = FormatSignature(
supported_dtypes=_NORM_DTYPES,
description="out = (x / RMS(x)) * weight; returns tensor",
)
descriptions = {
KernelBackend.CUDA_AOT: "RMS normalization (sgl_kernel wheel).",
KernelBackend.CUDA_JIT: "RMS normalization (sglang.jit_kernel).",
KernelBackend.AOT: "RMS normalization (sgl_kernel wheel).",
KernelBackend.JIT: "RMS normalization (sglang.jit_kernel).",
KernelBackend.TORCH: "RMS normalization (pure-torch reference).",
}
@@ -74,7 +74,7 @@ class RMSNormOp(BaseFusedOp):
out.copy_(result)
return out
def forward_cuda_aot(
def forward_aot(
self,
input: torch.Tensor,
weight: torch.Tensor,
@@ -86,7 +86,7 @@ class RMSNormOp(BaseFusedOp):
return sgl_kernel.rmsnorm(input, weight, eps, out, enable_pdl)
def forward_cuda_jit(
def forward_jit(
self,
input: torch.Tensor,
weight: torch.Tensor,
@@ -114,8 +114,8 @@ class FusedAddRMSNormOp(BaseFusedOp):
op = "layernorm.fused_add_rmsnorm"
priority = _NORM_PRIORITY
capabilities = {
KernelBackend.CUDA_AOT: _CUDA,
KernelBackend.CUDA_JIT: _CUDA,
KernelBackend.AOT: _CUDA,
KernelBackend.JIT: _CUDA,
}
format_signature = FormatSignature(
supported_dtypes=_NORM_DTYPES,
@@ -123,10 +123,10 @@ class FusedAddRMSNormOp(BaseFusedOp):
description="residual += x; x = RMSNorm(residual) * weight",
)
descriptions = {
KernelBackend.CUDA_AOT: (
KernelBackend.AOT: (
"Fused residual-add + RMS normalization (sgl_kernel wheel)."
),
KernelBackend.CUDA_JIT: (
KernelBackend.JIT: (
"Fused residual-add + RMS normalization (sglang.jit_kernel)."
),
KernelBackend.TORCH: (
@@ -150,7 +150,7 @@ class FusedAddRMSNormOp(BaseFusedOp):
normed = acc * torch.rsqrt(variance + eps)
input.copy_((normed * weight).to(input.dtype))
def forward_cuda_aot(
def forward_aot(
self,
input: torch.Tensor,
residual: torch.Tensor,
@@ -162,7 +162,7 @@ class FusedAddRMSNormOp(BaseFusedOp):
return sgl_kernel.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
def forward_cuda_jit(
def forward_jit(
self,
input: torch.Tensor,
residual: torch.Tensor,
@@ -180,13 +180,13 @@ class GemmaRMSNormOp(BaseFusedOp):
op = "layernorm.gemma_rmsnorm"
priority = _NORM_PRIORITY
capabilities = {KernelBackend.CUDA_AOT: _CUDA}
capabilities = {KernelBackend.AOT: _CUDA}
format_signature = FormatSignature(
supported_dtypes=_NORM_DTYPES,
description="out = (x / RMS(x)) * (weight + 1); returns tensor",
)
descriptions = {
KernelBackend.CUDA_AOT: "Gemma-style RMS normalization (sgl_kernel wheel).",
KernelBackend.AOT: "Gemma-style RMS normalization (sgl_kernel wheel).",
KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).",
}
@@ -209,7 +209,7 @@ class GemmaRMSNormOp(BaseFusedOp):
out.copy_(result)
return out
def forward_cuda_aot(
def forward_aot(
self,
input: torch.Tensor,
weight: torch.Tensor,
@@ -227,14 +227,14 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp):
op = "layernorm.gemma_fused_add_rmsnorm"
priority = _NORM_PRIORITY
capabilities = {KernelBackend.CUDA_AOT: _CUDA}
capabilities = {KernelBackend.AOT: _CUDA}
format_signature = FormatSignature(
supported_dtypes=_NORM_DTYPES,
in_place=True,
description="residual += x; x = GemmaRMSNorm(residual) * (weight + 1)",
)
descriptions = {
KernelBackend.CUDA_AOT: ("Gemma-style fused residual-add + RMS normalization."),
KernelBackend.AOT: ("Gemma-style fused residual-add + RMS normalization."),
KernelBackend.TORCH: (
"Gemma-style fused residual-add + RMS normalization "
"(pure-torch reference)."
@@ -257,7 +257,7 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp):
normed = acc * torch.rsqrt(variance + eps)
input.copy_((normed * (1.0 + weight.to(torch.float32))).to(input.dtype))
def forward_cuda_aot(
def forward_aot(
self,
input: torch.Tensor,
residual: torch.Tensor,
+4 -4
View File
@@ -14,7 +14,7 @@ if TYPE_CHECKING:
register_kernel(
KernelSpec(
op="mamba.causal_conv1d_fwd",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.mamba:causal_conv1d_fwd",
format_signature=FormatSignature(
in_place=True, description="causal depthwise conv1d forward (prefill)"
@@ -25,7 +25,7 @@ register_kernel(
register_kernel(
KernelSpec(
op="mamba.causal_conv1d_update",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.mamba:causal_conv1d_update",
format_signature=FormatSignature(
in_place=True, description="causal depthwise conv1d update (decode)"
@@ -47,7 +47,7 @@ def causal_conv1d_fwd(
pad_slot_id: int,
):
"""Causal depthwise conv1d forward (prefill)."""
return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.CUDA_AOT)(
return get_kernel("mamba.causal_conv1d_fwd", KernelBackend.AOT)(
x,
weight,
bias_,
@@ -71,7 +71,7 @@ def causal_conv1d_update(
pad_slot_id: int,
):
"""Causal depthwise conv1d update (decode)."""
return get_kernel("mamba.causal_conv1d_update", KernelBackend.CUDA_AOT)(
return get_kernel("mamba.causal_conv1d_update", KernelBackend.AOT)(
x,
conv_state,
weight,
+7 -7
View File
@@ -16,12 +16,12 @@ from sglang.kernels.spec import (
if TYPE_CHECKING:
import torch
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
register_kernel(
KernelSpec(
op="moe.moe_align_block_size",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel:moe_align_block_size",
format_signature=FormatSignature(
in_place=True,
@@ -33,9 +33,9 @@ register_kernel(
register_kernel(
KernelSpec(
op="moe.moe_align_block_size",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.moe_align:moe_align_block_size",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(
in_place=True,
description="MoE align-block-size (JIT variant, AOT signature)",
@@ -46,7 +46,7 @@ register_kernel(
register_kernel(
KernelSpec(
op="moe.topk_softmax",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel:topk_softmax",
format_signature=FormatSignature(
in_place=True,
@@ -69,7 +69,7 @@ def moe_align_block_size(
ignore_invalid_expert: bool = False,
) -> None:
"""Align and sort expert token ids into block-padded output buffers."""
kernel = get_kernel("moe.moe_align_block_size", KernelBackend.CUDA_AOT)
kernel = get_kernel("moe.moe_align_block_size", KernelBackend.AOT)
if ignore_invalid_expert:
return kernel(
topk_ids,
@@ -103,7 +103,7 @@ def topk_softmax(
correction_bias: Optional[torch.Tensor] = None,
) -> None:
"""Compute top-k softmax routing weights/ids for MoE."""
return get_kernel("moe.topk_softmax", KernelBackend.CUDA_AOT)(
return get_kernel("moe.topk_softmax", KernelBackend.AOT)(
topk_weights,
topk_ids,
gating_output,
@@ -16,12 +16,12 @@ from sglang.kernels.spec import (
if TYPE_CHECKING:
import torch
_CUDA = CapabilityRequirement(requires_cuda=True)
_CUDA = frozenset({CapabilityRequirement.CUDA})
register_kernel(
KernelSpec(
op="quantization.sgl_per_token_quant_fp8",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel:sgl_per_token_quant_fp8",
format_signature=FormatSignature(
supported_dtypes=("float8_e4m3fn",),
@@ -41,7 +41,7 @@ for _name in (
register_kernel(
KernelSpec(
op=f"quantization.{_name}",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target=f"sgl_kernel:{_name}",
format_signature=FormatSignature(
in_place=True,
@@ -55,9 +55,9 @@ del _name
register_kernel(
KernelSpec(
op="quantization.sgl_per_token_group_quant_8bit",
backend=KernelBackend.CUDA_JIT,
backend=KernelBackend.JIT,
target="sglang.jit_kernel.per_token_group_quant_8bit:per_token_group_quant_8bit",
capability=_CUDA,
capabilities=_CUDA,
format_signature=FormatSignature(
in_place=True,
description="per-token-group 8-bit quantization (JIT variant)",
@@ -73,7 +73,7 @@ def sgl_per_token_quant_fp8(
output_s: torch.Tensor,
) -> None:
"""Per-token FP8 quantization, writing into ``output_q`` / ``output_s``."""
return get_kernel("quantization.sgl_per_token_quant_fp8", KernelBackend.CUDA_AOT)(
return get_kernel("quantization.sgl_per_token_quant_fp8", KernelBackend.AOT)(
input, output_q, output_s
)
@@ -92,9 +92,7 @@ def sgl_per_token_group_quant_8bit(
enable_v2: Optional[bool] = None,
) -> None:
"""Per-token-group 8-bit quantization, writing into ``output_q`` / ``output_s``."""
return get_kernel(
"quantization.sgl_per_token_group_quant_8bit", KernelBackend.CUDA_AOT
)(
return get_kernel("quantization.sgl_per_token_group_quant_8bit", KernelBackend.AOT)(
input,
output_q,
output_s,
@@ -163,7 +161,7 @@ register_kernel(
"sglang.kernels.ops.quantization.nvfp4_gemm_swiglu_nvfp4_quant"
":nvfp4_gemm_swiglu_nvfp4_quant"
),
capability=CapabilityRequirement(requires_cuda=True, min_cuda_arch=(10, 0)),
capabilities=frozenset({CapabilityRequirement.cuda(min_sm=(10, 0))}),
description="Fused NVFP4 GEMM + SwiGLU + NVFP4 quant (CuTe DSL, SM100).",
)
)
@@ -14,7 +14,7 @@ if TYPE_CHECKING:
register_kernel(
KernelSpec(
op="sampling.top_k_renorm_probs",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.sampling:top_k_renorm_probs",
format_signature=FormatSignature(
description="renormalize probs by top-k thresholding; returns tensor"
@@ -25,7 +25,7 @@ register_kernel(
register_kernel(
KernelSpec(
op="sampling.top_p_renorm_probs",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.sampling:top_p_renorm_probs",
format_signature=FormatSignature(
description="renormalize probs by top-p thresholding; returns tensor"
@@ -39,18 +39,14 @@ def top_k_renorm_probs(
probs: torch.Tensor, top_k: Union[torch.Tensor, int]
) -> torch.Tensor:
"""Renormalize ``probs`` by top-k thresholding."""
return get_kernel("sampling.top_k_renorm_probs", KernelBackend.CUDA_AOT)(
probs, top_k
)
return get_kernel("sampling.top_k_renorm_probs", KernelBackend.AOT)(probs, top_k)
def top_p_renorm_probs(
probs: torch.Tensor, top_p: Union[torch.Tensor, float]
) -> torch.Tensor:
"""Renormalize ``probs`` by top-p thresholding."""
return get_kernel("sampling.top_p_renorm_probs", KernelBackend.CUDA_AOT)(
probs, top_p
)
return get_kernel("sampling.top_p_renorm_probs", KernelBackend.AOT)(probs, top_p)
__all__ = ["top_k_renorm_probs", "top_p_renorm_probs"]
@@ -11,7 +11,7 @@ from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
register_kernel(
KernelSpec(
op="spatial.get_sm_available",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.spatial:get_sm_available",
format_signature=FormatSignature(
description="number of SMs available on device"
@@ -22,7 +22,7 @@ register_kernel(
register_kernel(
KernelSpec(
op="spatial.create_greenctx_stream_by_value",
backend=KernelBackend.CUDA_AOT,
backend=KernelBackend.AOT,
target="sgl_kernel.spatial:create_greenctx_stream_by_value",
format_signature=FormatSignature(
description="create two green-context streams partitioned by SM count"
@@ -34,16 +34,16 @@ register_kernel(
def get_sm_available(device_id: Optional[int] = None) -> int:
"""Return the number of SMs available on ``device_id``."""
return get_kernel("spatial.get_sm_available", KernelBackend.CUDA_AOT)(device_id)
return get_kernel("spatial.get_sm_available", KernelBackend.AOT)(device_id)
def create_greenctx_stream_by_value(
SM_a: int, SM_b: int, device_id: Optional[int] = None
):
"""Create two green-context streams partitioned by ``SM_a`` / ``SM_b``."""
return get_kernel(
"spatial.create_greenctx_stream_by_value", KernelBackend.CUDA_AOT
)(SM_a, SM_b, device_id)
return get_kernel("spatial.create_greenctx_stream_by_value", KernelBackend.AOT)(
SM_a, SM_b, device_id
)
__all__ = ["get_sm_available", "create_greenctx_stream_by_value"]
+38 -13
View File
@@ -1,16 +1,24 @@
"""Fixed-path kernel resolution over the :data:`registry`.
"""Device-aware fixed-path kernel resolution over the :data:`registry`.
There is no priority ranking or heuristic backend selection. Each operator has
a fixed call path its :attr:`KernelSpec.target`:
There is no priority *ranking* or preference heuristic. Resolution of an op's
call path is deterministic:
- an op with a single registered backend resolves to it directly;
- an op with several registered backends must be resolved by naming the backend
explicitly (``backend=...``). The extra backends exist only as inventory, and
are never silently auto-picked.
- an op with several registered backends is filtered by the detected platform
(a hard :class:`~sglang.kernels.spec.CapabilityRequirement` check, not a
preference). If exactly one backend is usable on this device, it is the fixed
call path; if several remain usable, the caller must name the backend
explicitly (``backend=...``).
Because ``KernelBackend`` is now device-agnostic provenance, the *same* backend
(e.g. ``AOT``) may be registered for an op on more than one device; the
availability filter is what makes ``select_kernel`` pick the right one per
platform. Filtering by device is a hard eligibility gate, not the ranked
auto-selection that ``BaseFusedOp`` performs.
:func:`get_kernel` is the fast path used by the public ``ops.*`` wrappers: it
resolves the spec to its callable and caches the result so repeated calls do
not re-run resolution or re-import.
not re-run resolution or re-import (the platform is constant per process).
"""
from __future__ import annotations
@@ -19,7 +27,12 @@ from functools import lru_cache
from typing import Callable, Optional
from sglang.kernels.registry import registry
from sglang.kernels.spec import KernelBackend, KernelSpec
from sglang.kernels.spec import KernelBackend, KernelSpec, PlatformInfo
@lru_cache(maxsize=1)
def _platform() -> PlatformInfo:
return PlatformInfo.detect()
def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpec:
@@ -30,15 +43,16 @@ def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpe
op:
Operator id, ``"<group>.<name>"``.
backend:
Required only when ``op`` has more than one registered backend; selects
which one. For single-backend ops it is optional.
Required only when ``op`` has more than one backend *usable on the
current device*; selects which one. Otherwise optional.
Raises
------
KeyError
If ``op`` is unknown, or if ``backend`` is requested but not registered.
ValueError
If ``op`` has multiple backends and ``backend`` is not given.
If ``op`` has multiple device-eligible backends and ``backend`` is not
given, or if none are eligible on this platform.
"""
specs = registry.get(op)
if not specs:
@@ -53,9 +67,20 @@ def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpe
if len(specs) == 1:
return specs[0]
# Multiple backends: hard-filter by device eligibility.
platform = _platform()
eligible = [s for s in specs if s.is_available(platform)]
if len(eligible) == 1:
return eligible[0]
if not eligible:
raise ValueError(
f"op {op!r} has no backend usable on device {platform.device.value!r} "
f"(registered: {[s.backend.value for s in specs]})"
)
raise ValueError(
f"op {op!r} has multiple registered backends "
f"({[s.backend.value for s in specs]}); pass backend=... to choose one"
f"op {op!r} has multiple backends usable on device "
f"{platform.device.value!r} ({[s.backend.value for s in eligible]}); "
f"pass backend=... to choose one"
)
+96 -26
View File
@@ -8,34 +8,53 @@ box (see RFC #29630, Phase 2).
The concrete callable behind a :class:`KernelSpec` is resolved lazily through
``KernelSpec.load()``; nothing is imported until a kernel is actually called.
Backend vs. device (RFC #29630 follow-up): :class:`KernelBackend` names only the
*provenance* of an implementation (how it is built / where it comes from), not
the hardware it runs on. Both JIT and AOT sources already build for CUDA *and*
ROCm, and a wheel may ship only a per-op subset, so platform support is
per-``(op, backend)`` metadata carried by :class:`CapabilityRequirement`, not
derivable from the backend name.
"""
from __future__ import annotations
import importlib
from enum import Enum
from typing import Callable, Optional, Tuple
from typing import Callable, ClassVar, FrozenSet, Optional, Tuple, Union
import msgspec
class KernelBackend(str, Enum):
"""Implementation backend for a kernel.
"""Provenance of a kernel implementation (how it is built), not its device.
Values mirror the backends called out in RFC #29630: JIT CUDA, AOT
CUDA/C++ (the ``sgl_kernel`` wheel), Triton, CuTe DSL, FlashInfer, DeepGEMM,
and the pure-``torch`` fallback path.
``JIT`` (``sglang.jit_kernel``, compiles under nvcc *and* hipcc) and ``AOT``
(the ``sgl_kernel`` wheel, built for CUDA *and* ROCm) are both cross-device;
which devices a given op supports is expressed by its
:class:`CapabilityRequirement` list. Platform-specific libraries (e.g.
``aiter`` on AMD) are just additional provenance values.
"""
TORCH = "torch" # pure-torch reference (forward_native)
TORCH_COMPILE = "torch_compile" # torch.compile(forward_native)
TRITON = "triton"
CUDA_JIT = "cuda_jit" # sglang.jit_kernel
CUDA_AOT = "cuda_aot" # sgl_kernel wheel
JIT = "jit" # sglang.jit_kernel (nvcc / hipcc)
AOT = "aot" # sgl_kernel wheel (CUDA / ROCm builds)
CUTE_DSL = "cute_dsl"
FLASHINFER = "flashinfer"
DEEPGEMM = "deepgemm"
# TODO(RFC #29630): backends for other hardware (hip_c / npu / cpu-avx, ...)
AITER = "aiter" # AMD aiter library (device=HIP)
# TODO(RFC #29630): more provenance as needed (npu / cpu-avx, ...)
class DeviceType(str, Enum):
"""Accelerator device family a kernel can run on."""
CUDA = "cuda"
HIP = "hip"
CPU = "cpu"
# TODO(RFC #29630): NPU / XPU / ... as backends land.
class PlatformInfo(msgspec.Struct, frozen=True):
@@ -49,6 +68,13 @@ class PlatformInfo(msgspec.Struct, frozen=True):
cuda_arch_major: Optional[int] = None
cuda_arch_minor: Optional[int] = None
@property
def device(self) -> DeviceType:
try:
return DeviceType(self.device_type)
except (ValueError, TypeError):
return DeviceType.CPU
@property
def is_cuda(self) -> bool:
return self.device_type == "cuda"
@@ -85,23 +111,45 @@ class PlatformInfo(msgspec.Struct, frozen=True):
class CapabilityRequirement(msgspec.Struct, frozen=True):
"""Coarse hardware requirement used to filter out unusable backends.
"""One device (plus an optional CUDA-arch window) a backend can run on.
``min_cuda_arch`` / ``max_cuda_arch`` are ``(major, minor)`` tuples, e.g.
``(9, 0)`` for SM90. They only apply when the kernel requires CUDA.
A :class:`KernelSpec` / :class:`~sglang.kernels.fused_op.BaseFusedOp` backend
carries a *set* of these with **OR** semantics — any matching entry makes the
backend eligible, and an empty set means unrestricted (runs anywhere). A set
(not a tuple) because order and duplicates are meaningless here: ``{CUDA,
HIP}`` and ``{HIP, CUDA}`` describe the same thing. This replaces the old
``requires_cuda`` / ``requires_hip`` booleans (whose AND semantics could not
express "CUDA or HIP"); arch bounds now attach to the device they describe
(``min_cuda_arch`` / ``max_cuda_arch`` apply only when ``device == CUDA``).
The device-only cases are so common that they are exposed as class constants
(``CapabilityRequirement.CUDA`` / ``.HIP``); use :meth:`cuda` for an
arch-bounded CUDA requirement (e.g. ``CapabilityRequirement.cuda(
min_sm=(10, 0))`` for SM100+).
"""
requires_cuda: bool = False
requires_hip: bool = False
device: DeviceType
min_cuda_arch: Optional[Tuple[int, int]] = None
max_cuda_arch: Optional[Tuple[int, int]] = None
# Common device-only shortcuts, assigned after the class body (they are
# instances of the class itself). ClassVar keeps them out of msgspec fields.
CUDA: ClassVar[CapabilityRequirement]
HIP: ClassVar[CapabilityRequirement]
@classmethod
def cuda(
cls,
min_sm: Optional[Tuple[int, int]] = None,
max_sm: Optional[Tuple[int, int]] = None,
) -> CapabilityRequirement:
"""A CUDA requirement bounded to an SM-arch window (inclusive)."""
return cls(device=DeviceType.CUDA, min_cuda_arch=min_sm, max_cuda_arch=max_sm)
def is_satisfied_by(self, platform: PlatformInfo) -> bool:
if self.requires_hip and not platform.is_hip:
if self.device != platform.device:
return False
if self.requires_cuda and not platform.is_cuda:
return False
if platform.is_cuda and platform.cuda_arch_major is not None:
if self.device == DeviceType.CUDA and platform.cuda_arch_major is not None:
arch = (platform.cuda_arch_major, platform.cuda_arch_minor or 0)
if self.min_cuda_arch is not None and arch < self.min_cuda_arch:
return False
@@ -110,6 +158,28 @@ class CapabilityRequirement(msgspec.Struct, frozen=True):
return True
CapabilityRequirement.CUDA = CapabilityRequirement(device=DeviceType.CUDA)
CapabilityRequirement.HIP = CapabilityRequirement(device=DeviceType.HIP)
def capabilities_satisfied(
capabilities: Union[
FrozenSet[CapabilityRequirement],
Tuple[CapabilityRequirement, ...],
CapabilityRequirement,
],
platform: PlatformInfo,
) -> bool:
"""OR over ``capabilities`` (empty = unrestricted).
Accepts a set/tuple of requirements, or tolerates a single
:class:`CapabilityRequirement` (the pre-decouple API used one) by wrapping it.
"""
if isinstance(capabilities, CapabilityRequirement):
capabilities = (capabilities,)
return (not capabilities) or any(c.is_satisfied_by(platform) for c in capabilities)
class FormatSignature(msgspec.Struct, frozen=True):
"""A light description of a kernel's data contract.
@@ -131,15 +201,17 @@ class KernelSpec(msgspec.Struct, frozen=True):
Fully-qualified operator id, ``"<group>.<name>"`` (e.g.
``"layernorm.rmsnorm"``). This is the public lookup key.
backend:
Which :class:`KernelBackend` provides this implementation.
Which :class:`KernelBackend` (provenance) provides this implementation.
target:
Import path of the callable in ``"module:attr"`` form, resolved lazily
by :meth:`load` (e.g. ``"sgl_kernel:rmsnorm"``). ``attr`` may be a
dotted path into a module-level object, e.g.
``"sglang.kernels.ops.layernorm:_RMSNORM.forward_cuda_aot"`` for a
bound :class:`~sglang.kernels.fused_op.BaseFusedOp` backend method.
capability:
Hardware requirement used by the selector to skip unusable backends.
``"sglang.kernels.ops.layernorm:_RMSNORM.forward_aot"`` for a bound
:class:`~sglang.kernels.fused_op.BaseFusedOp` backend method.
capabilities:
Set of :class:`CapabilityRequirement` (OR semantics; empty = runs on
any device) used by the selector to skip backends unusable on the
detected platform.
format_signature:
Optional data-contract description for inventory/documentation.
description:
@@ -149,9 +221,7 @@ class KernelSpec(msgspec.Struct, frozen=True):
op: str
backend: KernelBackend
target: str
capability: CapabilityRequirement = msgspec.field(
default_factory=CapabilityRequirement
)
capabilities: FrozenSet[CapabilityRequirement] = frozenset()
format_signature: FormatSignature = msgspec.field(default_factory=FormatSignature)
description: str = ""
@@ -165,7 +235,7 @@ class KernelSpec(msgspec.Struct, frozen=True):
def is_available(self, platform: PlatformInfo) -> bool:
"""Whether this backend can run on ``platform`` (metadata-only check)."""
return self.capability.is_satisfied_by(platform)
return capabilities_satisfied(self.capabilities, platform)
def load(self) -> Callable:
"""Import and return the backing callable.
+1 -1
View File
@@ -701,7 +701,7 @@ class Envs:
# Kernels
# Force every sglang.kernels BaseFusedOp onto one backend (a KernelBackend
# value, e.g. "torch" / "torch_compile" / "triton" / "cuda_aot"); unset =
# value, e.g. "torch" / "torch_compile" / "triton" / "aot"); unset =
# auto-select by priority. "torch" flips all fused ops to their pure-torch
# reference implementations for numerical-bug bisection.
SGLANG_FORCE_FUSED_OP_BACKEND = EnvStr(None)