Add Arm64 CPU Phase 1A CI bootstrap (#22123)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

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

@@ -0,0 +1,118 @@
+name: PR Test (Arm64)
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+  workflow_dispatch:
+  workflow_call:
+    inputs:
+      ref:
+        description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.'
+        required: false
+        type: string
+        default: ''
+      run_all_tests:
+        description: "Run all tests (for releasing or testing purpose)"
+        required: false
+        type: boolean
+        default: false
+
+concurrency:
+  group: pr-test-arm64-${{ inputs.ref || github.ref }}
+  cancel-in-progress: false
+
+jobs:
+  check-changes:
+    runs-on: ubuntu-latest
+    outputs:
+      main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests}}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.ref }}
+
+      - name: Determine run mode
+        id: run-mode
+        run: |
+          if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then
+            echo "run_all_tests=true" >> $GITHUB_OUTPUT
+            echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})"
+          else
+            echo "run_all_tests=false" >> $GITHUB_OUTPUT
+            echo "Run mode: FILTERED (triggered by ${{ github.event_name }})"
+          fi
+
+      - name: Detect file changes
+        id: filter
+        uses: dorny/paths-filter@v3
+        if: steps.run-mode.outputs.run_all_tests != 'true'
+        with:
+          filters: |
+            main_package:
+              - "python/sglang/!(multimodal_gen)/**/!(*.md)"
+              - "python/pyproject_cpu.toml"
+              - "test/**/!(*.md)"
+              - "sgl-kernel/**/*.!(md|txt)"
+              - ".github/workflows/pr-test-arm64.yml"
+              - "docker/arm64.Dockerfile"
+
+  pr-gate:
+    needs: check-changes
+    if: needs.check-changes.outputs.main_package == 'true'
+    uses: ./.github/workflows/pr-gate.yml
+    secrets: inherit
+
+  build-test:
+    needs: [check-changes, pr-gate]
+    if: needs.check-changes.outputs.main_package == 'true'
+    runs-on: ubuntu-24.04-arm
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.ref }}
+
+      - name: Build container
+        run: |
+          PR_REPO=${{ github.event.pull_request.head.repo.clone_url }}
+          PR_HEAD_REF=${{ github.head_ref }}
+
+          docker build \
+            ${PR_REPO:+--build-arg SGLANG_REPO=$PR_REPO} \
+            ${PR_HEAD_REF:+--build-arg VER_SGLANG=$PR_HEAD_REF} \
+            . -f docker/arm64.Dockerfile -t sglang_arm64 --no-cache
+
+      - name: Run container
+        run: |
+          docker run -dt \
+            -v ${{ github.workspace }}:/sglang-checkout/ --ipc=host \
+            --name ci_sglang_arm64 \
+            sglang_arm64
+
+      - name: Arm sanity check
+        timeout-minutes: 5
+        run: |
+          docker exec -w /sglang-checkout/ ci_sglang_arm64 \
+            bash -c "source /opt/.venv/bin/activate && python3 -c 'import platform; import torch; import sgl_kernel; from sglang.srt.utils.common import is_host_cpu_arm64; assert platform.machine() in (\"aarch64\", \"arm64\"); assert is_host_cpu_arm64(); assert hasattr(torch.ops.sgl_kernel, \"decode_attention_cpu\"); assert hasattr(torch.ops.sgl_kernel, \"initialize\");'"
+
+      - name: Run unit tests
+        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"
+
+      - name: Change permission
+        timeout-minutes: 2
+        run: |
+          docker exec -u root ci_sglang_arm64 bash -c "
+            rm -rf /tmp/ci-home  &&
+            chown -R  $(id -u):$(id -g) /sglang-checkout/ 2>/dev/null || true
+          "
+
+      - name: Cleanup container
+        if: always()
+        run: |
+          docker rm -f ci_sglang_arm64 || true

docker/arm64.Dockerfile

@@ -0,0 +1,52 @@
+FROM ubuntu:24.04
+SHELL ["/bin/bash", "-c"]
+
+ARG SGLANG_REPO=https://github.com/sgl-project/sglang.git
+ARG VER_SGLANG=main
+
+RUN apt-get update && \
+    apt-get full-upgrade -y && \
+    DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
+    ca-certificates \
+    git \
+    curl \
+    wget \
+    vim \
+    gcc \
+    g++ \
+    make \
+    cmake \
+    libsqlite3-dev \
+    google-perftools \
+    libtbb-dev \
+    libnuma-dev \
+    numactl
+
+WORKDIR /opt
+
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    source $HOME/.local/bin/env && \
+    uv venv --python 3.12
+
+RUN echo -e '[[index]]\nname = "torch"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "torchvision"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "torchaudio"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "triton"\nurl = "https://download.pytorch.org/whl/cpu"' > .venv/uv.toml
+
+ENV UV_CONFIG_FILE=/opt/.venv/uv.toml
+ENV CMAKE_BUILD_PARALLEL_LEVEL=1
+
+WORKDIR /sgl-workspace
+RUN source $HOME/.local/bin/env && \
+    source /opt/.venv/bin/activate && \
+    git clone ${SGLANG_REPO} sglang && \
+    cd sglang && \
+    git checkout ${VER_SGLANG} && \
+    cd python && \
+    cp pyproject_cpu.toml pyproject.toml && \
+    uv pip install . && \
+    cd ../sgl-kernel && \
+    cp pyproject_cpu.toml pyproject.toml && \
+    uv pip install .
+
+ENV SGLANG_USE_CPU_ENGINE=1
+RUN echo 'source /opt/.venv/bin/activate' >> /root/.bashrc
+
+WORKDIR /sgl-workspace/sglang

python/sglang/srt/server_args.py

@@ -51,6 +51,7 @@ from sglang.srt.utils.common import (
     is_flashinfer_available,
     is_hip,
     is_hopper_with_cuda_12_3,
+    is_host_cpu_arm64,
     is_mps,
     is_musa,
     is_no_spec_infer_or_topk_one,
@@ -1203,7 +1204,9 @@ class ServerArgs:
     def _handle_cpu_backends(self):
         if self.device == "cpu":
             if self.attention_backend is None:
-                self.attention_backend = "intel_amx"
+                self.attention_backend = (
+                    "torch_native" if is_host_cpu_arm64() else "intel_amx"
+                )
             self.sampling_backend = "pytorch"
 
     def _handle_npu_backends(self):

sgl-kernel/csrc/cpu/CMakeLists.txt

@@ -75,6 +75,23 @@ endif()
 
 file(GLOB_RECURSE SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
 
+# These kernels still rely on x86-specific AMX/AVX512 implementations.
+# Keep them out of Arm64 bootstrap builds until native Arm paths land.
+set(SGLANG_CPU_X86_ONLY_SOURCES
+    ${CMAKE_CURRENT_SOURCE_DIR}/gemm_int4.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/moe.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/moe_fp8.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/moe_int4.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/moe_int8.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/qkv_proj.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/mamba/conv.cpp
+)
+
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
+    add_compile_definitions(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
+    list(REMOVE_ITEM SOURCES ${SGLANG_CPU_X86_ONLY_SOURCES})
+endif()
+
 if(NOT DEFINED ENV{SGLANG_CPU_FP8_CVT_FTZ})
     set(ENV{SGLANG_CPU_FP8_CVT_FTZ} "1")
 endif()

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -150,18 +150,6 @@ at::Tensor convert_scale_packed(at::Tensor& scale);
 // quant
 std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A);
 
-// gemm
-at::Tensor
-weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni);
-
-// gemm fusion
-at::Tensor fused_linear_sigmoid_mul(
-    at::Tensor& mat1,
-    at::Tensor& mat2,
-    const std::optional<at::Tensor>& bias,
-    bool is_vnni,
-    const at::Tensor& post_mul_mat);
-
 // igemm
 at::Tensor int8_scaled_mm_cpu(
     at::Tensor& mat1,
@@ -195,6 +183,7 @@ at::Tensor int8_scaled_mm_with_quant(
     at::ScalarType out_dtype,
     bool is_vnni);
 
+#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
 // int4 gemm
 at::Tensor int4_scaled_mm_cpu(
     at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
@@ -205,10 +194,24 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
     at::Tensor qzeros,   // awq: (*, K / group_size, N / 8) ||  gptq: (*, K / group_size, N / 8) , int32
     at::Tensor scales,   // awq: (*, K / group_size, N) ||  gptq: (*, K / group_size, N) , bfloat16
     int64_t quant_method_4bit);
+#endif
+
+// gemm
+at::Tensor
+weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni);
+
+// gemm fusion
+at::Tensor fused_linear_sigmoid_mul(
+    at::Tensor& mat1,
+    at::Tensor& mat2,
+    const std::optional<at::Tensor>& bias,
+    bool is_vnni,
+    const at::Tensor& post_mul_mat);
 
 // bmm
 void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
 
+#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
 // fused moe
 at::Tensor fused_experts_cpu(
     at::Tensor& hidden_states,
@@ -306,6 +309,7 @@ at::Tensor causal_conv1d_update_cpu(
     const std::optional<at::Tensor>& conv_state_indices,
     int64_t pad_slot_id,
     bool is_vnni);
+#endif
 
 // conv3d fast path for patch embedding
 at::Tensor conv3d_embed_weight_pack(const at::Tensor& weight);
@@ -495,15 +499,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.def("per_token_quant_int8_cpu(Tensor A) -> (Tensor, Tensor)");
   m.impl("per_token_quant_int8_cpu", torch::kCPU, &per_token_quant_int8_cpu);
 
-  // gemm
-  m.def("weight_packed_linear(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni) -> Tensor");
-  m.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
-
-  // gemm fusion
-  m.def(
-      "fused_linear_sigmoid_mul(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni, Tensor post_mul_mat) -> Tensor");
-  m.impl("fused_linear_sigmoid_mul", torch::kCPU, &fused_linear_sigmoid_mul);
-
   // igemm
   m.def(
       "int8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales1, Tensor scales2, Tensor? bias, ScalarType "
@@ -526,6 +521,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "is_vnni) -> Tensor");
   m.impl("int8_scaled_mm_with_quant", torch::kCPU, &int8_scaled_mm_with_quant);
 
+#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
   // int4 gemm
   m.def("int4_scaled_mm_cpu(Tensor x, Tensor w, Tensor w_zeros, Tensor w_scales, Tensor? bias) -> Tensor");
   m.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
@@ -535,11 +531,22 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor scales, int quant_method_4bit) -> (Tensor, "
       "Tensor, Tensor)");
   m.impl("convert_weight_packed_scale_zp", torch::kCPU, &convert_weight_packed_scale_zp);
+#endif
+
+  // gemm
+  m.def("weight_packed_linear(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni) -> Tensor");
+  m.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
+
+  // gemm fusion
+  m.def(
+      "fused_linear_sigmoid_mul(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni, Tensor post_mul_mat) -> Tensor");
+  m.impl("fused_linear_sigmoid_mul", torch::kCPU, &fused_linear_sigmoid_mul);
 
   // bmm
   m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()");
   m.impl("bmm_cpu", torch::kCPU, &bmm_cpu);
 
+#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
   // moe
   m.def(
       "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool "
@@ -585,6 +592,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor weight, Tensor? bias, bool silu_activation,"
       "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor");
   m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
+#endif
 
   // conv3d fast path for patch embedding
   m.def("conv3d_embed_weight_pack(Tensor weight) -> Tensor");

test/srt/cpu/test_server_args_backend.py

@@ -0,0 +1,35 @@
+import unittest
+from unittest.mock import patch
+
+from sglang.srt.server_args import ServerArgs
+
+
+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/srt/run_suite.py

@@ -44,6 +44,21 @@ suite_amd = {
     # by test/run_suite.py using the registry system.
 }
 
+# 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": [
@@ -66,6 +81,7 @@ suite_xeon = {
         TestFile("cpu/test_qkv_proj_with_rope.py"),
         TestFile("cpu/test_qwen3.py"),
         TestFile("cpu/test_rope.py"),
+        TestFile("cpu/test_server_args_backend.py"),
         TestFile("cpu/test_shared_expert.py"),
         TestFile("cpu/test_topk.py"),
     ],
@@ -83,6 +99,7 @@ suite_xpu = {
 }
 
 suites.update(suite_amd)
+suites.update(suite_arm64)
 suites.update(suite_xeon)
 suites.update(suite_xpu)