[AMD] add mori ep normal tbo unittest (#20941)

test/registered/amd/test_moriep_small.py

@@ -36,7 +36,7 @@ common_args = [
     "--mem-fraction-static",
     "0.6",
     "--chunked-prefill-size",
-    "131072",
+    "32768",
     "--max-running-requests",
     "128",
     "--context-length",
@@ -70,7 +70,7 @@ class TestPureDP(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
         cls.process = popen_launch_server(
@@ -114,7 +114,7 @@ class TestMTP(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
         cls.process = popen_launch_server(
@@ -167,7 +167,7 @@ class TestNormal(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
         cls.process = popen_launch_server(
@@ -214,7 +214,7 @@ class TestLowLatency(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
         # env["MORI_DISABLE_P2P"] = "1"
@@ -249,6 +249,54 @@ class TestLowLatency(CustomTestCase):
         self.assertGreaterEqual(metrics["accuracy"], 0.935)
 
 
+class TestTBOwithNormal(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = common_args + [
+            "--deepep-mode",
+            "normal",
+            "--enable-two-batch-overlap",
+        ]
+
+        env = dict(os.environ)
+        env["SGLANG_USE_AITER"] = "1"
+        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
+        env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(
+        self,
+    ):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+
+        self.assertGreaterEqual(metrics["accuracy"], 0.935)
+
+
 class TestTBOwithLowLatency(CustomTestCase):
 
     @classmethod
@@ -264,7 +312,7 @@ class TestTBOwithLowLatency(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
         # env["MORI_DISABLE_P2P"] = "1"
@@ -299,7 +347,65 @@ class TestTBOwithLowLatency(CustomTestCase):
         self.assertGreaterEqual(metrics["accuracy"], 0.935)
 
 
-class TestMTPwithTBO(CustomTestCase):
+class TestMTPwithTBONormal(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + mtp_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--enable-two-batch-overlap",
+            ]
+        )
+
+        env = dict(os.environ)
+        env["SGLANG_USE_AITER"] = "1"
+        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
+        env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(
+        self,
+    ):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.92)
+
+        server_info = requests.get(self.base_url + "/get_server_info")
+        avg_spec_accept_length = server_info.json()["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreaterEqual(avg_spec_accept_length, 2.8)
+
+
+class TestMTPwithTBOLowLatency(CustomTestCase):
 
     @classmethod
     def setUpClass(cls):
@@ -318,7 +424,7 @@ class TestMTPwithTBO(CustomTestCase):
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
         env["SGLANG_MORI_FP8_DISP"] = "False"
-        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
+        env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
         # env["MORI_DISABLE_P2P"] = "1"