/test/registered/kernels/test_fused_op.py

"""GPU-free unit tests for ``sglang.kernels``: BaseFusedOp + registry/selector/spec.

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
``test_fused_op_gpu_parity.py``.
"""

import unittest

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.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."""

    op = "test.toy_add"
    priority = (KernelBackend.TRITON, KernelBackend.TORCH)

    def forward_native(self, a, b):
        return a + b

    def forward_triton(self, a, b):
        # Marker so tests can tell which backend ran.
        return a + b + 1000


class _CudaOnlyToyOp(BaseFusedOp):
    """Toy op whose optimized backend requires CUDA (never eligible on CPU)."""

    op = "test.toy_cuda_only"
    priority = (KernelBackend.CUDA_AOT, KernelBackend.TORCH)
    capabilities = {KernelBackend.CUDA_AOT: CapabilityRequirement(requires_cuda=True)}

    def forward_native(self, a):
        return a * 2

    def forward_cuda_aot(self, a):
        raise AssertionError("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.CUDA_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.CUDA_JIT,
                KernelBackend.CUDA_AOT,
            },
        )
        # 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):
    """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)


class TestKernelSpecUnit(unittest.TestCase):
    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):
        spec = KernelSpec(
            op="g.n",
            backend=KernelBackend.TORCH,
            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")
        with self.assertRaises(ValueError):
            spec.load()


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)


class TestNativeReferenceImplementations(unittest.TestCase):
    """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)
        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(self):
        from sglang.kernels.ops.layernorm import _FUSED_ADD_RMSNORM

        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

        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))

    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)


if __name__ == "__main__":
    unittest.main()