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

[Kernel] Simplify sglang.kernels tests to idiomatic pytest style (RFC #29630) (#31546)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-17 15:06:26 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 44e3dd2713
commit 619609aa5a
3 changed files with 417 additions and 678 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/test_fused_op.py
+174 -236
View File
@@ -1,33 +1,25 @@
"""GPU-free unit tests for ``sglang.kernels``: BaseFusedOp + registry/selector/spec.
"""GPU-free BaseFusedOp + registry / spec unit tests (RFC #29630).
Part of RFC #29630, Phase 2. Covers the multi-backend operator contract
(structural backend detection, priority dispatch, forced backend, runtime
eligibility, tracing), the registry/selector units in isolation, and the
pure-torch reference implementations of the reworked layernorm / activation
ops. Runs in the CPU CI lane; every-backend-vs-native parity lives in
Every-backend-vs-native parity on real hardware lives in
``test_fused_op_gpu_parity.py``.
"""
import unittest
import math
import pytest
import torch
import sglang.kernels as K
from sglang.kernels.fused_op import BaseFusedOp
from sglang.kernels.registry import KernelRegistry
from sglang.kernels.spec import (
CapabilityRequirement,
KernelBackend,
KernelSpec,
)
from sglang.kernels.spec import CapabilityRequirement as Cap
from sglang.kernels.spec import KernelBackend, KernelSpec
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
class _ToyAddOp(BaseFusedOp):
"""Toy op: element-wise a + b, with a fake 'triton' backend."""
class _ToyAdd(BaseFusedOp):
op = "test.toy_add"
priority = (KernelBackend.TRITON, KernelBackend.TORCH)
@@ -35,258 +27,204 @@ class _ToyAddOp(BaseFusedOp):
return a + b
def forward_triton(self, a, b):
# Marker so tests can tell which backend ran.
return a + b + 1000
return a + b + 1000 # marker so tests can tell which backend ran
class _CudaOnlyToyOp(BaseFusedOp):
"""Toy op whose optimized backend requires CUDA (never eligible on CPU)."""
class _CudaOnlyToy(BaseFusedOp):
op = "test.toy_cuda_only"
priority = (KernelBackend.AOT, KernelBackend.TORCH)
capabilities = {KernelBackend.AOT: {CapabilityRequirement.CUDA}}
capabilities = {KernelBackend.AOT: {Cap.CUDA}}
def forward_native(self, a):
return a * 2
def forward_aot(self, a):
raise AssertionError("must not be selected on a CPU-only box")
raise AssertionError("CUDA backend must not be selected on a CPU-only box")
class TestBaseFusedOp(unittest.TestCase):
def tearDown(self):
K.set_fused_op_backend(None)
K.disable_fused_op_trace()
K.clear_fused_op_trace()
def test_structural_backend_detection(self):
backends = set(_ToyAddOp().available_backends())
self.assertEqual(
backends,
{KernelBackend.TORCH, KernelBackend.TORCH_COMPILE, KernelBackend.TRITON},
)
def test_native_always_available(self):
backends = _CudaOnlyToyOp().available_backends()
self.assertIn(KernelBackend.TORCH, backends)
self.assertIn(KernelBackend.TORCH_COMPILE, backends)
def test_priority_dispatch(self):
op = _ToyAddOp()
a, b = torch.tensor([1.0]), torch.tensor([2.0])
# TRITON is first in priority and always eligible (no capability).
self.assertEqual(op(a, b).item(), 1003.0)
def test_explicit_backend_overrides_priority(self):
op = _ToyAddOp()
a, b = torch.tensor([1.0]), torch.tensor([2.0])
self.assertEqual(op.forward(a, b, backend=KernelBackend.TORCH).item(), 3.0)
def test_capability_gates_runtime_eligibility(self):
# On a CPU-only box the CUDA backend is filtered out and auto-selection
# falls back to native instead of raising.
op = _CudaOnlyToyOp()
if K.PlatformInfo.detect().is_cuda:
self.skipTest("test requires a CPU-only environment")
self.assertEqual(op(torch.tensor([3.0])).item(), 6.0)
def test_forced_backend_global_switch(self):
op = _ToyAddOp()
a, b = torch.tensor([1.0]), torch.tensor([2.0])
K.set_fused_op_backend(KernelBackend.TORCH)
self.assertEqual(op(a, b).item(), 3.0)
K.set_fused_op_backend(None)
self.assertEqual(op(a, b).item(), 1003.0)
def test_forced_backend_env_var(self):
import sglang.kernels.fused_op as fused_op_module
from sglang.srt.environ import envs
op = _ToyAddOp()
a, b = torch.tensor([1.0]), torch.tensor([2.0])
with envs.SGLANG_FORCE_FUSED_OP_BACKEND.override("torch"):
# Reset the module cache so the env var is re-read.
fused_op_module._forced_backend = fused_op_module._UNRESOLVED
self.assertEqual(K.get_fused_op_backend(), KernelBackend.TORCH)
self.assertEqual(op(a, b).item(), 3.0)
fused_op_module._forced_backend = fused_op_module._UNRESOLVED
def test_unimplemented_backend_raises(self):
op = _ToyAddOp()
with self.assertRaises(NotImplementedError):
op.forward(
torch.tensor([1.0]),
torch.tensor([2.0]),
backend=KernelBackend.AOT,
)
def test_torch_compile_backend(self):
op = _ToyAddOp()
a, b = torch.tensor([1.0]), torch.tensor([2.0])
try:
result = op.forward(a, b, backend=KernelBackend.TORCH_COMPILE)
except Exception as e: # inductor toolchain missing in some CI images
self.skipTest(f"torch.compile unavailable: {e}")
self.assertEqual(result.item(), 3.0)
def test_trace_records_op_backend_and_shapes(self):
op = _ToyAddOp()
K.enable_fused_op_trace()
op(torch.zeros(2, 3), torch.zeros(2, 3))
records = K.get_fused_op_trace()
self.assertEqual(len(records), 1)
self.assertEqual(records[0].op, "test.toy_add")
self.assertEqual(records[0].backend, "triton")
self.assertEqual(
records[0].tensor_args,
("torch.float32[2, 3]", "torch.float32[2, 3]"),
)
def test_register_fused_op_specs(self):
op = K.registry.get("layernorm.rmsnorm")
backends = {s.backend for s in op}
self.assertEqual(
backends,
{
KernelBackend.TORCH,
KernelBackend.TORCH_COMPILE,
KernelBackend.JIT,
KernelBackend.AOT,
KernelBackend.AITER,
KernelBackend.TORCH_NPU,
},
)
# Dotted targets resolve to the bound backend methods.
native = K.registry.get_backend("layernorm.rmsnorm", KernelBackend.TORCH)
fn = native.load()
x = torch.randn(4, 64)
w = torch.randn(64)
self.assertTrue(torch.allclose(fn(x, w), _ref_rmsnorm(x, w, 1e-6)))
@pytest.fixture(autouse=True)
def _reset_global_state():
yield
K.set_fused_op_backend(None)
K.disable_fused_op_trace()
K.clear_fused_op_trace()
class TestKernelRegistryUnit(unittest.TestCase):
"""Isolated KernelRegistry behavior (fresh instance, no global state)."""
def _spec(self, op="g.n", backend=KernelBackend.TORCH, target="math:sqrt"):
return KernelSpec(op=op, backend=backend, target=target)
def test_register_and_get(self):
reg = KernelRegistry()
spec = self._spec()
reg.register(spec)
self.assertEqual(reg.get("g.n"), [spec])
self.assertTrue(reg.has("g.n"))
self.assertEqual(reg.ops(), ["g.n"])
def test_get_unknown_op_returns_empty(self):
reg = KernelRegistry()
self.assertEqual(reg.get("no.such"), [])
self.assertFalse(reg.has("no.such"))
def test_reregister_same_backend_replaces(self):
reg = KernelRegistry()
reg.register(self._spec(target="math:sqrt"))
reg.register(self._spec(target="math:floor"))
specs = reg.get("g.n")
self.assertEqual(len(specs), 1)
self.assertEqual(specs[0].target, "math:floor")
def test_get_backend_missing_raises(self):
reg = KernelRegistry()
reg.register(self._spec(backend=KernelBackend.TORCH))
with self.assertRaises(KeyError):
reg.get_backend("g.n", KernelBackend.TRITON)
with self.assertRaises(KeyError):
reg.get_backend("no.such", KernelBackend.TORCH)
def _t(*vals):
return torch.tensor(list(vals))
class TestKernelSpecUnit(unittest.TestCase):
def test_load_simple_target(self):
import math
def test_available_backends():
assert set(_ToyAdd().available_backends()) == {
KernelBackend.TORCH,
KernelBackend.TORCH_COMPILE,
KernelBackend.TRITON,
}
spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="math:sqrt")
self.assertIs(spec.load(), math.sqrt)
def test_load_dotted_target(self):
spec = KernelSpec(
op="g.n",
backend=KernelBackend.TORCH,
target="sglang.kernels.ops.layernorm:_RMSNORM.forward_native",
)
self.assertTrue(callable(spec.load()))
def test_priority_dispatch():
# TRITON is first in priority and always eligible (no capability).
assert _ToyAdd()(_t(1.0), _t(2.0)).item() == 1003.0
def test_load_bad_target_raises(self):
spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="no-colon")
with self.assertRaises(ValueError):
spec.load()
def test_explicit_backend_overrides_priority():
assert (
_ToyAdd().forward(_t(1.0), _t(2.0), backend=KernelBackend.TORCH).item() == 3.0
)
def test_capability_gates_eligibility():
if K.PlatformInfo.detect().is_cuda:
pytest.skip("requires a CPU-only environment")
# CUDA backend is filtered out; auto-selection falls back to native.
assert _CudaOnlyToy()(_t(3.0)).item() == 6.0
def test_forced_backend_global_switch():
op = _ToyAdd()
K.set_fused_op_backend(KernelBackend.TORCH)
assert op(_t(1.0), _t(2.0)).item() == 3.0
K.set_fused_op_backend(None)
assert op(_t(1.0), _t(2.0)).item() == 1003.0
def test_forced_backend_env_var():
import sglang.kernels.fused_op as m
from sglang.srt.environ import envs
op = _ToyAdd()
with envs.SGLANG_FORCE_FUSED_OP_BACKEND.override("torch"):
m._forced_backend = m._UNRESOLVED # drop cache so the env var is re-read
assert K.get_fused_op_backend() is KernelBackend.TORCH
assert op(_t(1.0), _t(2.0)).item() == 3.0
m._forced_backend = m._UNRESOLVED
def test_unimplemented_backend_raises():
with pytest.raises(NotImplementedError):
_ToyAdd().forward(_t(1.0), _t(2.0), backend=KernelBackend.AOT)
def test_trace_records_op_backend_and_shapes():
K.enable_fused_op_trace()
_ToyAdd()(torch.zeros(2, 3), torch.zeros(2, 3))
(rec,) = K.get_fused_op_trace()
assert rec.op == "test.toy_add"
assert rec.backend == "triton"
assert rec.tensor_args == ("torch.float32[2, 3]", "torch.float32[2, 3]")
def test_fused_op_registers_all_backends():
backends = {s.backend for s in K.registry.get("layernorm.rmsnorm")}
assert backends == {
KernelBackend.TORCH,
KernelBackend.TORCH_COMPILE,
KernelBackend.JIT,
KernelBackend.AOT,
KernelBackend.AITER,
KernelBackend.TORCH_NPU,
}
# --- KernelRegistry unit ---
def _spec(op="g.n", backend=KernelBackend.TORCH, target="math:sqrt"):
return KernelSpec(op=op, backend=backend, target=target)
def test_registry_register_and_get():
reg = KernelRegistry()
spec = _spec()
reg.register(spec)
assert reg.get("g.n") == [spec]
assert reg.has("g.n")
assert reg.ops() == ["g.n"]
assert reg.get("no.such") == []
def test_registry_reregister_replaces():
reg = KernelRegistry()
reg.register(_spec(target="math:sqrt"))
reg.register(_spec(target="math:floor"))
assert [s.target for s in reg.get("g.n")] == ["math:floor"]
def test_registry_get_backend_missing_raises():
reg = KernelRegistry()
reg.register(_spec())
with pytest.raises(KeyError):
reg.get_backend("g.n", KernelBackend.TRITON)
# --- KernelSpec.load ---
def test_spec_load_simple_and_dotted():
assert _spec(target="math:sqrt").load() is math.sqrt
dotted = _spec(target="sglang.kernels.ops.layernorm:_RMSNORM.forward_native")
assert callable(dotted.load())
def test_spec_load_bad_target_raises():
with pytest.raises(ValueError):
_spec(target="no-colon").load()
# --- native reference math of the reworked ops (CPU) ---
def _ref_rmsnorm(x, w, eps):
xf = x.to(torch.float32)
var = xf.pow(2).mean(dim=-1, keepdim=True)
return (xf * torch.rsqrt(var + eps) * w).to(x.dtype)
xf = x.float()
return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w).to(x.dtype)
class TestNativeReferenceImplementations(unittest.TestCase):
"""The forward_native math of the reworked ops, on CPU tensors."""
def test_rmsnorm_native():
from sglang.kernels.ops.layernorm import _RMSNORM
def setUp(self):
torch.manual_seed(0)
x, w = torch.randn(8, 128), torch.randn(128)
out = _RMSNORM.forward_native(x, w, 1e-6)
assert torch.allclose(out, _ref_rmsnorm(x, w, 1e-6))
buf = torch.empty_like(x)
assert _RMSNORM.forward_native(x, w, 1e-6, out=buf) is buf # out= is in place
assert torch.allclose(buf, out)
def test_rmsnorm_native(self):
from sglang.kernels.ops.layernorm import _RMSNORM
x = torch.randn(8, 128)
w = torch.randn(128)
out = _RMSNORM.forward_native(x, w, 1e-6)
self.assertTrue(torch.allclose(out, _ref_rmsnorm(x, w, 1e-6)))
# out= writes in place and returns out
buf = torch.empty_like(x)
self.assertIs(_RMSNORM.forward_native(x, w, 1e-6, out=buf), buf)
self.assertTrue(torch.allclose(buf, out))
def test_fused_add_rmsnorm_native():
from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM
def test_fused_add_rmsnorm_native(self):
from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM
x, residual, w = torch.randn(8, 128), torch.randn(8, 128), torch.randn(128)
x0, r0 = x.clone(), residual.clone()
assert _FUSED_ADD_RMSNORM.forward_native(x, residual, w, 1e-6) is None
acc = x0.float() + r0.float()
assert torch.allclose(residual, acc)
assert torch.allclose(
x, acc * torch.rsqrt(acc.pow(2).mean(-1, keepdim=True) + 1e-6) * w
)
x = torch.randn(8, 128)
residual = torch.randn(8, 128)
w = torch.randn(128)
x2, r2 = x.clone(), residual.clone()
self.assertIsNone(_FUSED_ADD_RMSNORM.forward_native(x, residual, w, 1e-6))
acc = x2.to(torch.float32) + r2.to(torch.float32)
self.assertTrue(torch.allclose(residual, acc))
ref = acc * torch.rsqrt(acc.pow(2).mean(-1, keepdim=True) + 1e-6) * w
self.assertTrue(torch.allclose(x, ref))
def test_gemma_rmsnorm_native(self):
from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM
@pytest.mark.parametrize(
"op_attr, approximate",
[
("_SILU_AND_MUL", None),
("_GELU_AND_MUL", "none"),
("_GELU_TANH_AND_MUL", "tanh"),
],
)
def test_gated_activation_native(op_attr, approximate):
import torch.nn.functional as F
x = torch.randn(8, 128)
w = torch.randn(128)
out = _GEMMA_RMSNORM.forward_native(x, w, 1e-6)
xf = x.to(torch.float32)
ref = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) * (1.0 + w)
self.assertTrue(torch.allclose(out, ref))
import sglang.kernels.ops.activation as act
def test_gated_activations_native(self):
import torch.nn.functional as F
from sglang.kernels.ops.activation import (
_GELU_AND_MUL,
_GELU_TANH_AND_MUL,
_SILU_AND_MUL,
)
x = torch.randn(8, 256)
gate, up = x[..., :128], x[..., 128:]
cases = [
(_SILU_AND_MUL, F.silu(gate) * up),
(_GELU_AND_MUL, F.gelu(gate, approximate="none") * up),
(_GELU_TANH_AND_MUL, F.gelu(gate, approximate="tanh") * up),
]
for op, ref in cases:
self.assertTrue(torch.allclose(op.forward_native(x), ref), op.op)
x = torch.randn(8, 256)
gate, up = x[..., :128], x[..., 128:]
ref = (
F.silu(gate) if approximate is None else F.gelu(gate, approximate=approximate)
) * up
assert torch.allclose(getattr(act, op_attr).forward_native(x), ref)
if __name__ == "__main__":
unittest.main()
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/test_fused_op_gpu_parity.py
+97 -99
View File
@@ -1,125 +1,123 @@
"""Generic every-backend-vs-native parity harness for BaseFusedOp operators.
"""Every-backend-vs-native parity for BaseFusedOp ops on real GPU (RFC #29630).
Part of RFC #29630, Phase 2. For each reworked fused op, enumerate its
available backends, run each one that is eligible on this platform, and
assert the output matches the pure-torch ``forward_native`` reference within
dtype tolerance. New backends added to an op are picked up automatically —
no per-kernel test boilerplate.
For each reworked fused op, run every backend eligible on this platform and
assert it matches the pure-torch ``forward_native`` reference within dtype
tolerance. New backends are picked up automatically.
"""
import unittest
import pytest
import torch
from sglang.kernels.spec import KernelBackend
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
_DEVICE = "cuda"
# torch_compile is native under the hood; exclude it from the sweep to keep
# CI time down (compilation dominates) — it is exercised in the CPU lane.
_SKIP_BACKENDS = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE}
_TOLERANCE = {
# torch_compile is native under the hood; skip it here (compile time dominates)
# -- it is exercised in the CPU lane.
_SKIP = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE}
_TOL = {
torch.float16: dict(atol=1e-2, rtol=1e-2),
torch.bfloat16: dict(atol=2e-2, rtol=2e-2),
}
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
class TestFusedOpGpuParity(CustomTestCase):
def setUp(self):
torch.manual_seed(0)
def _eligible_backends(self, op):
return [
b
for b in op.available_backends()
if b not in _SKIP_BACKENDS and op.backend_eligible(b)
]
def _eligible(op):
return [
b for b in op.available_backends() if b not in _SKIP and op.backend_eligible(b)
]
def _assert_close(self, got, ref, dtype, msg):
torch.testing.assert_close(got, ref, **_TOLERANCE[dtype], msg=msg)
def test_rmsnorm_backends_match_native(self):
from sglang.kernels.ops.layernorm import _RMSNORM
def _close(got, ref, dtype, msg):
torch.testing.assert_close(got, ref, **_TOL[dtype], msg=msg)
for dtype in (torch.float16, torch.bfloat16):
for shape in ((1, 4096), (128, 4096), (7, 2048)):
x = torch.randn(shape, dtype=dtype, device=_DEVICE)
w = torch.randn(shape[-1], dtype=dtype, device=_DEVICE)
ref = _RMSNORM.forward_native(x, w, 1e-6)
for backend in self._eligible_backends(_RMSNORM):
got = _RMSNORM.forward(x, w, 1e-6, backend=backend)
self._assert_close(
got, ref, dtype, f"rmsnorm {backend.value} {dtype} {shape}"
)
def test_fused_add_rmsnorm_backends_match_native(self):
from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("shape", [(1, 4096), (128, 4096), (7, 2048)])
def test_rmsnorm(dtype, shape):
from sglang.kernels.ops.layernorm import _RMSNORM
for dtype in (torch.float16, torch.bfloat16):
for shape in ((1, 4096), (128, 4096)):
x0 = torch.randn(shape, dtype=dtype, device=_DEVICE)
r0 = torch.randn(shape, dtype=dtype, device=_DEVICE)
w = torch.randn(shape[-1], dtype=dtype, device=_DEVICE)
x_ref, r_ref = x0.clone(), r0.clone()
_FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6)
for backend in self._eligible_backends(_FUSED_ADD_RMSNORM):
x, r = x0.clone(), r0.clone()
_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=backend)
label = f"fused_add_rmsnorm {backend.value} {dtype} {shape}"
self._assert_close(x, x_ref, dtype, label + " (normed)")
self._assert_close(r, r_ref, dtype, label + " (residual)")
def test_gemma_rmsnorm_backends_match_native(self):
from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM
for dtype in (torch.float16, torch.bfloat16):
x = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
w = torch.randn(2048, dtype=dtype, device=_DEVICE)
ref = _GEMMA_RMSNORM.forward_native(x, w, 1e-6)
for backend in self._eligible_backends(_GEMMA_RMSNORM):
got = _GEMMA_RMSNORM.forward(x, w, 1e-6, backend=backend)
self._assert_close(
got, ref, dtype, f"gemma_rmsnorm {backend.value} {dtype}"
)
def test_gemma_fused_add_rmsnorm_backends_match_native(self):
from sglang.kernels.ops.layernorm import _GEMMA_FUSED_ADD_RMSNORM
for dtype in (torch.float16, torch.bfloat16):
x0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
r0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
w = torch.randn(2048, dtype=dtype, device=_DEVICE)
x_ref, r_ref = x0.clone(), r0.clone()
_GEMMA_FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6)
for backend in self._eligible_backends(_GEMMA_FUSED_ADD_RMSNORM):
x, r = x0.clone(), r0.clone()
_GEMMA_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=backend)
label = f"gemma_fused_add_rmsnorm {backend.value} {dtype}"
self._assert_close(x, x_ref, dtype, label + " (normed)")
self._assert_close(r, r_ref, dtype, label + " (residual)")
def test_gated_activation_backends_match_native(self):
from sglang.kernels.ops.activation import (
_GELU_AND_MUL,
_GELU_TANH_AND_MUL,
_SILU_AND_MUL,
torch.manual_seed(0)
x = torch.randn(shape, dtype=dtype, device="cuda")
w = torch.randn(shape[-1], dtype=dtype, device="cuda")
ref = _RMSNORM.forward_native(x, w, 1e-6)
for b in _eligible(_RMSNORM):
_close(
_RMSNORM.forward(x, w, 1e-6, backend=b), ref, dtype, f"rmsnorm {b.value}"
)
for op in (_SILU_AND_MUL, _GELU_AND_MUL, _GELU_TANH_AND_MUL):
for dtype in (torch.float16, torch.bfloat16):
for shape in ((1, 8192), (128, 8192)):
x = torch.randn(shape, dtype=dtype, device=_DEVICE)
ref = op.forward_native(x)
for backend in self._eligible_backends(op):
got = op.forward(x, backend=backend)
self._assert_close(
got, ref, dtype, f"{op.op} {backend.value} {dtype} {shape}"
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("shape", [(1, 4096), (128, 4096)])
def test_fused_add_rmsnorm(dtype, shape):
from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM
torch.manual_seed(0)
x0 = torch.randn(shape, dtype=dtype, device="cuda")
r0 = torch.randn(shape, dtype=dtype, device="cuda")
w = torch.randn(shape[-1], dtype=dtype, device="cuda")
x_ref, r_ref = x0.clone(), r0.clone()
_FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6)
for b in _eligible(_FUSED_ADD_RMSNORM):
x, r = x0.clone(), r0.clone()
_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=b)
_close(x, x_ref, dtype, f"fused_add {b.value} (normed)")
_close(r, r_ref, dtype, f"fused_add {b.value} (residual)")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_gemma_rmsnorm(dtype):
from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM
torch.manual_seed(0)
x = torch.randn(64, 2048, dtype=dtype, device="cuda")
w = torch.randn(2048, dtype=dtype, device="cuda")
ref = _GEMMA_RMSNORM.forward_native(x, w, 1e-6)
for b in _eligible(_GEMMA_RMSNORM):
_close(
_GEMMA_RMSNORM.forward(x, w, 1e-6, backend=b),
ref,
dtype,
f"gemma {b.value}",
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_gemma_fused_add_rmsnorm(dtype):
from sglang.kernels.ops.layernorm import _GEMMA_FUSED_ADD_RMSNORM
torch.manual_seed(0)
x0 = torch.randn(64, 2048, dtype=dtype, device="cuda")
r0 = torch.randn(64, 2048, dtype=dtype, device="cuda")
w = torch.randn(2048, dtype=dtype, device="cuda")
x_ref, r_ref = x0.clone(), r0.clone()
_GEMMA_FUSED_ADD_RMSNORM.forward_native(x_ref, r_ref, w, 1e-6)
for b in _eligible(_GEMMA_FUSED_ADD_RMSNORM):
x, r = x0.clone(), r0.clone()
_GEMMA_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=b)
_close(x, x_ref, dtype, f"gemma_fused_add {b.value} (normed)")
_close(r, r_ref, dtype, f"gemma_fused_add {b.value} (residual)")
@pytest.mark.parametrize(
"op_attr", ["_SILU_AND_MUL", "_GELU_AND_MUL", "_GELU_TANH_AND_MUL"]
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("shape", [(1, 8192), (128, 8192)])
def test_gated_activation(op_attr, dtype, shape):
import sglang.kernels.ops.activation as act
torch.manual_seed(0)
op = getattr(act, op_attr)
x = torch.randn(shape, dtype=dtype, device="cuda")
ref = op.forward_native(x)
for b in _eligible(op):
_close(op.forward(x, backend=b), ref, dtype, f"{op.op} {b.value}")
if __name__ == "__main__":
unittest.main()
import sys
sys.exit(pytest.main([__file__]))
test/registered/kernels/test_kernels_namespace.py
+146 -343
View File
@@ -1,142 +1,22 @@
"""GPU-free import/registry tests for the ``sglang.kernels`` namespace.
Part of RFC #29630, Phase 2. These tests exercise the public namespace, the
kernel registry, and the heuristic selector without touching a GPU or importing
any kernel backend (``sgl_kernel`` / ``sglang.jit_kernel``). They run in the CPU
CI lane.
"""
"""GPU-free import / registry / selector tests for ``sglang.kernels`` (RFC #29630)."""
import importlib
import subprocess
import sys
import unittest
import pytest
import sglang.kernels as K
import sglang.kernels.fused_op as fo
import sglang.kernels.ops # noqa: F401 -- populate the registry
import sglang.kernels.selector as sel
from sglang.kernels import DeviceType, KernelBackend, PlatformInfo
from sglang.kernels.spec import CapabilityRequirement as Cap
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
# A must-contain subset of registered operators and their backends. The
# registry holds many more entries (every migrated Triton kernel), so this is
# checked as a subset, not an exact match.
EXPECTED_OPS = {
# BaseFusedOp-backed ops: native + torch_compile always available,
# plus the overridden CUDA backends.
"activation.silu_and_mul": {"aot", "jit", "aiter", "torch", "torch_compile"},
"activation.gelu_and_mul": {"aot", "jit", "torch", "torch_compile"},
"activation.gelu_tanh_and_mul": {
"aot",
"jit",
"torch",
"torch_compile",
},
"activation.relu2": {"jit", "torch", "torch_compile"},
"activation.gelu_quick": {"aot", "torch", "torch_compile"},
"layernorm.rmsnorm": {
"aot",
"jit",
"aiter",
"torch_npu",
"torch",
"torch_compile",
},
"layernorm.fused_add_rmsnorm": {
"aot",
"jit",
"aiter",
"torch_npu",
"torch",
"torch_compile",
},
"layernorm.gemma_rmsnorm": {
"aot",
"jit",
"torch_npu",
"torch",
"torch_compile",
},
"layernorm.gemma_fused_add_rmsnorm": {
"aot",
"jit",
"torch",
"torch_compile",
},
# curated dual/single-backend wrapper ops
"gemm.fp8_scaled_mm": {"aot"},
"gemm.dsv3_fused_a_gemm": {"aot", "jit"},
"gemm.dsv3_router_gemm": {"jit"},
"kvcache.reshape_and_cache_flash": {"triton"},
"moe.moe_align_block_size": {"aot", "jit"},
"moe.topk_softmax": {"aot"},
"quantization.sgl_per_token_quant_fp8": {"aot"},
# migrated from srt/layers/quantization (Phase 2.5)
"quantization.w8a8_block_fp8_matmul": {"triton"},
"quantization.per_token_quant_int8": {"triton"},
"quantization.awq_dequantize_triton": {"triton"},
"quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
"moe.pack_topk_ids": {"triton"},
"quantization.sgl_per_token_group_quant_8bit": {"aot", "jit"},
"quantization.sgl_per_token_group_quant_fp8": {"aot"},
"quantization.sgl_per_token_group_quant_int8": {"aot"},
# deferred-group wrappers, now populated
"sampling.top_k_renorm_probs": {"aot"},
"sampling.top_p_renorm_probs": {"aot"},
"spatial.get_sm_available": {"aot"},
"spatial.create_greenctx_stream_by_value": {"aot"},
"mamba.causal_conv1d_fwd": {"aot"},
"mamba.causal_conv1d_update": {"aot"},
"diffusion.apply_group_norm_silu": {"jit"},
"diffusion.residual_gate_add": {"jit"},
"diffusion.fused_inplace_qknorm_rope": {"jit"},
# representative migrated Triton kernels (inventory)
"grammar.apply_token_bitmask_inplace_triton": {"triton"},
"memory.alloc_extend_kernel": {"triton"},
"attention.decode_attention_fwd": {"triton"},
"embeddings.vocab_parallel_embedding": {"triton"},
"kvcache.create_flashinfer_kv_indices_triton": {"triton"},
"speculative.draft_topk1_postprocess": {"triton"},
"speculative.gather_spec_extras": {"triton"},
}
# Public wrapper callables that each populated group must expose.
EXPECTED_WRAPPERS = {
"sglang.kernels.ops.layernorm": [
"rmsnorm",
"fused_add_rmsnorm",
"gemma_rmsnorm",
"gemma_fused_add_rmsnorm",
],
"sglang.kernels.ops.activation": [
"silu_and_mul",
"gelu_and_mul",
"gelu_tanh_and_mul",
],
"sglang.kernels.ops.gemm": [
"fp8_scaled_mm",
"dsv3_fused_a_gemm",
"dsv3_router_gemm",
],
"sglang.kernels.ops.quantization": [
"sgl_per_token_quant_fp8",
"sgl_per_token_group_quant_8bit",
"sgl_per_token_group_quant_fp8",
"sgl_per_token_group_quant_int8",
],
"sglang.kernels.ops.moe": ["moe_align_block_size", "topk_softmax"],
"sglang.kernels.ops.kvcache": ["reshape_and_cache_flash"],
"sglang.kernels.ops.sampling": ["top_k_renorm_probs", "top_p_renorm_probs"],
"sglang.kernels.ops.spatial": [
"get_sm_available",
"create_greenctx_stream_by_value",
],
"sglang.kernels.ops.mamba": ["causal_conv1d_fwd", "causal_conv1d_update"],
"sglang.kernels.ops.diffusion": [
"apply_group_norm_silu",
"residual_gate_add",
"fused_inplace_qknorm_rope",
],
}
# All operator groups from the RFC's proposed shape must import as packages.
ALL_GROUPS = [
GROUPS = [
"activation",
"attention",
"communication",
@@ -155,241 +35,164 @@ ALL_GROUPS = [
"speculative",
]
# Representative ops checked as a subset (the registry holds many more).
EXPECTED = {
"activation.silu_and_mul": {"aot", "jit", "aiter", "torch", "torch_compile"},
"activation.relu2": {"jit", "torch", "torch_compile"},
"layernorm.rmsnorm": {"aot", "jit", "aiter", "torch_npu", "torch", "torch_compile"},
"layernorm.gemma_rmsnorm": {"aot", "jit", "torch_npu", "torch", "torch_compile"},
"gemm.fp8_scaled_mm": {"aot"},
"moe.moe_align_block_size": {"aot", "jit"},
"quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
"kvcache.reshape_and_cache_flash": {"triton"},
}
class TestKernelsNamespace(unittest.TestCase):
def setUp(self):
import importlib
_CPU = PlatformInfo(device_type="cpu")
_SM90 = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
_SM100 = PlatformInfo(device_type="cuda", cuda_arch_major=10, cuda_arch_minor=0)
_HIP = PlatformInfo(device_type="hip")
import sglang.kernels
import sglang.kernels.ops # populate the registry
self.K = sglang.kernels
self.importlib = importlib
def test_top_level_exports():
for name in (
"KernelSpec",
"KernelBackend",
"FormatSignature",
"CapabilityRequirement",
"PlatformInfo",
"registry",
"get_kernel",
"select_kernel",
):
assert hasattr(K, name), name
def test_top_level_exports(self):
for name in (
"KernelSpec",
"KernelBackend",
"FormatSignature",
"CapabilityRequirement",
"PlatformInfo",
"registry",
"get_kernel",
"select_kernel",
):
self.assertTrue(hasattr(self.K, name), f"missing export: {name}")
def test_all_groups_importable(self):
for group in ALL_GROUPS:
mod = self.importlib.import_module(f"sglang.kernels.ops.{group}")
self.assertTrue(hasattr(mod, "__all__"))
@pytest.mark.parametrize("group", GROUPS)
def test_group_importable(group):
assert hasattr(importlib.import_module(f"sglang.kernels.ops.{group}"), "__all__")
def test_registry_contents(self):
registry = self.K.registry
ops = set(registry.ops())
# EXPECTED_OPS is a must-contain subset (many more migrated kernels
# are also registered).
missing = set(EXPECTED_OPS) - ops
self.assertFalse(missing, f"missing registered ops: {sorted(missing)}")
for op, backends in EXPECTED_OPS.items():
got = {s.backend.value for s in registry.get(op)}
self.assertEqual(got, backends, f"backend mismatch for {op}")
self.assertGreaterEqual(len(ops), 80, "registry unexpectedly small")
def test_specs_are_well_formed(self):
for spec in self.K.registry.all_specs():
self.assertIn(".", spec.op)
self.assertEqual(spec.op, f"{spec.group}.{spec.name}")
# target must be an importable "module:attr" path
module_path, sep, attr = spec.target.partition(":")
self.assertEqual(sep, ":", f"bad target for {spec.op}: {spec.target}")
self.assertTrue(module_path and attr, spec.target)
@pytest.mark.parametrize("op, backends", list(EXPECTED.items()))
def test_registry_backends(op, backends):
assert {s.backend.value for s in K.registry.get(op)} == backends
def test_wrappers_exposed_and_callable(self):
for module_name, names in EXPECTED_WRAPPERS.items():
mod = self.importlib.import_module(module_name)
for name in names:
self.assertTrue(callable(getattr(mod, name)), f"{module_name}.{name}")
def test_single_backend_op_resolves_without_backend(self):
# An op with exactly one registered backend has a fixed call path.
for op, backends in EXPECTED_OPS.items():
if len(backends) == 1:
spec = self.K.select_kernel(op)
self.assertEqual(spec.backend.value, next(iter(backends)), op)
def test_specs_well_formed():
for spec in K.registry.all_specs():
assert spec.op == f"{spec.group}.{spec.name}"
mod, sep, attr = spec.target.partition(":")
assert sep == ":" and mod and attr, spec.target
def test_multi_backend_op_requires_explicit_backend(self):
# Device is a HARD eligibility filter, not a preference ranking: when
# more than one backend is usable on the current device, selection must
# be explicit (no hidden auto-ranking). Force a CUDA platform so the
# result is deterministic regardless of the test host.
import sglang.kernels.selector as sel
saved = sel._platform
try:
sel._platform = lambda: self.K.PlatformInfo(
device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0
)
# rmsnorm exposes torch/torch_compile/jit/aot, all eligible on CUDA.
with self.assertRaises(ValueError):
self.K.select_kernel("layernorm.rmsnorm")
# An explicit backend is always the fixed call path.
spec = self.K.select_kernel(
"layernorm.rmsnorm", backend=self.K.KernelBackend.JIT
)
self.assertEqual(spec.backend, self.K.KernelBackend.JIT)
finally:
sel._platform = saved
def test_single_backend_resolves_without_backend():
assert K.select_kernel("gemm.fp8_scaled_mm").backend is KernelBackend.AOT
def test_decoupled_backend_device_selection(self):
# Proves the decoupled backend/device model against production reality:
# - AOT (sgl_kernel) spans CUDA *and* HIP (OR-semantics capability);
# - JIT is CUDA-only; AITER is an opt-in HIP-only path on silu_and_mul;
# - gelu_and_mul has no AITER kernel (per-(op, backend) subset).
# Auto-selection matches production defaults: JIT on CUDA, AOT on HIP.
import sglang.kernels.fused_op as fo
from sglang.kernels.ops.activation import _GELU_AND_MUL, _SILU_AND_MUL
B = self.K.KernelBackend
hip = self.K.PlatformInfo(device_type="hip")
cuda = self.K.PlatformInfo(device_type="cuda", cuda_arch_major=9)
saved = fo._platform
try:
fo._platform = lambda: hip
# silu implements AITER (a HIP kernel); gelu does not (per-op subset).
self.assertIn(B.AITER, _SILU_AND_MUL.available_backends())
self.assertNotIn(B.AITER, _GELU_AND_MUL.available_backends())
self.assertTrue(_SILU_AND_MUL.backend_eligible(B.AOT)) # (cuda, hip)
self.assertTrue(_SILU_AND_MUL.backend_eligible(B.AITER)) # hip-only
self.assertFalse(_SILU_AND_MUL.backend_eligible(B.JIT)) # cuda-only
# HIP default = AOT (production default); AITER is opt-in below it.
self.assertEqual(_SILU_AND_MUL._resolve_backend(), B.AOT)
# gelu has no AITER but AOT spans HIP -> resolves to AOT, not torch.
self.assertEqual(_GELU_AND_MUL._resolve_backend(), B.AOT)
fo._platform = lambda: cuda
self.assertEqual(_SILU_AND_MUL._resolve_backend(), B.JIT) # CUDA default
finally:
fo._platform = saved
def test_unknown_op_or_backend_raises():
with pytest.raises(KeyError):
K.select_kernel("does_not.exist")
with pytest.raises(KeyError):
K.select_kernel("gemm.fp8_scaled_mm", backend=KernelBackend.TRITON)
def test_layernorm_cross_device_coverage(self):
# The rmsnorm ops illustrate that the *same* provenance covers different
# devices per op: AOT (sgl_kernel) is CUDA-only here (sgl_kernel does not
# build rmsnorm for ROCm), so HIP falls to AITER and NPU to torch_npu,
# each matching the production default for that device. gemma uses a
# rocm-triton JIT path on HIP -- a JIT provenance pinned to HIP, unlike
# the CUDA-only JIT on plain rmsnorm.
import sglang.kernels.fused_op as fo
from sglang.kernels.ops.layernorm import (
_FUSED_ADD_RMSNORM,
_GEMMA_RMSNORM,
_RMSNORM,
)
B = self.K.KernelBackend
cuda = self.K.PlatformInfo(device_type="cuda", cuda_arch_major=9)
hip = self.K.PlatformInfo(device_type="hip")
npu = self.K.PlatformInfo(device_type="npu")
saved = fo._platform
try:
for plat, expect in ((cuda, B.AOT), (hip, B.AITER), (npu, B.TORCH_NPU)):
fo._platform = lambda p=plat: p
self.assertEqual(_RMSNORM._resolve_backend(), expect)
self.assertEqual(_FUSED_ADD_RMSNORM._resolve_backend(), expect)
# gemma: AOT on CUDA, rocm-triton JIT on HIP, torch_npu on NPU.
for plat, expect in ((cuda, B.AOT), (hip, B.JIT), (npu, B.TORCH_NPU)):
fo._platform = lambda p=plat: p
self.assertEqual(_GEMMA_RMSNORM._resolve_backend(), expect)
# AOT rmsnorm is CUDA-only (not HIP) -- distinct from activation's AOT.
fo._platform = lambda: hip
self.assertFalse(_RMSNORM.backend_eligible(B.AOT))
self.assertTrue(_RMSNORM.backend_eligible(B.AITER))
finally:
fo._platform = saved
def test_multi_backend_requires_explicit_backend(monkeypatch):
# Device is a hard eligibility filter, not a ranking: >1 usable backend on
# the current device means selection must name one.
monkeypatch.setattr(sel, "_platform", lambda: _SM90)
with pytest.raises(ValueError):
K.select_kernel("layernorm.rmsnorm")
spec = K.select_kernel("layernorm.rmsnorm", backend=KernelBackend.JIT)
assert spec.backend is KernelBackend.JIT
assert spec.target == "sglang.kernels.ops.layernorm:_RMSNORM.forward_jit"
def test_selector_explicit_backend(self):
spec = self.K.select_kernel(
"layernorm.rmsnorm", backend=self.K.KernelBackend.JIT
)
self.assertEqual(
spec.target, "sglang.kernels.ops.layernorm:_RMSNORM.forward_jit"
)
def test_selector_unknown_op_raises(self):
with self.assertRaises(KeyError):
self.K.select_kernel("does_not.exist")
with self.assertRaises(KeyError):
self.K.select_kernel(
"gemm.fp8_scaled_mm", backend=self.K.KernelBackend.TRITON
)
@pytest.mark.parametrize("device, expect", [("cuda", "jit"), ("hip", "aot")])
def test_activation_default_backend(monkeypatch, device, expect):
# silu_and_mul default matches production: jit on CUDA, aot (sgl_kernel) on HIP.
from sglang.kernels.ops.activation import _SILU_AND_MUL
def test_capability_requirement_logic(self):
cap = self.K.CapabilityRequirement
dev = self.K.DeviceType
plat = self.K.PlatformInfo
cpu = plat(device_type="cpu")
sm90 = plat(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
sm100 = plat(device_type="cuda", cuda_arch_major=10, cuda_arch_minor=0)
hip = plat(device_type="hip")
monkeypatch.setattr(fo, "_platform", lambda: PlatformInfo(device_type=device))
assert _SILU_AND_MUL._resolve_backend().value == expect
self.assertFalse(cap(device=dev.CUDA).is_satisfied_by(cpu))
self.assertTrue(cap(device=dev.CUDA).is_satisfied_by(sm90))
self.assertFalse(cap(device=dev.CUDA).is_satisfied_by(hip))
self.assertTrue(cap(device=dev.HIP).is_satisfied_by(hip))
self.assertFalse(
cap(device=dev.CUDA, min_cuda_arch=(10, 0)).is_satisfied_by(sm90)
)
self.assertTrue(
cap(device=dev.CUDA, min_cuda_arch=(10, 0)).is_satisfied_by(sm100)
)
self.assertFalse(
cap(device=dev.CUDA, max_cuda_arch=(9, 0)).is_satisfied_by(sm100)
)
# OR semantics: a {cuda, hip} set is satisfied by either device.
cuda_or_hip = {cap.CUDA, cap.HIP}
self.assertTrue(self.K.capabilities_satisfied(cuda_or_hip, sm90))
self.assertTrue(self.K.capabilities_satisfied(cuda_or_hip, hip))
self.assertFalse(self.K.capabilities_satisfied(cuda_or_hip, cpu))
self.assertTrue(self.K.capabilities_satisfied((), cpu)) # empty = unrestricted
# single requirement is tolerated (pre-decouple API used one).
self.assertTrue(self.K.capabilities_satisfied(cap.CUDA, sm90))
@pytest.mark.parametrize(
"op_attr, device, expect",
[
("_RMSNORM", "cuda", "aot"),
("_RMSNORM", "hip", "aiter"),
("_RMSNORM", "npu", "torch_npu"),
("_GEMMA_RMSNORM", "cuda", "aot"),
("_GEMMA_RMSNORM", "hip", "jit"), # rocm-triton JIT pinned to HIP
("_GEMMA_RMSNORM", "npu", "torch_npu"),
],
)
def test_layernorm_default_backend(monkeypatch, op_attr, device, expect):
# Same AOT provenance, different device coverage per op: rmsnorm's AOT is
# CUDA-only, so HIP falls to aiter and NPU to torch_npu.
ln = importlib.import_module("sglang.kernels.ops.layernorm")
monkeypatch.setattr(fo, "_platform", lambda: PlatformInfo(device_type=device))
assert getattr(ln, op_attr)._resolve_backend().value == expect
# Class-constant shortcuts equal their explicit form; sets are unordered
# and dedup, so {CUDA, HIP} == {HIP, CUDA}.
self.assertEqual(cap.CUDA, cap(device=dev.CUDA))
self.assertEqual(cap.HIP, cap(device=dev.HIP))
self.assertEqual(cap.NPU, cap(device=dev.NPU))
self.assertEqual({cap.CUDA, cap.HIP}, {cap.HIP, cap.CUDA})
self.assertEqual(len({cap.CUDA, cap(device=dev.CUDA)}), 1)
# cuda(min_sm=...) factory: an SM100+ CUDA requirement.
self.assertEqual(
cap.cuda(min_sm=(10, 0)),
cap(device=dev.CUDA, min_cuda_arch=(10, 0)),
)
self.assertTrue(cap.cuda(min_sm=(10, 0)).is_satisfied_by(sm100))
self.assertFalse(cap.cuda(min_sm=(10, 0)).is_satisfied_by(sm90))
def test_platform_detect_does_not_raise(self):
plat = self.K.PlatformInfo.detect()
self.assertIn(plat.device_type, ("cpu", "cuda", "hip"))
def test_per_op_backend_subset():
# silu_and_mul ships an aiter (HIP) kernel; the gelu siblings deliberately
# do not -- ROCm coverage is a per-(op, backend) subset.
from sglang.kernels.ops.activation import _GELU_AND_MUL, _SILU_AND_MUL
def test_import_does_not_load_kernel_backends(self):
# Importing the namespace must stay metadata-only: no sgl_kernel or
# sglang.jit_kernel import, and no JIT compilation, on a CPU box.
code = (
"import sys; import sglang.kernels.ops; "
"backend = ('sgl_kernel' in sys.modules) or "
"any(m.startswith('sglang.jit_kernel') for m in sys.modules); "
"print('BACKEND_IMPORTED' if backend else 'CLEAN')"
)
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("CLEAN", result.stdout, result.stdout + result.stderr)
assert KernelBackend.AITER in _SILU_AND_MUL.available_backends()
assert KernelBackend.AITER not in _GELU_AND_MUL.available_backends()
@pytest.mark.parametrize(
"req, plat, ok",
[
(Cap.CUDA, _CPU, False),
(Cap.CUDA, _SM90, True),
(Cap.CUDA, _HIP, False),
(Cap.HIP, _HIP, True),
(Cap.cuda(min_sm=(10, 0)), _SM90, False),
(Cap.cuda(min_sm=(10, 0)), _SM100, True),
(Cap.cuda(max_sm=(9, 0)), _SM100, False),
],
)
def test_capability_is_satisfied_by(req, plat, ok):
assert req.is_satisfied_by(plat) is ok
def test_capabilities_or_semantics():
both = {Cap.CUDA, Cap.HIP}
assert K.capabilities_satisfied(both, _SM90)
assert K.capabilities_satisfied(both, _HIP)
assert not K.capabilities_satisfied(both, _CPU)
assert K.capabilities_satisfied((), _CPU) # empty = unrestricted
assert K.capabilities_satisfied(Cap.CUDA, _SM90) # single tolerated
def test_capability_shortcuts():
assert Cap.CUDA == Cap(device=DeviceType.CUDA)
assert Cap.HIP == Cap(device=DeviceType.HIP)
assert Cap.NPU == Cap(device=DeviceType.NPU)
assert {Cap.CUDA, Cap.HIP} == {Cap.HIP, Cap.CUDA}
assert Cap.cuda(min_sm=(10, 0)) == Cap(
device=DeviceType.CUDA, min_cuda_arch=(10, 0)
)
def test_platform_detect_does_not_raise():
assert PlatformInfo.detect().device_type in ("cpu", "cuda", "hip", "npu")
def test_import_stays_metadata_only():
# Importing the namespace must not pull in sgl_kernel / sglang.jit_kernel.
code = (
"import sys, sglang.kernels.ops; "
"print('DIRTY' if 'sgl_kernel' in sys.modules or any("
"m.startswith('sglang.jit_kernel') for m in sys.modules) else 'CLEAN')"
)
r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
assert r.returncode == 0, r.stderr
assert "CLEAN" in r.stdout
if __name__ == "__main__":
unittest.main()
sys.exit(pytest.main([__file__]))