[Kernel] Introduce sglang.kernels namespace and migrate scattered triton_ops kernels (RFC #29630, Phase 2) (#30044)

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

test/manual/lora/test_lora_tuning_config.py

@@ -6,7 +6,7 @@ import tempfile
 import unittest
 from unittest.mock import patch
 
-from sglang.srt.lora.triton_ops.lora_tuning_config import (
+from sglang.kernels.ops.gemm.lora_tuning_config import (
     DEFAULT_EXPAND_CONFIG,
     DEFAULT_SHRINK_CONFIG,
     get_lora_config_file_name,
@@ -15,7 +15,7 @@ from sglang.srt.lora.triton_ops.lora_tuning_config import (
     get_lora_shrink_config,
 )
 
-_MODULE = "sglang.srt.lora.triton_ops.lora_tuning_config"
+_MODULE = "sglang.kernels.ops.gemm.lora_tuning_config"
 
 # Shared fixture
 _TUNED_CONFIGS = {
@@ -88,7 +88,7 @@ class TestConfigSelection(unittest.TestCase):
 
     def setUp(self):
         get_lora_configs.cache_clear()
-        from sglang.srt.lora.triton_ops import lora_tuning_config
+        from sglang.kernels.ops.gemm import lora_tuning_config
 
         lora_tuning_config._logged_configs.clear()
 

test/manual/test_triton_attention_rocm_mla.py

@@ -3,10 +3,10 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.decode_attention import (
+from sglang.kernels.ops.attention.decode_attention import (
     decode_attention_fwd_grouped,
 )
-from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import (
+from sglang.kernels.ops.attention.rocm_mla_decode_rope import (
     decode_attention_fwd_grouped_rope,
 )
 from sglang.srt.layers.rotary_embedding import DeepseekScalingRotaryEmbedding

test/manual/test_trtllm_fp8_kv_kernel.py

@@ -6,7 +6,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
+from sglang.kernels.ops.kvcache.trtllm_fp8_kv_kernel import (
     fused_fp8_set_kv_buffer,
 )
 from sglang.test.test_utils import CustomTestCase

test/registered/attention/test_triton_attention_kernels.py

@@ -4,18 +4,18 @@ import unittest
 import torch
 import torch.nn.functional as F
 
-from sglang.srt.layers.attention.triton_ops.decode_attention import (
+from sglang.kernels.ops.attention.decode_attention import (
     decode_attention_fwd,
     decode_attention_fwd_grouped,
     decode_attention_fwd_normal,
 )
-from sglang.srt.layers.attention.triton_ops.extend_attention import (
+from sglang.kernels.ops.attention.extend_attention import (
     build_unified_kv_indices,
     extend_attention_fwd,
     extend_attention_fwd_unified,
     redundant_attention,
 )
-from sglang.srt.layers.attention.triton_ops.prefill_attention import (
+from sglang.kernels.ops.attention.prefill_attention import (
     context_attention_fwd,
 )
 from sglang.srt.utils import get_device
@@ -320,7 +320,7 @@ class TestTritonAttention(CustomTestCase):
             self._test_extend_attention_once(19, 12331, 12, 4, value)
 
     def test_extend_attention_block_sizes(self):
-        from sglang.srt.layers.attention.triton_ops import extend_attention as ea
+        from sglang.kernels.ops.attention import extend_attention as ea
 
         if not ea._is_hip:
             self.skipTest("HIP-only block-size selection")

test/registered/attention/test_trtllm_mha_graph_metadata.py

@@ -12,7 +12,7 @@ import pytest
 import torch
 
 import sglang.srt.layers.attention.trtllm_mha_backend as trtllm_mha_backend
-from sglang.srt.layers.attention.triton_ops.trtllm_mha_graph_metadata import (
+from sglang.kernels.ops.kvcache.trtllm_mha_graph_metadata import (
     Q_MODE_CUMSUM,
     Q_MODE_NONE,
     Q_MODE_STRIDED,

test/registered/attention/test_trtllm_mha_page_table.py

@@ -18,7 +18,7 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import (
+from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
     build_trtllm_mha_page_table,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/attention/test_verify_splitkv.py

@@ -20,10 +20,10 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.extend_attention import (
+from sglang.kernels.ops.attention.extend_attention import (
     extend_attention_fwd,
 )
-from sglang.srt.layers.attention.triton_ops.verify_splitkv import (
+from sglang.kernels.ops.attention.verify_splitkv import (
     can_handle,
     verify_splitkv_fwd,
 )

test/registered/attention/test_wave_attention_kernels.py

@@ -3,14 +3,14 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.decode_attention import (
+from sglang.kernels.ops.attention.decode_attention import (
     decode_attention_fwd_grouped as triton_decode_attention_fwd_grouped,
 )
-from sglang.srt.layers.attention.triton_ops.extend_attention import (
+from sglang.kernels.ops.attention.extend_attention import (
     extend_attention_fwd,
     redundant_attention,
 )
-from sglang.srt.layers.attention.triton_ops.prefill_attention import (
+from sglang.kernels.ops.attention.prefill_attention import (
     context_attention_fwd,
 )
 from sglang.srt.layers.attention.wave_ops.decode_attention import (

test/registered/kernels/test_dsa_metadata.py

@@ -2,7 +2,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
+from sglang.kernels.ops.attention.dsa_metadata import (
     fused_dsa_decode_metadata,
     fused_dsa_draft_extend_metadata,
     fused_dsa_target_verify_metadata,

test/registered/kernels/test_fused_op.py

@@ -0,0 +1,290 @@
+"""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()

test/registered/kernels/test_fused_op_gpu_parity.py

@@ -0,0 +1,125 @@
+"""Generic every-backend-vs-native parity harness for BaseFusedOp operators.
+
+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.
+"""
+
+import unittest
+
+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.float16: dict(atol=1e-2, rtol=1e-2),
+    torch.bfloat16: dict(atol=2e-2, rtol=2e-2),
+}
+
+
+@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 _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
+
+        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
+
+        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,
+        )
+
+        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}"
+                        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/test_gather_spec_extras.py

@@ -7,7 +7,7 @@ import unittest
 
 import torch
 
-from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras
+from sglang.kernels.ops.speculative.gather_spec_extras import gather_spec_extras
 from sglang.test.test_utils import CustomTestCase
 
 _OUTPUT_NAMES = ("topk_p", "topk_index", "bonus_tokens", "hidden_states")

test/registered/kernels/test_kernels_namespace.py

@@ -0,0 +1,256 @@
+"""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.
+"""
+
+import subprocess
+import sys
+import unittest
+
+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": {"cuda_aot", "cuda_jit", "torch", "torch_compile"},
+    "activation.gelu_and_mul": {"cuda_aot", "cuda_jit", "torch", "torch_compile"},
+    "activation.gelu_tanh_and_mul": {
+        "cuda_aot",
+        "cuda_jit",
+        "torch",
+        "torch_compile",
+    },
+    "layernorm.rmsnorm": {"cuda_aot", "cuda_jit", "torch", "torch_compile"},
+    "layernorm.fused_add_rmsnorm": {
+        "cuda_aot",
+        "cuda_jit",
+        "torch",
+        "torch_compile",
+    },
+    "layernorm.gemma_rmsnorm": {"cuda_aot", "torch", "torch_compile"},
+    "layernorm.gemma_fused_add_rmsnorm": {"cuda_aot", "torch", "torch_compile"},
+    # curated dual/single-backend wrapper ops
+    "gemm.fp8_scaled_mm": {"cuda_aot"},
+    "gemm.dsv3_fused_a_gemm": {"cuda_aot", "cuda_jit"},
+    "gemm.dsv3_router_gemm": {"cuda_jit"},
+    "kvcache.reshape_and_cache_flash": {"triton"},
+    "moe.moe_align_block_size": {"cuda_aot", "cuda_jit"},
+    "moe.topk_softmax": {"cuda_aot"},
+    "quantization.sgl_per_token_quant_fp8": {"cuda_aot"},
+    "quantization.sgl_per_token_group_quant_8bit": {"cuda_aot", "cuda_jit"},
+    "quantization.sgl_per_token_group_quant_fp8": {"cuda_aot"},
+    "quantization.sgl_per_token_group_quant_int8": {"cuda_aot"},
+    # deferred-group wrappers, now populated
+    "sampling.top_k_renorm_probs": {"cuda_aot"},
+    "sampling.top_p_renorm_probs": {"cuda_aot"},
+    "spatial.get_sm_available": {"cuda_aot"},
+    "spatial.create_greenctx_stream_by_value": {"cuda_aot"},
+    "mamba.causal_conv1d_fwd": {"cuda_aot"},
+    "mamba.causal_conv1d_update": {"cuda_aot"},
+    "diffusion.apply_group_norm_silu": {"cuda_jit"},
+    "diffusion.residual_gate_add": {"cuda_jit"},
+    "diffusion.fused_inplace_qknorm_rope": {"cuda_jit"},
+    # representative migrated Triton kernels (inventory)
+    "grammar.apply_token_bitmask_inplace_triton": {"triton"},
+    "memory.alloc_extend_kernel": {"triton"},
+    "attention.decode_attention_fwd": {"triton"},
+    "kvcache.create_flashinfer_kv_indices_triton": {"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",
+    "diffusion",
+    "gemm",
+    "grammar",
+    "kvcache",
+    "layernorm",
+    "mamba",
+    "memory",
+    "moe",
+    "quantization",
+    "sampling",
+    "spatial",
+    "speculative",
+]
+
+
+class TestKernelsNamespace(unittest.TestCase):
+    def setUp(self):
+        import importlib
+
+        import sglang.kernels
+        import sglang.kernels.ops  # populate the registry
+
+        self.K = sglang.kernels
+        self.importlib = importlib
+
+    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__"))
+
+    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)
+
+    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_multi_backend_op_requires_explicit_backend(self):
+        # No hidden ranking: a multi-backend op must be resolved explicitly.
+        multi = [op for op, b in EXPECTED_OPS.items() if len(b) > 1]
+        self.assertTrue(multi)  # sanity: we do have multi-backend ops
+        for op in multi:
+            with self.assertRaises(ValueError):
+                self.K.select_kernel(op)
+
+    def test_selector_explicit_backend(self):
+        spec = self.K.select_kernel(
+            "layernorm.rmsnorm", backend=self.K.KernelBackend.CUDA_JIT
+        )
+        self.assertEqual(
+            spec.target, "sglang.kernels.ops.layernorm:_RMSNORM.forward_cuda_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):
+        cap = self.K.CapabilityRequirement
+        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)
+
+        self.assertFalse(cap(requires_cuda=True).is_satisfied_by(cpu))
+        self.assertTrue(cap(requires_cuda=True).is_satisfied_by(sm90))
+        self.assertFalse(
+            cap(requires_cuda=True, min_cuda_arch=(10, 0)).is_satisfied_by(sm90)
+        )
+        self.assertTrue(
+            cap(requires_cuda=True, min_cuda_arch=(10, 0)).is_satisfied_by(sm100)
+        )
+        self.assertFalse(
+            cap(requires_cuda=True, max_cuda_arch=(9, 0)).is_satisfied_by(sm100)
+        )
+
+    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; "
+            "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)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/lora/test_chunked_sgmv_backend.py

@@ -5,19 +5,25 @@ from typing import List, Optional, Tuple
 
 import torch
 
-from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor
-from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
-from sglang.srt.lora.triton_ops import (
+from sglang.kernels.ops.gemm.chunked_embedding_lora_a import (
     chunked_embedding_lora_a_forward,
+)
+from sglang.kernels.ops.gemm.chunked_sgmv_expand import (
+    _chunked_lora_expand_kernel,
     chunked_sgmv_lora_expand_forward,
+)
+from sglang.kernels.ops.gemm.chunked_sgmv_shrink import (
+    _chunked_lora_shrink_kernel,
     chunked_sgmv_lora_shrink_forward,
+)
+from sglang.kernels.ops.gemm.kv_b_lora_absorbed import (
     step_a_q_fwd,
     step_a_v_fwd,
     step_b_q_fwd,
     step_b_v_fwd,
 )
-from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel
-from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel
+from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor
+from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
 from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.test.ci.ci_register import register_cuda_ci

test/registered/lora/test_fused_moe_lora_kernel.py

@@ -9,7 +9,7 @@ import torch
 # IMPORT PREBUILT KERNEL
 # ==============================================================================
 from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
-from sglang.srt.lora.triton_ops import fused_moe_lora
+from sglang.kernels.ops.moe.fused_moe_lora_kernel import fused_moe_lora
 from sglang.srt.utils import set_random_seed
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/lora/test_virtual_experts_kernels.py

@@ -32,7 +32,7 @@ from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small")
 
-from sglang.srt.lora.triton_ops.virtual_experts import (
+from sglang.kernels.ops.moe.virtual_experts import (
     _align_block_size_jit,
     _align_block_size_torch,
     _fused_virtual_topk_ids,

test/registered/unit/constrained/test_token_filter_ops.py

@@ -19,7 +19,7 @@ register_cpu_ci(2.0, "base-a-test-cpu")
 # Conditionally import Triton path
 _has_cuda = torch.cuda.is_available()
 if _has_cuda:
-    from sglang.srt.constrained.triton_ops.token_filter_ops import (
+    from sglang.kernels.ops.grammar.token_filter_ops import (
         set_token_filter_triton,
     )
 

test/registered/unit/mem_cache/test_store_cache_4d.py

@@ -111,7 +111,7 @@ class TestStoreCache4D(unittest.TestCase):
         dtype: torch.dtype = torch.bfloat16,
         loc_dtype: torch.dtype = torch.int64,
     ):
-        from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
+        from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
 
         # Two independent target buffers — one for the kernel, one for the
         # legacy reference path.
@@ -219,7 +219,7 @@ class TestStoreCache4D(unittest.TestCase):
     def test_store_cache_4d_empty_loc(self):
         """N=0 must be a no-op: no kernel launch, no exception, no buffer
         mutation."""
-        from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
+        from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
 
         k_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
         v_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
@@ -259,7 +259,7 @@ class TestStoreCache4DAssertions(unittest.TestCase):
         """Wrapper requires `stride[-1] == 1` and `stride[-2] == head_dim`
         (the trailing two dims must be contiguous). A permutation that
         breaks this should trigger AssertionError."""
-        from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
+        from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
 
         # Build a 4-D view, then permute the last two dims → trailing
         # contiguity violated.
@@ -278,7 +278,7 @@ class TestStoreCache4DAssertions(unittest.TestCase):
     def test_rejects_dtype_mismatch(self):
         """All four tensors must share a dtype; the caller is responsible
         for any cast before the call."""
-        from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
+        from sglang.kernels.ops.kvcache.cache_move import store_cache_4d
 
         k_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
         v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")

test/registered/unit/mem_cache/test_triton_kernel_layout.py

@@ -57,7 +57,7 @@ class TestTritonKernelLayoutParity(unittest.TestCase):
         return q, logical_kv_k, logical_kv_v, kv_indptr, kv_indices, seq_len
 
     def _run_decode(self, q, k_buf, v_buf, kv_indptr, kv_indices, page_size):
-        from sglang.srt.layers.attention.triton_ops.decode_attention import (
+        from sglang.kernels.ops.attention.decode_attention import (
             decode_attention_fwd,
         )
 
@@ -115,7 +115,7 @@ class TestTritonKernelLayoutParity(unittest.TestCase):
 
     def test_extend_3d_vs_4d_ps1_byte_identical(self):
         """Same parity check for extend kernel."""
-        from sglang.srt.layers.attention.triton_ops.extend_attention import (
+        from sglang.kernels.ops.attention.extend_attention import (
             extend_attention_fwd,
         )