diff --git a/test/registered/kernels/test_fused_op.py b/test/registered/kernels/test_fused_op.py index b0729d1a2..83b9b37d6 100644 --- a/test/registered/kernels/test_fused_op.py +++ b/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 -(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__])) diff --git a/test/registered/kernels/test_fused_op_gpu_parity.py b/test/registered/kernels/test_fused_op_gpu_parity.py index 1178748f7..920474439 100644 --- a/test/registered/kernels/test_fused_op_gpu_parity.py +++ b/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 -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__])) diff --git a/test/registered/kernels/test_kernels_namespace.py b/test/registered/kernels/test_kernels_namespace.py index acc516a31..79f19da23 100644 --- a/test/registered/kernels/test_kernels_namespace.py +++ b/test/registered/kernels/test_kernels_namespace.py @@ -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__]))