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

This commit is contained in:
billishyahao
2026-03-20 02:52:25 +08:00
committed by GitHub
parent 9e629d31fd
commit 2ee7f41e25

View File

@@ -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"