Add Arm64 INT8 MoE test coverage (#25007)

test/registered/cpu/arm64/test_moe.py

@@ -0,0 +1,118 @@
+"""Arm64 MoE test.
+
+Tests fused_experts_cpu with W8A8 INT8 quantization, which is supported
+on Arm64 via aarch64/moe.cpp (PR #16045). Additional quantization paths
+(BF16, INT4) will be added here as Arm kernels land.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
+
+import itertools
+import math
+import os
+import platform
+import sys
+import unittest
+
+import torch
+
+# Add parent dir (test/srt/cpu/) to path for utils import
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from sglang.srt.layers.amx_utils import CPUQuantMethod
+from sglang.test.test_utils import CustomTestCase
+
+kernel = torch.ops.sgl_kernel
+IS_ARM64 = platform.machine().lower() in ("aarch64", "arm64")
+
+torch.manual_seed(128)
+
+from utils import (
+    precision,
+    torch_w8a8_per_column_fused_moe,
+)
+
+
+class TestFusedExpertsInt8(CustomTestCase):
+    M = [1, 6, 32, 64]
+    N = [256, 512]
+    K = [256, 512]
+    E = [8]
+    topk = [4]
+
+    def _int8_moe(self, M, N, K, E, topk):
+        dtype = torch.bfloat16
+        # Arm64 INT8 MoE currently uses the unpacked, out-of-place path.
+        prepack = not IS_ARM64
+
+        int8_factor_for_scale = 1e-2
+        int8_max = 127
+        int8_min = -128
+
+        a = torch.randn((M, K), dtype=dtype) / math.sqrt(K)
+
+        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)
+
+        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
+
+        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 = not IS_ARM64
+        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]
+        if IS_ARM64:
+            atol = rtol = 0.03
+        elif 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,
+            self.N,
+            self.K,
+            self.E,
+            self.topk,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                E=params[3],
+                topk=params[4],
+            ):
+                self._int8_moe(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_activation.py

@@ -9,6 +9,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 

test/registered/cpu/test_decode.py

@@ -7,6 +7,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 

test/registered/cpu/test_norm.py

@@ -9,6 +9,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 

test/registered/cpu/test_qwen3.py

@@ -7,6 +7,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 

test/registered/cpu/test_rope.py

@@ -16,6 +16,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 

test/registered/cpu/test_server_args_backend.py

@@ -8,6 +8,7 @@ from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 
 class TestServerArgsCPUBackend(unittest.TestCase):

test/registered/cpu/test_topk.py

@@ -14,6 +14,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
+register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 
 # This is used by the Deepseek-V2 model

test/run_suite.py

@@ -29,7 +29,7 @@ HW_MAPPING = {
 # (label-gated; pr-test-extra.yml). Tests are tagged per-commit regardless;
 # pr-test-extra.yml's `run-ci-extra` PR label decides whether extra-* dispatches.
 PER_COMMIT_SUITES = {
-    HWBackend.CPU: ["base-a-test-cpu", "base-b-test-cpu"],
+    HWBackend.CPU: ["base-a-test-cpu", "base-b-test-cpu", "base-b-test-cpu-arm64"],
     HWBackend.AMD: [
         "stage-a-test-1-gpu-small-amd",
         "stage-b-test-1-gpu-small-amd",

test/srt/cpu/arm64/test_moe.py

@@ -0,0 +1,118 @@
+"""Arm64 MoE test.
+
+Tests fused_experts_cpu with W8A8 INT8 quantization, which is supported
+on Arm64 via aarch64/moe.cpp (PR #16045). Additional quantization paths
+(BF16, INT4) will be added here as Arm kernels land.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="per-commit-cpu-arm64")
+
+import itertools
+import math
+import os
+import platform
+import sys
+import unittest
+
+import torch
+
+# Add parent dir (test/srt/cpu/) to path for utils import
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from sglang.srt.layers.amx_utils import CPUQuantMethod
+from sglang.test.test_utils import CustomTestCase
+
+kernel = torch.ops.sgl_kernel
+IS_ARM64 = platform.machine().lower() in ("aarch64", "arm64")
+
+torch.manual_seed(128)
+
+from utils import (
+    precision,
+    torch_w8a8_per_column_fused_moe,
+)
+
+
+class TestFusedExpertsInt8(CustomTestCase):
+    M = [1, 6, 32, 64]
+    N = [256, 512]
+    K = [256, 512]
+    E = [8]
+    topk = [4]
+
+    def _int8_moe(self, M, N, K, E, topk):
+        dtype = torch.bfloat16
+        # Arm64 INT8 MoE currently uses the unpacked, out-of-place path.
+        prepack = not IS_ARM64
+
+        int8_factor_for_scale = 1e-2
+        int8_max = 127
+        int8_min = -128
+
+        a = torch.randn((M, K), dtype=dtype) / math.sqrt(K)
+
+        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)
+
+        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
+
+        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 = not IS_ARM64
+        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]
+        if IS_ARM64:
+            atol = rtol = 0.03
+        elif 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,
+            self.N,
+            self.K,
+            self.E,
+            self.topk,
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                E=params[3],
+                topk=params[4],
+            ):
+                self._int8_moe(*params)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/srt/run_suite.py

@@ -45,20 +45,6 @@ suite_amd = {
 }
 
 # Keep the Arm64 bootstrap suite limited to hosted-runner-safe unit kernels.
-# `test_extend.py`, `test_mamba.py`, and `test_mla.py` still hit the
-# x86-specific BF16 BRGEMM/VNNI path on Arm and need dedicated fallbacks.
-suite_arm64 = {
-    "per-commit-cpu-arm64": [
-        TestFile("cpu/test_activation.py"),
-        TestFile("cpu/test_decode.py"),
-        TestFile("cpu/test_norm.py"),
-        TestFile("cpu/test_qwen3.py"),
-        TestFile("cpu/test_rope.py"),
-        TestFile("cpu/test_server_args_backend.py"),
-        TestFile("cpu/test_topk.py"),
-    ],
-}
-
 # Add Intel Xeon tests
 suite_xeon = {
     "per-commit-cpu": [
@@ -93,7 +79,6 @@ suite_xeon = {
 suite_xpu = {}
 
 suites.update(suite_amd)
-suites.update(suite_arm64)
 suites.update(suite_xeon)
 suites.update(suite_xpu)