diff --git a/.github/workflows/pr-test-arm64.yml b/.github/workflows/pr-test-arm64.yml new file mode 100644 index 000000000..4525ed388 --- /dev/null +++ b/.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 diff --git a/docker/arm64.Dockerfile b/docker/arm64.Dockerfile new file mode 100644 index 000000000..5173e46be --- /dev/null +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bd167be0d..eb89090c9 100644 --- a/python/sglang/srt/server_args.py +++ b/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): diff --git a/sgl-kernel/csrc/cpu/CMakeLists.txt b/sgl-kernel/csrc/cpu/CMakeLists.txt index ca7e133bc..7b4275ecd 100755 --- a/sgl-kernel/csrc/cpu/CMakeLists.txt +++ b/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() diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 51700a281..9ec3ea450 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -150,18 +150,6 @@ at::Tensor convert_scale_packed(at::Tensor& scale); // quant std::tuple per_token_quant_int8_cpu(at::Tensor& A); -// gemm -at::Tensor -weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, bool is_vnni); - -// gemm fusion -at::Tensor fused_linear_sigmoid_mul( - at::Tensor& mat1, - at::Tensor& mat2, - const std::optional& 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 bias); @@ -205,10 +194,24 @@ std::tuple 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& bias, bool is_vnni); + +// gemm fusion +at::Tensor fused_linear_sigmoid_mul( + at::Tensor& mat1, + at::Tensor& mat2, + const std::optional& 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& 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& 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"); diff --git a/test/srt/cpu/test_server_args_backend.py b/test/srt/cpu/test_server_args_backend.py new file mode 100644 index 000000000..9780bbbb2 --- /dev/null +++ b/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() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index d053a895a..00af66144 100644 --- a/test/srt/run_suite.py +++ b/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)