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

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

test/registered/kernels/test_fused_op.py

@@ -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
+Every-backend-vs-native parity on real hardware lives in
-(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
 ``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 (
+from sglang.kernels.spec import CapabilityRequirement as Cap
-    CapabilityRequirement,
+from sglang.kernels.spec import KernelBackend, KernelSpec
-    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):
+class _ToyAdd(BaseFusedOp):
-    """Toy op: element-wise a + b, with a fake 'triton' backend."""
-
     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  # marker so tests can tell which backend ran
-        return a + b + 1000
 
 
-class _CudaOnlyToyOp(BaseFusedOp):
+class _CudaOnlyToy(BaseFusedOp):
-    """Toy op whose optimized backend requires CUDA (never eligible on CPU)."""
-
     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):
+@pytest.fixture(autouse=True)
-    def tearDown(self):
+def _reset_global_state():
+    yield
     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())
+def _t(*vals):
-        self.assertEqual(
+    return torch.tensor(list(vals))
-            backends,
+
-            {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE, KernelBackend.TRITON},
+
+def test_available_backends():
+    assert set(_ToyAdd().available_backends()) == {
+        KernelBackend.TORCH,
+        KernelBackend.TORCH_COMPILE,
+        KernelBackend.TRITON,
+    }
+
+
+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_explicit_backend_overrides_priority():
+    assert (
+        _ToyAdd().forward(_t(1.0), _t(2.0), backend=KernelBackend.TORCH).item() == 3.0
     )
 
-    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):
+def test_capability_gates_eligibility():
-        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")
+        pytest.skip("requires a CPU-only environment")
-        self.assertEqual(op(torch.tensor([3.0])).item(), 6.0)
+    # 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(self):
+
-        op = _ToyAddOp()
+def test_forced_backend_global_switch():
-        a, b = torch.tensor([1.0]), torch.tensor([2.0])
+    op = _ToyAdd()
     K.set_fused_op_backend(KernelBackend.TORCH)
-        self.assertEqual(op(a, b).item(), 3.0)
+    assert op(_t(1.0), _t(2.0)).item() == 3.0
     K.set_fused_op_backend(None)
-        self.assertEqual(op(a, b).item(), 1003.0)
+    assert op(_t(1.0), _t(2.0)).item() == 1003.0
 
-    def test_forced_backend_env_var(self):
+
-        import sglang.kernels.fused_op as fused_op_module
+def test_forced_backend_env_var():
+    import sglang.kernels.fused_op as m
     from sglang.srt.environ import envs
 
-        op = _ToyAddOp()
+    op = _ToyAdd()
-        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.
+        m._forced_backend = m._UNRESOLVED  # drop cache so the env var is re-read
-            fused_op_module._forced_backend = fused_op_module._UNRESOLVED
+        assert K.get_fused_op_backend() is KernelBackend.TORCH
-            self.assertEqual(K.get_fused_op_backend(), KernelBackend.TORCH)
+        assert op(_t(1.0), _t(2.0)).item() == 3.0
-            self.assertEqual(op(a, b).item(), 3.0)
+    m._forced_backend = m._UNRESOLVED
-        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):
+def test_unimplemented_backend_raises():
-        op = _ToyAddOp()
+    with pytest.raises(NotImplementedError):
-        a, b = torch.tensor([1.0]), torch.tensor([2.0])
+        _ToyAdd().forward(_t(1.0), _t(2.0), backend=KernelBackend.AOT)
-        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()
+def test_trace_records_op_backend_and_shapes():
     K.enable_fused_op_trace()
-        op(torch.zeros(2, 3), torch.zeros(2, 3))
+    _ToyAdd()(torch.zeros(2, 3), torch.zeros(2, 3))
-        records = K.get_fused_op_trace()
+    (rec,) = K.get_fused_op_trace()
-        self.assertEqual(len(records), 1)
+    assert rec.op == "test.toy_add"
-        self.assertEqual(records[0].op, "test.toy_add")
+    assert rec.backend == "triton"
-        self.assertEqual(records[0].backend, "triton")
+    assert rec.tensor_args == ("torch.float32[2, 3]", "torch.float32[2, 3]")
-        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")
+def test_fused_op_registers_all_backends():
-        backends = {s.backend for s in op}
+    backends = {s.backend for s in K.registry.get("layernorm.rmsnorm")}
-        self.assertEqual(
+    assert backends == {
-            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)))
 
 
-class TestKernelRegistryUnit(unittest.TestCase):
+# --- KernelRegistry unit ---
-    """Isolated KernelRegistry behavior (fresh instance, no global state)."""
 
-    def _spec(self, op="g.n", backend=KernelBackend.TORCH, target="math:sqrt"):
+
+def _spec(op="g.n", backend=KernelBackend.TORCH, target="math:sqrt"):
     return KernelSpec(op=op, backend=backend, target=target)
 
-    def test_register_and_get(self):
+
+def test_registry_register_and_get():
     reg = KernelRegistry()
-        spec = self._spec()
+    spec = _spec()
     reg.register(spec)
-        self.assertEqual(reg.get("g.n"), [spec])
+    assert reg.get("g.n") == [spec]
-        self.assertTrue(reg.has("g.n"))
+    assert reg.has("g.n")
-        self.assertEqual(reg.ops(), ["g.n"])
+    assert reg.ops() == ["g.n"]
+    assert reg.get("no.such") == []
 
-    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):
+def test_registry_reregister_replaces():
     reg = KernelRegistry()
-        reg.register(self._spec(target="math:sqrt"))
+    reg.register(_spec(target="math:sqrt"))
-        reg.register(self._spec(target="math:floor"))
+    reg.register(_spec(target="math:floor"))
-        specs = reg.get("g.n")
+    assert [s.target for s in reg.get("g.n")] == ["math:floor"]
-        self.assertEqual(len(specs), 1)
-        self.assertEqual(specs[0].target, "math:floor")
 
-    def test_get_backend_missing_raises(self):
+
+def test_registry_get_backend_missing_raises():
     reg = KernelRegistry()
-        reg.register(self._spec(backend=KernelBackend.TORCH))
+    reg.register(_spec())
-        with self.assertRaises(KeyError):
+    with pytest.raises(KeyError):
         reg.get_backend("g.n", KernelBackend.TRITON)
-        with self.assertRaises(KeyError):
-            reg.get_backend("no.such", KernelBackend.TORCH)
 
 
-class TestKernelSpecUnit(unittest.TestCase):
+# --- KernelSpec.load ---
-    def test_load_simple_target(self):
-        import math
 
-        spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="math:sqrt")
-        self.assertIs(spec.load(), math.sqrt)
 
-    def test_load_dotted_target(self):
+def test_spec_load_simple_and_dotted():
-        spec = KernelSpec(
+    assert _spec(target="math:sqrt").load() is math.sqrt
-            op="g.n",
+    dotted = _spec(target="sglang.kernels.ops.layernorm:_RMSNORM.forward_native")
-            backend=KernelBackend.TORCH,
+    assert callable(dotted.load())
-            target="sglang.kernels.ops.layernorm:_RMSNORM.forward_native",
-        )
-        self.assertTrue(callable(spec.load()))
 
-    def test_load_bad_target_raises(self):
+
-        spec = KernelSpec(op="g.n", backend=KernelBackend.TORCH, target="no-colon")
+def test_spec_load_bad_target_raises():
-        with self.assertRaises(ValueError):
+    with pytest.raises(ValueError):
-            spec.load()
+        _spec(target="no-colon").load()
+
+
+# --- native reference math of the reworked ops (CPU) ---
 
 
 def _ref_rmsnorm(x, w, eps):
-    xf = x.to(torch.float32)
+    xf = x.float()
-    var = xf.pow(2).mean(dim=-1, keepdim=True)
+    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w).to(x.dtype)
-    return (xf * torch.rsqrt(var + eps) * w).to(x.dtype)
 
 
-class TestNativeReferenceImplementations(unittest.TestCase):
+def test_rmsnorm_native():
-    """The forward_native math of the reworked ops, on CPU tensors."""
-
-    def setUp(self):
-        torch.manual_seed(0)
-
-    def test_rmsnorm_native(self):
     from sglang.kernels.ops.layernorm import _RMSNORM
 
-        x = torch.randn(8, 128)
+    x, w = torch.randn(8, 128), torch.randn(128)
-        w = torch.randn(128)
     out = _RMSNORM.forward_native(x, w, 1e-6)
-        self.assertTrue(torch.allclose(out, _ref_rmsnorm(x, w, 1e-6)))
+    assert 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)
+    assert _RMSNORM.forward_native(x, w, 1e-6, out=buf) is buf  # out= is in place
-        self.assertTrue(torch.allclose(buf, out))
+    assert torch.allclose(buf, out)
 
-    def test_fused_add_rmsnorm_native(self):
+
+def test_fused_add_rmsnorm_native():
     from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM
 
-        x = torch.randn(8, 128)
+    x, residual, w = torch.randn(8, 128), torch.randn(8, 128), torch.randn(128)
-        residual = torch.randn(8, 128)
+    x0, r0 = x.clone(), residual.clone()
-        w = torch.randn(128)
+    assert _FUSED_ADD_RMSNORM.forward_native(x, residual, w, 1e-6) is None
-        x2, r2 = x.clone(), residual.clone()
+    acc = x0.float() + r0.float()
-        self.assertIsNone(_FUSED_ADD_RMSNORM.forward_native(x, residual, w, 1e-6))
+    assert torch.allclose(residual, acc)
-        acc = x2.to(torch.float32) + r2.to(torch.float32)
+    assert torch.allclose(
-        self.assertTrue(torch.allclose(residual, acc))
+        x, acc * torch.rsqrt(acc.pow(2).mean(-1, keepdim=True) + 1e-6) * w
-        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
 
-        x = torch.randn(8, 128)
+@pytest.mark.parametrize(
-        w = torch.randn(128)
+    "op_attr, approximate",
-        out = _GEMMA_RMSNORM.forward_native(x, w, 1e-6)
+    [
-        xf = x.to(torch.float32)
+        ("_SILU_AND_MUL", None),
-        ref = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) * (1.0 + w)
+        ("_GELU_AND_MUL", "none"),
-        self.assertTrue(torch.allclose(out, ref))
+        ("_GELU_TANH_AND_MUL", "tanh"),
-
+    ],
-    def test_gated_activations_native(self):
+)
+def test_gated_activation_native(op_attr, approximate):
     import torch.nn.functional as F
 
-        from sglang.kernels.ops.activation import (
+    import sglang.kernels.ops.activation as act
-            _GELU_AND_MUL,
-            _GELU_TANH_AND_MUL,
-            _SILU_AND_MUL,
-        )
 
     x = torch.randn(8, 256)
     gate, up = x[..., :128], x[..., 128:]
-        cases = [
+    ref = (
-            (_SILU_AND_MUL, F.silu(gate) * up),
+        F.silu(gate) if approximate is None else F.gelu(gate, approximate=approximate)
-            (_GELU_AND_MUL, F.gelu(gate, approximate="none") * up),
+    ) * up
-            (_GELU_TANH_AND_MUL, F.gelu(gate, approximate="tanh") * up),
+    assert torch.allclose(getattr(act, op_attr).forward_native(x), ref)
-        ]
-        for op, ref in cases:
-            self.assertTrue(torch.allclose(op.forward_native(x), ref), op.op)
 
 
 if __name__ == "__main__":
-    unittest.main()
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/kernels/test_fused_op_gpu_parity.py

@@ -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
+For each reworked fused op, run every backend eligible on this platform and
-available backends, run each one that is eligible on this platform, and
+assert it matches the pure-torch ``forward_native`` reference within dtype
-assert the output matches the pure-torch ``forward_native`` reference within
+tolerance. New backends are picked up automatically.
-dtype tolerance. New backends added to an op are picked up automatically —
-no per-kernel test boilerplate.
 """
 
-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; skip it here (compile time dominates)
-# torch_compile is native under the hood; exclude it from the sweep to keep
+# -- it is exercised in the CPU lane.
-# CI time down (compilation dominates) — it is exercised in the CPU lane.
+_SKIP = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE}
-_SKIP_BACKENDS = {KernelBackend.TORCH, KernelBackend.TORCH_COMPILE}
+_TOL = {
-_TOLERANCE = {
     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):
+def _eligible(op):
     return [
-            b
+        b for b in op.available_backends() if b not in _SKIP and op.backend_eligible(b)
-            for b in op.available_backends()
-            if b not in _SKIP_BACKENDS 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):
+def _close(got, ref, dtype, msg):
+    torch.testing.assert_close(got, ref, **_TOL[dtype], msg=msg)
+
+
+@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):
+    torch.manual_seed(0)
-            for shape in ((1, 4096), (128, 4096), (7, 2048)):
+    x = torch.randn(shape, dtype=dtype, device="cuda")
-                x = torch.randn(shape, dtype=dtype, device=_DEVICE)
+    w = torch.randn(shape[-1], dtype=dtype, device="cuda")
-                w = torch.randn(shape[-1], dtype=dtype, device=_DEVICE)
     ref = _RMSNORM.forward_native(x, w, 1e-6)
-                for backend in self._eligible_backends(_RMSNORM):
+    for b in _eligible(_RMSNORM):
-                    got = _RMSNORM.forward(x, w, 1e-6, backend=backend)
+        _close(
-                    self._assert_close(
+            _RMSNORM.forward(x, w, 1e-6, backend=b), ref, dtype, f"rmsnorm {b.value}"
-                        got, ref, dtype, f"rmsnorm {backend.value} {dtype} {shape}"
         )
 
-    def test_fused_add_rmsnorm_backends_match_native(self):
+
+@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
 
-        for dtype in (torch.float16, torch.bfloat16):
+    torch.manual_seed(0)
-            for shape in ((1, 4096), (128, 4096)):
+    x0 = torch.randn(shape, dtype=dtype, device="cuda")
-                x0 = torch.randn(shape, dtype=dtype, device=_DEVICE)
+    r0 = torch.randn(shape, dtype=dtype, device="cuda")
-                r0 = torch.randn(shape, dtype=dtype, device=_DEVICE)
+    w = torch.randn(shape[-1], dtype=dtype, device="cuda")
-                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):
+    for b in _eligible(_FUSED_ADD_RMSNORM):
         x, r = x0.clone(), r0.clone()
-                    _FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=backend)
+        _FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=b)
-                    label = f"fused_add_rmsnorm {backend.value} {dtype} {shape}"
+        _close(x, x_ref, dtype, f"fused_add {b.value} (normed)")
-                    self._assert_close(x, x_ref, dtype, label + " (normed)")
+        _close(r, r_ref, dtype, f"fused_add {b.value} (residual)")
-                    self._assert_close(r, r_ref, dtype, label + " (residual)")
 
-    def test_gemma_rmsnorm_backends_match_native(self):
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_gemma_rmsnorm(dtype):
     from sglang.kernels.ops.layernorm import _GEMMA_RMSNORM
 
-        for dtype in (torch.float16, torch.bfloat16):
+    torch.manual_seed(0)
-            x = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
+    x = torch.randn(64, 2048, dtype=dtype, device="cuda")
-            w = torch.randn(2048, dtype=dtype, device=_DEVICE)
+    w = torch.randn(2048, dtype=dtype, device="cuda")
     ref = _GEMMA_RMSNORM.forward_native(x, w, 1e-6)
-            for backend in self._eligible_backends(_GEMMA_RMSNORM):
+    for b in _eligible(_GEMMA_RMSNORM):
-                got = _GEMMA_RMSNORM.forward(x, w, 1e-6, backend=backend)
+        _close(
-                self._assert_close(
+            _GEMMA_RMSNORM.forward(x, w, 1e-6, backend=b),
-                    got, ref, dtype, f"gemma_rmsnorm {backend.value} {dtype}"
+            ref,
+            dtype,
+            f"gemma {b.value}",
         )
 
-    def test_gemma_fused_add_rmsnorm_backends_match_native(self):
+
+@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
 
-        for dtype in (torch.float16, torch.bfloat16):
+    torch.manual_seed(0)
-            x0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
+    x0 = torch.randn(64, 2048, dtype=dtype, device="cuda")
-            r0 = torch.randn(64, 2048, dtype=dtype, device=_DEVICE)
+    r0 = torch.randn(64, 2048, dtype=dtype, device="cuda")
-            w = torch.randn(2048, dtype=dtype, device=_DEVICE)
+    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 backend in self._eligible_backends(_GEMMA_FUSED_ADD_RMSNORM):
+    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=backend)
+        _GEMMA_FUSED_ADD_RMSNORM.forward(x, r, w, 1e-6, backend=b)
-                label = f"gemma_fused_add_rmsnorm {backend.value} {dtype}"
+        _close(x, x_ref, dtype, f"gemma_fused_add {b.value} (normed)")
-                self._assert_close(x, x_ref, dtype, label + " (normed)")
+        _close(r, r_ref, dtype, f"gemma_fused_add {b.value} (residual)")
-                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,
-        )
 
-        for op in (_SILU_AND_MUL, _GELU_AND_MUL, _GELU_TANH_AND_MUL):
+@pytest.mark.parametrize(
-            for dtype in (torch.float16, torch.bfloat16):
+    "op_attr", ["_SILU_AND_MUL", "_GELU_AND_MUL", "_GELU_TANH_AND_MUL"]
-                for shape in ((1, 8192), (128, 8192)):
+)
-                    x = torch.randn(shape, dtype=dtype, device=_DEVICE)
+@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 backend in self._eligible_backends(op):
+    for b in _eligible(op):
-                        got = op.forward(x, backend=backend)
+        _close(op.forward(x, backend=b), ref, dtype, f"{op.op} {b.value}")
-                        self._assert_close(
-                            got, ref, dtype, f"{op.op} {backend.value} {dtype} {shape}"
-                        )
 
 
 if __name__ == "__main__":
-    unittest.main()
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/kernels/test_kernels_namespace.py

@@ -1,142 +1,22 @@
-"""GPU-free import/registry tests for the ``sglang.kernels`` namespace.
+"""GPU-free import / registry / selector tests for ``sglang.kernels`` (RFC #29630)."""
-
-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.
-"""
 
+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
+GROUPS = [
-# 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 = [
     "activation",
     "attention",
     "communication",
@@ -155,18 +35,25 @@ 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):
+_CPU = PlatformInfo(device_type="cpu")
-    def setUp(self):
+_SM90 = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
-        import importlib
+_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
+def test_top_level_exports():
-        self.importlib = importlib
-
-    def test_top_level_exports(self):
     for name in (
         "KernelSpec",
         "KernelBackend",
@@ -177,219 +64,135 @@ class TestKernelsNamespace(unittest.TestCase):
         "get_kernel",
         "select_kernel",
     ):
-            self.assertTrue(hasattr(self.K, name), f"missing export: {name}")
+        assert hasattr(K, name), 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__"))
 
-    def test_registry_contents(self):
+@pytest.mark.parametrize("group", GROUPS)
-        registry = self.K.registry
+def test_group_importable(group):
-        ops = set(registry.ops())
+    assert hasattr(importlib.import_module(f"sglang.kernels.ops.{group}"), "__all__")
-        # 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)
 
-    def test_wrappers_exposed_and_callable(self):
+@pytest.mark.parametrize("op, backends", list(EXPECTED.items()))
-        for module_name, names in EXPECTED_WRAPPERS.items():
+def test_registry_backends(op, backends):
-            mod = self.importlib.import_module(module_name)
+    assert {s.backend.value for s in K.registry.get(op)} == backends
-            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_multi_backend_op_requires_explicit_backend(self):
+def test_specs_well_formed():
-        # Device is a HARD eligibility filter, not a preference ranking: when
+    for spec in K.registry.all_specs():
-        # more than one backend is usable on the current device, selection must
+        assert spec.op == f"{spec.group}.{spec.name}"
-        # be explicit (no hidden auto-ranking). Force a CUDA platform so the
+        mod, sep, attr = spec.target.partition(":")
-        # result is deterministic regardless of the test host.
+        assert sep == ":" and mod and attr, spec.target
-        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_decoupled_backend_device_selection(self):
+def test_single_backend_resolves_without_backend():
-        # Proves the decoupled backend/device model against production reality:
+    assert K.select_kernel("gemm.fp8_scaled_mm").backend is KernelBackend.AOT
-        #  - 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).
+def test_unknown_op_or_backend_raises():
-        # Auto-selection matches production defaults: JIT on CUDA, AOT on HIP.
+    with pytest.raises(KeyError):
-        import sglang.kernels.fused_op as fo
+        K.select_kernel("does_not.exist")
+    with pytest.raises(KeyError):
+        K.select_kernel("gemm.fp8_scaled_mm", backend=KernelBackend.TRITON)
+
+
+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"
+
+
+@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
+
+    monkeypatch.setattr(fo, "_platform", lambda: PlatformInfo(device_type=device))
+    assert _SILU_AND_MUL._resolve_backend().value == expect
+
+
+@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
+
+
+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
 
-        B = self.K.KernelBackend
+    assert KernelBackend.AITER in _SILU_AND_MUL.available_backends()
-        hip = self.K.PlatformInfo(device_type="hip")
+    assert KernelBackend.AITER not in _GELU_AND_MUL.available_backends()
-        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_layernorm_cross_device_coverage(self):
+
-        # The rmsnorm ops illustrate that the *same* provenance covers different
+@pytest.mark.parametrize(
-        # devices per op: AOT (sgl_kernel) is CUDA-only here (sgl_kernel does not
+    "req, plat, ok",
-        # 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
+        (Cap.CUDA, _CPU, False),
-        # rocm-triton JIT path on HIP -- a JIT provenance pinned to HIP, unlike
+        (Cap.CUDA, _SM90, True),
-        # the CUDA-only JIT on plain rmsnorm.
+        (Cap.CUDA, _HIP, False),
-        import sglang.kernels.fused_op as fo
+        (Cap.HIP, _HIP, True),
-        from sglang.kernels.ops.layernorm import (
+        (Cap.cuda(min_sm=(10, 0)), _SM90, False),
-            _FUSED_ADD_RMSNORM,
+        (Cap.cuda(min_sm=(10, 0)), _SM100, True),
-            _GEMMA_RMSNORM,
+        (Cap.cuda(max_sm=(9, 0)), _SM100, False),
-            _RMSNORM,
+    ],
+)
+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)
     )
 
-        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_selector_explicit_backend(self):
+def test_platform_detect_does_not_raise():
-        spec = self.K.select_kernel(
+    assert PlatformInfo.detect().device_type in ("cpu", "cuda", "hip", "npu")
-            "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
-            )
 
-    def test_capability_requirement_logic(self):
+def test_import_stays_metadata_only():
-        cap = self.K.CapabilityRequirement
+    # Importing the namespace must not pull in sgl_kernel / sglang.jit_kernel.
-        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")
-
-        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))
-
-        # 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_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; "
+        "import sys, sglang.kernels.ops; "
-            "backend = ('sgl_kernel' in sys.modules) or "
+        "print('DIRTY' if 'sgl_kernel' in sys.modules or any("
-            "any(m.startswith('sglang.jit_kernel') for m in sys.modules); "
+        "m.startswith('sglang.jit_kernel') for m in sys.modules) else 'CLEAN')"
-            "print('BACKEND_IMPORTED' if backend else 'CLEAN')"
     )
-        result = subprocess.run(
+    r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
-            [sys.executable, "-c", code],
+    assert r.returncode == 0, r.stderr
-            capture_output=True,
+    assert "CLEAN" in r.stdout
-            text=True,
-        )
-        self.assertEqual(result.returncode, 0, result.stderr)
-        self.assertIn("CLEAN", result.stdout, result.stdout + result.stderr)
 
 
 if __name__ == "__main__":
-    unittest.main()
+    sys.exit(pytest.main([__file__]))