From 4b14f622e142eedf199b7ae7308ff0c4cc450352 Mon Sep 17 00:00:00 2001
From: Shangming Cai <csmthu@gmail.com>
Date: Sat, 10 Jan 2026 14:44:54 +0800
Subject: [PATCH] [CI] Add PD Disaggregation aarch64 test (#16572)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
---
 docker/Dockerfile                       |  9 +--
 test/srt/run_suite.py                   |  1 +
 test/srt/test_disaggregation_aarch64.py | 93 +++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 8 deletions(-)
 create mode 100644 test/srt/test_disaggregation_aarch64.py

diff --git a/docker/Dockerfile b/docker/Dockerfile
index c8cb55c59..46a54b103 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -273,6 +273,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install \
     datamodel_code_generator \
+    mooncake-transfer-engine==0.3.8.post1 \
     pre-commit \
     pytest \
     black \
@@ -286,14 +287,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     cubloaty \
     google-cloud-storage
 
-# Install mooncake
-RUN --mount=type=cache,target=/root/.cache/pip \
-    if [ "$GRACE_BLACKWELL" = "1" ]; then \
-        python3 -m pip install mooncake-transfer-engine==0.3.7.post2; \
-    else \
-        python3 -m pip install mooncake-transfer-engine==0.3.8.post1; \
-    fi
-
 # Build and install sgl-model-gateway (install Rust, build, then remove to save space)
 RUN --mount=type=cache,target=/root/.cache/pip \
     curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs | sh -s -- -y \
diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py
index 4676be856..b0a6a38f1 100644
--- a/test/srt/run_suite.py
+++ b/test/srt/run_suite.py
@@ -41,6 +41,7 @@ suites = {
     # ],
     "per-commit-4-gpu-gb200": [
         TestFile("test_deepseek_v3_cutedsl_4gpu.py", 1800),
+        TestFile("test_disaggregation_aarch64.py", 300),
     ],
     "per-commit-4-gpu-deepep": [
         TestFile("ep/test_deepep_small.py", 531),
diff --git a/test/srt/test_disaggregation_aarch64.py b/test/srt/test_disaggregation_aarch64.py
new file mode 100644
index 000000000..8a6048adc
--- /dev/null
+++ b/test/srt/test_disaggregation_aarch64.py
@@ -0,0 +1,93 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+
+class TestDisaggregationMooncakeAARCH64Accuracy(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        os.environ["SGLANG_MOONCAKE_CUSTOM_MEM_POOL"] = "true"
+        os.environ["MC_FORCE_MNNVL"] = "true"
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        # Non blocking start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        cls.wait_server_ready(cls.prefill_url + "/health")
+        cls.wait_server_ready(cls.decode_url + "/health")
+
+        cls.launch_lb()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.environ.pop("SGLANG_MOONCAKE_CUSTOM_MEM_POOL")
+        os.environ.pop("MC_FORCE_MNNVL")
+        super().tearDownClass()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "2",
+        ]
+        prefill_args += cls.transfer_backend + cls.rdma_devices
+        cls.process_prefill = popen_launch_pd_server(
+            cls.model,
+            cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp",
+            "2",
+            "--base-gpu-id",
+            "2",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+        )
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host=f"http://{self.base_host}",
+            port=int(self.lb_port),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"Evaluation metrics: {metrics}")
+
+        self.assertGreater(metrics["accuracy"], 0.62)
+
+
+if __name__ == "__main__":
+    unittest.main()