From 7e3e6161592f2f98a51b0b8436fad55087ab4087 Mon Sep 17 00:00:00 2001 From: Mandepudi Rani Chowdary <66872302+ranimandepudi@users.noreply.github.com> Date: Tue, 9 Jun 2026 21:36:57 -0500 Subject: [PATCH] Add Arm64 INT8 MoE test coverage (#25007) --- .github/workflows/pr-test-arm64.yml | 2 +- test/registered/cpu/arm64/test_moe.py | 118 ++++++++++++++++++ test/registered/cpu/test_activation.py | 1 + test/registered/cpu/test_decode.py | 1 + test/registered/cpu/test_norm.py | 1 + test/registered/cpu/test_qwen3.py | 1 + test/registered/cpu/test_rope.py | 1 + .../cpu/test_server_args_backend.py | 1 + test/registered/cpu/test_topk.py | 1 + test/run_suite.py | 2 +- test/srt/cpu/arm64/test_moe.py | 118 ++++++++++++++++++ test/srt/run_suite.py | 15 --- 12 files changed, 245 insertions(+), 17 deletions(-) create mode 100644 test/registered/cpu/arm64/test_moe.py create mode 100644 test/srt/cpu/arm64/test_moe.py diff --git a/.github/workflows/pr-test-arm64.yml b/.github/workflows/pr-test-arm64.yml index 4525ed388..d49b45c2a 100644 --- a/.github/workflows/pr-test-arm64.yml +++ b/.github/workflows/pr-test-arm64.yml @@ -102,7 +102,7 @@ jobs: timeout-minutes: 36 run: | docker exec -w /sglang-checkout/ ci_sglang_arm64 \ - bash -c "source /opt/.venv/bin/activate && cd ./test/srt && python3 run_suite.py --suite per-commit-cpu-arm64 --timeout-per-file 1500" + bash -c "source /opt/.venv/bin/activate && cd ./test && python3 run_suite.py --hw cpu --suite base-b-test-cpu-arm64 --timeout-per-file 1500" - name: Change permission timeout-minutes: 2 diff --git a/test/registered/cpu/arm64/test_moe.py b/test/registered/cpu/arm64/test_moe.py new file mode 100644 index 000000000..9d6c3c9e6 --- /dev/null +++ b/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() diff --git a/test/registered/cpu/test_activation.py b/test/registered/cpu/test_activation.py index 70843a9dc..7ca963359 100644 --- a/test/registered/cpu/test_activation.py +++ b/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) diff --git a/test/registered/cpu/test_decode.py b/test/registered/cpu/test_decode.py index 494fe73f7..b6685c31c 100644 --- a/test/registered/cpu/test_decode.py +++ b/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) diff --git a/test/registered/cpu/test_norm.py b/test/registered/cpu/test_norm.py index 87e525b02..cdcc2b369 100644 --- a/test/registered/cpu/test_norm.py +++ b/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) diff --git a/test/registered/cpu/test_qwen3.py b/test/registered/cpu/test_qwen3.py index 21c62a550..602144fec 100644 --- a/test/registered/cpu/test_qwen3.py +++ b/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) diff --git a/test/registered/cpu/test_rope.py b/test/registered/cpu/test_rope.py index febe0b710..d829a0622 100644 --- a/test/registered/cpu/test_rope.py +++ b/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) diff --git a/test/registered/cpu/test_server_args_backend.py b/test/registered/cpu/test_server_args_backend.py index 336a57d35..02c9af632 100644 --- a/test/registered/cpu/test_server_args_backend.py +++ b/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): diff --git a/test/registered/cpu/test_topk.py b/test/registered/cpu/test_topk.py index c594a37ab..54b78c74e 100644 --- a/test/registered/cpu/test_topk.py +++ b/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 diff --git a/test/run_suite.py b/test/run_suite.py index 1f935973a..accc5191a 100644 --- a/test/run_suite.py +++ b/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", diff --git a/test/srt/cpu/arm64/test_moe.py b/test/srt/cpu/arm64/test_moe.py new file mode 100644 index 000000000..a2f0532f0 --- /dev/null +++ b/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() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 029853c16..d21e9952e 100644 --- a/test/srt/run_suite.py +++ b/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)