From 56e5c07424be69f498d85d60694633d99fe152f9 Mon Sep 17 00:00:00 2001
From: b8zhong <b8zhong@uwaterloo.ca>
Date: Wed, 10 Dec 2025 00:03:43 -0800
Subject: [PATCH] fix b200 fa4 ci (#14788)

---
 test/srt/run_suite.py              |  2 +-
 test/srt/test_flash_attention_4.py | 16 ++++++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py
index ebe766925..ae97062c5 100644
--- a/test/srt/run_suite.py
+++ b/test/srt/run_suite.py
@@ -169,7 +169,7 @@ suites = {
     ],
     "per-commit-4-gpu-b200": [
         TestFile("test_deepseek_v3_fp4_4gpu.py", 1800),
-        TestFile("test_flash_attention_4.py", 300),
+        TestFile("test_flash_attention_4.py", 90),
         TestFile("test_fp8_blockwise_gemm.py", 280),
         TestFile("test_gpt_oss_4gpu.py", 700),
         TestFile("test_llama31_fp4.py", 90),
diff --git a/test/srt/test_flash_attention_4.py b/test/srt/test_flash_attention_4.py
index 44623a132..9d81ccd85 100644
--- a/test/srt/test_flash_attention_4.py
+++ b/test/srt/test_flash_attention_4.py
@@ -1,5 +1,6 @@
 import unittest
 from types import SimpleNamespace
+from urllib.parse import urlparse
 
 from sglang.srt.utils import get_device_sm, kill_process_tree
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
@@ -18,8 +19,6 @@ class TestFlashAttention4(unittest.TestCase):
         cls.base_url = DEFAULT_URL_FOR_TEST
         other_args = [
             "--trust-remote-code",
-            "--mem-fraction-static",
-            "0.8",
             "--prefill-attention-backend",
             "fa4",
             "--decode-attention-backend",
@@ -37,19 +36,20 @@ class TestFlashAttention4(unittest.TestCase):
         kill_process_tree(cls.process.pid)
 
     def test_gsm8k(self):
+        parsed_url = urlparse(self.base_url)
         args = SimpleNamespace(
-            num_shots=4,
+            num_shots=5,
             data_path=None,
-            num_questions=100,
+            num_questions=1319,
             max_new_tokens=512,
-            parallel=128,
-            host="http://127.0.0.1",
-            port=int(self.base_url.split(":")[-1]),
+            parallel=200,
+            host=f"{parsed_url.scheme}://{parsed_url.hostname}",
+            port=parsed_url.port,
         )
         metrics = run_eval_few_shot_gsm8k(args)
         print(metrics)
 
-        self.assertGreater(metrics["accuracy"], 0.75)
+        self.assertGreater(metrics["accuracy"], 0.89)
 
 
 if __name__ == "__main__":