Migrate Intel CPU cases to the test/registered (#25139)

.github/workflows/pr-test-xeon.yml

@@ -115,7 +115,7 @@ jobs:
         timeout-minutes: 36
         run: |
           docker exec -w /sglang-checkout/ ci_sglang_xeon \
-            bash -c "source /opt/.venv/bin/activate && cd ./test/srt && python3 run_suite.py --suite per-commit-cpu --timeout-per-file 1500"
+            bash -c "source /opt/.venv/bin/activate && cd ./test && python3 run_suite.py --hw cpu --suite stage-b-test-cpu"
 
       - name: Change permission
         timeout-minutes: 2

.pre-commit-config.yaml

@@ -92,6 +92,7 @@ repos:
         entry: python3 scripts/ci/check_registered_tests.py
         language: system
         files: ^test/registered/.*\.py$
+        exclude: ^test/registered/.*/utils\.py$
         pass_filenames: false
       - id: check-no-docs-changes
         name: reject changes under legacy docs/

scripts/ci/check_registered_tests.py

@@ -22,11 +22,11 @@ def main() -> int:
     ci_register = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(ci_register)
 
-    # Same filter as run_suite.py: skip conftest.py and __init__.py
+    # Same filter as run_suite.py: skip conftest.py, __init__.py, and utils.py
     files = sorted(
         f
         for f in glob.glob("test/registered/**/*.py", recursive=True)
-        if os.path.basename(f) not in ("conftest.py", "__init__.py")
+        if os.path.basename(f) not in ("conftest.py", "__init__.py", "utils.py")
     )
     if not files:
         return 0

test/registered/cpu/test_activation.py

@@ -0,0 +1,59 @@
+import itertools
+import unittest
+
+import torch
+from utils import GeluAndMul, SiluAndMul, precision
+
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestActivation(CustomTestCase):
+    M = [128, 129, 257]
+    N = [22016, 22018]
+    dtype = [torch.float16, torch.bfloat16]
+
+    def _silu_and_mul_test(self, m, n, dtype):
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+        x = torch.randn([m, n], dtype=dtype)
+
+        out = torch.ops.sgl_kernel.silu_and_mul_cpu(x)
+        ref_out = SiluAndMul(x)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def _gelu_and_mul_test(self, m, n, dtype):
+        x = torch.randn([m, n], dtype=dtype)
+
+        out = torch.ops.sgl_kernel.gelu_and_mul_cpu(x)
+        ref_out = GeluAndMul(x, approximate="none")
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def _gelu_tanh_and_mul_test(self, m, n, dtype):
+        x = torch.randn([m, n], dtype=dtype)
+
+        out = torch.ops.sgl_kernel.gelu_tanh_and_mul_cpu(x)
+        ref_out = GeluAndMul(x, approximate="tanh")
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def test_activation(self):
+        for params in itertools.product(self.M, self.N, self.dtype):
+            with self.subTest(m=params[0], n=params[1], dtype=params[2]):
+                self._silu_and_mul_test(*params)
+                self._gelu_and_mul_test(*params)
+                self._gelu_tanh_and_mul_test(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_binding.py

@@ -0,0 +1,30 @@
+import re
+import unittest
+
+import torch
+
+kernel = torch.ops.sgl_kernel
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestBinding(CustomTestCase):
+    def test_binding(self):
+        start_id = 1
+        n_cpu = 6
+
+        expected_cores = list(map(str, range(start_id, start_id + n_cpu)))
+        cpu_ids = ",".join(expected_cores)
+        output = kernel.init_cpu_threads_env(cpu_ids)
+
+        bindings = re.findall(r"OMP tid: \d+, core (\d+)", output)
+        self.assertEqual(len(bindings), n_cpu)
+
+        self.assertEqual(bindings, expected_cores)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_bmm.py

@@ -0,0 +1,98 @@
+import itertools
+import unittest
+
+# TODO: use interface in cpu.py
+import torch
+import torch.nn as nn
+from utils import precision
+
+from sglang.srt.layers.quantization.fp8_utils import input_to_float8
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class Mod(nn.Module):
+    def __init__(self, input_channel, output_channel, has_bias):
+        super(Mod, self).__init__()
+        self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class TestBmm(CustomTestCase):
+    M = [1, 2, 11, 111]
+    N = [128 + 32, 512]
+    K = [512 + 32, 128 + 32]
+    B = [1, 16, 17]
+    chunk = [True, False]
+
+    def _get_bmm_inputs(self, B, M, N, K, chunk, dtype):
+        if chunk:
+            mat1 = (
+                torch.randn(M, B, K + 64, dtype=dtype).narrow(2, 0, K).transpose_(0, 1)
+            )
+            mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
+            mat3 = (
+                torch.randn(M, B, N + 64, dtype=dtype).narrow(2, 0, N).transpose_(0, 1)
+            )
+        else:
+            mat1 = torch.randn(M, B, K, dtype=dtype).transpose_(0, 1)
+            mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
+            mat3 = torch.randn(M, B, N, dtype=dtype).transpose_(0, 1)
+        return mat1, mat2, mat3
+
+    def _bf16_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
+        mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
+        ref = torch.bmm(mat1, mat2)
+        mat2_t = mat2.transpose_(1, 2)
+        mat3.zero_()
+        torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2, False, None)
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
+
+        packed_B = torch.ops.sgl_kernel.convert_weight_packed(mat2_t)
+        mat3.zero_()
+        torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B, True, None)
+        torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
+
+    def _fp8_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
+        mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
+        mat2_q, mat2_s = input_to_float8(mat2)
+        ref = torch.bmm(mat1, mat2_q.to(torch.bfloat16)) * mat2_s
+        mat2_q_t = mat2_q.transpose_(1, 2).contiguous()
+        mat3.zero_()
+        atol = rtol = precision[ref.dtype]
+        torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2_q_t, False, mat2_s)
+        torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
+
+        packed_B_q = torch.ops.sgl_kernel.convert_weight_packed(mat2_q_t)
+        mat3.zero_()
+        torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B_q, True, mat2_s)
+        torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
+
+    def test_bmm(self):
+        for params in itertools.product(
+            self.B,
+            self.M,
+            self.N,
+            self.K,
+            self.chunk,
+        ):
+            with self.subTest(
+                B=params[0],
+                M=params[1],
+                N=params[2],
+                K=params[3],
+                chunk=params[4],
+            ):
+                self._bf16_bmm(*params)
+                self._fp8_bmm(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_causal_conv1d.py

@@ -0,0 +1,330 @@
+import unittest
+from typing import Optional
+
+import sgl_kernel  # noqa: F401
+import torch
+import torch.nn.functional as F
+from utils import parametrize, precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+causal_conv1d_weight_pack = torch.ops.sgl_kernel.causal_conv1d_weight_pack
+causal_conv1d_fwd = torch.ops.sgl_kernel.causal_conv1d_fwd_cpu
+causal_conv1d_update = torch.ops.sgl_kernel.causal_conv1d_update_cpu
+
+
+torch.manual_seed(1234)
+
+PAD_SLOT_ID = -1
+
+
+def causal_conv1d_ref(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    initial_states: Optional[torch.Tensor] = None,
+    return_final_states: bool = False,
+    final_states_out: Optional[torch.Tensor] = None,
+    activation: Optional[str] = "silu",
+):
+    """
+    x: (batch, dim, seqlen)
+    weight: (dim, width)
+    bias: (dim,)
+    initial_states: (batch, dim, width - 1)
+    final_states_out: (batch, dim, width - 1)
+
+    out: (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    dtype_in = x.dtype
+    x = x.to(weight.dtype)
+    seqlen = x.shape[-1]
+    dim, width = weight.shape
+    if initial_states is None:
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim)
+    else:
+        x = torch.cat([initial_states, x], dim=-1)
+        out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim)
+    out = out[..., :seqlen]
+    if return_final_states:
+        final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(
+            dtype_in
+        )  # (batch, dim, width - 1)
+        if final_states_out is not None:
+            final_states_out.copy_(final_states)
+        else:
+            final_states_out = final_states
+    out = (out if activation is None else F.silu(out)).to(dtype=dtype_in)
+    return (out, None) if not return_final_states else (out, final_states_out)
+
+
+def causal_conv1d_update_ref(
+    x, conv_state, weight, bias=None, activation=None, cache_seqlens=None
+):
+    """
+    x: (batch, dim) or (batch, dim, seqlen)
+    conv_state: (batch, dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    cache_seqlens: (batch,), dtype int32.
+        If not None, the conv_state is treated as a circular buffer.
+        The conv_state will be updated by copying x to the
+        conv_state starting at the index
+        @cache_seqlens % state_len before performing the convolution.
+
+    out: (batch, dim) or (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+
+    x = x.unsqueeze(-1)
+    batch, dim, seqlen = x.shape
+    width = weight.shape[1]
+    state_len = conv_state.shape[-1]
+
+    x_new = torch.cat([conv_state, x], dim=-1)
+    conv_state.copy_(x_new[:, :, -state_len:])
+    out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[
+        :, :, -seqlen:
+    ]
+
+    out = out.squeeze(-1)
+    return out if activation is None else F.silu(out)
+
+
+class TestCausalConv1d(CustomTestCase):
+    activation = "silu"
+
+    @parametrize(
+        batch=[1, 1024],
+        dim=[96, 512],
+        seqlen=[2, 36],
+        width=[4],
+        has_bias=[True, False],
+        has_initial_state=[True, False],
+    )
+    def test_causal_conv1d(
+        self,
+        batch,
+        dim,
+        seqlen,
+        width,
+        has_bias,
+        has_initial_state,
+        dtype=torch.bfloat16,
+        prepack=True,
+    ):
+        x = torch.randn(batch, seqlen, dim).to(dtype).transpose_(-1, -2)
+        weight = torch.randn(dim, width).to(dtype)
+        bias = torch.randn(dim).to(dtype) if has_bias else None
+
+        if has_initial_state:
+            initial_states = torch.randn(batch, dim, width - 1, dtype=dtype)
+            has_initial_state_tensor = torch.ones(batch, dtype=torch.bool)
+        else:
+            initial_states = None
+            has_initial_state_tensor = None
+
+        packed_weight = causal_conv1d_weight_pack(weight) if prepack else weight
+
+        out_ref, final_states_ref = causal_conv1d_ref(
+            x,
+            weight,
+            bias,
+            initial_states,
+            return_final_states=has_initial_state,
+            activation=self.activation,
+        )
+
+        out = causal_conv1d_fwd(
+            x,
+            packed_weight,
+            bias,
+            initial_states,
+            None,
+            None,
+            has_initial_state_tensor,
+            self.activation in ["silu"],
+            PAD_SLOT_ID,
+            prepack,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(
+            final_states_ref, initial_states, atol=atol, rtol=rtol
+        )
+
+    @parametrize(
+        batch=[11],
+        dim=[96],
+        max_seqlen=[66],
+        width=[4],
+    )
+    def test_causal_conv1d_varlen(
+        self,
+        batch,
+        dim,
+        max_seqlen,
+        width,
+        has_bias=False,
+        dtype=torch.bfloat16,
+        prepack=False,
+    ):
+        total_entries = batch + 3
+
+        seqlens = torch.randint(1, max_seqlen, (batch + 1,))
+        seqlens[0] = 0
+        # 1 or 2 must test
+        seqlens[-2] = 2
+
+        query_start_loc = torch.cumsum(seqlens, dim=0).to(torch.int32)
+
+        seqlen = query_start_loc[-1].item()
+        x = torch.randn(seqlen, dim, dtype=dtype).transpose_(-1, -2)
+        weight = torch.randn(dim, width, dtype=dtype)
+        bias = torch.randn(dim, dtype=dtype) if has_bias else None
+
+        final_states = torch.randn(total_entries, dim, width - 1, dtype=dtype)
+        final_states_ref = final_states.clone()
+
+        has_initial_states = torch.randint(0, 2, (batch,), dtype=torch.bool).fill_(
+            False
+        )
+        state_indices = torch.randperm(total_entries, dtype=torch.int32)[:batch]
+
+        out_ref = []
+        out_ref_b = []
+
+        return_final_states = final_states is not None
+        splits = torch.split(x, seqlens[1:].tolist(), dim=1)
+        for i, x_s in enumerate(splits):
+            out_ref_b.append(
+                causal_conv1d_ref(
+                    x_s.unsqueeze(0),
+                    weight,
+                    bias,
+                    activation=self.activation,
+                    return_final_states=return_final_states,
+                    final_states_out=(
+                        final_states_ref[state_indices[i]].unsqueeze(0)
+                        if return_final_states
+                        else None
+                    ),
+                    initial_states=(
+                        final_states_ref[state_indices[i]].unsqueeze(0)
+                        if has_initial_states[i]
+                        else None
+                    ),
+                )
+            )
+        out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2))
+        out_ref_tensor = torch.cat(out_ref, dim=0).squeeze(0)
+
+        out = causal_conv1d_fwd(
+            x,
+            weight,
+            bias,
+            final_states,
+            query_start_loc,
+            state_indices,
+            has_initial_states,
+            self.activation in ["silu"],
+            PAD_SLOT_ID,
+            prepack,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref_tensor, out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(final_states_ref, final_states, atol=atol, rtol=rtol)
+
+    @parametrize(
+        batch=[11],
+        dim=[32, 64, 96],
+        width=[4],
+    )
+    def test_causal_conv1d_update(
+        self, batch, dim, width, has_bias=False, dtype=torch.bfloat16, prepack=True
+    ):
+        x = torch.randn(batch, dim).to(dtype)
+        conv_state = torch.randn(batch, dim, width - 1, dtype=dtype)
+        weight = torch.randn(dim, width).to(dtype)
+        bias = torch.randn(dim).to(dtype) if has_bias else None
+
+        packed_weight = causal_conv1d_weight_pack(weight) if prepack else weight
+
+        conv_state_ref = conv_state.clone()
+        out_ref = causal_conv1d_update_ref(
+            x, conv_state_ref, weight, bias, activation=self.activation
+        )
+
+        cache_seqlens = None
+        conv_state_indices = None
+        out = causal_conv1d_update(
+            x,
+            conv_state,
+            packed_weight,
+            bias,
+            self.activation in ["silu"],
+            cache_seqlens,
+            conv_state_indices,
+            PAD_SLOT_ID,
+            prepack,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(conv_state_ref, conv_state, atol=atol, rtol=rtol)
+
+    @parametrize(
+        batch=[7],
+        dim=[96],
+        width=[4],
+    )
+    def test_causal_conv1d_update_with_batch_gather(
+        self, batch, dim, width, has_bias=False, dtype=torch.bfloat16, prepack=True
+    ):
+        total_entries = batch + 3
+
+        x = torch.randn(batch, dim).to(dtype=dtype)
+
+        conv_state_indices = torch.randperm(total_entries)[:batch].to(dtype=torch.int32)
+        conv_state = torch.randn(total_entries, dim, width - 1, dtype=dtype)
+
+        weight = torch.randn(dim, width).to(dtype=dtype)
+        bias = torch.randn(dim).to(dtype=dtype) if has_bias else None
+        conv_state_ref = conv_state[conv_state_indices, :]
+
+        packed_weight = causal_conv1d_weight_pack(weight) if prepack else weight
+
+        out_ref = causal_conv1d_update_ref(
+            x, conv_state_ref, weight, bias, activation=self.activation
+        )
+
+        cache_seqlens = None
+        out = causal_conv1d_update(
+            x,
+            conv_state,
+            packed_weight,
+            bias,
+            self.activation in ["silu"],
+            cache_seqlens,
+            conv_state_indices,
+            PAD_SLOT_ID,
+            prepack,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(
+            conv_state_ref, conv_state[conv_state_indices, :], atol=atol, rtol=rtol
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_cpu_graph.py

@@ -0,0 +1,91 @@
+"""
+Usage:
+python3 -m unittest test_cpu_graph.TestCPUGraph.test_mmlu_torch_compile_cpu
+"""
+
+import copy
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import get_cpu_ids_by_node, kill_process_tree
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    intel_amx_benchmark,
+    is_in_ci,
+    popen_launch_server,
+)
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestCPUGraph(CustomTestCase):
+
+    @intel_amx_benchmark(
+        extra_args=[
+            "--batch-size",
+            "1",
+            "--mem-fraction-static",
+            "0.05",
+            "--enable-torch-compile",
+            "--torch-compile-max-bs",
+            "2",
+            "--cuda-graph-bs",
+            "2",
+        ],
+        min_throughput=7,
+    )
+    def test_latency_torch_compile_cpu(self):
+        return DEFAULT_MLA_MODEL_NAME_FOR_TEST
+
+    def test_mmlu_torch_compile_cpu(self):
+        model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        base_url = DEFAULT_URL_FOR_TEST
+        cpu_ids_by_node = get_cpu_ids_by_node()
+        n_numa_node = len(cpu_ids_by_node)
+        env = copy.deepcopy(os.environ)
+        env["SGLANG_CPU_OMP_THREADS_BIND"] = "all"
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--attention-backend",
+                "intel_amx",
+                "--mem-fraction-static",
+                "0.05",
+                "--disable-radix",
+                "--trust-remote-code",
+                "--disable-overlap-schedule",
+                "--enable-torch-compile",
+                "--cuda-graph-bs",
+                "2",
+                "--tp",
+                f"{n_numa_node}",
+            ],
+            env=env,
+        )
+
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                model=model,
+                eval_name="mmlu",
+                num_examples=64,
+                num_threads=32,
+            )
+
+            metrics = run_eval(args)
+            if is_in_ci():
+                self.assertGreater(metrics["score"], 0.45)
+        finally:
+            kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_decode.py

@@ -0,0 +1,172 @@
+import unittest
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestDecodeAttention(CustomTestCase):
+    def _run_sdpa_forward_decode(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            seq_len_q = 1
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_out = (
+                scaled_dot_product_attention(
+                    per_req_query.unsqueeze(0),
+                    per_req_key.unsqueeze(0),
+                    per_req_value.unsqueeze(0),
+                    enable_gqa=enable_gqa,
+                    scale=scaling,
+                    is_causal=causal,
+                )
+                .squeeze(0)
+                .movedim(query.dim() - 2, 0)
+            )
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+
+        return output
+
+    def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
+        # This represents the number of tokens already in the sequence
+        seq_len = 1024
+        total_tokens = B * seq_len
+        sm_scale = 1.0 / (D**0.5)
+        logit_cap = 0.0
+        num_kv_splits = 8
+        enable_gqa = H_Q != H_KV
+
+        # q represents the new token being generated, one per batch
+        q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
+
+        # k_buffer and v_buffer represent all previous tokens
+        k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
+        v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)
+
+        key = torch.randn(B, H_KV, D, dtype=dtype)
+        value = torch.randn(B, H_KV, D_V, dtype=dtype)
+        loc = torch.randint(0, 10, (B,)).to(torch.int64)
+
+        # set kv cache
+        k_buffer[loc] = key
+        v_buffer[loc] = value
+
+        # o will have the same shape as q
+        o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
+        o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
+
+        req_to_token = (
+            torch.arange(total_tokens, device=device)
+            .reshape(B, seq_len)
+            .to(torch.int32)
+        )
+        b_req_idx = torch.arange(B, device=device).to(torch.int64)
+        b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
+
+        attn_logits = torch.empty(
+            (B, H_Q, num_kv_splits, D_V + 1),
+            dtype=torch.float32,
+            device=device,
+        )
+
+        # k_buffer, v_buffer, query, key and value supports non-contiguous tensors
+        k_buffer = k_buffer.transpose(0, 1).contiguous().transpose(0, 1)
+        v_buffer = v_buffer.transpose(0, 1).contiguous().transpose(0, 1)
+        q = q.transpose(0, 1).contiguous().transpose(0, 1)
+        key = key.transpose(0, 1).contiguous().transpose(0, 1)
+        value = value.transpose(0, 1).contiguous().transpose(0, 1)
+        torch.ops.sgl_kernel.decode_attention_cpu(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            key,
+            value,
+            loc,
+            attn_logits,
+            req_to_token,
+            b_req_idx,
+            b_seq_len,
+            sm_scale,
+            logit_cap,
+        )
+
+        self._run_sdpa_forward_decode(
+            q,
+            o_grouped,
+            k_buffer,
+            v_buffer,
+            req_to_token,
+            b_req_idx,
+            b_seq_len,
+            scaling=sm_scale,
+            enable_gqa=enable_gqa,
+        )
+
+        cos_sim = torch.nn.functional.cosine_similarity(
+            o.flatten(), o_grouped.flatten(), dim=0
+        )
+        self.assertGreater(cos_sim.item(), 0.99)
+        torch.testing.assert_close(o, o_grouped, atol=3e-2, rtol=1e-6)
+
+    def _test_grouped_decode_attention(self, device="cpu"):
+        configs = [
+            (2, 16, 16, 64, 64),
+            (2, 16, 1, 16, 16),
+            (2, 32, 8, 33, 55),
+            (2, 16, 1, 64, 64),
+            (2, 64, 1, 13, 13),
+            (2, 128, 1, 80, 80),
+            (2, 128, 2, 512, 512),
+            (1, 16, 1, 576, 512),
+            (1, 16, 16, 576, 512),
+            (1, 22, 1, 576, 512),
+            (1, 40, 8, 128, 128),
+        ]
+
+        for B, H_Q, H_KV, D, D_V in configs:
+            for dtype in [torch.bfloat16, torch.float16]:
+                self._test_grouped_decode_attention_once(
+                    B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
+                )
+
+    def test_grouped_decode_attention(self):
+        self._test_grouped_decode_attention("cpu")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_extend.py

@@ -0,0 +1,225 @@
+import unittest
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestExtendAttention(CustomTestCase):
+
+    def _run_sdpa_forward_extend(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+
+        assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
+        assert seq_lens.shape[0] == extend_seq_lens.shape[0]
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+
+            extend_seq_len_q = extend_seq_lens[seq_idx]
+            prefill_seq_len_q = extend_prefix_lens[seq_idx]
+
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + extend_seq_len_q
+            end_kv = start_kv + seq_len_kv
+
+            per_req_query = query[:, start_q:end_q, :]
+            per_req_query_redudant = torch.empty(
+                (per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
+                dtype=per_req_query.dtype,
+                device=per_req_query.device,
+            )
+
+            per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_out_redudant = (
+                scaled_dot_product_attention(
+                    per_req_query_redudant.unsqueeze(0),
+                    per_req_key.unsqueeze(0),
+                    per_req_value.unsqueeze(0),
+                    enable_gqa=enable_gqa,
+                    scale=scaling,
+                    is_causal=causal,
+                )
+                .squeeze(0)
+                .movedim(query.dim() - 2, 0)
+            )
+            output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
+            start_q, start_kv = end_q, end_kv
+        return output
+
+    def _test_extend_attention_once(
+        self,
+        B,
+        N_CTX,
+        H_Q,
+        H_KV,
+        D,
+        DV,
+        mla=False,
+        *,
+        b_seq_len_prefix=None,
+        b_seq_len_extend=None,
+    ):
+        dtype = torch.bfloat16
+
+        if b_seq_len_prefix is None:
+            b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
+            if mla:
+                b_seq_len_prefix.zero_()
+        else:
+            b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
+
+        if b_seq_len_extend is None:
+            b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
+        else:
+            b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
+
+        b_seq_len = b_seq_len_prefix + b_seq_len_extend
+        max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
+
+        b_req_idx = torch.arange(B, dtype=torch.int32)
+        req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
+        b_start_loc = torch.zeros((B,), dtype=torch.int32)
+        b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
+        b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
+        b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
+
+        for i in range(B):
+            req_to_tokens[i, : b_seq_len[i]] = torch.arange(
+                b_start_loc[i], b_start_loc[i] + b_seq_len[i]
+            )
+
+        total_token_num = torch.sum(b_seq_len).item()
+        extend_token_num = torch.sum(b_seq_len_extend).item()
+
+        H_BUF = 1 if mla else H_KV
+        k_buffer = torch.randn((total_token_num, H_BUF, D), dtype=dtype)
+        v_buffer = torch.randn((total_token_num, H_BUF, DV), dtype=dtype)
+
+        k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype)
+        v_extend = torch.empty((extend_token_num, H_KV, DV), dtype=dtype)
+        q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
+
+        for i in range(B):
+            extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
+            extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
+            extend_start = b_start_loc_extend[i]
+            extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
+            k_extend[extend_start:extend_end] = k_buffer[
+                extend_start_in_buffer:extend_end_in_buffer
+            ]
+            v_extend[extend_start:extend_end] = v_buffer[
+                extend_start_in_buffer:extend_end_in_buffer
+            ]
+            q_extend[extend_start:extend_end] = (
+                torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * 20
+            )
+
+        # q_extend, k_extend, v_extend, k_buffer and v_buffer supports non-contiguous tensors
+        q_extend = q_extend.transpose(0, 1).contiguous().transpose(0, 1)
+        k_extend = k_extend.transpose(0, 1).contiguous().transpose(0, 1)
+        v_extend = v_extend.transpose(0, 1).contiguous().transpose(0, 1)
+        k_buffer = k_buffer.transpose(0, 1).contiguous().transpose(0, 1)
+        v_buffer = v_buffer.transpose(0, 1).contiguous().transpose(0, 1)
+
+        b_seq_len_extend = b_seq_len - b_seq_len_prefix
+        b_start_loc_extend = torch.zeros_like(b_seq_len)
+        b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
+        max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
+
+        sm_scale = 1.0 / (D**0.5)
+        logit_cap = 0.0
+
+        # handle index type
+        b_req_idx = b_req_idx.to(torch.int64)
+        b_seq_len = b_seq_len.to(torch.int64)
+
+        enable_gqa = H_Q != H_KV
+        o_ref = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
+        self._run_sdpa_forward_extend(
+            q_extend,
+            o_ref,
+            k_buffer,
+            v_buffer,
+            req_to_tokens,
+            b_req_idx,
+            b_seq_len,
+            b_seq_len_prefix,
+            b_seq_len_extend,
+            scaling=sm_scale,
+            enable_gqa=enable_gqa,
+            causal=True,
+        )
+
+        o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
+        torch.ops.sgl_kernel.extend_attention_cpu(
+            q_extend,
+            k_extend,
+            v_extend,
+            o_extend,
+            k_buffer,
+            v_buffer,
+            req_to_tokens,
+            b_req_idx,
+            b_seq_len,
+            b_seq_len_extend,
+            b_start_loc_extend,
+            max_len_extend,
+            sm_scale,
+            logit_cap,
+        )
+
+        torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
+
+    def test_extend_attention(self):
+        for is_mla in [True, False]:
+            self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
+            self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
+            self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
+            self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
+
+    def test_extend_attention_large_seq_causal_mask(self):
+        self._test_extend_attention_once(
+            B=1,
+            N_CTX=5001,
+            H_Q=8,
+            H_KV=2,
+            D=64,
+            DV=64,
+            b_seq_len_prefix=[0],
+            b_seq_len_extend=[5000],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_flash_attn.py

@@ -0,0 +1,245 @@
+import unittest
+
+import sgl_kernel  # noqa: F401
+import torch
+import torch.nn.functional as F
+from utils import parametrize, precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+flash_attn_varlen_func = torch.ops.sgl_kernel.flash_attn_varlen_func
+
+torch.manual_seed(1234)
+
+
+def flash_attn_varlen_ref(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    is_causal,
+    enable_gqa,
+):
+    cu_q = cu_seqlens_q.tolist()
+    cu_k = cu_seqlens_k.tolist()
+    batch = len(cu_k) - 1
+
+    # [T, H, D] -> [1, H, T, D]
+    q, k, v = [x.unsqueeze(0).transpose(1, 2) for x in [q, k, v]]
+
+    B, H, T, D = q.shape
+    out = torch.empty(B, H, T, v.size(-1), dtype=q.dtype)
+    for b in range(batch):
+        start_q, end_q = cu_q[b], cu_q[b + 1]
+        start_k, end_k = cu_k[b], cu_k[b + 1]
+
+        out[:, :, start_q:end_q, :] = F.scaled_dot_product_attention(
+            q[:, :, start_q:end_q, :],
+            k[:, :, start_k:end_k, :],
+            v[:, :, start_k:end_k, :],
+            is_causal=is_causal,
+            enable_gqa=enable_gqa,
+        )
+
+    # [1, H, T, D] -> [T, H, D]
+    return out.transpose(1, 2).squeeze(0)
+
+
+# faster version ref kernel for non varlen case
+def flash_attn_non_varlen_ref(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    is_causal,
+    enable_gqa,
+):
+    cu_q = cu_seqlens_q.tolist()
+    cu_k = cu_seqlens_k.tolist()
+    batch = len(cu_k) - 1
+
+    B_T, H, D = q.shape
+    T = B_T // batch
+
+    # [T, H, D] -> [1, H, T, D]
+    q, k, v = [x.reshape(batch, T, H, D).transpose(1, 2) for x in [q, k, v]]
+
+    out = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        is_causal=is_causal,
+        enable_gqa=enable_gqa,
+    )
+    # [B, H, T, D] -> [B * T, H, D]
+    return out.transpose(1, 2).reshape(batch * T, H, D)
+
+
+class TestFlashAttn(CustomTestCase):
+
+    @parametrize(
+        batch=[4],
+        max_seqlen_q=[35, 96],
+        max_seqlen_k=[35, 96],
+        num_heads=[16],
+        num_heads_kv=[16, 2],
+        head_dim=[32, 48],  # test when D is not 32x
+        head_dim_v=[32],
+        is_causal=[True, False],
+    )
+    def test_flash_attn_varlen(
+        self,
+        batch,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads,
+        num_heads_kv,
+        head_dim,
+        head_dim_v,
+        is_causal,
+    ):
+        dtype = torch.bfloat16
+
+        # random seqlens for k and kv
+        seqlens_q = torch.randint(1, max_seqlen_q, (batch,), dtype=torch.int32)
+        seqlens_k = torch.randint(1, max_seqlen_k, (batch,), dtype=torch.int32)
+        cu_seqlens_q = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_k = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_q[1:] = torch.cumsum(seqlens_q, 0)
+        cu_seqlens_k[1:] = torch.cumsum(seqlens_k, 0)
+
+        sum_seqlen_q = seqlens_q.sum().item()
+        sum_seqlen_k = seqlens_k.sum().item()
+        q = torch.randn(sum_seqlen_q, num_heads, head_dim).to(dtype)
+        k = torch.randn(sum_seqlen_k, num_heads_kv, head_dim).to(dtype)
+        v = torch.randn(sum_seqlen_k, num_heads_kv, head_dim_v).to(dtype)
+
+        out_ref = flash_attn_varlen_ref(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            is_causal=is_causal,
+            enable_gqa=num_heads != num_heads_kv,
+        )
+
+        out = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqlens_q.max().item(),
+            seqlens_k.max().item(),
+            is_causal,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+
+    # test with large size to capture overflow issue
+    @parametrize(
+        batch=[4097],
+        max_seqlen_q=[4097],
+        max_seqlen_k=[4097],
+        num_heads=[4],
+        num_heads_kv=[4],
+        head_dim=[32],
+        head_dim_v=[32],
+        is_causal=[False],
+    )
+    def test_flash_attn_large_size(
+        self,
+        batch,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads,
+        num_heads_kv,
+        head_dim,
+        head_dim_v,
+        is_causal,
+    ):
+        dtype = torch.bfloat16
+
+        # test the non varlen case
+        seqlens_q = torch.full((batch,), max_seqlen_q, dtype=torch.int32)
+        seqlens_k = torch.full((batch,), max_seqlen_k, dtype=torch.int32)
+
+        cu_seqlens_q = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_k = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_q[1:] = torch.cumsum(seqlens_q, 0)
+        cu_seqlens_k[1:] = torch.cumsum(seqlens_k, 0)
+
+        sum_seqlen_q = seqlens_q.sum().item()
+        sum_seqlen_k = seqlens_k.sum().item()
+        q = torch.randn(sum_seqlen_q, num_heads, head_dim).to(dtype)
+        k = torch.randn(sum_seqlen_k, num_heads_kv, head_dim).to(dtype)
+        v = torch.randn(sum_seqlen_k, num_heads_kv, head_dim_v).to(dtype)
+
+        out_ref = flash_attn_non_varlen_ref(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            is_causal=is_causal,
+            enable_gqa=num_heads != num_heads_kv,
+        )
+
+        out = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqlens_q.max().item(),
+            seqlens_k.max().item(),
+            is_causal,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+
+    def _test_flash_attn_large_seq_causal_mask_once(self, seqlens):
+        dtype = torch.bfloat16
+        num_heads = 8
+        num_heads_kv = 2
+        head_dim = 64
+
+        seqlens_t = torch.tensor(seqlens, dtype=torch.int32)
+        cu_seqlens = torch.zeros(len(seqlens) + 1, dtype=torch.int32)
+        cu_seqlens[1:] = torch.cumsum(seqlens_t, 0)
+        total = cu_seqlens[-1].item()
+        max_seqlen = seqlens_t.max().item()
+
+        q = torch.randn(total, num_heads, head_dim, dtype=dtype)
+        k = torch.randn(total, num_heads_kv, head_dim, dtype=dtype)
+        v = torch.randn(total, num_heads_kv, head_dim, dtype=dtype)
+
+        out_ref = flash_attn_varlen_ref(
+            q, k, v, cu_seqlens, cu_seqlens, is_causal=True, enable_gqa=True
+        )
+        out = flash_attn_varlen_func(
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, True
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+
+    def test_flash_attn_large_seq_causal_mask(self):
+        # Non-varlen path: single sequence, has_varlen_sequences returns False
+        # → dispatches to flash_attn_kernel_impl.
+        self._test_flash_attn_large_seq_causal_mask_once([5000])
+        # Varlen path: sequences with different lengths, has_varlen_sequences
+        # returns True → dispatches to flash_attn_varlen_kernel_impl
+        self._test_flash_attn_large_seq_causal_mask_once([5000, 4999])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_gemm.py

@@ -0,0 +1,335 @@
+import itertools
+import unittest
+
+# TODO: use interface in cpu.py
+import torch
+import torch.nn as nn
+from utils import (
+    convert_weight,
+    native_w8a8_per_token_matmul,
+    per_token_quant_int8,
+    precision,
+    unpack_and_dequant_awq,
+    unpack_and_dequant_gptq,
+)
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class Mod(nn.Module):
+    def __init__(self, input_channel, output_channel, has_bias):
+        super(Mod, self).__init__()
+        self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class TestGemm(CustomTestCase):
+    M = [1, 101]
+    N = [16, 32 * 13]
+    K = [32 * 16]
+    has_bias = [False, True]
+
+    M_int8 = [2, 128]
+    N_int8 = [32 * 12]
+    K_int8 = [32 * 17]
+
+    M_fp8 = [1, 11]
+    N_fp8 = [128, 224]
+    K_fp8 = [512, 576]
+
+    M_awq = [1, 32]
+    N_awq = [4096]
+    K_awq = [4096]
+
+    M_gptq = [1, 32]
+    N_gptq = [4096]
+    K_gptq = [4096]
+
+    def _bf16_gemm(self, M, N, K, has_bias):
+
+        mat1 = torch.randn(M, K, dtype=torch.bfloat16)
+        mat2 = torch.randn(N, K, dtype=torch.bfloat16)
+
+        ref = torch.matmul(mat1.float(), mat2.float().t())
+        if has_bias:
+            bias = torch.randn(N, dtype=torch.float32)
+            ref.add_(bias.bfloat16())
+
+        ref = ref.bfloat16()
+
+        out = torch.ops.sgl_kernel.weight_packed_linear(
+            mat1, mat2, bias if has_bias else None, False
+        )
+
+        packed_mat2 = torch.ops.sgl_kernel.convert_weight_packed(mat2)
+        out2 = torch.ops.sgl_kernel.weight_packed_linear(
+            mat1, packed_mat2, bias if has_bias else None, True
+        )
+
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(ref, out2, atol=atol, rtol=rtol)
+
+    def test_bf16_gemm(self):
+        for params in itertools.product(
+            self.M,
+            self.N,
+            self.K,
+            self.has_bias,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                has_bias=params[3],
+            ):
+                self._bf16_gemm(*params)
+
+    def _bf16_gemm_with_small_oc(self, M, N, K, has_bias, use_post_sigmul):
+        use_post_sigmul = use_post_sigmul and N == 1
+        mat_mul = (
+            None if not use_post_sigmul else torch.randn(M, 2 * K, dtype=torch.bfloat16)
+        )
+        mat1 = torch.randn(M, K, dtype=torch.bfloat16)
+        mat2 = torch.randn(N, K, dtype=torch.bfloat16)
+
+        ref = torch.nn.functional.linear(mat1, mat2)
+        if has_bias:
+            bias = torch.randn(N, dtype=torch.float32)
+            ref.add_(bias)
+        if use_post_sigmul:
+            ref = torch.nn.functional.sigmoid(ref) * mat_mul
+            out = torch.ops.sgl_kernel.fused_linear_sigmoid_mul(
+                mat1,
+                torch.ops.sgl_kernel.convert_weight_packed(mat2),
+                bias if has_bias else None,
+                True,
+                mat_mul if use_post_sigmul else None,
+            )
+        else:
+            out = torch.ops.sgl_kernel.weight_packed_linear(
+                mat1,
+                torch.ops.sgl_kernel.convert_weight_packed(mat2),
+                bias if has_bias else None,
+                True,
+            )
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
+
+    def test_bf16_gemm_with_small_oc(self):
+        for params in itertools.product(
+            [1, 8, 32, 1024], [12, 1], self.K, self.has_bias, [False, True]
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                has_bias=params[3],
+                use_post_sigmul=params[4],
+            ):
+                self._bf16_gemm_with_small_oc(*params)
+
+    def _int8_gemm(self, M, N, K, has_bias):
+        dtype = torch.bfloat16
+        A = torch.randn((M, K), dtype=dtype) / 10
+        Aq, As = per_token_quant_int8(A)
+
+        factor_for_scale = 1e-2
+        int8_max = 127
+        int8_min = -128
+
+        B = (torch.rand((N, K), dtype=torch.float32) - 0.5) * 2
+        Bq = (B * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
+        Bs = torch.rand(N) * factor_for_scale
+
+        bias = torch.randn(N) if has_bias else None
+        ref_out = native_w8a8_per_token_matmul(Aq, Bq, As, Bs, bias, dtype)
+
+        atol = rtol = precision[ref_out.dtype]
+
+        Aq2, As2 = torch.ops.sgl_kernel.per_token_quant_int8_cpu(A)
+        out = torch.ops.sgl_kernel.int8_scaled_mm_cpu(
+            Aq2, Bq, As2, Bs, bias if has_bias else None, torch.bfloat16, False
+        )
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+        # test the fused version
+        fused_out = torch.ops.sgl_kernel.int8_scaled_mm_with_quant(
+            A, Bq, Bs, bias if has_bias else None, torch.bfloat16, False
+        )
+        torch.testing.assert_close(ref_out, fused_out, atol=atol, rtol=rtol)
+
+    def test_int8_gemm(self):
+        for params in itertools.product(
+            self.M_int8,
+            self.N_int8,
+            self.K_int8,
+            self.has_bias,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                has_bias=params[3],
+            ):
+                self._int8_gemm(*params)
+
+    def _fp8_gemm(self, M, N, K, has_bias):
+        prepack = True
+        chunk = False
+        scale_block_size_N = 64
+        scale_block_size_K = 128
+        assert scale_block_size_N <= N
+        assert scale_block_size_K <= K
+        A_dtype = torch.bfloat16
+
+        model = Mod(K, N, has_bias).eval()
+        if chunk:
+            data = torch.randn(M, K + 6, dtype=A_dtype).narrow(1, 0, K)
+        else:
+            data = torch.randn(M, K, dtype=A_dtype)
+
+        weight = model.linear.weight  # (N, K)
+
+        if has_bias:
+            bias = model.linear.bias
+
+        fp8_weight, scales, dq_weight = convert_weight(
+            weight, [scale_block_size_N, scale_block_size_K], A_dtype
+        )
+
+        if has_bias:
+            ref = torch.matmul(data.to(A_dtype), dq_weight.T) + bias.to(A_dtype)
+        else:
+            ref = torch.matmul(data.to(A_dtype), dq_weight.T)
+
+        if prepack:
+            fp8_weight = torch.ops.sgl_kernel.convert_weight_packed(fp8_weight)
+
+        opt = torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
+            data,
+            fp8_weight,
+            scales,
+            [scale_block_size_N, scale_block_size_K],
+            bias if has_bias else None,
+            data.dtype,
+            prepack,
+        )
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, opt, atol=atol, rtol=rtol)
+
+    def test_fp8_gemm(self):
+        for params in itertools.product(
+            self.M_fp8,
+            self.N_fp8,
+            self.K_fp8,
+            self.has_bias,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                has_bias=params[3],
+            ):
+                self._fp8_gemm(*params)
+
+    def _int4_awq_gemm(self, M, N, K, group_size, has_bias):
+        awq_weight = torch.randint(-128, 128, (K, N // 8)).to(torch.int)
+        awq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
+        awq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16)
+        bf16_weight, _ = unpack_and_dequant_awq(
+            awq_weight, awq_zero, awq_scales, 4, 128
+        )
+        if has_bias:
+            bias = torch.rand(bf16_weight.shape[0]).to(torch.float)
+        else:
+            bias = None
+        x = torch.rand(M, bf16_weight.size(-1)).to(torch.bfloat16)
+        ref_res = torch.nn.functional.linear(
+            x, bf16_weight, bias=bias.to(torch.bfloat16) if has_bias else None
+        )
+
+        packed_weight, packed_zero, packed_scales = (
+            torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
+                awq_weight, awq_zero, awq_scales, 0
+            )
+        )
+        target_res = torch.ops.sgl_kernel.int4_scaled_mm_cpu(
+            x,
+            packed_weight,
+            packed_zero,
+            packed_scales,
+            bias,
+        )
+
+        atol = rtol = precision[ref_res.dtype]
+        torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol)
+
+    def test_int4_awq_gemm(self):
+        for params in itertools.product(
+            self.M_awq, self.N_awq, self.K_awq, [128], self.has_bias
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                group_size=params[3],
+                has_bias=params[4],
+            ):
+                self._int4_awq_gemm(*params)
+
+    def _int4_gptq_gemm(self, M, N, K, group_size, has_bias):
+        torch.manual_seed(127)
+        gptq_weight = torch.randint(-128, 128, (K // 8, N)).to(torch.int)
+        gptq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
+        gptq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16) // 10
+
+        bf16_weight = unpack_and_dequant_gptq(gptq_weight, gptq_zero, gptq_scales)
+        if has_bias:
+            bias = torch.rand(bf16_weight.shape[0]).to(torch.float)
+        else:
+            bias = None
+        x = torch.rand(M, bf16_weight.size(-1)).to(torch.bfloat16)
+        ref_res = torch.nn.functional.linear(
+            x, bf16_weight, bias=bias.to(torch.bfloat16) if has_bias else None
+        )
+
+        packed_weight, packed_zero, packed_scales = (
+            torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
+                gptq_weight, gptq_zero, gptq_scales, 1
+            )
+        )
+        target_res = torch.ops.sgl_kernel.int4_scaled_mm_cpu(
+            x,
+            packed_weight,
+            packed_zero,
+            packed_scales,
+            bias,
+        )
+
+        atol = rtol = precision[ref_res.dtype]
+        torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol)
+
+    def test_int4_gptq_gemm(self):
+        for params in itertools.product(
+            self.M_gptq, self.N_gptq, self.K_gptq, [128], self.has_bias
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                group_size=params[3],
+                has_bias=params[4],
+            ):
+                self._int4_gptq_gemm(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_intel_amx_attention_backend_a.py

@@ -0,0 +1,76 @@
+"""
+Usage:
+python3 -m unittest test_intel_amx_attention_backend.TestIntelAMXAttnBackend.test_latency_default_model
+"""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    intel_amx_benchmark,
+    is_in_ci,
+    popen_launch_server,
+)
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestIntelAMXAttnBackend(CustomTestCase):
+
+    @intel_amx_benchmark(
+        extra_args=["--batch-size", "4", "--mem-fraction-static", "0.3"],
+        min_throughput=10,
+    )
+    def test_latency_mla_model(self):
+        return DEFAULT_MLA_MODEL_NAME_FOR_TEST
+
+    @intel_amx_benchmark(
+        extra_args=["--batch-size", "4", "--mem-fraction-static", "0.1"],
+        min_throughput=40,
+    )
+    def test_latency_default_model(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST
+
+    def test_mmlu(self):
+        model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        base_url = DEFAULT_URL_FOR_TEST
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--attention-backend",
+                "intel_amx",
+                "--mem-fraction-static",
+                "0.3",
+                "--disable-radix",
+                "--trust-remote-code",
+                "--disable-overlap-schedule",
+            ],
+        )
+
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                model=model,
+                eval_name="mmlu",
+                num_examples=64,
+                num_threads=32,
+            )
+            metrics = run_eval(args)
+            if is_in_ci():
+                self.assertGreater(metrics["score"], 0.45)
+        finally:
+            kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_intel_amx_attention_backend_b.py

@@ -0,0 +1,38 @@
+"""
+For intel_amx attention backend FP8 tests
+Usage:
+python3 -m unittest test_intel_amx_attention_backend_1.TestIntelAMXAttnBackendQuant.test_latency_fp8_qwen
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
+    DEFAULT_MODEL_NAME_FOR_TEST_QWEN_FP8,
+    CustomTestCase,
+    intel_amx_benchmark,
+)
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestIntelAMXAttnBackendQuant(CustomTestCase):
+
+    @intel_amx_benchmark(
+        extra_args=["--batch-size", "4", "--mem-fraction-static", "0.1"],
+        min_throughput=150,
+    )
+    def test_latency_fp8_qwen(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST_QWEN_FP8
+
+    @intel_amx_benchmark(
+        extra_args=["--batch-size", "4", "--mem-fraction-static", "0.1"],
+        min_throughput=50,
+    )
+    def test_latency_fp8_moe_model(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_intel_amx_attention_backend_c.py

@@ -0,0 +1,56 @@
+"""
+For intel_amx attention backend w8a8 tests
+Usage:
+python3 -m unittest test_intel_amx_attention_backend_2.TestIntelAMXAttnBackendQuant.test_latency_w8a8_default_model
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_W8A8,
+    DEFAULT_MODEL_NAME_FOR_TEST_W8A8_WITH_MOE,
+    CustomTestCase,
+    intel_amx_benchmark,
+)
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestIntelAMXAttnBackendQuant(CustomTestCase):
+
+    @intel_amx_benchmark(
+        extra_args=[
+            "--batch-size",
+            "4",
+            "--quantization",
+            "w8a8_int8",
+            "--mem-fraction-static",
+            "0.1",
+        ],
+        min_throughput=100,
+    )
+    def test_latency_w8a8_default_model(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST_W8A8
+
+    @intel_amx_benchmark(
+        extra_args=[
+            "--batch-size",
+            "4",
+            "--quantization",
+            "w8a8_int8",
+            "--mem-fraction-static",
+            "0.9",
+            "--max-total-tokens",
+            "65536",
+            "--tp",
+            "6",
+        ],
+        min_throughput=100,
+    )
+    def test_latency_w8a8_moe_model(self):
+        return DEFAULT_MODEL_NAME_FOR_TEST_W8A8_WITH_MOE
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_mamba.py

@@ -0,0 +1,397 @@
+import unittest
+
+import torch
+import torch.nn.functional as F
+from torch.nn.functional import softplus
+from utils import precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6):
+    """This function is intended to align with the l2norm implementation in the FLA library."""
+    inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+    return x * inv_norm
+
+
+def torch_chunk_gated_delta_rule(
+    query,
+    key,
+    value,
+    g,
+    beta,
+    chunk_size=64,
+    initial_state=None,
+    output_final_state=False,
+    use_qk_l2norm_in_kernel=False,
+):
+    initial_dtype = query.dtype
+    if use_qk_l2norm_in_kernel:
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
+    query, key, value, beta, g = [
+        x.transpose(1, 2).contiguous().to(torch.float32)
+        for x in (query, key, value, beta, g)
+    ]
+
+    batch_size, sequence_length, num_heads, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    pad_size = (chunk_size - num_heads % chunk_size) % chunk_size
+    query = F.pad(query, (0, 0, 0, pad_size))
+    key = F.pad(key, (0, 0, 0, pad_size))
+    value = F.pad(value, (0, 0, 0, pad_size))
+    beta = F.pad(beta, (0, pad_size))
+    g = F.pad(g, (0, pad_size))
+    tot_heads = num_heads + pad_size
+    scale = 1 / (query.shape[-1] ** 0.5)
+    query = query * scale
+
+    v_beta = value * beta.unsqueeze(-1)
+    k_beta = key * beta.unsqueeze(-1)
+    # reshape to chunks
+    query, key, value, k_beta, v_beta = [
+        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
+        for x in (query, key, value, k_beta, v_beta)
+    ]
+    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+    mask = torch.triu(
+        torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
+        diagonal=0,
+    )
+
+    # chunk decay
+    g = g.cumsum(dim=-1)
+    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        row = attn[..., i, :i].clone()
+        sub = attn[..., :i, :i].clone()
+        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    value = attn @ v_beta
+    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+    last_recurrent_state = (
+        torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
+        if initial_state is None
+        else initial_state.to(value)
+    )
+    core_attn_out = torch.zeros_like(value)
+    mask = torch.triu(
+        torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
+        diagonal=1,
+    )
+
+    # for each chunk
+    for i in range(0, tot_heads // chunk_size):
+        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
+        v_new = v_i - v_prime
+        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
+        core_attn_out[:, :, i] = attn_inter + attn @ v_new
+        last_recurrent_state = (
+            last_recurrent_state * g[:, :, i, -1, None, None].exp()
+            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(
+                -1, -2
+            )
+            @ v_new
+        )
+
+    if not output_final_state:
+        last_recurrent_state = None
+    core_attn_out = core_attn_out.reshape(
+        core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]
+    )
+    core_attn_out = core_attn_out[:, :, :num_heads]
+    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+    return core_attn_out, last_recurrent_state
+
+
+def chunk_gated_delta_rule_update(
+    query,  # [B, T, HK, K]
+    key,  # [B, T, HK, K]
+    value,  # [B, T, HV, V]
+    g,  # [B, T, HV]
+    beta,  # [B, T, HV]
+    cu_seqlens,  # [N+1]
+    initial_state,  # [N, HV, K, V]
+    use_qk_l2norm_in_kernel,  # True
+):
+    num_heads = query.shape[2]
+    num_value_heads = value.shape[2]
+    batch_size = initial_state.shape[0]
+    if num_value_heads // num_heads > 1:
+        query = query.repeat_interleave(num_value_heads // num_heads, dim=2)
+        key = key.repeat_interleave(num_value_heads // num_heads, dim=2)
+    output = torch.empty_like(value)
+    final_state = torch.empty_like(initial_state)
+    start_q = 0
+    for i in range(batch_size):
+        end_q = cu_seqlens[i + 1]
+        core_attn_outi, last_recurrent_state = torch_chunk_gated_delta_rule(
+            query=query[:, start_q:end_q, :, :],
+            key=key[:, start_q:end_q, :, :],
+            value=value[:, start_q:end_q, :, :],
+            g=g[:, start_q:end_q, :],
+            beta=beta[:, start_q:end_q, :],
+            initial_state=initial_state[i],
+            output_final_state=True,
+            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        )
+        output[:, start_q:end_q, :, :] = core_attn_outi
+        final_state[i] = last_recurrent_state
+        start_q = end_q
+    return output, final_state
+
+
+def torch_recurrent_gated_delta_rule(
+    query,
+    key,
+    value,
+    g,
+    beta,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+):
+    initial_dtype = query.dtype
+    if use_qk_l2norm_in_kernel:
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
+    query, key, value, beta, g = [
+        x.transpose(1, 2).contiguous().to(torch.float32)
+        for x in (query, key, value, beta, g)
+    ]
+
+    batch_size, num_heads, sequence_length, k_head_dim = key.shape
+    v_head_dim = value.shape[-1]
+    scale = 1 / (query.shape[-1] ** 0.5)
+    query = query * scale
+
+    core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(
+        value
+    )
+    last_recurrent_state = (
+        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
+        if initial_state is None
+        else initial_state.to(value)
+    )
+
+    for i in range(sequence_length):
+        q_t = query[:, :, i]
+        k_t = key[:, :, i]
+        v_t = value[:, :, i]
+        g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
+        beta_t = beta[:, :, i].unsqueeze(-1)
+
+        last_recurrent_state = last_recurrent_state * g_t
+        kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)
+        delta = (v_t - kv_mem) * beta_t
+        last_recurrent_state = last_recurrent_state + k_t.unsqueeze(
+            -1
+        ) * delta.unsqueeze(-2)
+        core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)
+
+    if not output_final_state:
+        last_recurrent_state = None
+    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+    return core_attn_out, last_recurrent_state
+
+
+def sigmoid_gating_delta_rule_update(
+    query,
+    key,
+    value,
+    A_log,
+    a,
+    dt_bias,
+    b,
+    initial_state,
+    output_final_state,
+    use_qk_l2norm_in_kernel=False,
+):
+    beta = b.sigmoid()
+    g = -A_log.float().exp() * softplus(a.float() + dt_bias)
+    return torch_recurrent_gated_delta_rule(
+        query,
+        key,
+        value,
+        g.unsqueeze(0),
+        beta.unsqueeze(0),
+        initial_state,
+        output_final_state,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+    )
+
+
+def torch_gdn_gating(A_log, a, b, dt_bias):
+    return -A_log.float().exp() * softplus(a.float() + dt_bias).unsqueeze(
+        0
+    ), b.sigmoid().unsqueeze(0)
+
+
+class TestMambaAttention(CustomTestCase):
+    def test_chunk_gated_delta_rule(self):
+        B, L, HK, HV, EK, EV, N = 1, 100, 3, 6, 64, 64, 4
+        seqlens = torch.randint(1, L, (N + 1,))
+        seqlens[0] = 0
+        cu_seqlens_ = torch.cumsum(seqlens, dim=0).to(torch.int32)
+        T = cu_seqlens_[-1].item()
+        query_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05
+        key_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05
+        value_ = torch.rand((B, T, HV, EV), dtype=torch.bfloat16) * 0.05
+        g_ = torch.rand((B, T, HV), dtype=torch.float32) * 0.05
+        beta_ = torch.rand((B, T, HV), dtype=torch.bfloat16) * 0.05
+        initial_state_ = torch.rand((N, HV, EK, EV), dtype=torch.float32) * 0.05
+
+        for use_qk_l2norm_in_kernel in [True, False]:
+            core_attn_out_ref, last_recurrent_state_ref = chunk_gated_delta_rule_update(
+                query=query_,
+                key=key_,
+                value=value_,
+                g=g_,
+                beta=beta_,
+                cu_seqlens=cu_seqlens_,
+                initial_state=initial_state_,
+                use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+            )
+
+            query = query_.clone()
+            key = key_.clone()
+            value = value_.clone()
+            g = g_.clone()
+            beta = beta_.clone()
+            cu_seqlens = cu_seqlens_.clone()
+            initial_state = initial_state_.clone()
+
+            core_attn_out, last_recurrent_state = (
+                torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
+                    query=query,
+                    key=key,
+                    value=value,
+                    g=g,
+                    beta=beta,
+                    initial_state=initial_state,
+                    output_final_state=True,
+                    cu_seqlens=cu_seqlens,
+                    head_first=False,
+                    use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                )
+            )
+            atol = rtol = precision[core_attn_out.dtype]
+            torch.testing.assert_close(
+                core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
+            )
+            torch.testing.assert_close(
+                last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
+            )
+
+    def test_fused_gdn_gating(self):
+        dims = [6, 32]
+        for dim in dims:
+            for A_log_dtype in [torch.float32, torch.bfloat16]:
+                A_log = torch.rand(dim, dtype=A_log_dtype)
+                a = torch.rand(1024, dim, dtype=torch.bfloat16)
+                b = torch.rand(1024, dim, dtype=torch.bfloat16)
+                dt_bias = torch.rand(dim, dtype=torch.bfloat16)
+
+                g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
+                g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
+                    A_log, a, b, dt_bias
+                )
+                atol = rtol = precision[g.dtype]
+                atol2 = rtol2 = precision[beta.dtype]
+                torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
+                torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
+
+    def test_fused_sigmoid_gating_delta_rule_update(self):
+        batch_size = 1
+        num_value_heads = 32
+        head_k_dim = 128
+        head_v_dim = 128
+        num_heads = 16
+        seq_len = 1
+        attn_tp_size = 1
+        key_dim = head_k_dim * num_heads
+        value_dim = head_v_dim * num_value_heads
+        mixed_qkv_dim = (key_dim * 2 + value_dim) // attn_tp_size
+        mixed_qkv = torch.rand(
+            seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16
+        )
+        query, key, value = torch.split(
+            mixed_qkv,
+            [
+                key_dim // attn_tp_size,
+                key_dim // attn_tp_size,
+                value_dim // attn_tp_size,
+            ],
+            dim=-1,
+        )
+        query = query.view(1, seq_len, num_heads, head_k_dim)
+        key = key.view(1, seq_len, num_heads, head_k_dim)
+        value = value.view(1, seq_len, num_value_heads, head_v_dim)
+        A_log = torch.rand(num_value_heads, dtype=torch.float32)
+        a = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
+        b = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16)
+        dt_bias = torch.rand(num_value_heads, dtype=torch.bfloat16)
+        ssm_states = torch.rand(
+            513, num_value_heads, head_k_dim, head_v_dim, dtype=torch.float32
+        )
+        cache_indices = torch.randint(0, 513, (batch_size,), dtype=torch.int32)
+        query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
+        use_qk_l2norm_in_kernel = True
+        query_ref = query.clone()
+        key_ref = key.clone()
+        if num_value_heads // num_heads > 1:
+            query_ref = query_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
+            key_ref = key_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
+        for A_log_dtype in [torch.float32, torch.bfloat16]:
+            A_log = A_log.to(A_log_dtype)
+            core_attn_out_ref, last_recurrent_state_ref = (
+                sigmoid_gating_delta_rule_update(
+                    query_ref.transpose(0, 1),
+                    key_ref.transpose(0, 1),
+                    value.transpose(0, 1),
+                    A_log,
+                    a,
+                    dt_bias,
+                    b,
+                    initial_state=ssm_states[cache_indices],
+                    output_final_state=True,
+                    use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                )
+            )
+            core_attn_out = (
+                torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
+                    A_log=A_log,
+                    dt_bias=dt_bias,
+                    q=query,
+                    k=key,
+                    v=value,
+                    a=a,
+                    b=b,
+                    initial_state_source=ssm_states,
+                    initial_state_indices=cache_indices,
+                    cu_seqlens=query_start_loc,
+                    use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                    softplus_beta=1.0,
+                    softplus_threshold=20.0,
+                )
+            )
+            last_recurrent_state = ssm_states[cache_indices]
+            atol = rtol = precision[core_attn_out.dtype]
+            torch.testing.assert_close(
+                core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
+            )
+            torch.testing.assert_close(
+                last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_mla.py

@@ -0,0 +1,158 @@
+import unittest
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+from utils import precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestMLA(CustomTestCase):
+    def _run_sdpa_forward_decode(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        key: torch.Tensor,
+        loc: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+    ):
+        # set kv cache
+        k_cache[loc] = key
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            seq_len_q = 1
+            seq_len_kv = seq_lens[seq_idx]
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            per_req_out = (
+                scaled_dot_product_attention(
+                    per_req_query.unsqueeze(0),
+                    per_req_key.unsqueeze(0),
+                    per_req_value.unsqueeze(0),
+                    enable_gqa=enable_gqa,
+                    scale=scaling,
+                    is_causal=causal,
+                )
+                .squeeze(0)
+                .movedim(query.dim() - 2, 0)
+            )
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+
+        return output
+
+    def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, seq_len):
+        dtype = torch.bfloat16
+
+        total_tokens = B * seq_len
+        sm_scale = 1.0 / (D**0.5)
+        logit_cap = 0.0
+        num_kv_splits = 8
+        enable_gqa = H_Q != H_KV
+
+        # q represents the new token being generated, one per batch
+        q = torch.randn(B, H_Q, D, dtype=dtype)
+
+        # k_buffer and v_buffer represent all previous tokens
+        k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype)
+        v_buffer = k_buffer.narrow(2, 0, D_V)
+
+        key = torch.randn(B, H_KV, D, dtype=dtype)
+        value = key.narrow(2, 0, D_V)
+        # make sure no duplicates in loc
+        loc = torch.randperm(total_tokens)[:B].to(torch.int64)
+
+        k_buffer2 = k_buffer.clone()
+        v_buffer2 = k_buffer2.narrow(2, 0, D_V)
+
+        # o will have the same shape as q
+        o = torch.zeros(B, H_Q, D_V, dtype=dtype)
+        o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype)
+
+        req_to_token = torch.arange(total_tokens).reshape(B, seq_len).to(torch.int32)
+        b_req_idx = torch.arange(B).to(torch.int64)
+        b_seq_len = torch.full((B,), seq_len).to(torch.int64)
+
+        attn_logits = torch.empty(
+            (B, H_Q, num_kv_splits, D_V + 1),
+            dtype=torch.float32,
+        )
+
+        torch.ops.sgl_kernel.decode_attention_cpu(
+            q,
+            k_buffer2,
+            v_buffer2,
+            o,
+            key,
+            value,
+            loc,
+            attn_logits,
+            req_to_token,
+            b_req_idx,
+            b_seq_len,
+            sm_scale,
+            logit_cap,
+        )
+
+        self._run_sdpa_forward_decode(
+            q,
+            o_grouped,
+            k_buffer,
+            v_buffer,
+            key,
+            loc,
+            req_to_token,
+            b_req_idx,
+            b_seq_len,
+            scaling=sm_scale,
+            enable_gqa=enable_gqa,
+        )
+
+        cos_sim = torch.nn.functional.cosine_similarity(
+            o.flatten(), o_grouped.flatten(), dim=0
+        )
+        atol = rtol = precision[q.dtype]
+        self.assertGreater(cos_sim.item(), 0.99)
+        torch.testing.assert_close(o, o_grouped, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k_buffer, k_buffer2, atol=atol, rtol=rtol)
+        torch.testing.assert_close(v_buffer, v_buffer2, atol=atol, rtol=rtol)
+
+    def test_grouped_decode_attention(self):
+        configs = [
+            (1, 22, 1, 576, 512, 8 * 111),
+            (4, 22, 1, 576, 512, 8 * 128),
+            (40, 22, 1, 576, 512, 8 * 133),
+        ]
+
+        for B, H_Q, H_KV, D, D_V, seqlen in configs:
+            self._test_grouped_decode_attention_once(B, H_Q, H_KV, D, D_V, seqlen)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_moe.py

@@ -0,0 +1,355 @@
+import itertools
+import math
+import unittest
+
+# TODO: use interface in cpu.py
+import torch
+
+from sglang.srt.layers.amx_utils import CPUQuantMethod
+
+kernel = torch.ops.sgl_kernel
+
+torch.manual_seed(128)
+
+from utils import (
+    BLOCK_K,
+    BLOCK_N,
+    factor_for_scale,
+    fp8_max,
+    fp8_min,
+    native_fp8_fused_moe,
+    precision,
+    scaled_weight,
+    torch_naive_fused_moe,
+    torch_w8a8_per_column_fused_moe,
+    unpack_and_dequant_awq,
+)
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+def fused_moe(a, w1, w2, score, topk, renormalize, prepack):
+
+    G = 1
+    topk_group = 1
+
+    B, D = a.shape
+    topk_weights = torch.empty(B, topk, dtype=torch.float32)
+    topk_ids = torch.empty(B, topk, dtype=torch.int32)
+    topk_weights, topk_ids = kernel.grouped_topk_cpu(
+        a, score, topk, renormalize, G, topk_group, 0, None, None
+    )
+
+    packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
+    packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
+
+    inplace = True
+    return kernel.fused_experts_cpu(
+        a,
+        packed_w1,
+        packed_w2,
+        topk_weights,
+        topk_ids,
+        inplace,
+        CPUQuantMethod.UNQUANT,
+        None,
+        None,
+        None,
+        None,
+        None,
+        prepack,
+    )
+
+
+class TestFusedExperts(CustomTestCase):
+    M = [2, 114]
+    N = [32]
+    K = [32]
+    E = [4]
+    topk = [2]
+    renormalize = [False, True]
+
+    M_int8 = [1, 39]
+    N_int8 = [128]
+    K_int8 = [256]
+    E_int8 = [8]
+    topk_int8 = [3]
+
+    M_fp8 = [2, 121]
+    N_fp8 = [352, 512]
+    K_fp8 = [256, 320]
+    E_fp8 = [8]
+    topk_fp8 = [4]
+
+    M_int4 = [1, 6]
+    N_int4 = [512]
+    K_int4 = [256]
+    E_int4 = [8]
+    topk_int4 = [4]
+
+    def _bf16_moe(self, m, n, k, e, topk, renormalize):
+        dtype = torch.bfloat16
+        prepack = True
+
+        a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
+        w1 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
+        w2 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
+        score = torch.randn((m, e), device="cpu", dtype=dtype)
+
+        torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize)
+        fused_output = fused_moe(a, w1, w2, score, topk, renormalize, prepack)
+
+        atol = rtol = precision[torch_output.dtype]
+        torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
+
+    def test_bf16_moe(self):
+        for params in itertools.product(
+            self.M,
+            self.N,
+            self.K,
+            self.E,
+            self.topk,
+            self.renormalize,
+        ):
+            with self.subTest(
+                m=params[0],
+                n=params[1],
+                k=params[2],
+                e=params[3],
+                topk=params[4],
+                renormalize=params[5],
+            ):
+                self._bf16_moe(*params)
+
+    def _int8_moe(self, M, N, K, E, topk):
+        dtype = torch.bfloat16
+        prepack = True
+
+        # Initialize int8 quantization parameters
+        int8_factor_for_scale = 1e-2
+        int8_max = 127
+        int8_min = -128
+
+        # Input tensor
+        # M * K
+        a = torch.randn((M, K), dtype=dtype) / math.sqrt(K)
+
+        # Generate int8 weights
+        w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
+        w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
+
+        w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
+        w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
+
+        # Generate scale for each column (per-column quantization)
+        w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * int8_factor_for_scale
+        w2_s = torch.rand(E, K, device=w2_fp32.device) * int8_factor_for_scale
+
+        # Calculate routing
+        score = torch.randn((M, E), dtype=dtype)
+        score = torch.softmax(score, dim=-1, dtype=torch.float32)
+        topk_weight, topk_ids = torch.topk(score, topk)
+
+        ref_out = torch_w8a8_per_column_fused_moe(
+            a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk
+        )
+
+        inplace = True
+        packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
+        packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
+        out = kernel.fused_experts_cpu(
+            a,
+            packed_w1,
+            packed_w2,
+            topk_weight,
+            topk_ids.to(torch.int32),
+            inplace,
+            CPUQuantMethod.INT8_W8A8,
+            w1_s,
+            w2_s,
+            None,
+            None,
+            None,
+            prepack,
+        )
+
+        atol = rtol = precision[ref_out.dtype]
+        # Increase the tolerance for large input shapes
+        if M > 35:
+            atol = rtol = 0.02
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def test_int8_moe(self):
+        for params in itertools.product(
+            self.M_int8,
+            self.N_int8,
+            self.K_int8,
+            self.E_int8,
+            self.topk_int8,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                E=params[3],
+                topk=params[4],
+            ):
+                self._int8_moe(*params)
+
+    def _fp8_moe(self, M, N, K, E, topk):
+        dtype = torch.bfloat16
+
+        a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
+
+        w1_fp32 = torch.randn(E, 2 * N, K)
+        w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        w2_fp32 = torch.randn(E, K, N)
+        w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        w1s = (
+            torch.randn(E, math.ceil(2 * N / BLOCK_N), math.ceil(K / BLOCK_K))
+            * factor_for_scale
+        )
+        w2s = (
+            torch.randn(E, math.ceil(K / BLOCK_N), math.ceil(N / BLOCK_K))
+            * factor_for_scale
+        )
+
+        w1_scaled = scaled_weight(w1, w1s)
+        w2_scaled = scaled_weight(w2, w2s)
+
+        score = torch.randn((M, E), dtype=dtype)
+        score = torch.softmax(score, dim=-1, dtype=torch.float32)
+        topk_weight, topk_ids = torch.topk(score, topk)
+
+        w1 = kernel.convert_weight_packed(w1)
+        w2 = kernel.convert_weight_packed(w2)
+
+        ref_out = native_fp8_fused_moe(
+            a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk
+        )
+        out = kernel.fused_experts_cpu(
+            a,
+            w1,
+            w2,
+            topk_weight,
+            topk_ids.to(torch.int32),
+            False,
+            CPUQuantMethod.FP8_W8A16,
+            w1s,
+            w2s,
+            None,
+            None,
+            [BLOCK_N, BLOCK_K],
+            True,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
+
+    def test_fp8_moe(self):
+        for params in itertools.product(
+            self.M_fp8,
+            self.N_fp8,
+            self.K_fp8,
+            self.E_fp8,
+            self.topk_fp8,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                E=params[3],
+                topk=params[4],
+            ):
+                self._fp8_moe(*params)
+
+    def _int4_moe(self, M, N, K, E, topk, group_size=128):
+        dtype = torch.bfloat16
+
+        a = torch.rand(M, K, dtype=dtype) / math.sqrt(K)
+
+        awq_w13_weight = torch.randint(-127, 128, (E, K, 2 * N // 8)).to(torch.int)
+        awq_w13_zero = torch.randint(0, 10, (E, K // group_size, 2 * N // 8)).to(
+            torch.int
+        )
+        awq_w13_scales = torch.rand(E, int(K // group_size), 2 * N).to(torch.bfloat16)
+
+        awq_w2_weight = torch.randint(-127, 128, (E, N, K // 8)).to(torch.int)
+        awq_w2_zero = torch.randint(0, 10, (E, N // group_size, K // 8)).to(torch.int)
+        awq_w2_scales = torch.rand(E, int(N // group_size), K).to(torch.bfloat16)
+        bf16_w13_weight = []
+        bf16_w2_weight = []
+        for i in range(E):
+            bf16_w13_weight_i, _ = unpack_and_dequant_awq(
+                awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i], 4, 128
+            )
+            bf16_w2_weight_i, _ = unpack_and_dequant_awq(
+                awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i], 4, 128
+            )
+            bf16_w13_weight.append(bf16_w13_weight_i)
+            bf16_w2_weight.append(bf16_w2_weight_i)
+        bf16_w13_weight = torch.stack(bf16_w13_weight).detach()
+        bf16_w2_weight = torch.stack(bf16_w2_weight).detach()
+
+        score = torch.rand((M, E), dtype=dtype)
+
+        ref_out = torch_naive_fused_moe(
+            a, bf16_w13_weight, bf16_w2_weight, score, topk, False
+        )
+        score = torch.softmax(score, dim=-1, dtype=torch.float32)
+        topk_weight, topk_ids = torch.topk(score, topk)
+        awq_w13_weight_pack, awq_w13_zero_pack, awq_w13_scales_pack = (
+            torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
+                awq_w13_weight, awq_w13_zero, awq_w13_scales, 0
+            )
+        )
+        awq_w2_weight_pack, awq_w2_zero_pack, awq_w2_scales_pack = (
+            torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
+                awq_w2_weight, awq_w2_zero, awq_w2_scales, 0
+            )
+        )
+
+        out = kernel.fused_experts_cpu(
+            a,
+            awq_w13_weight_pack,
+            awq_w2_weight_pack,
+            topk_weight,
+            topk_ids.to(torch.int32),
+            False,
+            CPUQuantMethod.INT4_W4A8,
+            awq_w13_scales_pack,
+            awq_w2_scales_pack,
+            awq_w13_zero_pack,
+            awq_w2_zero_pack,
+            None,
+            True,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
+
+    def test_int4_moe(self):
+        for params in itertools.product(
+            self.M_int4,
+            self.N_int4,
+            self.K_int4,
+            self.E_int4,
+            self.topk_int4,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                E=params[3],
+                topk=params[4],
+            ):
+                self._int4_moe(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_norm.py

@@ -0,0 +1,435 @@
+import itertools
+import unittest
+from typing import Optional, Tuple, Union
+
+import torch
+from utils import make_non_contiguous, parametrize, precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestNorm(CustomTestCase):
+
+    def _forward_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float = 1e-6,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        if residual is not None:
+            x = x + residual.to(torch.float32)
+            residual = x.to(orig_dtype)
+
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + variance_epsilon)
+        x = x.to(orig_dtype) * weight
+        if residual is None:
+            return x
+        else:
+            return x, residual
+
+    def _norm(self, x, eps):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+    def _gemma3_rmsnorm_native(
+        self, x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float = 1e-6
+    ):
+        output = self._norm(x.float(), variance_epsilon)
+        output = output * (1.0 + weight.float())
+        return output.type_as(x)
+
+    def _gemma_rmsnorm_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float = 1e-6,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        if residual is not None:
+            x = x + residual
+            residual = x
+
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + variance_epsilon)
+        x = x * (1.0 + weight.float())
+        x = x.to(orig_dtype)
+        return x if residual is None else (x, residual)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_norm(self, m, n, dtype):
+
+        x = torch.randn([m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, variance_epsilon)
+        ref_out = self._forward_native(x, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+        ref_x = x.clone()
+        residual = torch.randn([m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(
+            x, residual, weight, variance_epsilon
+        )
+
+        ref_x, ref_residual = self._forward_native(
+            ref_x, weight, variance_epsilon, ref_residual
+        )
+
+        torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+    @parametrize(
+        l=[1, 2],
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_norm_3d(self, l, m, n, dtype):
+
+        x = torch.randn([l, m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, variance_epsilon)
+        ref_out = self._forward_native(x, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+        ref_x = x.clone()
+        residual = torch.randn([l, m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(
+            x, residual, weight, variance_epsilon
+        )
+
+        ref_x, ref_residual = self._forward_native(
+            ref_x, weight, variance_epsilon, ref_residual
+        )
+
+        torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_l2norm(self, m, n, dtype):
+
+        x = torch.randn([m, n], dtype=dtype)
+        hidden_size = x.size(-1)
+        fake_ones_weight = torch.ones(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        out = torch.ops.sgl_kernel.l2norm_cpu(x, variance_epsilon)
+        ref_out = self._forward_native(x, fake_ones_weight, variance_epsilon)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_gemma_rmsnorm(self, m, n, dtype):
+
+        x = torch.randn([m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        out = torch.ops.sgl_kernel.gemma_rmsnorm_cpu(x, weight, variance_epsilon)
+        ref_out = self._gemma_rmsnorm_native(x, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+        ref_x = x.clone()
+        residual = torch.randn([m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu(
+            x, residual, weight, variance_epsilon
+        )
+
+        ref_x, ref_residual = self._gemma_rmsnorm_native(
+            ref_x, weight, variance_epsilon, ref_residual
+        )
+
+        torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_gemma3_rmsnorm(self, m, n, dtype):
+        x_list = [
+            torch.randn([m, n], dtype=dtype),
+            torch.randn([1, m, 2, n], dtype=dtype),
+        ]
+        for x in x_list:
+            x = make_non_contiguous(x)
+            hidden_size = x.size(-1)
+            weight = torch.randn(hidden_size, dtype=dtype)
+            variance_epsilon = 1e-6
+            out = torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, weight, variance_epsilon)
+            ref_out = self._gemma3_rmsnorm_native(x, weight, variance_epsilon)
+
+            atol = rtol = precision[ref_out.dtype]
+            torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def _gemma4_rmsnorm_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float = 1e-6,
+        scale_shift: float = 0.0,
+        with_scale: bool = True,
+    ):
+        output = self._norm(x.float(), variance_epsilon)
+        if with_scale:
+            output = output * (weight.float() + scale_shift)
+        return output.type_as(x)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_gemma4_rmsnorm(self, m, n, dtype):
+        for scale_shift, with_scale in [
+            (0.0, True),
+            (1.0, True),
+            (0.0, False),
+            (1.0, False),
+        ]:
+            x_list = [
+                torch.randn([m, n], dtype=dtype),
+                torch.randn([4, m, n], dtype=dtype),
+            ]
+            # Add non-block-contiguous 3D input
+            base = torch.randn([4, 2 * m, n], dtype=dtype)
+            x_list.append(base[:, :m, :])
+
+            for x in x_list:
+                x = make_non_contiguous(x)
+                hidden_size = x.size(-1)
+                weight = torch.randn(hidden_size, dtype=dtype)
+                variance_epsilon = 1e-6
+
+                out = torch.ops.sgl_kernel.gemma4_rmsnorm_cpu(
+                    x, weight, variance_epsilon, scale_shift, with_scale
+                )
+                ref_out = self._gemma4_rmsnorm_native(
+                    x, weight, variance_epsilon, scale_shift, with_scale
+                )
+
+                atol = rtol = precision[ref_out.dtype]
+                torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+
+class TestFusedRMSNormGated(CustomTestCase):
+    M = [4096, 1024]
+    N = [4096, 4096 + 13]
+    dtype = [torch.float16, torch.bfloat16]
+
+    def _forward_native(
+        self,
+        hidden_states: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float = 1e-6,
+        gate: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        # Norm before gate
+        hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
+        hidden_states = weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states * torch.nn.functional.silu(gate.to(torch.float32))
+
+        return hidden_states.to(input_dtype)
+
+    def _norm_test(self, m, n, dtype):
+
+        x = torch.randn([m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        batch_size = x.size(0)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+        gate = torch.randn([batch_size, hidden_size], dtype=dtype)
+
+        out = torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
+            x, weight, gate, variance_epsilon
+        )
+        ref_out = self._forward_native(x, weight, variance_epsilon, gate)
+
+        atol = rtol = precision[ref_out.dtype] * 2
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+    def test_norm(self):
+        for params in itertools.product(self.M, self.N, self.dtype):
+            with self.subTest(m=params[0], n=params[1], dtype=params[2]):
+                self._norm_test(*params)
+
+
+class TestLayerNorm(CustomTestCase):
+
+    def _forward_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float,
+        residual: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        if residual is not None:
+            x = x + residual.to(torch.float32)
+            residual = x.to(orig_dtype)
+
+        variance, mean = torch.var_mean(x, dim=-1, keepdim=True, correction=0)
+        x = (x - mean) * torch.rsqrt(variance + variance_epsilon)
+        x = x * weight.to(torch.float32)
+        if bias is not None:
+            x = x + bias.to(torch.float32)
+        x = x.to(orig_dtype)
+        return x if residual is None else (x, residual)
+
+    @parametrize(
+        m=[4096, 1024],
+        n=[4096, 4109],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_norm_input_2d(self, m: int, n: int, dtype: torch.dtype) -> None:
+        x = torch.randn([m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        bias = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon)
+        ref_ln_out = self._forward_native(x, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_ln_out.dtype]
+        torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
+
+        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon)
+        ref_ln_out = self._forward_native(
+            x, weight, variance_epsilon, residual=None, bias=bias
+        )
+        torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
+
+        residual = torch.randn([m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+            x, residual, weight, None, variance_epsilon
+        )
+        ref_add_ln_out, ref_residual = self._forward_native(
+            x, weight, variance_epsilon, residual=ref_residual
+        )
+
+        torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+        residual = torch.randn([m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+            x, residual, weight, bias, variance_epsilon
+        )
+        ref_add_ln_out, ref_residual = self._forward_native(
+            x, weight, variance_epsilon, residual=ref_residual, bias=bias
+        )
+
+        torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+    @parametrize(
+        l=[4096, 1024],
+        m=[1, 4],
+        n=[4096, 4109, 2304],
+        dtype=[torch.float16, torch.bfloat16],
+    )
+    def test_norm_input_3d(self, l: int, m: int, n: int, dtype: torch.dtype) -> None:
+        x = torch.randn([l, m, n], dtype=dtype)
+        x = make_non_contiguous(x)
+        hidden_size = x.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        bias = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon)
+        ref_ln_out = self._forward_native(x, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_ln_out.dtype]
+        torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
+
+        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon)
+        ref_ln_out = self._forward_native(
+            x, weight, variance_epsilon, residual=None, bias=bias
+        )
+        torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
+
+        residual = torch.randn([l, m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+            x, residual, weight, None, variance_epsilon
+        )
+        ref_add_ln_out, ref_residual = self._forward_native(
+            x, weight, variance_epsilon, ref_residual
+        )
+
+        torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+        residual = torch.randn([l, m, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+            x, residual, weight, bias, variance_epsilon
+        )
+        ref_add_ln_out, ref_residual = self._forward_native(
+            x, weight, variance_epsilon, residual=ref_residual, bias=bias
+        )
+
+        torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_qkv_proj_with_rope.py

@@ -0,0 +1,443 @@
+import unittest
+
+import torch
+from utils import (
+    convert_weight,
+    native_w8a8_per_token_matmul,
+    per_token_quant_int8,
+    precision,
+)
+
+from sglang.srt.layers.quantization.fp8_utils import input_to_float8
+from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+convert_weight_packed = torch.ops.sgl_kernel.convert_weight_packed
+qkv_proj_with_rope = torch.ops.sgl_kernel.qkv_proj_with_rope
+qkv_proj_with_rope_fused_weight = torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight
+torch.manual_seed(1234)
+# constants
+kv_lora_rank = 512
+qk_head_dim = 192
+qk_nope_head_dim = 128
+qk_rope_head_dim = 64
+rotary_dim = qk_rope_head_dim
+num_heads = 22
+q_lora_rank = 1536
+hidden_size = 7168
+B = 1
+eps = 1e-6
+
+
+def layernorm(x, weight, variance_epsilon=1e-6, residual=None):
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + variance_epsilon)
+    return (x * weight).to(orig_dtype)
+
+
+def rotary_emb(q_pe, k_pe, pos, cos_sin_cache):
+    orig_dtype = q_pe.dtype
+    q_pe = q_pe.float()
+    k_pe = k_pe.float()
+    cos_sin_cache = cos_sin_cache.float()
+
+    query_rot = q_pe[..., :rotary_dim]
+    key_rot = k_pe[..., :rotary_dim]
+    cos_sin = cos_sin_cache[pos]
+    cos, sin = cos_sin.chunk(2, dim=-1)
+    query_rot = apply_rotary_emb(query_rot, cos, sin, False)
+    key_rot = apply_rotary_emb(key_rot, cos, sin, False)
+    return query_rot.to(orig_dtype), key_rot.to(orig_dtype)
+
+
+def native_torch(
+    q_input,
+    hidden_states,
+    q_a_proj_weight,
+    norm_weight1,
+    q_b_proj_weight,
+    w_kc,
+    kv_a_proj_weight,
+    norm_weight2,
+    pos,
+    cos_sin_cache,
+):
+
+    q = torch.matmul(hidden_states, q_a_proj_weight.t())
+    q = layernorm(q, norm_weight1)
+    q = torch.matmul(q, q_b_proj_weight.t()).view(-1, num_heads, qk_head_dim)
+
+    q_nope, q_pe = q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
+    q_nope_out = torch.bmm(q_nope.transpose(0, 1), w_kc)
+
+    q_input[..., :kv_lora_rank] = q_nope_out.transpose(0, 1)
+    latent_cache = torch.matmul(hidden_states, kv_a_proj_weight.t())
+    v_input = latent_cache[..., :kv_lora_rank]
+    v_input = layernorm(v_input.contiguous(), norm_weight2).unsqueeze(1)
+    k_input = latent_cache.unsqueeze(1)
+    k_input[..., :kv_lora_rank] = v_input
+    k_pe = k_input[..., kv_lora_rank:]
+
+    q_pe, k_pe = rotary_emb(q_pe, k_pe, pos, cos_sin_cache)
+    q_input[..., kv_lora_rank:] = q_pe
+    k_input[..., kv_lora_rank:] = k_pe
+
+    return q_input, k_input, v_input
+
+
+def native_torch_int8(
+    q_input,
+    hidden_states,
+    w1_q,
+    w1_s,
+    norm_weight1,
+    w2_q,
+    w2_s,
+    w_kc,
+    w3_q,
+    w3_s,
+    norm_weight2,
+    pos,
+    cos_sin_cache,
+):
+
+    a_q, a_s = per_token_quant_int8(hidden_states)
+    q = native_w8a8_per_token_matmul(a_q, w1_q, a_s, w1_s, None, torch.bfloat16)
+    q = layernorm(q, norm_weight1)
+
+    a_q, a_s = per_token_quant_int8(q)
+    q = native_w8a8_per_token_matmul(a_q, w2_q, a_s, w2_s, None, torch.bfloat16).view(
+        -1, num_heads, qk_head_dim
+    )
+
+    q_nope, q_pe = q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
+    q_nope_out = torch.bmm(q_nope.transpose(0, 1), w_kc)
+
+    q_input[..., :kv_lora_rank] = q_nope_out.transpose(0, 1)
+    a_q, a_s = per_token_quant_int8(hidden_states)
+    latent_cache = native_w8a8_per_token_matmul(
+        a_q, w3_q, a_s, w3_s, None, torch.bfloat16
+    )
+    v_input = latent_cache[..., :kv_lora_rank]
+    v_input = layernorm(v_input.contiguous(), norm_weight2).unsqueeze(1)
+    k_input = latent_cache.unsqueeze(1)
+    k_input[..., :kv_lora_rank] = v_input
+    k_pe = k_input[..., kv_lora_rank:]
+
+    q_pe, k_pe = rotary_emb(q_pe, k_pe, pos, cos_sin_cache)
+    q_input[..., kv_lora_rank:] = q_pe
+    k_input[..., kv_lora_rank:] = k_pe
+
+    return q_input, k_input, v_input
+
+
+class TestQKVProjWithROPE(CustomTestCase):
+    def test_bf16_qkv_proj_with_rope(self):
+        dtype = torch.bfloat16
+        hidden_states = torch.randn(B, hidden_size, dtype=dtype) / hidden_size
+        q_input = torch.empty(
+            B, num_heads, kv_lora_rank + qk_rope_head_dim, dtype=dtype
+        )
+        q_a_proj_weight = torch.randn(q_lora_rank, hidden_size, dtype=dtype) * 0.1
+        norm_weight1 = torch.randn(q_lora_rank, dtype=dtype)
+        q_b_proj_weight = (
+            torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1
+        )
+        w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1
+        kv_a_proj_weight = (
+            torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1
+        )
+        fused_weight = torch.cat([q_a_proj_weight, kv_a_proj_weight], dim=0)
+        norm_weight2 = torch.randn(kv_lora_rank, dtype=dtype)
+        pos = torch.randint(10, 100, (B,))
+        cos_sin_cache = torch.randn(100, rotary_dim, dtype=dtype)
+        q_ref, k_ref, v_ref = native_torch(
+            q_input,
+            hidden_states,
+            q_a_proj_weight,
+            norm_weight1,
+            q_b_proj_weight,
+            w_kc.transpose(1, 2),
+            kv_a_proj_weight,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+        )
+        qa_packed = convert_weight_packed(q_a_proj_weight)
+        qb_packed = convert_weight_packed(q_b_proj_weight)
+        kva_packed = convert_weight_packed(kv_a_proj_weight)
+        wkc_packed = convert_weight_packed(w_kc)
+        fused_weight_packed = convert_weight_packed(fused_weight)
+
+        q_out, k_out, v_out = qkv_proj_with_rope(
+            hidden_states,
+            qa_packed,
+            qb_packed,
+            kva_packed,
+            wkc_packed,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            False,
+            False,
+            None,
+            None,
+            None,
+            None,
+            True,
+            None,
+        )
+        fused_q_out, fused_k_out, fused_v_out = qkv_proj_with_rope_fused_weight(
+            hidden_states,
+            fused_weight_packed,
+            qb_packed,
+            wkc_packed,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            False,
+            False,
+            None,
+            None,
+            None,
+            True,
+            None,
+            q_lora_rank,
+            kv_lora_rank,
+            qk_rope_head_dim,
+        )
+        atol = rtol = precision[q_ref.dtype]
+        torch.testing.assert_close(q_ref, q_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k_ref, k_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(v_ref, v_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(fused_q_out, q_out)
+        torch.testing.assert_close(fused_k_out, k_out)
+        torch.testing.assert_close(fused_v_out, v_out)
+
+    def test_int8_qkv_proj_with_rope(self):
+        dtype = torch.bfloat16
+        hidden_states = torch.randn(B, hidden_size, dtype=dtype) / hidden_size
+        q_input = torch.empty(
+            B, num_heads, kv_lora_rank + qk_rope_head_dim, dtype=dtype
+        )
+        q_a_proj_weight = torch.randn(q_lora_rank, hidden_size, dtype=dtype) * 0.1
+        norm_weight1 = torch.randn(q_lora_rank, dtype=dtype)
+        q_b_proj_weight = (
+            torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1
+        )
+        w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1
+        kv_a_proj_weight = (
+            torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1
+        )
+        norm_weight2 = torch.randn(kv_lora_rank, dtype=dtype)
+        pos = torch.randint(10, 100, (B,))
+        cos_sin_cache = torch.randn(100, rotary_dim, dtype=dtype)
+
+        w1_q, w1_s = per_token_quant_int8(q_a_proj_weight)
+        w2_q, w2_s = per_token_quant_int8(q_b_proj_weight)
+        w3_q, w3_s = per_token_quant_int8(kv_a_proj_weight)
+        q_ref, k_ref, v_ref = native_torch_int8(
+            q_input,
+            hidden_states,
+            w1_q,
+            w1_s,
+            norm_weight1,
+            w2_q,
+            w2_s,
+            w_kc.transpose(1, 2),
+            w3_q,
+            w3_s,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+        )
+        w1_q_packed = convert_weight_packed(w1_q)
+        w2_q_packed = convert_weight_packed(w2_q)
+        w3_q_packed = convert_weight_packed(w3_q)
+        wkc_packed = convert_weight_packed(w_kc)
+        q_out, k_out, v_out = qkv_proj_with_rope(
+            hidden_states,
+            w1_q_packed,
+            w2_q_packed,
+            w3_q_packed,
+            wkc_packed,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            True,
+            False,
+            w1_s,
+            w2_s,
+            w3_s,
+            None,
+            True,
+            None,
+        )
+        fused_weight = torch.cat([w1_q, w3_q], dim=0)
+        fused_weight_s = torch.cat([w1_s, w3_s], dim=0)
+        w_fused_q_packed = convert_weight_packed(fused_weight)
+        fused_q_out, fused_k_out, fused_v_out = qkv_proj_with_rope_fused_weight(
+            hidden_states,
+            w_fused_q_packed,
+            w2_q_packed,
+            wkc_packed,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            True,
+            False,
+            fused_weight_s,
+            w2_s,
+            None,
+            True,
+            None,
+            q_lora_rank,
+            kv_lora_rank,
+            qk_rope_head_dim,
+        )
+        atol = rtol = precision[q_ref.dtype]
+        torch.testing.assert_close(q_ref, q_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(k_ref, k_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(v_ref, v_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(fused_q_out, q_out)
+        torch.testing.assert_close(fused_k_out, k_out)
+        torch.testing.assert_close(fused_v_out, v_out)
+
+    def test_fp8_qkv_proj_with_rope(self):
+        dtype = torch.bfloat16
+        hidden_states = torch.randn(B, hidden_size, dtype=dtype) / hidden_size
+        q_input = torch.empty(
+            B, num_heads, kv_lora_rank + qk_rope_head_dim, dtype=dtype
+        )
+        q_a_proj_weight = torch.randn(q_lora_rank, hidden_size, dtype=dtype) * 0.1
+        norm_weight1 = torch.randn(q_lora_rank, dtype=dtype)
+        q_b_proj_weight = (
+            torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1
+        )
+        w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1
+        w_kc_q, w_kc_s = input_to_float8(w_kc)
+        kv_a_proj_weight = (
+            torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1
+        )
+        norm_weight2 = torch.randn(kv_lora_rank, dtype=dtype)
+        pos = torch.randint(10, 100, (B,))
+        cos_sin_cache = torch.randn(100, rotary_dim, dtype=dtype)
+
+        scale_block_size_N = 128
+        scale_block_size_K = 128
+        fp8_q_a_proj_weight, q_a_proj_weight_scale_inv, q_a_proj_weight_dq = (
+            convert_weight(
+                q_a_proj_weight,
+                [scale_block_size_N, scale_block_size_K],
+                torch.bfloat16,
+            )
+        )
+        fp8_q_b_proj_weight, q_b_proj_weight_scale_inv, q_b_proj_weight_dq = (
+            convert_weight(
+                q_b_proj_weight,
+                [scale_block_size_N, scale_block_size_K],
+                torch.bfloat16,
+            )
+        )
+        (
+            fp8_kv_a_proj_with_mqa_weight,
+            kv_a_proj_with_mqa_weight_scale_inv,
+            kv_a_proj_with_mqa_weight_dq,
+        ) = convert_weight(
+            kv_a_proj_weight, [scale_block_size_N, scale_block_size_K], torch.bfloat16
+        )
+        w_kc_dq = w_kc_q.to(torch.bfloat16) * w_kc_s
+        q_ref, k_ref, v_ref = native_torch(
+            q_input,
+            hidden_states,
+            q_a_proj_weight_dq,
+            norm_weight1,
+            q_b_proj_weight_dq,
+            w_kc_dq.transpose(1, 2),
+            kv_a_proj_with_mqa_weight_dq,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+        )
+        fp8_q_a_proj_weight_packed = convert_weight_packed(fp8_q_a_proj_weight)
+        fp8_q_b_proj_weight_packed = convert_weight_packed(fp8_q_b_proj_weight)
+        fp8_kv_a_proj_with_mqa_weight_packed = convert_weight_packed(
+            fp8_kv_a_proj_with_mqa_weight
+        )
+        w_kc_q = convert_weight_packed(w_kc_q)
+        q_out, k_out, v_out = qkv_proj_with_rope(
+            hidden_states,
+            fp8_q_a_proj_weight_packed,
+            fp8_q_b_proj_weight_packed,
+            fp8_kv_a_proj_with_mqa_weight_packed,
+            w_kc_q,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            False,
+            True,
+            q_a_proj_weight_scale_inv.float(),
+            q_b_proj_weight_scale_inv.float(),
+            kv_a_proj_with_mqa_weight_scale_inv.float(),
+            w_kc_s,
+            True,
+            [scale_block_size_N, scale_block_size_K],
+        )
+
+        fused_weight = torch.cat(
+            [fp8_q_a_proj_weight, fp8_kv_a_proj_with_mqa_weight], dim=0
+        )
+        fused_weight_s = torch.cat(
+            [q_a_proj_weight_scale_inv, kv_a_proj_with_mqa_weight_scale_inv], dim=0
+        )
+        fused_weight_packed = convert_weight_packed(fused_weight)
+        fused_q_out, fused_k_out, fused_v_out = qkv_proj_with_rope_fused_weight(
+            hidden_states,
+            fused_weight_packed,
+            fp8_q_b_proj_weight_packed,
+            w_kc_q,
+            norm_weight1,
+            norm_weight2,
+            pos,
+            cos_sin_cache,
+            eps,
+            False,
+            True,
+            fused_weight_s.float(),
+            q_b_proj_weight_scale_inv.float(),
+            w_kc_s,
+            True,
+            [scale_block_size_N, scale_block_size_K],
+            q_lora_rank,
+            kv_lora_rank,
+            qk_rope_head_dim,
+        )
+        atol = rtol = precision[q_ref.dtype]
+        # Due to the change in multiplication order, the error is amplified.
+        # In the model, with fewer layers, this doesn't cause issues, but in
+        # tests with more layers, we need to enlarge the tolerance to pass the tests.
+        torch.testing.assert_close(q_ref, q_out, atol=1e-1, rtol=1e-1)
+        torch.testing.assert_close(k_ref, k_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(v_ref, v_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(fused_q_out, q_out)
+        torch.testing.assert_close(fused_k_out, k_out)
+        torch.testing.assert_close(fused_v_out, v_out)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_qwen3.py

@@ -0,0 +1,152 @@
+import unittest
+
+import torch
+from utils import precision
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+def fix_query_key_value_ordering_reshape_cat(
+    mixed_qkvz, mixed_ba, num_k_heads, num_v_heads, attn_tp_size, head_k_dim, head_v_dim
+):
+    new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
+        num_k_heads // attn_tp_size,
+        (
+            head_k_dim
+            + head_k_dim
+            + (head_v_dim + head_v_dim) * num_v_heads // num_k_heads
+        ),
+    )
+    new_tensor_shape_ba = mixed_ba.size()[:-1] + (
+        num_k_heads // attn_tp_size,
+        2 * num_v_heads // num_k_heads,
+    )
+
+    mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
+    mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
+
+    split_arg_list_qkvz = [
+        head_k_dim,
+        head_k_dim,
+        (num_v_heads // num_k_heads * head_v_dim),
+        (num_v_heads // num_k_heads * head_v_dim),
+    ]
+    split_arg_list_ba = [
+        num_v_heads // num_k_heads,
+        num_v_heads // num_k_heads,
+    ]
+    # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)]
+    # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng]
+    query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2)
+    b, a = torch.split(mixed_ba, split_arg_list_ba, dim=2)
+
+    # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
+    value = value.reshape(value.size(0), -1, head_v_dim)
+    z = z.reshape(z.size(0), -1, head_v_dim)
+    b = b.reshape(b.size(0), num_v_heads // attn_tp_size)
+    a = a.reshape(a.size(0), num_v_heads // attn_tp_size)
+    query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
+    mixed_qkv = torch.cat((query, key, value), dim=-1)
+
+    return mixed_qkv, z, b, a
+
+
+def fix_query_key_value_ordering_reshape_cat_contiguous(
+    mixed_qkvz: torch.Tensor,
+    mixed_ba: torch.Tensor,
+    key_dim: int,
+    value_dim: int,
+    num_v_heads: int,
+    head_v_dim: int,
+    attn_tp_size: int,
+):
+    """
+    Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
+    """
+    k_tp = key_dim // attn_tp_size
+    v_tp = value_dim // attn_tp_size
+    nv_tp = num_v_heads // attn_tp_size
+
+    # Directly split, no head group reshape
+    query, key, value, z = mixed_qkvz.split([k_tp, k_tp, v_tp, v_tp], dim=-1)
+    b, a = mixed_ba.split([nv_tp, nv_tp], dim=-1)
+
+    # value / z reshape to (seq, num_v_heads/tp, head_v_dim)
+    value = value.reshape(value.size(0), -1, head_v_dim)
+    z = z.reshape(z.size(0), -1, head_v_dim)
+    query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
+    mixed_qkv = torch.cat((query, key, value), dim=-1)
+    return mixed_qkv, z, b, a
+
+
+class TestQwen3(CustomTestCase):
+    def test_fused_qkvzba_split_reshape_cat(self):
+        mixed_qkvz = torch.rand(1024, 12288, dtype=torch.bfloat16)
+        mixed_ba = torch.rand(1024, 64, dtype=torch.bfloat16)
+        head_k_dim = 128
+        head_v_dim = 128
+        num_v_heads = 32
+        num_k_heads = 16
+        attn_tp_size = 1
+        mixed_qkv_ref, z_ref, b_ref, a_ref = fix_query_key_value_ordering_reshape_cat(
+            mixed_qkvz,
+            mixed_ba,
+            num_k_heads,
+            num_v_heads,
+            attn_tp_size,
+            head_k_dim,
+            head_v_dim,
+        )
+        num_heads_qk = num_k_heads // attn_tp_size
+        num_heads_v = num_v_heads // attn_tp_size
+        mixed_qkv, z, b, a = torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_cpu(
+            mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
+        )
+        atol = rtol = precision[mixed_qkv.dtype]
+        torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
+
+    def test_fused_qkvzba_split_reshape_cat_contiguous(self):
+        mixed_qkvz = torch.rand(1, 12288, dtype=torch.bfloat16)
+        mixed_ba = torch.rand(1, 64, dtype=torch.bfloat16)
+        head_k_dim = 128
+        head_v_dim = 128
+        num_v_heads = 32
+        num_k_heads = 16
+        attn_tp_size = 1
+        key_dim = head_k_dim * num_k_heads
+        value_dim = head_v_dim * num_v_heads
+        mixed_qkv_ref, z_ref, b_ref, a_ref = (
+            fix_query_key_value_ordering_reshape_cat_contiguous(
+                mixed_qkvz,
+                mixed_ba,
+                key_dim,
+                value_dim,
+                num_v_heads,
+                head_v_dim,
+                attn_tp_size,
+            )
+        )
+        num_heads_qk = num_k_heads // attn_tp_size
+        num_heads_v = num_v_heads // attn_tp_size
+        mixed_qkv, z, b, a = (
+            torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu(
+                mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
+            )
+        )
+        atol = rtol = precision[mixed_qkv.dtype]
+        torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
+        torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_rope.py

@@ -0,0 +1,285 @@
+import unittest
+
+import torch
+from utils import precision
+
+from sglang.srt.layers.rotary_embedding import (
+    MRotaryEmbedding,
+    RotaryEmbedding,
+)
+from sglang.srt.layers.rotary_embedding.rope_variant import (
+    DeepseekScalingRotaryEmbedding,
+    apply_rotary_pos_emb_native,
+)
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestROPE(CustomTestCase):
+    def test_mrope(self):
+        torch.manual_seed(100)
+        head_size = 128
+        seq_len = 512
+        num_heads = 16
+        num_kv_heads = 1
+        rotary_dim = 128
+        max_pos = 262144
+        base = 5000000
+        is_neox_style = True
+        dtype = torch.bfloat16
+        mrope_section = [24, 20, 20]
+        mrope_interleaved = True
+        positions_mrope = torch.randint(0, max_pos, (3, seq_len))
+        positions_text = torch.randint(0, max_pos, (seq_len,))
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+        test_config = [
+            # (dtype, is_neox_stype, mrope_interleaved, positions, mrope_section)
+            (torch.bfloat16, False, True, positions_mrope, mrope_section),
+            (torch.bfloat16, False, False, positions_mrope, mrope_section),
+            (torch.bfloat16, False, False, positions_text, None),
+            (torch.bfloat16, True, True, positions_mrope, mrope_section),
+            (torch.bfloat16, True, False, positions_mrope, mrope_section),
+            (torch.bfloat16, True, False, positions_text, None),
+        ]
+        for (
+            dtype,
+            is_neox_style,
+            mrope_interleaved,
+            positions,
+            mrope_section,
+        ) in test_config:
+            rope = MRotaryEmbedding(
+                head_size,
+                rotary_dim,
+                max_pos,
+                base,
+                is_neox_style,
+                dtype,
+                mrope_section,
+                mrope_interleaved,
+            )
+            enable_autocast = True
+
+            with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast):
+                q = torch.randn(seq_len, num_heads * head_size, dtype=dtype)
+                q_clone = q.clone()
+                k = torch.randn(seq_len, num_kv_heads * head_size, dtype=dtype)
+                k_clone = k.clone()
+
+                # ref kernel
+                q_ref, k_ref = rope.forward_native(
+                    query=q,
+                    key=k,
+                    positions=positions,
+                )
+                # fused rope kernel
+                q_sgl, k_sgl = torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu(
+                    positions,
+                    q_clone,
+                    k_clone,
+                    rope.head_size,
+                    rope.cos_sin_cache,
+                    rope.mrope_section,
+                    rope.mrope_interleaved,
+                    is_neox_style,
+                )
+                atol = rtol = precision[q_ref.dtype]
+                torch.testing.assert_close(q_ref, q_sgl, atol=atol, rtol=rtol)
+                torch.testing.assert_close(k_ref, k_sgl, atol=atol, rtol=rtol)
+
+    def test_deepseek_v2_rope(self):
+        num_head = 16
+        seq_len = 1024
+        q_head_dim = 192
+        qk_nope_head_dim = 128
+        qk_rope_head_dim = 64
+        max_pos = 256
+        k_dim = 576
+        rotary_dim = 64
+        is_neox_style = False
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+        # Create cos_sin_cache
+        freqs = torch.rand(max_pos, qk_rope_head_dim // 2)
+        cos = freqs.cos() * 0.7
+        sin = freqs.sin() * 0.7
+        cos_sin_cache = torch.cat((cos, sin), dim=-1).to(torch.bfloat16)
+        positions = torch.randint(0, max_pos, (seq_len,))
+
+        rope = DeepseekScalingRotaryEmbedding(
+            qk_rope_head_dim,
+            rotary_dim,
+            max_pos,
+            16,  # not used since cos_sin_cache is provided
+            is_neox_style,
+            1.0,
+            torch.bfloat16,
+            device="cpu",
+        )
+        rope.register_buffer("cos_sin_cache", cos_sin_cache)
+
+        for dtype in [torch.bfloat16]:
+            enable_autocast = True
+
+            with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast):
+                q = torch.randn(seq_len, num_head, q_head_dim, dtype=dtype)
+                q_clone = q.clone()
+                k = torch.randn(seq_len, 1, k_dim, dtype=dtype)
+                k_clone = k.clone()
+                _, q_pe = q.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
+                _, q_pe_clone = q_clone.split(
+                    [qk_nope_head_dim, qk_rope_head_dim], dim=-1
+                )
+                k_pe = k[:, :, k_dim - qk_rope_head_dim :]
+                k_pe_clone = k_clone[:, :, k_dim - qk_rope_head_dim :]
+
+                # ref kernel
+                q_pe, k_pe = rope.forward_native(
+                    query=q_pe,
+                    key=k_pe,
+                    positions=positions,
+                )
+
+                # fused rope kernel
+                q_pe_clone, k_pe_clone = torch.ops.sgl_kernel.rotary_embedding_cpu(
+                    positions,
+                    q_pe_clone,
+                    k_pe_clone,
+                    rope.head_size,
+                    cos_sin_cache,
+                    False,
+                )
+
+                atol = rtol = precision[q_pe.dtype]
+                torch.testing.assert_close(q_pe, q_pe_clone, atol=atol, rtol=rtol)
+                torch.testing.assert_close(k_pe, k_pe_clone, atol=atol, rtol=rtol)
+                torch.testing.assert_close(k_pe, k_pe_clone)
+
+    def test_origin_rope(self):
+        def single_test(
+            head_size: int,
+            rotary_dim: int,
+            max_position_embeddings: int,
+            base: int,
+            dims: int,
+            is_neox_style: bool,
+            dtype: torch.dtype,
+            device: str,
+            batch_size: int,
+            seq_len: int,
+            num_q_heads: int,
+            num_kv_heads: int,
+        ):
+            set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+            torch.manual_seed(100)
+            rope_ref = RotaryEmbedding(
+                head_size,
+                rotary_dim,
+                max_position_embeddings,
+                base,
+                is_neox_style,
+                dtype,
+            ).to(device)
+            pos_ids = torch.arange(seq_len, device=device).repeat(batch_size)
+            query = torch.randn(
+                batch_size * seq_len,
+                num_q_heads * head_size,
+                dtype=dtype,
+                device=device,
+            )
+            key = torch.randn(
+                batch_size * seq_len,
+                num_kv_heads * head_size,
+                dtype=dtype,
+                device=device,
+            )
+            if dims == 4:
+                query = query.view(batch_size, seq_len, num_q_heads, head_size)
+                key = key.view(batch_size, seq_len, num_kv_heads, head_size)
+            query_ref, key_ref = query.clone(), key.clone()
+            query_cpu, key_cpu = query.clone(), key.clone()
+
+            query_ref_out, key_ref_out = rope_ref.forward_native(
+                pos_ids, query_ref, key_ref
+            )
+            query_cpu_out, key_cpu_out = torch.ops.sgl_kernel.rotary_embedding_cpu(
+                pos_ids,
+                query_cpu,
+                key_cpu,
+                rope_ref.head_size,
+                rope_ref.cos_sin_cache.to(query.dtype),
+                rope_ref.is_neox_style,
+            )
+            torch.testing.assert_close(
+                query_ref_out, query_cpu_out, atol=1e-2, rtol=1e-2
+            )
+            torch.testing.assert_close(key_ref_out, key_cpu_out, atol=1e-2, rtol=1e-2)
+
+        test_config = [
+            (64, 64, 32, 8000, True, torch.bfloat16, "cpu", 32, 32, 1, 1),
+            (256, 128, 4096, 10000, True, torch.bfloat16, "cpu", 2, 512, 32, 8),
+            (512, 128, 311, 10000, True, torch.bfloat16, "cpu", 3, 39, 4, 2),
+            (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 32, 8),
+            (128, 128, 2048, 10000, False, torch.bfloat16, "cpu", 2, 512, 16, 4),
+            (512, 128, 311, 10000, False, torch.bfloat16, "cpu", 3, 39, 4, 2),
+        ]
+
+        for (
+            head_size,
+            rotary_dim,
+            max_position_embeddings,
+            base,
+            is_neox_style,
+            dtype,
+            device,
+            batch_size,
+            seq_len,
+            num_q_heads,
+            num_kv_heads,
+        ) in test_config:
+            for dim in [2, 4]:
+                single_test(
+                    head_size,
+                    rotary_dim,
+                    max_position_embeddings,
+                    base,
+                    dim,
+                    is_neox_style,
+                    dtype,
+                    device,
+                    batch_size,
+                    seq_len,
+                    num_q_heads,
+                    num_kv_heads,
+                )
+
+    def test_apply_rotary_pos_emb(self):
+        num_tokens = 1024
+        num_heads = 8
+        head_size = 72
+        qkv = torch.randn(num_tokens, num_heads * head_size * 3).to(torch.bfloat16)
+        query, key, _ = qkv.split(
+            [num_heads * head_size, num_heads * head_size, num_heads * head_size],
+            dim=-1,
+        )
+        query = query.view(num_tokens, num_heads, head_size)
+        key = key.view(num_tokens, num_heads, head_size)
+        for sincos_dtype in [torch.float32, torch.bfloat16]:
+            cos = torch.rand(num_tokens, head_size).to(sincos_dtype)
+            sin = torch.rand(num_tokens, head_size).to(sincos_dtype)
+            q_out_ref, k_out_ref = apply_rotary_pos_emb_native(query, key, cos, sin)
+            q_out_sgl, k_out_sgl = torch.ops.sgl_kernel.apply_rotary_pos_emb_cpu(
+                query, key, cos, sin
+            )
+            torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2)
+            torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_server_args_backend.py

@@ -0,0 +1,38 @@
+import unittest
+from unittest.mock import patch
+
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+
+class TestServerArgsCPUBackend(unittest.TestCase):
+    def _make_server_args(self, attention_backend=None):
+        server_args = ServerArgs.__new__(ServerArgs)
+        server_args.device = "cpu"
+        server_args.attention_backend = attention_backend
+        server_args.sampling_backend = None
+        return server_args
+
+    @patch("sglang.srt.server_args.is_host_cpu_arm64", return_value=True)
+    def test_arm_cpu_defaults_to_torch_native(self, _mock_is_arm64):
+        server_args = self._make_server_args()
+
+        ServerArgs._handle_cpu_backends(server_args)
+
+        self.assertEqual(server_args.attention_backend, "torch_native")
+        self.assertEqual(server_args.sampling_backend, "pytorch")
+
+    @patch("sglang.srt.server_args.is_host_cpu_arm64", return_value=False)
+    def test_x86_cpu_defaults_to_intel_amx(self, _mock_is_arm64):
+        server_args = self._make_server_args()
+
+        ServerArgs._handle_cpu_backends(server_args)
+
+        self.assertEqual(server_args.attention_backend, "intel_amx")
+        self.assertEqual(server_args.sampling_backend, "pytorch")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_shared_expert.py

@@ -0,0 +1,235 @@
+import itertools
+import math
+import unittest
+
+import torch
+from utils import (
+    BLOCK_K,
+    BLOCK_N,
+    factor_for_scale,
+    fp8_max,
+    fp8_min,
+    per_token_quant_int8,
+    precision,
+    scaled_weight,
+    torch_naive_moe,
+    torch_w8a8_per_column_moe,
+)
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+class TestSharedExpert(CustomTestCase):
+    M = [2, 121]
+    N = [32, 32 * 4]
+    K = [32, 32 * 2]
+    routed_scaling_factor = [16]
+    apply_scaling_factor = [True, False]
+
+    M_fp8 = [2, 12]
+    N_fp8 = [512]
+    K_fp8 = [256]
+
+    def _bf16_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
+        dtype = torch.bfloat16
+
+        hidden_states = torch.randn(m, k, dtype=dtype) / k
+        w1 = torch.randn(2 * n, k, dtype=dtype)
+        w2 = torch.randn(k, n, dtype=dtype)
+        fused_output = (
+            torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
+        )
+        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
+
+        # fused moe mutates content in hs
+        hidden_states2 = hidden_states.clone()
+
+        # bfloat16
+        ref = torch_naive_moe(
+            hidden_states,
+            w1,
+            w2,
+            fused_output,
+            routed_scaling_factor,
+            output_dtype=dtype,
+        )
+        out = torch.ops.sgl_kernel.shared_expert_cpu(
+            hidden_states2,
+            w1,
+            w2,
+            fused_output,
+            routed_scaling_factor,
+            True,
+            False,
+            False,
+            None,
+            None,
+            None,
+            False,
+        )
+
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
+
+    def test_bf16_shared_expert(self):
+        for params in itertools.product(
+            self.M,
+            self.N,
+            self.K,
+            self.routed_scaling_factor,
+            self.apply_scaling_factor,
+        ):
+            with self.subTest(
+                m=params[0],
+                n=params[1],
+                k=params[2],
+                routed_scaling_factor=params[3],
+                apply_scaling_factor=params[4],
+            ):
+                self._bf16_shared_expert(*params)
+
+    def _int8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
+        dtype = torch.bfloat16
+
+        hidden_states = torch.randn(m, k, dtype=dtype) / k
+        w1 = torch.randn(2 * n, k, dtype=dtype)
+        w2 = torch.randn(k, n, dtype=dtype)
+        fused_output = (
+            torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
+        )
+        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
+
+        # fused moe mutates content in hs
+        hidden_states2 = hidden_states.clone()
+
+        w1_q, w1_s = per_token_quant_int8(w1)
+        w2_q, w2_s = per_token_quant_int8(w2)
+        ref = torch_w8a8_per_column_moe(
+            hidden_states,
+            w1_q,
+            w2_q,
+            w1_s,
+            w2_s,
+            fused_output,
+            routed_scaling_factor,
+        )
+        out = torch.ops.sgl_kernel.shared_expert_cpu(
+            hidden_states2,
+            w1_q,
+            w2_q,
+            fused_output,
+            routed_scaling_factor,
+            True,
+            True,
+            False,
+            w1_s,
+            w2_s,
+            None,
+            False,
+        )
+
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
+
+    def test_int8_shared_expert(self):
+        for params in itertools.product(
+            self.M,
+            self.N,
+            self.K,
+            self.routed_scaling_factor,
+            self.apply_scaling_factor,
+        ):
+            with self.subTest(
+                m=params[0],
+                n=params[1],
+                k=params[2],
+                routed_scaling_factor=params[3],
+                apply_scaling_factor=params[4],
+            ):
+                self._int8_shared_expert(*params)
+
+    def _fp8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
+        dtype = torch.bfloat16
+
+        hidden_states = torch.randn(m, k, dtype=dtype) / math.sqrt(k)
+
+        w1_fp32 = torch.randn(1, 2 * n, k)
+        w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        w2_fp32 = torch.randn(1, k, n)
+        w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+        w1s = torch.randn(1, 2 * n // BLOCK_N, k // BLOCK_K) * factor_for_scale
+        w2s = torch.randn(1, k // BLOCK_N, n // BLOCK_K) * factor_for_scale
+
+        w1_scaled = scaled_weight(w1, w1s).view(2 * n, k)
+        w2_scaled = scaled_weight(w2, w2s).view(k, n)
+
+        # change back to 2D
+        w1, w2 = w1.squeeze(0), w2.squeeze(0)
+        w1s, w2s = w1s.squeeze(0), w2s.squeeze(0)
+        w1_scaled, w2_scaled = w1_scaled.squeeze(0), w2_scaled.squeeze(0)
+
+        fused_output = (
+            torch.randn(m, k, dtype=dtype) / math.sqrt(k)
+            if apply_scaling_factor
+            else None
+        )
+        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
+        hidden_states2 = hidden_states.clone()
+
+        # ref with bfloat16
+        ref = torch_naive_moe(
+            hidden_states,
+            w1_scaled,
+            w2_scaled,
+            fused_output,
+            routed_scaling_factor,
+            output_dtype=dtype,
+        )
+
+        w1 = torch.ops.sgl_kernel.convert_weight_packed(w1)  # [2N, K]
+        w2 = torch.ops.sgl_kernel.convert_weight_packed(w2)  # [K, N]
+        out = torch.ops.sgl_kernel.shared_expert_cpu(
+            hidden_states2,
+            w1,
+            w2,
+            fused_output,
+            routed_scaling_factor,
+            True,
+            False,
+            True,
+            w1s,
+            w2s,
+            [BLOCK_N, BLOCK_K],
+            True,
+        )
+
+        atol = rtol = precision[ref.dtype]
+        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
+
+    def test_fp8_shared_expert(self):
+        for params in itertools.product(
+            self.M_fp8,
+            self.N_fp8,
+            self.K_fp8,
+            self.routed_scaling_factor,
+            self.apply_scaling_factor,
+        ):
+            with self.subTest(
+                m=params[0],
+                n=params[1],
+                k=params[2],
+                routed_scaling_factor=params[3],
+                apply_scaling_factor=params[4],
+            ):
+                self._fp8_shared_expert(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_topk.py

@@ -0,0 +1,223 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.moe.topk import (
+    biased_grouped_topk_impl as native_biased_grouped_topk,
+)
+from sglang.srt.layers.moe.topk import fused_topk_torch_native as native_fused_topk
+from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk
+from sglang.srt.models.llama4 import Llama4MoE
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="stage-b-test-cpu")
+
+torch.manual_seed(1234)
+
+
+# This is used by the Deepseek-V2 model
+class TestGroupedTopK(CustomTestCase):
+    def _run_single_test(self, M, E, G, topk, topk_group, renormalize, dtype):
+        torch.manual_seed(1234)
+
+        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
+        hidden_states = torch.randn(M, 100, dtype=dtype)
+        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M
+
+        ref_topk_weights, ref_topk_ids = native_grouped_topk(
+            hidden_states.float(),
+            gating_output.float(),
+            topk,
+            renormalize,
+            G,
+            topk_group,
+        )
+
+        # fused version
+        topk_weights, topk_ids = torch.ops.sgl_kernel.grouped_topk_cpu(
+            hidden_states,
+            gating_output,
+            topk,
+            renormalize,
+            G,
+            topk_group,
+            0,
+            None,
+            None,
+        )
+
+        res = torch.zeros(M, E, dtype=torch.float)
+        ref = torch.zeros(M, E, dtype=torch.float)
+        res.scatter_(1, topk_ids.long(), topk_weights)
+        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
+        torch.testing.assert_close(res, ref)
+
+    def test_grouped_topk(self):
+        for renormalize in [True, False]:
+            self._run_single_test(123, 8, 2, 2, 1, renormalize, torch.bfloat16)
+            self._run_single_test(123, 16, 4, 3, 2, renormalize, torch.bfloat16)
+            self._run_single_test(123, 32, 4, 3, 2, renormalize, torch.bfloat16)
+            self._run_single_test(1123, 32, 4, 3, 2, renormalize, torch.bfloat16)
+            self._run_single_test(123, 64, 1, 6, 1, renormalize, torch.bfloat16)
+            self._run_single_test(123, 256, 8, 4, 8, renormalize, torch.bfloat16)
+            self._run_single_test(123, 160, 8, 6, 2, renormalize, torch.bfloat16)
+
+
+# DeepSeek V2/V3/R1 uses biased_grouped_top
+class TestBiasedGroupedTopK(CustomTestCase):
+    def _run_single_test(
+        self,
+        M,
+        E,
+        G,
+        topk,
+        topk_group,
+        renormalize,
+        gating_dtype,
+        bias_dtype,
+        routed_scaling_factor,
+    ):
+        torch.manual_seed(1024)
+
+        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
+        hidden_states = torch.randn(M, 100, dtype=torch.bfloat16)
+        gating_output = torch.randn(M, E, dtype=gating_dtype) * 2 * M
+        correction_bias = torch.randn(E, dtype=bias_dtype)
+
+        ref_topk_weights, ref_topk_ids = native_biased_grouped_topk(
+            hidden_states.float(),
+            gating_output.float(),
+            correction_bias.float(),
+            topk,
+            renormalize,
+            G,
+            topk_group,
+        )
+        ref_topk_weights = (
+            ref_topk_weights * routed_scaling_factor
+            if routed_scaling_factor is not None
+            else ref_topk_weights
+        )
+        # fused version
+        topk_weights, topk_ids = torch.ops.sgl_kernel.biased_grouped_topk_cpu(
+            hidden_states,
+            gating_output,
+            correction_bias,
+            topk,
+            renormalize,
+            G,
+            topk_group,
+            0,
+            routed_scaling_factor,
+            None,
+        )
+
+        res = torch.zeros(M, E, dtype=torch.float)
+        ref = torch.zeros(M, E, dtype=torch.float)
+        res.scatter_(1, topk_ids.long(), topk_weights)
+        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
+        torch.testing.assert_close(res, ref)
+
+    def test_biased_grouped_topk(self):
+        for renormalize in [False]:
+            for bias_dtype in [torch.float32, torch.bfloat16]:
+                for gating_dtype in [torch.float32, torch.bfloat16]:
+                    for routed_scaling_factor in [None, 1.125]:
+                        for E_num in [128, 192, 256, 384]:
+                            self._run_single_test(
+                                34,
+                                E_num,
+                                8,
+                                8,
+                                2,
+                                renormalize,
+                                gating_dtype,
+                                bias_dtype,
+                                routed_scaling_factor,
+                            )
+
+
+class TestTopK(CustomTestCase):
+    def _run_single_test(self, M, E, topk, renormalize, dtype):
+        torch.manual_seed(1998)
+
+        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
+        hidden_states = torch.randn(M, 100, dtype=dtype)
+        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M
+
+        ref_topk_weights, ref_topk_ids = native_fused_topk(
+            hidden_states.float(),
+            gating_output.float(),
+            topk,
+            renormalize,
+        )
+
+        # fused version
+        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
+            hidden_states, gating_output, topk, renormalize
+        )
+
+        res = torch.zeros(M, E, dtype=torch.float)
+        ref = torch.zeros(M, E, dtype=torch.float)
+        res.scatter_(1, topk_ids.long(), topk_weights)
+        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
+        torch.testing.assert_close(res, ref)
+
+    def test_topk(self):
+        for renormalize in [True, False]:
+            self._run_single_test(123, 8, 2, renormalize, torch.bfloat16)
+            self._run_single_test(123, 16, 3, renormalize, torch.bfloat16)
+            self._run_single_test(123, 32, 3, renormalize, torch.bfloat16)
+            self._run_single_test(123, 32, 3, renormalize, torch.bfloat16)
+            self._run_single_test(123, 64, 6, renormalize, torch.bfloat16)
+            self._run_single_test(123, 256, 4, renormalize, torch.bfloat16)
+            self._run_single_test(123, 160, 6, renormalize, torch.bfloat16)
+
+
+class TestCustomTopK(CustomTestCase):
+    def _run_single_test(
+        self, M, E, topk, renormalize, dtype, native_custom_f, fused_custom_f
+    ):
+        torch.manual_seed(16)
+
+        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
+        hidden_states = torch.randn(M, 100, dtype=dtype)
+        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M
+
+        ref_topk_weights, ref_topk_ids = native_custom_f(
+            hidden_states.float(),
+            gating_output.float(),
+            topk,
+            renormalize,
+        )
+
+        # fused version
+        topk_weights, topk_ids = fused_custom_f(
+            hidden_states, gating_output, topk, renormalize
+        )
+
+        res = torch.zeros(M, E, dtype=torch.float)
+        ref = torch.zeros(M, E, dtype=torch.float)
+        res.scatter_(1, topk_ids.long(), topk_weights)
+        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
+        torch.testing.assert_close(res, ref)
+
+    def test_custom_topk(self):
+        test_custom_functions = [
+            (Llama4MoE.custom_routing_function, torch.ops.sgl_kernel.topk_sigmoid_cpu)
+        ]
+        for native_custom_f, fused_custom_f in test_custom_functions:
+            self._run_single_test(
+                123, 8, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
+            )
+            self._run_single_test(
+                123, 16, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
+            )
+            self._run_single_test(
+                123, 32, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/utils.py

@@ -0,0 +1,440 @@
+import itertools
+import math
+
+import torch
+import torch.nn.functional as F
+
+precision = {
+    torch.bfloat16: 1e-2,
+    torch.float16: 1e-3,
+    torch.float32: 1e-5,
+}
+
+
+BLOCK_N, BLOCK_K = 64, 128
+factor_for_scale = 1e-3
+fp8_max, fp8_min = 400, -400
+
+
+def parametrize(**params):
+    def decorator(func):
+        def wrapper(self):
+            for combo in itertools.product(*params.values()):
+                kwargs = dict(zip(params.keys(), combo))
+                with self.subTest(**kwargs):
+                    func(self, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def SiluAndMul(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+def GeluAndMul(x: torch.Tensor, approximate="tanh") -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.gelu(x[..., :d], approximate=approximate) * x[..., d:]
+
+
+def per_token_quant_int8(x):
+    x = x.float()
+    absmax = x.abs().max(dim=-1).values
+    absmax = absmax.clamp_min(1e-10).unsqueeze(-1)
+    scale_x = absmax / 127
+    x_q = x.mul(127 / absmax)
+    x_q = torch.round(x_q).to(torch.int8)
+
+    return x_q, scale_x
+
+
+def convert_weight(weight, scale_block_size, A_dtype):
+    N, K = weight.size()
+    fp8_max = 448.0
+    scale_block_size_N, scale_block_size_K = scale_block_size  # (128, 128)
+
+    pad_N = (scale_block_size_N - (N % scale_block_size_N)) % scale_block_size_N
+    pad_K = (scale_block_size_K - (K % scale_block_size_K)) % scale_block_size_K
+
+    if pad_N > 0 or pad_K > 0:
+        weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
+
+    weight_blocks = weight.view(
+        math.ceil(N / scale_block_size_N),
+        scale_block_size_N,
+        math.ceil(K / scale_block_size_K),
+        scale_block_size_K,
+    )  # (8, 128, 8, 128)
+    weight_blocks = weight_blocks.permute(0, 2, 1, 3).contiguous()  # (8, 8, 128, 128)
+
+    # Step 2: compute per-block max abs values → scale
+    abs_max = weight_blocks.abs().amax(dim=(-2, -1), keepdim=True)  # (8, 8, 1, 1)
+    scales = abs_max / fp8_max
+    scales = torch.where(
+        scales == 0, torch.ones_like(scales), scales
+    )  # avoid division by zero
+
+    q_fp8 = (weight_blocks / scales).to(torch.float8_e4m3fn)
+    q_fp8_reshape = q_fp8.permute(0, 2, 1, 3).contiguous()
+
+    if pad_N > 0 or pad_K > 0:
+        q_fp8_reshape = q_fp8_reshape.view(N + pad_N, K + pad_K)
+        q_fp8_reshape = q_fp8_reshape[:N, :K].contiguous()
+    else:
+        q_fp8_reshape = q_fp8_reshape.view(N, K)
+
+    dq_weight = q_fp8.float() * scales
+    dq_weight = dq_weight.permute(0, 2, 1, 3).contiguous()  # (8, 128, 8, 128)
+
+    if pad_N > 0 or pad_K > 0:
+        w_dq = dq_weight.view(N + pad_N, K + pad_K).to(A_dtype)
+        w_dq = w_dq[:N, :K].contiguous()
+    else:
+        w_dq = dq_weight.view(N, K).to(A_dtype)
+
+    scales = scales.view(
+        math.ceil(N / scale_block_size_N), math.ceil(K / scale_block_size_K)
+    )
+
+    return q_fp8_reshape, scales, w_dq
+
+
+def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16):
+    """Matrix multiplication function that supports per-token input quantization and per-column weight quantization"""
+    A = A.to(torch.float32)
+    B = B.to(torch.float32)
+
+    assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
+    assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
+
+    # Reshape input
+    M = A.numel() // A.shape[-1]
+    B = B.t()  # Transpose weight matrix
+    N, K = B.shape
+    origin_C_shape = A.shape[:-1] + (K,)
+    A = A.reshape(M, N)
+
+    # As is per-token [M, 1], Bs is per-column [1, K]
+    C = torch.matmul(A, B)  # [M, K]
+    C = As * C * Bs.view(1, -1)  # Broadcast per-column scale
+
+    if bias is not None:
+        C.add_(bias.view(1, -1))
+
+    return C.reshape(origin_C_shape).to(output_dtype)
+
+
+def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16):
+
+    a = a.to(torch.float32)
+    w1 = w1.to(torch.float32)
+    w2 = w2.to(torch.float32)
+    b = b.to(torch.float32) if b is not None else None
+
+    ic1 = torch.matmul(a, w1.transpose(0, 1))
+    ic2 = SiluAndMul(ic1)
+    ic3 = torch.matmul(ic2, w2.transpose(0, 1))
+
+    out = ic3 if b is None else ic3 + b * routed_scaling_factor
+
+    return out.to(output_dtype)
+
+
+def torch_w8a8_per_column_moe(
+    a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16
+):
+
+    a = a.to(torch.float32)
+    b = b.to(torch.float32) if b is not None else None
+
+    # Perform per-token quantization
+    a_q, a_s = per_token_quant_int8(a)
+
+    ic1 = native_w8a8_per_token_matmul(
+        a_q, w1_q, a_s, w1_s, bias=None, output_dtype=torch.float32
+    )
+    ic2 = SiluAndMul(ic1)
+
+    a1_q, a1_s = per_token_quant_int8(ic2)
+    ic3 = native_w8a8_per_token_matmul(
+        a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32
+    )
+
+    out = ic3 if b is None else ic3 + b * routed_scaling_factor
+
+    return out.to(output_dtype)
+
+
+def scaled_weight(weight, scales):
+    E, N, K = weight.shape
+    pad_N = (BLOCK_N - (N % BLOCK_N)) % BLOCK_N
+    pad_K = (BLOCK_K - (K % BLOCK_K)) % BLOCK_K
+
+    if pad_N > 0 or pad_K > 0:
+        weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
+
+    weight_block = (
+        weight.view(E, math.ceil(N / BLOCK_N), BLOCK_N, math.ceil(K / BLOCK_K), BLOCK_K)
+        .permute(0, 1, 3, 2, 4)
+        .float()
+        .contiguous()
+    )
+
+    weight_scaled = (
+        (
+            weight_block
+            * scales.view(E, math.ceil(N / BLOCK_N), math.ceil(K / BLOCK_K), 1, 1)
+        )
+        .permute(0, 1, 3, 2, 4)
+        .contiguous()
+    )
+    if pad_N > 0 or pad_K > 0:
+        weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
+        weight_scaled = weight_scaled[..., :N, :K].contiguous()
+    else:
+        weight_scaled = weight_scaled.view(E, N, K)
+    return weight_scaled
+
+
+def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize):
+    B, D = a.shape
+    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+    out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+
+    if renormalize:
+        topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
+
+    topk_weight = topk_weight.view(-1)
+    topk_ids = topk_ids.view(-1)
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            out[mask] = SiluAndMul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
+                0, 1
+            )
+    return (
+        out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
+    ).sum(dim=1)
+
+
+def torch_w8a8_per_column_fused_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk):
+    """This function performs fused moe with per-column int8 quantization using native torch."""
+
+    B, D = a.shape
+    # Perform per-token quantization
+    a_q, a_s = per_token_quant_int8(a)
+    # Repeat tokens to match topk
+    a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+    # Also repeat the scale
+    a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1)  # [B*topk, 1]
+
+    out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
+
+    # Calculate routing
+    topk_weight = topk_weight.view(-1)
+    topk_ids = topk_ids.view(-1)
+    # Process each expert
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            # First MLP layer: note that a_s is now per-token
+            inter_out = native_w8a8_per_token_matmul(
+                a_q[mask],
+                w1[i],
+                a_s[mask],
+                w1_s[i],
+                bias=None,
+                output_dtype=torch.float32,
+            )
+            # Activation function
+            act_out = SiluAndMul(inter_out)
+            # Quantize activation output with per-token
+            act_out_q, act_out_s = per_token_quant_int8(act_out)
+            # Second MLP layer
+            out[mask] = native_w8a8_per_token_matmul(
+                act_out_q,
+                w2[i],
+                act_out_s,
+                w2_s[i],
+                bias=None,
+                output_dtype=torch.float32,
+            )
+    # Apply routing weights and sum
+    return (
+        (out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
+        .sum(dim=1)
+        .to(a.dtype)
+    )
+
+
+def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk):
+    B, D = a.shape
+    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
+    out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
+
+    # Calculate routing
+    topk_weight = topk_weight.view(-1)
+    topk_ids = topk_ids.view(-1)
+
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            ic0 = torch.matmul(a[mask], w1[i].transpose(0, 1))
+            ic1 = SiluAndMul(ic0)
+            out[mask] = torch.matmul(ic1, w2[i].transpose(0, 1))
+
+    return (
+        (out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
+        .sum(dim=1)
+        .to(a.dtype)
+    )
+
+
+def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
+    """
+    Make a tensor non-contiguous by slicing it via last dimension.
+    """
+    last_dim = x.shape[-1]
+    return x[..., : last_dim // 2] if x.is_contiguous() else x
+
+
+def awq_reverse_reorder_int_tensor(int_tensor, bits: int):
+    assert bits == 4
+
+    int_tensor = int_tensor.T.contiguous()
+    compress_ratio = 32 // bits
+    assert int_tensor.shape[-1] % compress_ratio == 0
+
+    order_map = [0, 2, 4, 6, 1, 3, 5, 7]
+    order_tensor = torch.tensor(
+        order_map, dtype=torch.int32, device=int_tensor.device
+    ).reshape(1, -1)
+    order_tensor = order_tensor.repeat(int_tensor.shape[1] // compress_ratio, 1)
+    order_tensor = order_tensor + torch.arange(
+        0,
+        int_tensor.shape[1],
+        compress_ratio,
+        dtype=torch.int32,
+        device=int_tensor.device,
+    ).reshape(-1, 1)
+    order_tensor = order_tensor.reshape(-1)
+
+    reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
+    reverse_order_tensor = reverse_order_tensor[order_tensor]
+    int_tensor = int_tensor[:, reverse_order_tensor]
+    return int_tensor
+
+
+def unpack_and_dequant_awq(
+    awq_qweight: torch.Tensor,
+    awq_qzeros: torch.Tensor,
+    awq_scales: torch.Tensor,
+    bits: int,
+    group_size: int,
+):
+    """
+    Args:
+        awq_qweight (`torch.LongTensor`):
+            Expected shape: (in_features, out_features // (32 // bits))
+        awq_qzeros (`torch.LongTensor`):
+            Expected shape: (in_features // group_size, out_features // (32 // bits))
+        awq_scales (`torch.LongTensor`):
+            Expected shape: (in_features // group_size, out_features)
+
+    Returns:
+        fp16_weight (`torch.LongTensor`):
+            With shape (in_features, out_features).
+        zeros (`torch.LongTensor`):
+            With shape (in_features // group_size, out_features).
+    """
+    assert bits == 4
+
+    qzeros = awq_qzeros
+    qweight = awq_qweight
+    qweight = qweight.T.contiguous()
+
+    scales = awq_scales
+    scales = scales.reshape(-1, 1, scales.shape[-1])
+
+    infeatures = awq_qweight.shape[0]
+
+    wf = torch.tensor(
+        list(range(0, 32, bits)), dtype=torch.int32, device=qzeros.device
+    ).unsqueeze(0)
+    zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2), wf.unsqueeze(0)).to(
+        torch.int16 if bits == 8 else torch.int8
+    )
+
+    torch.bitwise_and(zeros, (2**bits) - 1, out=zeros)
+
+    zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2])
+
+    weight = torch.bitwise_right_shift(
+        torch.unsqueeze(qweight, 1), wf.unsqueeze(-1)
+    ).to(torch.int16 if bits == 8 else torch.int8)
+    torch.bitwise_and(weight, (2**bits) - 1, out=weight)
+    weight = weight.reshape(-1, group_size, weight.shape[2])
+
+    weight = weight.view(-1, weight.shape[-1])
+    zeros = zeros.view(-1, zeros.shape[-1])
+
+    zeros = zeros.T.contiguous()
+    zeros = awq_reverse_reorder_int_tensor(zeros, bits)
+    weight = awq_reverse_reorder_int_tensor(weight, bits)
+
+    # Dequantize weights.
+    scales = awq_scales
+    zeros = zeros.contiguous()
+    scale_zeros = zeros * scales
+
+    g_idx = torch.tensor(
+        [i // group_size for i in range(infeatures)], dtype=torch.int32
+    )
+    scale_mat = scales[g_idx]
+    scale_zeros_mat = scale_zeros[g_idx].to(torch.bfloat16)
+
+    qdq_weight_T = weight * scale_mat - scale_zeros_mat.to(torch.bfloat16)
+
+    fp16_weight = qdq_weight_T.T
+
+    return fp16_weight, zeros
+
+
+def unpack_4bit_to_32bit_signed(qweight, qzeros):
+    # Unpack 4-bit values and interpret them as signed integers
+    unpacked_weights = torch.zeros(
+        (qweight.shape[0] * 8, qweight.shape[1]),
+        dtype=torch.int8,
+        device=qweight.device,
+        requires_grad=False,
+    )
+    unpacked_zeros = torch.zeros(
+        (qzeros.shape[0], qzeros.shape[1] * 8),
+        dtype=torch.int8,
+        device=qzeros.device,
+        requires_grad=False,
+    )
+
+    for row in range(unpacked_weights.shape[0]):
+        i = row % 8
+        unpacked_weights[row, :] = (qweight[row // 8, :] >> (4 * i)) & 0xF
+
+    for col in range(unpacked_zeros.shape[1]):
+        i = col % 8
+        unpacked_zeros[:, col] = (qzeros[:, col // 8] >> (4 * i)) & 0xF
+
+    return unpacked_weights, unpacked_zeros + 1
+
+
+def unpack_and_dequant_gptq(qweight, qzeros, scales):
+    unpacked_qweight, unpacked_qzeros = unpack_4bit_to_32bit_signed(qweight, qzeros)
+    group_size = unpacked_qweight.shape[0] // scales.shape[0]
+    scales = scales.repeat_interleave(group_size, dim=0)
+    unpacked_qzeros = unpacked_qzeros.repeat_interleave(group_size, dim=0)
+    unpacked_qweight = (unpacked_qweight - unpacked_qzeros) * scales
+
+    return unpacked_qweight.T

test/run_suite.py

@@ -23,7 +23,7 @@ HW_MAPPING = {
 
 # Per-commit test suites (run on every PR)
 PER_COMMIT_SUITES = {
-    HWBackend.CPU: ["stage-a-test-cpu"],
+    HWBackend.CPU: ["stage-a-test-cpu", "stage-b-test-cpu"],
     HWBackend.AMD: [
         "stage-a-test-1-gpu-small-amd",
         "stage-b-test-1-gpu-small-amd",
@@ -239,7 +239,9 @@ def run_a_suite(args):
         for f in glob.glob(
             os.path.join(script_dir, "registered", "**", "*.py"), recursive=True
         )
-        if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
+        if not f.endswith("/conftest.py")
+        and not f.endswith("/__init__.py")
+        and not f.endswith("/cpu/utils.py")
     ]
 
     # JIT kernel tests and benchmarks (live alongside kernel source)