Add TP2 MOE benchmarks for AMD. (#5909)

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

@@ -141,11 +141,31 @@ jobs:
           mkdir -p dummy-grok && wget https://sharkpublic.blob.core.windows.net/sharkpublic/sglang/dummy_grok.json -O dummy-grok/config.json
           docker cp ./dummy-grok ci_sglang:/
 
-      - name: Evaluate Benchmark
+      - name: Benchmark dummy grok (TP=2)
         timeout-minutes: 20
         run: |
           docker exec -w /sglang-checkout/test/srt -e SGLANG_IS_IN_CI=1 ci_sglang python3 models/test_dummy_grok_models.py
 
+      - name: Benchmark single latency (TP=2)
+        timeout-minutes: 20
+        run: |
+          docker exec -w /sglang-checkout/test/srt -e SGLANG_IS_IN_CI=1 -e SGLANG_AMD_CI=1 ci_sglang python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_moe_tp2_bs1
+
+      - name: Benchmark single latency + torch.compile (TP=2)
+        timeout-minutes: 20
+        run: |
+          docker exec -w /sglang-checkout/test/srt -e SGLANG_IS_IN_CI=1 ci_sglang python3 -m unittest test_bench_one_batch.TestBenchOneBatch.test_torch_compile_tp2_bs1
+
+      - name: Benchmark offline throughput (TP=2)
+        timeout-minutes: 20
+        run: |
+          docker exec -w /sglang-checkout/test/srt -e SGLANG_IS_IN_CI=1 -e SGLANG_AMD_CI=1 ci_sglang python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_default
+
+      - name: Benchmark offline throughput (w/o RadixAttention) (TP=2)
+        timeout-minutes: 20
+        run: |
+          docker exec -w /sglang-checkout/test/srt -e SGLANG_IS_IN_CI=1 -e SGLANG_AMD_CI=1 ci_sglang python3 -m unittest test_bench_serving.TestBenchServing.test_moe_offline_throughput_without_radix_cache
+
   finish:
     if: always()
     needs: [

test/srt/test_bench_one_batch.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from sglang.test.test_utils import (
@@ -45,7 +46,10 @@ class TestBenchOneBatch(CustomTestCase):
                 f"### test_moe_tp2_bs1 (Mixtral-8x7B)\n"
                 f"output_throughput: {output_throughput:.2f} token/s\n"
             )
-            self.assertGreater(output_throughput, 125)
+            if os.getenv("SGLANG_AMD_CI") == "1":
+                self.assertGreater(output_throughput, 85)
+            else:
+                self.assertGreater(output_throughput, 125)
 
     def test_torch_compile_tp2_bs1(self):
         output_throughput = run_bench_offline_throughput(

test/srt/test_bench_serving.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from sglang.test.test_utils import (
@@ -180,7 +181,10 @@ class TestBenchServing(CustomTestCase):
                 f"### test_moe_offline_throughput_default\n"
                 f'Output throughput: {res["output_throughput"]:.2f} token/s\n'
             )
-            self.assertGreater(res["output_throughput"], 2200)
+            if os.getenv("SGLANG_AMD_CI") == "1":
+                self.assertGreater(res["output_throughput"], 2100)
+            else:
+                self.assertGreater(res["output_throughput"], 2200)
 
     def test_moe_offline_throughput_without_radix_cache(self):
         res = run_bench_serving(
@@ -195,7 +199,10 @@ class TestBenchServing(CustomTestCase):
                 f"### test_moe_offline_throughput_without_radix_cache\n"
                 f'Output throughput: {res["output_throughput"]:.2f} token/s\n'
             )
-            self.assertGreater(res["output_throughput"], 2200)
+            if os.getenv("SGLANG_AMD_CI") == "1":
+                self.assertGreater(res["output_throughput"], 2100)
+            else:
+                self.assertGreater(res["output_throughput"], 2200)
 
 
 if __name__ == "__main__":