diff --git a/test/registered/amd/test_moriep_small.py b/test/registered/amd/test_moriep_small.py index eb808b1f2..ffe56e926 100644 --- a/test/registered/amd/test_moriep_small.py +++ b/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"