chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Migrate 4-GPU/8-GPU workflow jobs to stage-c and add CI registry decorators (#17299)

Browse Source
This commit is contained in:
Alison Shao
2026-01-31 22:37:22 -08:00
committed by GitHub
parent 95180484e9
commit a0bae4c343
39 changed files with 189 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/ep/test_deepep_large.py
-215
View File
@@ -1,215 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
@unittest.skip("Skip for saving ci time")
class TestDeepseek(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"8",
"--enable-dp-attention",
"--dp",
"8",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--moe-runner-backend",
"deep_gemm",
"--enable-two-batch-overlap",
"--ep-num-redundant-experts",
"32",
"--ep-dispatch-algorithm",
"dynamic",
"--eplb-algorithm",
"deepseek",
"--cuda-graph-bs",
"256",
"--max-running-requests",
"2048",
"--disable-radix-cache",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1200,
parallel=1200,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
class TestDeepseekMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"8",
"--enable-dp-attention",
"--dp",
"8",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--moe-runner-backend",
"deep_gemm",
"--enable-two-batch-overlap",
"--ep-num-redundant-experts",
"32",
"--ep-dispatch-algorithm",
"dynamic",
"--eplb-algorithm",
"deepseek",
"--cuda-graph-bs",
"64", # TODO: increase it to 128 when TBO is supported in draft_extend
"--max-running-requests",
"512",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--disable-radix-cache",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1200,
parallel=1200,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(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"###test_gsm8k:\n"
f"accuracy={metrics['accuracy']=:.3f}\n"
f"{avg_spec_accept_length=:.3f}\n"
)
self.assertGreater(avg_spec_accept_length, 1.85)
class TestDeepseekV32TBO(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--enable-two-batch-overlap",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"256",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1200,
parallel=1200,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if __name__ == "__main__":
unittest.main()
test/srt/ep/test_deepep_small.py
-575
View File
@@ -1,575 +0,0 @@
import os
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestPureDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static",
"0.5",
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
class TestHybridDPTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"2",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
class TestTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"128",
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
@unittest.skip("covered in test_deepep_large.py")
class TestNoGatherdBuffer(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"32",
"--max-running-requests",
"512",
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
class TestTBO(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--moe-dense-tp-size",
"1",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "1",
},
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
class TestTBOWithTPAttn(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "1",
},
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case.
# @unittest.skip("covered in TestMTPWithTPAttnAndTBO")
class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--moe-dense-tp-size",
"1",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "1",
},
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
@unittest.skip("covered in TestMTPWithTBO")
class TestMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"2",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--speculative-algo",
"EAGLE",
"--speculative-draft-model-path",
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
"--speculative-num-steps",
"2",
"--speculative-eagle-topk",
"3",
"--speculative-num-draft-tokens",
"3",
"--cuda-graph-max-bs",
"32",
"--max-running-requests",
"64",
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
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"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
f"accuracy={metrics['accuracy']=:.3f}\n"
f"{avg_spec_accept_length=:.3f}\n"
)
self.assertGreater(avg_spec_accept_length, 2.1)
class TestMTPWithTBO(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--enable-dp-attention",
"--dp-size",
"4",
"--enable-two-batch-overlap",
"--moe-a2a-backend",
"deepep",
"--trust-remote-code",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"2",
"--speculative-eagle-topk",
"3",
"--speculative-num-draft-tokens",
"3",
"--speculative-draft-model-path",
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
"--chunked-prefill-size",
"256",
"--cuda-graph-max-bs",
"32",
"--max-running-requests",
"128",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "1",
},
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
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"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
f"accuracy={metrics['accuracy']=:.3f}\n"
f"{avg_spec_accept_length=:.3f}\n"
)
self.assertGreater(avg_spec_accept_length, 2.1)
@unittest.skip("skipped due to bug when using MTP & TBO & attn_tp_size > 1")
class TestMTPWithTPAttnAndTBO(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--moe-dense-tp-size",
"1",
"--enable-two-batch-overlap",
"--moe-a2a-backend",
"deepep",
"--trust-remote-code",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"2",
"--speculative-eagle-topk",
"3",
"--speculative-num-draft-tokens",
"3",
"--speculative-draft-model-path",
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
"--chunked-prefill-size",
"256",
"--cuda-graph-max-bs",
"32",
"--max-running-requests",
"128",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "1",
},
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
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"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
f"accuracy={metrics['accuracy']=:.3f}\n"
f"{avg_spec_accept_length=:.3f}\n"
)
self.assertGreater(avg_spec_accept_length, 2.1)
if __name__ == "__main__":
unittest.main()
test/srt/ep/test_mooncake_ep_small.py
-122
View File
@@ -1,122 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
ib_devices = get_rdma_devices_args()
class TestTP(CustomTestCase):
extra_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--elastic-ep-backend",
"mooncake",
"--mooncake-ib-device",
ib_devices,
"--moe-a2a-backend",
"mooncake",
"--deepep-mode",
"low_latency",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--enable-two-batch-overlap",
"--disable-custom-all-reduce",
"--enable-eplb",
"--ep-num-redundant-experts",
"72",
"--chunked-prefill-size",
"512",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static",
"0.5",
*cls.extra_args,
],
)
@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(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
class TestPureDP(TestTP):
extra_args = [
"--enable-dp-attention",
"--dp",
"4",
]
pkill_process_1 = "sglang::scheduler_DP1_TP1_EP1"
pkill_process_2 = "sglang::scheduler_DP3_TP3_EP3"
def test_gsm8k_fault_1(self):
"""
Kill one rank and the system should remain operational.
"""
os.system(f"pkill -f {self.pkill_process_1}")
super().test_gsm8k()
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
def test_gsm8k_fault_2(self):
"""
Kill another rank and the system should remain operational.
"""
os.system(f"pkill -f {self.pkill_process_2}")
super().test_gsm8k()
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
class TestHybridDPTP(TestPureDP):
extra_args = [
"--enable-dp-attention",
"--dp",
"2",
]
pkill_process_1 = "sglang::scheduler_DP1_TP2_EP2"
pkill_process_2 = "sglang::scheduler_DP1_TP3_EP3"
if __name__ == "__main__":
unittest.main()
test/srt/ep/test_moriep_small.py
-170
View File
@@ -1,170 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST_NEXTN,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestPureDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp-size",
"8",
"--ep-size",
"8",
"--dp-size",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"mori",
"--trust-remote-code",
"--load-balance-method",
"round_robin",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--mem-fraction-static",
"0.6",
"--chunked-prefill-size",
"131072",
"--max-running-requests",
"128",
"--context-length",
"12288",
"--attention-backend",
"aiter",
]
env = dict(os.environ)
env["SGLANG_USE_AITER"] = "1"
env["SGLANG_MORI_FP8_DISP"] = "True"
env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
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.assertGreater(metrics["accuracy"], 0.935)
class TestMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp-size",
"8",
"--ep-size",
"8",
"--dp-size",
"8",
"--enable-dp-attention",
"--moe-a2a-backend",
"mori",
"--trust-remote-code",
"--load-balance-method",
"round_robin",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--mem-fraction-static",
"0.6",
"--chunked-prefill-size",
"131072",
"--max-running-requests",
"128",
"--context-length",
"12288",
"--attention-backend",
"aiter",
"--speculative-algo",
"EAGLE",
"--speculative-draft-model-path",
DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST_NEXTN,
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--cuda-graph-max-bs",
"32",
]
env = dict(os.environ)
env["SGLANG_USE_AITER"] = "1"
env["SGLANG_MORI_FP8_DISP"] = "True"
env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "16384"
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.assertGreater(metrics["accuracy"], 0.935)
if __name__ == "__main__":
unittest.main()
test/srt/models/test_dummy_grok_models.py
-34
View File
@@ -1,34 +0,0 @@
import unittest
from sglang.test.test_utils import CustomTestCase, is_in_ci, run_bench_one_batch
class TestDummyGrok1(CustomTestCase):
def test_dummy_grok_1(self):
_, output_throughput, _ = run_bench_one_batch(
None,
[
"--model",
"/dummy-grok",
"--tokenizer-path",
"Xenova/grok-1-tokenizer",
"--batch-size",
"2",
"--tp",
"2",
"--quantization",
"fp8",
"--load-format",
"dummy",
"--json-model-override-args",
'{"num_hidden_layers": 2}',
],
)
if is_in_ci():
self.assertGreater(output_throughput, 0)
if __name__ == "__main__":
unittest.main()
test/srt/models/test_kimi_k2_models.py
-70
View File
@@ -1,70 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
class TestKimiK2Thinking(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "moonshotai/Kimi-K2-Thinking"
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--tool-call-parser",
"kimi_k2",
"--reasoning-parser",
"kimi_k2",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
):
requests.get(self.base_url + "/flush_cache")
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=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (Kimi-K2-Thinking)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.95)
if __name__ == "__main__":
unittest.main()
test/srt/models/test_mimo_models.py
-47
View File
@@ -1,47 +0,0 @@
import unittest
from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin
from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase):
gsm8k_accuracy_thres = 0.75
gsm8k_num_questions = 1319
gsm8k_parallel = 1319
model = "XiaomiMiMo/MiMo-V2-Flash"
other_args = [
"--tp",
"4",
"--dp",
"2",
"--enable-dp-attention",
"--trust-remote-code",
"--attention-backend",
"fa3",
"--max-running-requests",
"128",
"--cuda-graph-max-bs",
"64",
"--mem-fraction-static",
"0.75",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--enable-multi-layer-eagle",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
]
bs_1_speed_thres = 170
accept_length_thres = 3.2
if __name__ == "__main__":
unittest.main()
test/srt/models/test_ministral3_models.py
-26
View File
@@ -1,26 +0,0 @@
import unittest
from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin
from sglang.test.kits.mmmu_vlm_kit import MMMUMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase
MODEL = "mistralai/Ministral-3-3B-Instruct-2512"
class TestMinistral3TextOnly(GSM8KMixin, DefaultServerBase):
gsm8k_accuracy_thres = 0.6
model = MODEL
other_args = ["--trust-remote-code"]
class TestMinistral3MMMU(MMMUMixin, MMMUServerBase):
accuracy = 0.3
model = MODEL
other_args = ["--trust-remote-code"]
mmmu_args = ["--limit=0.1"]
"""`--limit=0.1`: 10 percent of each task - this is fine for testing since the nominal result isn't interesting - this run is just to prevent relative regressions."""
if __name__ == "__main__":
unittest.main()
test/srt/models/test_qwen3_next_models.py
-127
View File
@@ -1,127 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.kl_test_utils import (
test_input_output_logprobs_match_decode_cache_hit_helper,
test_input_output_logprobs_match_prefill_cache_hit_helper,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
ACC_THRESHOLDS = {
QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93},
}
def send_request_helper(base_url: str, text: str):
response = requests.post(
base_url + "/generate",
json={
"text": text,
"sampling_params": {
"max_new_tokens": 1,
},
},
)
return response.json()
class TestQwen3Next(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = QWEN3_NEXT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
"128",
],
)
@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(args)
print(f"{metrics=}")
self.assertGreaterEqual(
metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
)
def test_input_output_logprobs_match_prefill_cache_hit(self):
test_input_output_logprobs_match_prefill_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS,
self.model,
max_samples=32,
max_new_tokens=512,
)
def test_input_output_logprobs_match_decode_cache_hit(self):
test_input_output_logprobs_match_decode_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS,
self.model,
max_samples=32,
max_new_tokens=512,
)
def test_prefix_cache_branching(self):
print("running test_prefix_cache_branching")
requests.get(self.base_url + "/flush_cache")
branching_pos = 257
text_prefix = "hi" * branching_pos
suffix_list = ["this" * 256, "here" * 256, "that" * 256]
cache_hit_list = [False, False, True]
# First request only prefill the entire sequence
# Second request won't have cache hit, but will cache the branching point
# Third request will have cache hit on the branching point
for i, (suffix, cache_hit) in enumerate(
zip(suffix_list, cache_hit_list, strict=True)
):
result = send_request_helper(self.base_url, text_prefix + suffix)
cached_tokens = result["meta_info"]["cached_tokens"]
if cache_hit:
expected_cached_tokens = branching_pos // 64 * 64
assert (
cached_tokens == expected_cached_tokens
), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}"
else:
assert (
cached_tokens == 0
), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0"
print("test_prefix_cache_branching passed")
if __name__ == "__main__":
unittest.main()
test/srt/models/test_qwen3_next_models_mtp.py
-212
View File
@@ -1,212 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.kl_test_utils import (
test_input_output_logprobs_match_decode_cache_hit_helper,
test_input_output_logprobs_match_prefill_cache_hit_helper,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
ACC_THRESHOLDS = {
QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93},
}
# MTP has higher KL divergence threshold
ACC_THRESHOLDS_MTP = {
QWEN3_NEXT_MODEL: {"kl_div": 0.008, "gsm8k": 0.93},
}
def send_request_helper(base_url: str, text: str):
response = requests.post(
base_url + "/generate",
json={
"text": text,
"sampling_params": {
"max_new_tokens": 1,
},
},
)
return response.json()
class TestQwen3NextMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = QWEN3_NEXT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-fraction-static",
"0.8",
"--tp",
"4",
"--chunked-prefill-size",
"2048",
"--mamba-scheduler-strategy",
"no_buffer",
],
)
@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(args)
print(f"{metrics=}")
self.assertGreaterEqual(
metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
)
def test_input_output_logprobs_match_prefill_cache_hit(self):
test_input_output_logprobs_match_prefill_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS,
self.model,
max_samples=32,
max_new_tokens=512,
)
def test_input_output_logprobs_match_decode_cache_hit(self):
test_input_output_logprobs_match_decode_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS,
self.model,
max_samples=32,
max_new_tokens=512,
)
class TestQwen3NextMTPTopk(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = QWEN3_NEXT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
"5",
"--speculative-eagle-topk",
"4",
"--speculative-num-draft-tokens",
"8",
"--mem-fraction-static",
"0.8",
"--tp",
"4",
"--chunked-prefill-size",
"2048",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
"128",
],
)
@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(args)
print(f"{metrics=}")
self.assertGreaterEqual(
metrics["accuracy"], ACC_THRESHOLDS_MTP[self.model]["gsm8k"]
)
def test_input_output_logprobs_match_prefill_cache_hit(self):
test_input_output_logprobs_match_prefill_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS_MTP,
self.model,
max_samples=32,
max_new_tokens=512,
)
def test_input_output_logprobs_match_decode_cache_hit(self):
test_input_output_logprobs_match_decode_cache_hit_helper(
self.base_url,
ACC_THRESHOLDS_MTP,
self.model,
max_samples=32,
max_new_tokens=512,
)
def test_prefix_cache_branching(self):
print("running test_prefix_cache_branching")
requests.get(self.base_url + "/flush_cache")
branching_pos = 257
text_prefix = "hi" * branching_pos
suffix_list = ["this" * 256, "here" * 256, "that" * 256]
cache_hit_list = [False, False, True]
# First request only prefill the entire sequence
# Second request won't have cache hit, but will cache the branching point
# Third request will have cache hit on the branching point
for i, (suffix, cache_hit) in enumerate(
zip(suffix_list, cache_hit_list, strict=True)
):
result = send_request_helper(self.base_url, text_prefix + suffix)
cached_tokens = result["meta_info"]["cached_tokens"]
if cache_hit:
expected_cached_tokens = branching_pos // 64 * 64
assert (
cached_tokens == expected_cached_tokens
), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}"
else:
assert (
cached_tokens == 0
), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0"
print("test_prefix_cache_branching passed")
if __name__ == "__main__":
unittest.main()
test/srt/models/test_qwen3_next_models_pcg.py
-70
View File
@@ -1,70 +0,0 @@
"""
Qwen3 Next piecewise CUDA graph tests.
DISABLED: See https://github.com/sgl-project/sglang/issues/17039
PCG tests for Qwen3 Next have intermittent failures (5-10% probability).
Investigation ongoing by @YuweiAn.
"""
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
ACC_THRESHOLDS = {
QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93},
}
@unittest.skip("Disabled: intermittent failures, see #17039")
class TestQwen3NextPiecewiseCudaGraph(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = QWEN3_NEXT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp",
"4",
"--enable-piecewise-cuda-graph",
"--piecewise-cuda-graph-compiler",
"eager",
],
)
@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(args)
print(f"{metrics=}")
self.assertGreaterEqual(
metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
)
if __name__ == "__main__":
unittest.main()
test/srt/quant/test_w4a8_deepseek_v3.py
-246
View File
@@ -1,246 +0,0 @@
import os
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_amd_ci,
is_in_ci,
popen_launch_server,
try_cached_model,
write_github_step_summary,
)
class TestDeepseekV3W4afp8(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = ["--trust-remote-code", "--tp", "8", "--ep-size", "8"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1200,
parallel=1200,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
class TestDeepseekV3W4Afp8Mtp(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--ep-size",
"8",
"--cuda-graph-bs",
"256",
"--disable-radix-cache",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"2",
"--speculative-num-draft-tokens",
"4",
]
if not is_in_amd_ci():
other_args += ["--mem-frac", "0.7"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@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=}")
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=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3 mtp)\n"
f'{metrics["accuracy"]=:.3f}\n'
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], 0.935)
self.assertGreater(avg_spec_accept_length, 2.9)
class TestDeepseekV3W4Afp8DeepepNormal(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--ep-size",
"8",
"--cuda-graph-bs",
"256",
"--disable-radix-cache",
"--moe-a2a-backend",
"deepep",
"--deepep-mode",
"normal",
"--dp",
"8",
"--enable-dp-attention",
"--moe-runner-backend",
"cutlass",
]
if not is_in_amd_ci():
other_args += ["--mem-frac", "0.7"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@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"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
class TestDeepseekV3W4Afp8DeepepAutoMtp(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--ep-size",
"8",
"--cuda-graph-bs",
"256",
"--disable-radix-cache",
"--moe-a2a-backend",
"deepep",
"--deepep-mode",
"auto",
"--dp",
"8",
"--enable-dp-attention",
"--moe-runner-backend",
"cutlass",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
if not is_in_amd_ci():
other_args += ["--mem-frac", "0.7"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env={
**os.environ,
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
},
)
@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"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
if __name__ == "__main__":
unittest.main()
test/srt/run_suite.py
+7 -60
View File
@@ -7,65 +7,13 @@ import tabulate
from sglang.test.ci.ci_utils import TestFile, run_unittest_files
# NOTE: please sort the test cases alphabetically by the test file name
# NOTE: per-commit-4-gpu, per-commit-8-gpu-h200, per-commit-8-gpu-h20, per-commit-4-gpu-b200,
# per-commit-4-gpu-gb200, per-commit-4-gpu-deepep, and per-commit-8-gpu-h200-deepep suites
# have been migrated to stage-c suites in test/registered/ using the CI registry system.
suites = {
"per-commit-4-gpu": [
TestFile("models/test_qwen3_next_models.py", 350),
TestFile("models/test_qwen3_next_models_mtp.py", 500),
TestFile("test_gpt_oss_4gpu.py", 300),
TestFile("test_multi_instance_release_memory_occupation.py", 64),
TestFile("test_pp_single_node.py", 500),
TestFile("test_epd_disaggregation.py", 150),
],
"per-commit-8-gpu-h200": [
TestFile("test_deepseek_v3_basic.py", 275),
TestFile("test_deepseek_v3_mtp.py", 275),
TestFile("test_disaggregation_hybrid_attention.py", 400),
TestFile("models/test_kimi_k2_models.py", 200),
TestFile("test_deepseek_v32_basic.py", 360),
TestFile("test_deepseek_v32_mtp.py", 360),
TestFile("models/test_mimo_models.py", 200),
],
"per-commit-8-gpu-h20": [
TestFile("quant/test_w4a8_deepseek_v3.py", 520),
TestFile("test_disaggregation_different_tp.py", 600),
TestFile("test_disaggregation_pp.py", 180),
TestFile("test_disaggregation_dp_attention.py", 155),
],
"per-commit-4-gpu-b200": [
TestFile("test_deepseek_v3_fp4_4gpu.py", 1500),
TestFile("test_fp8_blockwise_gemm.py", 280),
TestFile("test_gpt_oss_4gpu.py", 300),
TestFile("test_nvfp4_gemm.py", 360),
TestFile("test_deepseek_v32_fp4_4gpu.py", 600),
],
# "per-commit-8-gpu-b200": [
# TestFile("test_mistral_large3_basic.py", 275), # Moved to nightly - large model
# ],
"per-commit-4-gpu-gb200": [
TestFile("test_deepseek_v3_cutedsl_4gpu.py", 1800),
TestFile("test_disaggregation_aarch64.py", 300),
],
"per-commit-4-gpu-deepep": [
TestFile("ep/test_deepep_small.py", 531),
TestFile("ep/test_mooncake_ep_small.py", 660),
],
"per-commit-8-gpu-h200-deepep": [
TestFile("ep/test_deepep_large.py", 563),
],
# quantization_test suite migrated to test/registered/quant/
"__not_in_ci__": [
TestFile("test_release_memory_occupation.py", 200), # Temporarily disabled
TestFile("models/test_dummy_grok_models.py"),
TestFile("test_profile_v2.py"),
TestFile("models/test_ministral3_models.py"),
TestFile("test_mistral_large3_basic.py"),
TestFile("test_prefill_delayer.py"),
TestFile("test_fla_layernorm_guard.py"),
TestFile(
"models/test_qwen3_next_models_pcg.py"
), # Disabled: intermittent failures, see #17039
TestFile("ep/test_moriep_small.py"),
],
# All CUDA tests migrated to test/registered/
"__not_in_ci__": [],
}
# Add AMD tests
@@ -89,9 +37,8 @@ suite_amd = {
# TestFile("test_wave_attention_backend.py", 150), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
# The time estimation for `test_int4fp8_moe.py` assumes `mistralai/Mixtral-8x7B-Instruct-v0.1` is already cached (running on 1xMI300X).
],
"per-commit-4-gpu-amd": [
TestFile("test_pp_single_node.py", 150),
],
# per-commit-4-gpu-amd migrated to test/registered/distributed/ using the CI registry system
"per-commit-4-gpu-amd": [],
# NOTE: AMD nightly suites (nightly-amd, nightly-amd-vlm, nightly-amd-8-gpu)
# have been migrated to test/registered/amd/nightly/ and are now managed
# by test/run_suite.py using the registry system.
test/srt/test_deepseek_v32_basic.py
-137
View File
@@ -1,137 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
class TestDeepseekV32DP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 50)
class TestDeepseekV32TP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 70)
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v32_fp4_4gpu.py
-79
View File
@@ -1,79 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3.2-NVFP4"
SERVER_LAUNCH_TIMEOUT = 1200
class TestDeepseekV32FP4(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--attention-backend",
"nsa",
"--moe-runner-backend",
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--tool-call-parser",
"deepseekv32",
"--reasoning-parser",
"deepseek-v3",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v32_mtp.py
-189
View File
@@ -1,189 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
FULL_DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
class TestDeepseekV32DPMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-frac",
"0.7",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
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=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32 mtp)\n"
f'{metrics["accuracy"]=:.3f}\n'
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], 0.94)
self.assertGreater(avg_spec_accept_length, 2.7)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{acc_length=:.2f} {speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32 mtp)\n"
f"{acc_length=:.2f}\n"
f"{speed=:.2f} token/s\n"
)
self.assertGreater(acc_length, 2.7)
self.assertGreater(speed, 75)
class TestDeepseekV32TPMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-frac",
"0.7",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
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=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v32 mtp)\n"
f'{metrics["accuracy"]=:.3f}\n'
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], 0.94)
self.assertGreater(avg_spec_accept_length, 2.7)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{acc_length=:.2f} {speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v32 mtp)\n"
f"{acc_length=:.2f}\n"
f"{speed=:.2f} token/s\n"
)
self.assertGreater(acc_length, 2.7)
self.assertGreater(speed, 130)
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v3_basic.py
-81
View File
@@ -1,81 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_amd_ci,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
class TestDeepseekV3Basic(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3)\n" f"{speed=:.2f} token/s\n"
)
if is_in_amd_ci():
self.assertGreater(speed, 12)
else:
self.assertGreater(speed, 75)
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v3_cutedsl_4gpu.py
-161
View File
@@ -1,161 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
class TestDeepseekR1Nvfp4CuteDSLDeepEP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--disable-radix-cache",
"--mem-fraction-static",
"0.8",
"--max-prefill-tokens",
"16384",
"--max-running-requests",
"256",
"--chunked-prefill-size",
"1024",
"--tp",
"4",
"--dp",
"4",
"--ep",
"4",
"--moe-dense-tp-size",
"1",
"--enable-dp-attention",
"--quantization",
"modelopt_fp4",
"--attention-backend",
"trtllm_mla",
"--moe-a2a-backend",
"deepep",
"--moe-runner-backend",
"flashinfer_cutedsl",
"--deepep-mode",
"low_latency",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env={
**os.environ,
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
"SGLANG_MOE_NVFP4_DISPATCH": "0",
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=512,
parallel=512,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["accuracy"], 0.92)
class TestDummyWithSBO(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--disable-radix-cache",
"--mem-fraction-static",
"0.05",
"--max-prefill-tokens",
"16384",
"--max-running-requests",
"256",
"--chunked-prefill-size",
"1024",
"--cuda-graph-bs",
"64",
"--tp",
"4",
"--dp",
"4",
"--ep",
"4",
"--moe-dense-tp-size",
"1",
"--enable-dp-attention",
"--quantization",
"modelopt_fp4",
"--attention-backend",
"trtllm_mla",
"--moe-a2a-backend",
"deepep",
"--moe-runner-backend",
"flashinfer_cutedsl",
"--deepep-mode",
"low_latency",
"--json-model-override-args",
'{"num_hidden_layers": 1, "first_k_dense_replace": 0, "n_routed_experts": 24}',
"--enable-single-batch-overlap",
"--load-format",
"dummy",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env={
**os.environ,
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
"SGLANG_MOE_NVFP4_DISPATCH": "0",
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=0,
data_path=None,
num_questions=512,
parallel=512,
max_new_tokens=16,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v3_fp4_4gpu.py
-265
View File
@@ -1,265 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4"
SERVER_LAUNCH_TIMEOUT = 1200
class TestDeepseekV3FP4(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"4",
"--attention-backend",
"trtllm_mla",
"--moe-runner-backend",
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3-fp4)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 75)
class TestDeepseekV3FP4PiecewiseCudaGraph(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"4",
"--attention-backend",
"trtllm_mla",
"--moe-runner-backend",
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--enable-piecewise-cuda-graph",
"--kv-cache-dtype",
"fp8_e4m3",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
):
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
_, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3-fp4)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 120)
class TestDeepseekV3FP4CutlassMoE(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"4",
"--ep",
"4",
"--attention-backend",
"trtllm_mla",
"--moe-runner-backend",
"flashinfer_cutlass",
"--quantization",
"modelopt_fp4",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
env={
**os.environ,
"SGLANG_MOE_NVFP4_DISPATCH": "1", # Enable nvfp4 all gather
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3-fp4-cutlass-moe)\n"
f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
class TestDeepseekV3FP4SymmetricMemory(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"4",
"--attention-backend",
"trtllm_mla",
"--moe-runner-backend",
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--model-loader-extra-config",
'{"enable_multithread_load": true,"num_threads": 64}',
"--enable-symm-mem",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
if __name__ == "__main__":
unittest.main()
test/srt/test_deepseek_v3_mtp.py
-107
View File
@@ -1,107 +0,0 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_amd_ci,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
class TestDeepseekV3MTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
if not is_in_amd_ci():
other_args += ["--mem-frac", "0.7"]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
requests.get(self.base_url + "/flush_cache")
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=}")
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=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3 mtp)\n"
f'{metrics["accuracy"]=:.3f}\n'
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], 0.935)
self.assertGreater(avg_spec_accept_length, 2.8)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{acc_length=:.2f} {speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v3 mtp)\n"
f"{acc_length=:.2f}\n"
f"{speed=:.2f} token/s\n"
)
self.assertGreater(acc_length, 2.8)
if is_in_amd_ci():
self.assertGreater(speed, 15)
else:
self.assertGreater(speed, 130)
if __name__ == "__main__":
unittest.main()
test/srt/test_disaggregation_aarch64.py
-93
View File
@@ -1,93 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
class TestDisaggregationMooncakeAARCH64Accuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
os.environ["SGLANG_MOONCAKE_CUSTOM_MEM_POOL"] = "true"
os.environ["MC_FORCE_MNNVL"] = "true"
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def tearDownClass(cls):
os.environ.pop("SGLANG_MOONCAKE_CUSTOM_MEM_POOL")
os.environ.pop("MC_FORCE_MNNVL")
super().tearDownClass()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"2",
"--base-gpu-id",
"2",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.62)
if __name__ == "__main__":
unittest.main()
test/srt/test_disaggregation_different_tp.py
-303
View File
@@ -1,303 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.environ import envs
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
try_cached_model,
)
class TestDisaggregationMooncakePrefillLargerTP(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Temporarily disable JIT DeepGEMM
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"4",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"2",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.60)
class TestDisaggregationMooncakeDecodeLargerTP(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Temporarily disable JIT DeepGEMM
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"4",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.60)
class TestDisaggregationMooncakeMHAPrefillLargerTP(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Temporarily disable JIT DeepGEMM
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"4",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"2",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.60)
class TestDisaggregationMooncakeMHADecodeLargerTP(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Temporarily disable JIT DeepGEMM
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"4",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.60)
if __name__ == "__main__":
unittest.main()
test/srt/test_disaggregation_dp_attention.py
-98
View File
@@ -1,98 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.environ import envs
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
try_cached_model,
)
class TestDisaggregationDPAttention(PDDisaggregationServerBase):
PREFILL_DP_SIZE = 4
DECODE_DP_SIZE = 4
@classmethod
def setUpClass(cls):
super().setUpClass()
# Temporarily disable JIT DeepGEMM
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
str(cls.PREFILL_DP_SIZE),
"--dp",
str(cls.PREFILL_DP_SIZE),
"--enable-dp-attention",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
str(cls.DECODE_DP_SIZE),
"--dp",
str(cls.DECODE_DP_SIZE),
"--enable-dp-attention",
"--base-gpu-id",
str(cls.PREFILL_DP_SIZE),
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1400,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.60)
if __name__ == "__main__":
unittest.main()
test/srt/test_disaggregation_hybrid_attention.py
-235
View File
@@ -1,235 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
is_in_ci,
popen_launch_pd_server,
)
@unittest.skipIf(is_in_ci(), "Temporarily disable the flaky test.")
class TestDisaggregationHybridAttentionMamba(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"4",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"4",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.93)
class TestDisaggregationHybridAttentionMambaExtraBuffer(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"4",
"--mamba-scheduler-strategy",
"extra_buffer",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"4",
"--base-gpu-id",
"4",
"--mamba-scheduler-strategy",
"extra_buffer",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.93)
@unittest.skipIf(
is_in_ci(),
"Temporarily disable the flaky test: tcp fallback is not stable currently.",
)
class TestDisaggregationHybridAttentionMambaDPDecode(PDDisaggregationServerBase):
"""Test with prefill tp=2 and decode tp=2/dp=2 with dp-attention enabled."""
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"2",
"--dp",
"2",
"--enable-dp-attention",
"--enable-dp-lm-head",
"--base-gpu-id",
"2",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.93)
if __name__ == "__main__":
unittest.main()
test/srt/test_disaggregation_pp.py
-240
View File
@@ -1,240 +0,0 @@
import time
import unittest
from types import SimpleNamespace
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
try_cached_model,
)
class TestDisaggregationPrefillPPAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.24)
# Wait a little bit so that the memory check happens.
time.sleep(5)
class TestDisaggregationPrefillPPDynamicChunkAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
"--enable-dynamic-chunking",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.24)
# Wait a little bit so that the memory check happens.
time.sleep(5)
class TestDisaggregationDecodePPAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--pp-size",
"2",
"--base-gpu-id",
"4",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.24)
# Wait a little bit so that the memory check happens.
time.sleep(5)
if __name__ == "__main__":
unittest.main()
test/srt/test_epd_disaggregation.py
-426
View File
@@ -1,426 +0,0 @@
import os
import threading
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
is_in_ci,
popen_launch_server,
)
@unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
class TestEPDDisaggregationOneEncoder(PDDisaggregationServerBase):
"""Test EPD disaggregation with single encode server"""
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
cls.encode_port = f"{int(cls.lb_port) + 300}"
cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
print(
f"Setting up EPD (one encoder): encode={cls.encode_port}, "
f"prefill={cls.prefill_port}, decode={cls.decode_port}"
)
# Start servers in order: encode -> prefill/decode
cls.start_encode()
prefill_thread = threading.Thread(target=cls.start_prefill)
decode_thread = threading.Thread(target=cls.start_decode)
prefill_thread.start()
decode_thread.start()
prefill_thread.join()
decode_thread.join()
# Wait for all servers to be ready
cls.wait_server_ready(cls.encode_url + "/health")
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
cls.api_key = "sk-123456"
os.environ["OPENAI_API_KEY"] = cls.api_key
os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
@classmethod
def start_encode(cls):
"""Start encode server for multimodal processing"""
encode_args = [
"--trust-remote-code",
"--encoder-only",
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--tp",
"1",
"--port",
cls.encode_port,
"--enable-prefix-mm-cache",
]
cls.process_encode = popen_launch_server(
cls.model,
base_url=cls.encode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=encode_args,
)
@classmethod
def start_prefill(cls):
"""Start prefill server with language model only"""
prefill_args = [
"--trust-remote-code",
"--language-only",
"--encoder-urls",
cls.encode_url,
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--base-gpu-id",
"1",
"--port",
cls.prefill_port,
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_server(
cls.model,
base_url=cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
"""Start decode server"""
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"2",
"--port",
cls.decode_port,
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_server(
cls.model,
base_url=cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
@classmethod
def tearDownClass(cls):
"""Clean up all processes"""
for process in [
cls.process_lb,
cls.process_decode,
cls.process_prefill,
cls.process_encode,
]:
if process:
try:
kill_process_tree(process.pid)
except Exception as e:
print(f"Error killing process: {e}")
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
"""
Evaluate a VLM on the MMMU validation set with lmms-eval.
Reference: test_vlm_models.py
Args:
model_version: Model version/checkpoint to evaluate
output_path: Path to save evaluation results
limit: Number of samples to evaluate (default: "50" for CI time constraints)
"""
model = "openai_compatible"
tp = 1
tasks = "mmmu_val"
batch_size = 32
log_suffix = "openai_compatible"
os.makedirs(output_path, exist_ok=True)
model_args = f'model_version="{model_version}",' f"tp={tp}"
cmd = [
"python3",
"-m",
"lmms_eval",
"--model",
model,
"--model_args",
model_args,
"--tasks",
tasks,
"--batch_size",
str(batch_size),
"--log_samples",
"--log_samples_suffix",
log_suffix,
"--output_path",
str(output_path),
"--limit",
limit,
]
_run_lmms_eval_with_retry(cmd, timeout=3600)
def test_mmmu(self):
"""Test MMMU evaluation with EPD disaggregation"""
import glob
import json
output_path = "./logs/epd_one_encoder_mmmu"
self.run_mmmu_eval(self.model, output_path)
# Get the result file
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
if not result_files:
result_files = glob.glob(f"{output_path}/*.json")
if not result_files:
self.fail(f"No JSON result files found in {output_path}")
result_file_path = result_files[0]
with open(result_file_path, "r") as f:
result = json.load(f)
print(f"MMMU result: {result}")
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
print(f"MMMU accuracy: {mmmu_accuracy:.4f}")
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
self.assertGreater(mmmu_accuracy, 0.40)
class TestEPDDisaggregationMultiEncoders(PDDisaggregationServerBase):
"""
Test EPD disaggregation with multiple encode servers for load balancing.
Both encode servers run on GPU 0 (different ports) for testing load distribution.
"""
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
cls.encode_port1 = f"{int(cls.lb_port) + 300}"
cls.encode_port2 = f"{int(cls.lb_port) + 301}"
cls.encode_url1 = f"http://{cls.base_host}:{cls.encode_port1}"
cls.encode_url2 = f"http://{cls.base_host}:{cls.encode_port2}"
print(
f"Setting up EPD (multiple encoders): encode1={cls.encode_port1}, "
f"encode2={cls.encode_port2}, prefill={cls.prefill_port}, decode={cls.decode_port}"
)
# Start two encode servers on GPU 0/1
encode1_thread = threading.Thread(
target=cls.start_encode_server, args=(cls.encode_port1, 0)
)
encode2_thread = threading.Thread(
target=cls.start_encode_server, args=(cls.encode_port2, 1)
)
encode1_thread.start()
encode2_thread.start()
encode1_thread.join()
encode2_thread.join()
prefill_thread = threading.Thread(target=cls.start_prefill)
decode_thread = threading.Thread(target=cls.start_decode)
prefill_thread.start()
decode_thread.start()
prefill_thread.join()
decode_thread.join()
cls.wait_server_ready(cls.encode_url1 + "/health")
cls.wait_server_ready(cls.encode_url2 + "/health")
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
# Set OpenAI API key and base URL environment variables. Needed for lmms-eval to work.
cls.api_key = "sk-123456"
os.environ["OPENAI_API_KEY"] = cls.api_key
os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
@classmethod
def start_encode_server(cls, port, gpu_id):
"""Start an encode server on specific port and GPU"""
encode_args = [
"--trust-remote-code",
"--encoder-only",
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--tp",
"1",
"--port",
port,
"--enable-prefix-mm-cache",
]
# Only set base-gpu-id if not using GPU 0
if gpu_id != 0:
encode_args.extend(["--base-gpu-id", str(gpu_id)])
process = popen_launch_server(
cls.model,
base_url=f"http://{cls.base_host}:{port}",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=encode_args,
)
if port == cls.encode_port1:
cls.process_encode1 = process
else:
cls.process_encode2 = process
@classmethod
def start_prefill(cls):
"""Start prefill server with multiple encode URLs"""
prefill_args = [
"--trust-remote-code",
"--language-only",
"--encoder-urls",
cls.encode_url1,
cls.encode_url2,
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--base-gpu-id",
"2",
"--port",
cls.prefill_port,
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_server(
cls.model,
base_url=cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
"""Start decode server"""
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"3",
"--port",
cls.decode_port,
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_server(
cls.model,
base_url=cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
@classmethod
def tearDownClass(cls):
"""Clean up all processes"""
for process in [
cls.process_lb,
cls.process_decode,
cls.process_prefill,
cls.process_encode1,
cls.process_encode2,
]:
if process:
try:
kill_process_tree(process.pid)
except Exception as e:
print(f"Error killing process: {e}")
def run_mmmu_eval(self, model_version: str, output_path: str, limit: str = "50"):
"""
Evaluate a VLM on the MMMU validation set with lmms-eval.
Reference: test_vlm_models.py
Args:
model_version: Model version/checkpoint to evaluate
output_path: Path to save evaluation results
limit: Number of samples to evaluate (default: "50" for CI time constraints)
"""
model = "openai_compatible"
tp = 1
tasks = "mmmu_val"
batch_size = 32
log_suffix = "openai_compatible"
os.makedirs(output_path, exist_ok=True)
model_args = f'model_version="{model_version}",' f"tp={tp}"
cmd = [
"python3",
"-m",
"lmms_eval",
"--model",
model,
"--model_args",
model_args,
"--tasks",
tasks,
"--batch_size",
str(batch_size),
"--log_samples",
"--log_samples_suffix",
log_suffix,
"--output_path",
str(output_path),
"--limit",
limit,
]
_run_lmms_eval_with_retry(cmd, timeout=3600)
def test_mmmu(self):
"""Test MMMU evaluation with EPD disaggregation (multiple encoders)"""
import glob
import json
output_path = "./logs/epd_multi_encoder_mmmu"
self.run_mmmu_eval(self.model, output_path)
# Get the result file
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
if not result_files:
result_files = glob.glob(f"{output_path}/*.json")
if not result_files:
self.fail(f"No JSON result files found in {output_path}")
result_file_path = result_files[0]
with open(result_file_path, "r") as f:
result = json.load(f)
print(f"MMMU result (multi encoder): {result}")
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
print(f"MMMU accuracy (multi encoder): {mmmu_accuracy:.4f}")
# for qwen2.5-vl-3b-instruct, the accuracy is 0.40
self.assertGreater(mmmu_accuracy, 0.40)
if __name__ == "__main__":
unittest.main()
test/srt/test_fla_layernorm_guard.py
-384
View File
@@ -1,384 +0,0 @@
from __future__ import annotations
import socket
from dataclasses import dataclass
import pytest
import torch
import torch.nn.functional as F
from sglang.srt.layers.attention.fla.layernorm_gated import (
_layer_norm_fwd as layer_norm_fwd,
)
from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn, rms_norm_ref
# Optional dependency in sglang repo; skip collection cleanly if absent.
custom_all_reduce_utils = pytest.importorskip(
"sglang.srt.distributed.device_communicators.custom_all_reduce_utils"
)
parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state")
update_environment_variables = custom_all_reduce_utils.update_environment_variables
init_distributed_environment = parallel_state.init_distributed_environment
initialize_model_parallel = parallel_state.initialize_model_parallel
NUM_GPUS = 2
def _find_free_port() -> int:
# Avoid hard-coded port collisions when pytest runs tests in parallel.
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("localhost", 0))
s.listen(1)
return int(s.getsockname()[1])
def _skip_if_no_cuda_or_not_enough_gpus(required_gpus: int = NUM_GPUS) -> None:
if not torch.cuda.is_available():
pytest.skip("CUDA device not available")
if torch.cuda.device_count() < required_gpus:
pytest.skip(f"Need >= {required_gpus} GPUs, got {torch.cuda.device_count()}")
def _skip_if_dtype_unsupported(dtype: torch.dtype) -> None:
if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 not supported on this CUDA device")
def _setup_sglang_distributed(
local_rank: int,
world_size: int,
master_port: int,
dtype: torch.dtype,
) -> torch.device:
# Match sglang test style: set per-rank CUDA device + default dtype/device.
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
if hasattr(torch, "set_default_device"):
torch.set_default_device(device)
if hasattr(torch, "set_default_dtype"):
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": str(master_port),
}
)
init_distributed_environment(
world_size=world_size, rank=local_rank, local_rank=local_rank
)
initialize_model_parallel(tensor_model_parallel_size=world_size)
return device
def layer_norm_ref(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
z: torch.Tensor | None = None,
eps: float = 1e-6,
group_size: int | None = None,
norm_before_gate: bool = True,
is_rms_norm: bool = False,
) -> torch.Tensor:
"""Reference implementation for both LayerNorm and RMSNorm (supports optional gate + group norm)."""
if is_rms_norm:
return rms_norm_ref(
x,
weight,
bias,
z=z,
eps=eps,
group_size=group_size,
norm_before_gate=norm_before_gate,
upcast=True,
)
dtype = x.dtype
x_f = x.float()
w_f = weight.float()
b_f = bias.float() if bias is not None else None
z_f = z.float() if z is not None else None
if z_f is not None and not norm_before_gate:
x_f = x_f * F.silu(z_f)
if group_size is None:
mean = x_f.mean(dim=-1, keepdim=True)
var = (x_f - mean).square().mean(dim=-1, keepdim=True)
rstd = torch.rsqrt(var + eps)
out = (x_f - mean) * rstd * w_f
if b_f is not None:
out = out + b_f
else:
hidden = x_f.shape[-1]
assert hidden % group_size == 0
ng = hidden // group_size
xg = x_f.view(*x_f.shape[:-1], ng, group_size)
mean = xg.mean(dim=-1, keepdim=True)
var = (xg - mean).square().mean(dim=-1, keepdim=True)
rstd = torch.rsqrt(var + eps)
xg = (xg - mean) * rstd
out = xg.reshape(*x_f.shape[:-1], hidden) * w_f
if b_f is not None:
out = out + b_f
if z_f is not None and norm_before_gate:
out = out * F.silu(z_f)
return out.to(dtype)
@dataclass(frozen=True)
class FwdCase:
name: str
with_gate: bool
norm_before_gate: bool
group_size: int | None
is_rms_norm: bool
CASES: list[FwdCase] = [
FwdCase(
"layernorm",
with_gate=False,
norm_before_gate=True,
group_size=None,
is_rms_norm=False,
),
FwdCase(
"rmsnorm",
with_gate=False,
norm_before_gate=True,
group_size=None,
is_rms_norm=True,
),
FwdCase(
"layernorm_gate_pre",
with_gate=True,
norm_before_gate=True,
group_size=None,
is_rms_norm=False,
),
FwdCase(
"layernorm_gate_post",
with_gate=True,
norm_before_gate=False,
group_size=None,
is_rms_norm=False,
),
FwdCase(
"rmsnorm_gate_pre",
with_gate=True,
norm_before_gate=True,
group_size=None,
is_rms_norm=True,
),
FwdCase(
"group_layernorm",
with_gate=False,
norm_before_gate=True,
group_size=128,
is_rms_norm=False,
),
FwdCase(
"group_rmsnorm",
with_gate=False,
norm_before_gate=True,
group_size=128,
is_rms_norm=True,
),
]
@pytest.mark.parametrize("num_tokens", [128])
@pytest.mark.parametrize("hidden_size", [256])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("case", CASES, ids=lambda c: c.name)
def test_layernorm_guard_fwd_spawn(
num_tokens: int,
hidden_size: int,
dtype: torch.dtype,
case: FwdCase,
device: str = "cuda",
):
_skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
_skip_if_dtype_unsupported(dtype)
if case.group_size is not None and hidden_size % case.group_size != 0:
pytest.skip(
f"hidden_size {hidden_size} not divisible by group_size {case.group_size}"
)
master_port = _find_free_port()
world_size = NUM_GPUS
torch.multiprocessing.spawn(
_layernorm_guard_fwd_worker,
args=(
world_size,
master_port,
num_tokens,
hidden_size,
dtype,
case,
device,
),
nprocs=world_size,
join=True,
)
def _layernorm_guard_fwd_worker(
local_rank: int,
world_size: int,
master_port: int,
num_tokens: int,
hidden_size: int,
dtype: torch.dtype,
case: FwdCase,
device: str,
):
device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
with torch.inference_mode():
torch.manual_seed(42 + local_rank)
torch.cuda.manual_seed_all(42 + local_rank)
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
z = (
torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
if case.with_gate
else None
)
weight = torch.randn(hidden_size, dtype=dtype, device=device)
bias = (
None
if case.is_rms_norm
else torch.randn(hidden_size, dtype=dtype, device=device)
)
eps = 1e-6
out, mean, rstd = layer_norm_fwd(
x,
weight,
bias,
eps,
z=z,
group_size=case.group_size,
norm_before_gate=case.norm_before_gate,
is_rms_norm=case.is_rms_norm,
)
ref_out = layer_norm_ref(
x,
weight,
bias,
z=z,
eps=eps,
group_size=case.group_size,
norm_before_gate=case.norm_before_gate,
is_rms_norm=case.is_rms_norm,
)
assert out.shape == x.shape
assert out.dtype == x.dtype
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
# mean/rstd shape checks (same spirit as original vLLM tests)
if case.group_size is None:
if not case.is_rms_norm:
assert mean.shape == (num_tokens,)
assert rstd.shape == (num_tokens,)
else:
ngroups = hidden_size // case.group_size
if not case.is_rms_norm:
assert mean.shape == (ngroups * num_tokens,)
assert rstd.shape == (ngroups * num_tokens,)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
def test_layernorm_guard_misc_spawn(dtype: torch.dtype, device: str = "cuda"):
_skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
_skip_if_dtype_unsupported(dtype)
master_port = _find_free_port()
world_size = NUM_GPUS
torch.multiprocessing.spawn(
_layernorm_guard_misc_worker,
args=(world_size, master_port, dtype, device),
nprocs=world_size,
join=True,
)
def _layernorm_guard_misc_worker(
local_rank: int,
world_size: int,
master_port: int,
dtype: torch.dtype,
device: str,
):
device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
with torch.inference_mode():
torch.manual_seed(123 + local_rank)
torch.cuda.manual_seed_all(123 + local_rank)
# 1) rows_per_block-like sizes
hidden_size = 1024
weight = torch.randn(hidden_size, dtype=dtype, device=device)
bias = torch.randn(hidden_size, dtype=dtype, device=device)
eps = 1e-6
for num_tokens in [513]:
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
out, _, _ = layer_norm_fwd(x, weight, bias, eps, z=None, is_rms_norm=False)
ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
# 2) strided input (slice then contiguous)
num_tokens = 128
x_large = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
x = x_large[:, :hidden_size]
x_contig = x.contiguous()
out, _, _ = layer_norm_fwd(
x_contig, weight, bias, eps, z=None, is_rms_norm=False
)
ref = layer_norm_ref(x_contig, weight, bias, z=None, eps=eps, is_rms_norm=False)
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
# 3) provided output buffer
num_tokens = 256
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
out_buf = torch.empty_like(x)
out, _, _ = layer_norm_fwd(
x, weight, bias, eps, z=None, out=out_buf, is_rms_norm=False
)
assert out.data_ptr() == out_buf.data_ptr()
ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
# 4) multidimensional input via autograd fn
for shape in [(4, 16, 1024)]:
hs = shape[-1]
x = torch.randn(*shape, dtype=dtype, device=device)
w = torch.randn(hs, dtype=dtype, device=device)
b = torch.randn(hs, dtype=dtype, device=device)
out = layernorm_fn(x, w, b, z=None, eps=eps)
ref = layer_norm_ref(x, w, b, z=None, eps=eps, is_rms_norm=False)
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
if __name__ == "__main__":
pytest.main([__file__])
test/srt/test_fp8_blockwise_gemm.py
-78
View File
@@ -1,78 +0,0 @@
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
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
popen_launch_server,
try_cached_model,
)
MODEL_PATH = "Qwen/Qwen3-4B-Instruct-2507-FP8"
class FP8BlockwiseGemmBase:
backend = None
@classmethod
def setUpClass(cls):
if cls.backend is None:
raise NotImplementedError("Subclass must set 'backend' attribute")
cls.model = try_cached_model(MODEL_PATH)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--fp8-gemm-backend",
cls.backend,
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
parsed_url = urlparse(self.base_url)
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
max_new_tokens=512,
parallel=200,
host=f"{parsed_url.scheme}://{parsed_url.hostname}",
port=parsed_url.port,
)
metrics = run_eval_few_shot_gsm8k(args)
print(metrics)
self.assertGreaterEqual(metrics["accuracy"], 0.41)
class TestFP8BlockwiseGemmTriton(FP8BlockwiseGemmBase, unittest.TestCase):
backend = "triton"
class TestFP8BlockwiseGemmDeepGemm(FP8BlockwiseGemmBase, unittest.TestCase):
backend = "deep_gemm"
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFP8BlockwiseGemmFlashinferTrtllm(FP8BlockwiseGemmBase, unittest.TestCase):
backend = "flashinfer_trtllm"
@unittest.skipIf(get_device_sm() != 90, "Test requires CUDA SM 90")
class TestFP8BlockwiseGemmFlashinferDeepGemm(FP8BlockwiseGemmBase, unittest.TestCase):
backend = "flashinfer_deepgemm"
if __name__ == "__main__":
unittest.main()
test/srt/test_gpt_oss_4gpu.py
-36
View File
@@ -1,36 +0,0 @@
import unittest
from sglang.test.gpt_oss_common import BaseTestGptOss
class TestGptOss4Gpu(BaseTestGptOss):
def test_bf16_120b(self):
self.run_test(
model_variant="120b",
quantization="bf16",
expected_score_of_reasoning_effort={
"low": 0.60,
},
other_args=["--tp", "4", "--cuda-graph-max-bs", "200"],
)
def test_mxfp4_120b(self):
self.run_test(
model_variant="120b",
quantization="mxfp4",
expected_score_of_reasoning_effort={
"low": 0.60,
},
other_args=[
"--tp",
"4",
"--cuda-graph-max-bs",
"200",
"--mem-fraction-static",
"0.93",
],
)
if __name__ == "__main__":
unittest.main()
test/srt/test_mistral_large3_basic.py
-89
View File
@@ -1,89 +0,0 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_cuda_ci(est_time=600, suite="nightly-8-gpu-b200", nightly=True)
MISTRAL_LARGE3_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512"
class TestMistralLarge3Basic(CustomTestCase):
@classmethod
def setUpClass(cls):
# Set environment variable to disable JIT DeepGemm
os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"] = "0"
cls.model = MISTRAL_LARGE3_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--attention-backend",
"trtllm_mla",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
"--chat-template",
"mistral",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
# Clean up environment variable
if "SGLANG_ENABLE_JIT_DEEPGEMM" in os.environ:
del os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"]
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (mistral-large-3)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.90)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (mistral-large-3)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 50)
if __name__ == "__main__":
unittest.main()
test/srt/test_multi_instance_release_memory_occupation.py
-253
View File
@@ -1,253 +0,0 @@
import multiprocessing
import os
import time
import traceback
import unittest
from multiprocessing import Process
from typing import Iterable, Tuple
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from transformers import AutoModelForCausalLM
from sglang.srt.entrypoints.engine import Engine as SglangEngine
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
CustomTestCase,
find_available_port,
)
TEST_SUITE = dict(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
mem_fraction_static=0.83,
dp_size=2,
tp_size=2,
)
class EngineWrapper:
"""
A wrapper around Sglang engine to mock multi instance cases such as RL traing.
"""
def __init__(
self, model_path, random_seed, mem_fraction_static, device_mesh_cpu, base_gpu_id
):
self._device_mesh_cpu = device_mesh_cpu
self._tp_rank = device_mesh_cpu.get_local_rank()
self._rank = device_mesh_cpu.get_rank()
self._tp_size = device_mesh_cpu.size()
tp_size_per_node = self._tp_size
node_rank = self._tp_rank // tp_size_per_node
first_rank_in_node = self._tp_rank % tp_size_per_node == 0
engine_kwargs = dict(
model_path=model_path,
random_seed=random_seed,
mem_fraction_static=mem_fraction_static,
base_gpu_id=base_gpu_id,
enable_memory_saver=True,
tp_size=self._tp_size,
node_rank=node_rank,
nnodes=1,
)
self._engine = None
if first_rank_in_node:
os.environ["SGLANG_BLOCK_NONZERO_RANK_CHILDREN"] = "0"
self._engine = SglangEngine(**engine_kwargs)
dist.barrier(group=self._device_mesh_cpu.get_group())
def update_weights_from_tensor(
self, named_tensors: Iterable[Tuple[str, torch.Tensor]]
):
if self._tp_rank == 0:
self._engine.update_weights_from_tensor(list(named_tensors))
self._engine.flush_cache()
dist.barrier(group=self._device_mesh_cpu.get_group())
def release_memory_occupation(self, tags):
if self._tp_rank == 0:
self._engine.release_memory_occupation(tags)
dist.barrier(group=self._device_mesh_cpu.get_group())
def resume_memory_occupation(self, tags):
if self._tp_rank == 0:
self._engine.resume_memory_occupation(tags)
dist.barrier(group=self._device_mesh_cpu.get_group())
def shutdown(self):
if self._tp_rank == 0:
self._engine.shutdown()
dist.barrier(group=self._device_mesh_cpu.get_group())
def get_gpu_memory_gb(gpu_id=0):
return torch.cuda.device_memory_used() / 1024**3
class TestMultiInstanceReleaseMemoryOccupation(CustomTestCase):
@classmethod
def setUpClass(cls):
multiprocessing.set_start_method("spawn")
def test_multi_instance_release_memory_occupation(self):
master_port = find_available_port(23456)
dp_size = TEST_SUITE["dp_size"]
tp_size = TEST_SUITE["tp_size"]
world_size = dp_size * tp_size
processes = []
output_reader, output_writer = multiprocessing.Pipe(duplex=False)
for rank in range(world_size):
p = Process(
target=_run_sglang_subprocess,
kwargs=dict(
rank=rank,
dp_size=dp_size,
tp_size=tp_size,
model_path=TEST_SUITE["model_path"],
master_port=master_port,
output_writer=output_writer,
mem_fraction_static=TEST_SUITE["mem_fraction_static"],
),
)
p.start()
processes.append(p)
for _ in range(world_size):
self.assertTrue(
output_reader.recv(), f"Subprocess fail. Check the logs above."
)
for p in processes:
p.join()
def _run_sglang_subprocess(
rank: int,
dp_size: int,
tp_size: int,
model_path: str,
master_port: int,
output_writer,
mem_fraction_static: float,
):
engine = None
try:
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(master_port)
dist.init_process_group(
rank=rank,
device_id=torch.device(f"cuda:{rank}"),
world_size=dp_size * tp_size,
)
torch.cuda.set_device(rank)
base_gpu_id = rank // tp_size * tp_size
mesh_kwargs = dict(
mesh_shape=(dp_size, tp_size, 1), mesh_dim_names=["dp", "tp", "pp"]
)
inference_device_mesh_device = init_device_mesh("cuda", **mesh_kwargs)
inference_device_mesh_cpu = init_device_mesh("cpu", **mesh_kwargs)
print(
f"subprocess[{rank=},{base_gpu_id=},{rank=},{tp_size=}] {inference_device_mesh_device=} {inference_device_mesh_cpu=}"
)
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before starting Engine: {_mem_usage}")
engine = EngineWrapper(
model_path=model_path,
random_seed=42,
mem_fraction_static=mem_fraction_static,
device_mesh_cpu=inference_device_mesh_cpu["tp"],
base_gpu_id=base_gpu_id,
)
print(f"subprocess[{rank=}] {engine=}", flush=True)
# 1 - release kv cache
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before releasing Sgl KV cache: {_mem_usage}")
engine.release_memory_occupation(tags=["kv_cache"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing KV cache must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 2 - release sglang weights
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before releasing Sgl weights: {_mem_usage}")
engine.release_memory_occupation(tags=["weights"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing weights must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 3 - load hf model
_mem_usage = get_gpu_memory_gb(rank)
print(
f"GPU{rank} Memory usage after releasing Sgl weights and kv cache: {_mem_usage}"
)
hf_model = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map=f"cuda:{rank}",
trust_remote_code=True,
).cuda()
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage > _mem_usage
), f"Memory usage after loading hf model must be increased! before: {_mem_usage} vs after: {_curr_usage}"
# 4 - resume sglang weights and update the weights
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after loading hf model: {_mem_usage}")
engine.resume_memory_occupation(tags=["weights"])
engine.update_weights_from_tensor(
named_tensors=list(hf_model.named_parameters())
)
# 5 - release hf model
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after resuming Sgl weights: {_mem_usage}")
del hf_model
hf_model = None
torch.cuda.empty_cache()
time.sleep(3)
torch.cuda.empty_cache()
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing hf model must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 6 - resume slgang kv cache
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after releasing hf model: {_mem_usage}")
engine.resume_memory_occupation(tags=["kv_cache"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage > _mem_usage
), f"Memory usage after resuming kv cache must be increased! before: {_mem_usage} vs after: {_curr_usage}"
# 7 - Final checking!
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after resuming Sgl KV cache: {_mem_usage}")
execution_ok = True
except Exception as e:
print(f"subprocess[{rank=}] has error: {e}", flush=True)
traceback.print_exc()
execution_ok = False
output_writer.send(execution_ok)
output_writer.close()
if engine:
engine.shutdown()
if __name__ == "__main__":
unittest.main()
test/srt/test_nvfp4_gemm.py
-82
View File
@@ -1,82 +0,0 @@
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
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
popen_launch_server,
try_cached_model,
)
MODEL_PATH = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
class FP4GemmBase:
backend = None
@classmethod
def setUpClass(cls):
if cls.backend is None:
raise NotImplementedError("Subclass must set 'backend' attribute")
cls.model = try_cached_model(MODEL_PATH)
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--quantization",
"modelopt_fp4",
"--fp4-gemm-backend",
cls.backend,
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
parsed_url = urlparse(self.base_url)
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1319,
max_new_tokens=512,
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.64)
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFP4GemmAuto(FP4GemmBase, unittest.TestCase):
backend = "auto"
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFP4GemmFlashinferCutlass(FP4GemmBase, unittest.TestCase):
backend = "flashinfer_cutlass"
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFP4GemmFlashinferCudnn(FP4GemmBase, unittest.TestCase):
backend = "flashinfer_cudnn"
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
class TestFP4GemmFlashinferTrtllm(FP4GemmBase, unittest.TestCase):
backend = "flashinfer_trtllm"
if __name__ == "__main__":
unittest.main()
test/srt/test_pp_single_node.py
-424
View File
@@ -1,424 +0,0 @@
"""
Usage:
python3 -m unittest test_pp_single_node.TestPPAccuracy.test_gsm8k
python3 -m unittest test_pp_single_node.TestQwenPPAccuracy.test_pp_consistency
python3 -m unittest test_pp_single_node.TestFixedBugs.test_chunked_prefill_with_small_bs
python3 -m unittest test_pp_single_node.TestQwenVLPPAccuracy.test_mmmu
"""
import time
import unittest
from types import SimpleNamespace
import requests
from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP,
DEFAULT_MODEL_NAME_FOR_TEST_VL_PP,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
run_bench_one_batch_server,
)
class TestPPAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = "http://127.0.0.1:23333"
cls.process = popen_launch_server(
DEFAULT_MODEL_NAME_FOR_TEST,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
2,
"--pp-size",
2,
"--chunked-prefill-size",
256,
],
)
@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.assertGreater(metrics["accuracy"], 0.74)
# Wait a little bit so that the memory check happens.
time.sleep(4)
def test_logprob(self):
response = requests.post(
f"{self.base_url}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 16,
},
"return_logprob": True,
"top_logprobs_num": 5,
"logprob_start_len": 0,
},
)
response_json = response.json()
input_token_logprobs = response_json["meta_info"]["input_token_logprobs"]
output_token_logprobs = response_json["meta_info"]["output_token_logprobs"]
output_top_logprobs = response_json["meta_info"]["output_top_logprobs"]
assert len(input_token_logprobs) == 6
assert len(output_token_logprobs) == 16
assert len(output_top_logprobs) == 16
class TestDPAttentionDP2PP2(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"2",
"--pp-size",
"2",
"--enable-dp-attention",
"--dp",
"2",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_mgsm_en(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mgsm_en",
num_examples=None,
num_threads=1024,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.8)
class TestQwenVLPPAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_VL_PP
cls.base_url = "http://127.0.0.1:23333"
cls.process = popen_launch_server(
DEFAULT_MODEL_NAME_FOR_TEST_VL_PP,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
4,
"--chunked-prefill-size",
8192,
"--enable-multimodal",
],
)
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.assertGreater(metrics["accuracy"], 0.65)
# Wait a little bit so that the memory check happens.
time.sleep(4)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
def test_mmmu(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmmu",
num_examples=None,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.26)
class TestQwenPPAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = "http://127.0.0.1:23334" # different ports to avoid conflicts
cls.model_name = "Qwen/Qwen3-8B" # replace with your Qwen Model if needed
def run_gsm8k_test(self, pp_size):
process = popen_launch_server(
self.model_name,
self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--pp-size",
pp_size,
"--chunked-prefill-size",
256,
],
)
try:
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)
time.sleep(5)
return metrics
finally:
kill_process_tree(process.pid)
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
def test_pp_consistency(self):
baseline = self.run_gsm8k_test(pp_size=1)
pp_metrics = self.run_gsm8k_test(pp_size=2)
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
self.assertGreaterEqual(baseline["accuracy"], 0.74)
self.assertGreaterEqual(
pp_metrics["accuracy"],
baseline["accuracy"] - 0.02,
msg=(
f"PP accuracy dropped more than 1% compared to baseline. "
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
),
)
class TestQwenPPTieWeightsAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = "http://127.0.0.1:23335" # different ports to avoid conflicts
cls.model_name = (
"Qwen/Qwen3-0.6B" # qwen3 < 8B all have tie_word_embeddings = True
)
def run_gsm8k_test(self, pp_size):
process = popen_launch_server(
self.model_name,
self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--pp-size",
pp_size,
"--chunked-prefill-size",
256,
],
)
try:
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)
time.sleep(5)
return metrics
finally:
kill_process_tree(process.pid)
def test_pp_consistency(self):
baseline = self.run_gsm8k_test(pp_size=1)
pp_metrics = self.run_gsm8k_test(pp_size=2)
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
self.assertGreaterEqual(baseline["accuracy"], 0.38)
self.assertGreaterEqual(
pp_metrics["accuracy"],
baseline["accuracy"] - 0.02,
msg=(
f"PP accuracy dropped more than 1% compared to baseline. "
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
),
)
class TestQwenMoePPAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = "http://127.0.0.1:23336" # different ports to avoid conflicts
cls.model_name = "Qwen/Qwen3-30B-A3B" # replace with your Qwen Model if needed
def run_gsm8k_test(self, pp_size):
process = popen_launch_server(
self.model_name,
self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--pp-size",
pp_size,
"--chunked-prefill-size",
256,
],
)
try:
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)
time.sleep(5)
return metrics
finally:
kill_process_tree(process.pid)
def test_pp_consistency(self):
baseline = self.run_gsm8k_test(pp_size=1)
pp_metrics = self.run_gsm8k_test(pp_size=2)
print(f"[Qwen PP Comparison] Baseline: {baseline} | PP: {pp_metrics}")
self.assertGreaterEqual(baseline["accuracy"], 0.74)
self.assertGreaterEqual(
pp_metrics["accuracy"],
baseline["accuracy"] - 0.02,
msg=(
f"PP accuracy dropped more than 1% compared to baseline. "
f"Baseline: {baseline['accuracy']:.2%}, PP: {pp_metrics['accuracy']:.2%}"
),
)
class TestFixedBugs(unittest.TestCase):
def test_chunked_prefill_with_small_bs(self):
model = DEFAULT_MODEL_NAME_FOR_TEST
server_args = ServerArgs(model_path=model)
bench_args = OneBatchBenchArgs(
batch_size=(1,),
input_len=(1,),
output_len=(1,),
base_url=DEFAULT_URL_FOR_TEST,
)
other_server_args = [
"--tp-size",
2,
"--pp-size",
2,
"--chunked-prefill",
256,
"--max-running-requests",
2,
]
run_bench_one_batch_server(
model,
DEFAULT_URL_FOR_TEST,
server_args,
bench_args,
other_server_args,
)
@unittest.skipIf(
is_in_ci(), "Skipping GLM41V PP accuracy test before it gets more stable"
)
class TestGLM41VPPAccuracy(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
2,
"--chunked-prefill-size",
8192,
"--enable-multimodal",
"--reasoning-parser",
"glm45",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_mmmu(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmmu",
num_examples=None,
num_threads=32,
response_answer_regex="<\|begin_of_box\|>(.*)<\|end_of_box\|>",
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.45)
if __name__ == "__main__":
unittest.main()
test/srt/test_prefill_delayer.py
-559
View File
@@ -1,559 +0,0 @@
import asyncio
import os
import re
import time
import unittest
from collections import defaultdict
from dataclasses import dataclass
from types import SimpleNamespace
from typing import List, Optional
import openai
import requests
import torch
import torch.multiprocessing as mp
from sglang.bench_serving import run_benchmark
from sglang.srt.managers.prefill_delayer import PrefillDelayer
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
get_benchmark_args,
popen_launch_server,
)
WORLD_SIZE = os.environ.get("SGLANG_TEST_WORLD_SIZE", "8")
# ============================ Unit Tests ============================
@dataclass
class NegotiateCall:
prefillable: List[bool]
token_usage: List[float]
@dataclass
class NegotiateTestCase:
name: str
max_delay_passes: int
token_usage_low_watermark: Optional[float]
calls: List[NegotiateCall]
expected_allow: bool
expected_reason: str
def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
torch.distributed.init_process_group(
backend="gloo",
init_method=f"tcp://127.0.0.1:{port}",
world_size=world_size,
rank=rank,
)
cpu_group = torch.distributed.new_group(backend="gloo")
for case in test_cases:
delayer = PrefillDelayer(
dp_size=world_size,
attn_tp_size=1,
cpu_group=cpu_group,
server_args=SimpleNamespace(
enable_dp_attention=True,
disaggregation_mode="null",
disable_overlap_schedule=False,
),
max_delay_passes=case.max_delay_passes,
token_usage_low_watermark=case.token_usage_low_watermark,
)
for call in case.calls:
result = delayer._negotiate_should_allow_prefill(
local_prefillable=call.prefillable[rank],
token_usage=call.token_usage[rank],
)
results_queue.put((rank, case.name, result.output_allow, result.output_reason))
torch.distributed.destroy_process_group()
_NEGOTIATE_TEST_CASES = [
NegotiateTestCase(
name="all_prefillable",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
)
],
expected_allow=True,
expected_reason="no_wait",
),
NegotiateTestCase(
name="all_prefillable_with_previous_wait",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
),
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
),
],
expected_allow=True,
expected_reason="wait_success",
),
NegotiateTestCase(
name="none_prefillable",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[False, False, False, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
)
],
expected_allow=True,
expected_reason="",
),
NegotiateTestCase(
name="mixed_delay",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
)
],
expected_allow=False,
expected_reason="delay",
),
NegotiateTestCase(
name="mixed_watermark_force_allow",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.5, 0.9, 0.9, 0.9],
)
],
expected_allow=True,
expected_reason="token_watermark",
),
NegotiateTestCase(
name="mixed_watermark_disabled",
max_delay_passes=100,
token_usage_low_watermark=None,
calls=[
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.5, 0.9, 0.9, 0.9],
)
],
expected_allow=False,
expected_reason="delay",
),
NegotiateTestCase(
name="mixed_watermark_not_prefillable",
max_delay_passes=100,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[False, False, True, False],
token_usage=[0.5, 0.9, 0.9, 0.9],
)
],
expected_allow=False,
expected_reason="delay",
),
NegotiateTestCase(
name="mixed_timeout",
max_delay_passes=3,
token_usage_low_watermark=0.8,
calls=[
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
),
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
),
NegotiateCall(
prefillable=[True, False, True, False],
token_usage=[0.9, 0.9, 0.9, 0.9],
),
],
expected_allow=True,
expected_reason="wait_timeout",
),
]
class TestPrefillDelayerNegotiate(unittest.TestCase):
def test_negotiate(self):
world_size = 4
test_cases = _NEGOTIATE_TEST_CASES
ctx = mp.get_context("spawn")
results_queue = ctx.Queue()
port = 29500 + os.getpid() % 1000
processes = []
for rank in range(world_size):
p = ctx.Process(
target=_run_negotiate_test,
args=(rank, world_size, test_cases, results_queue, port),
)
p.start()
processes.append(p)
for p in processes:
p.join()
results = defaultdict(dict)
for _ in range(world_size * len(test_cases)):
rank, case_name, output_allow, output_reason = results_queue.get()
results[case_name][rank] = (output_allow, output_reason)
for case in test_cases:
for rank in range(world_size):
output_allow, output_reason = results[case.name][rank]
self.assertEqual(
(output_allow, output_reason),
(case.expected_allow, case.expected_reason),
f"Case {case.name} rank {rank}",
)
# ============================ E2E Tests ============================
class TestPrefillDelayerThroughputOnlineServing(CustomTestCase):
def test_throughput_comparison(self):
_run_throughput_comparison(
self,
test_name="online_serving",
other_launch_args=[
# Not really needed, only to test support non-FCFS algorithms
"--schedule-policy",
"lpm",
],
other_benchmark_args=dict(
num_prompts=500,
random_input_len=30000,
random_output_len=256,
request_rate=32,
),
min_improvement_pct=5,
)
class TestPrefillDelayerThroughputOfflineGen(CustomTestCase):
def test_throughput_comparison(self):
_run_throughput_comparison(
self,
test_name="offline_gen",
other_launch_args=["--max-total-tokens", "200000"],
other_benchmark_args=dict(
num_prompts=800,
random_input_len=30000,
random_output_len=500,
),
token_usage_low_watermark=0.8,
min_improvement_pct=20,
)
def _run_throughput_comparison(
test_case,
test_name: str,
other_launch_args,
other_benchmark_args,
min_improvement_pct: float,
token_usage_low_watermark: float = None,
):
common_kwargs = dict(
debug_name=test_name,
other_launch_args=other_launch_args,
other_benchmark_args=other_benchmark_args,
token_usage_low_watermark=token_usage_low_watermark,
)
res_enabled = _run_throughput_test(prefill_delayer=True, **common_kwargs)
res_disabled = _run_throughput_test(prefill_delayer=False, **common_kwargs)
_assert_throughput_improvement(
test_case,
test_name=test_name,
res_enabled=res_enabled,
res_disabled=res_disabled,
min_improvement_pct=min_improvement_pct,
)
def _run_throughput_test(
debug_name: str,
prefill_delayer: bool,
other_launch_args,
other_benchmark_args,
token_usage_low_watermark: float = None,
):
model = "Qwen/Qwen3-0.6B"
base_url = DEFAULT_URL_FOR_TEST
process = _launch_server(
prefill_delayer=prefill_delayer,
model=model,
base_url=base_url,
other_args=other_launch_args,
token_usage_low_watermark=token_usage_low_watermark,
)
try:
args = get_benchmark_args(
base_url=base_url,
dataset_name="random",
tokenizer=model,
**other_benchmark_args,
)
res = run_benchmark(args)
_print_prefill_delayer_metrics(base_url, expect_metrics=prefill_delayer)
finally:
kill_process_tree(process.pid)
print(f"=== {debug_name} ({prefill_delayer=}) ===")
res["total_throughput"] = res["input_throughput"] + res["output_throughput"]
print(f"Input throughput: {res['input_throughput']:.2f} token/s")
print(f"Output throughput: {res['output_throughput']:.2f} token/s")
print(f"Total throughput: {res['total_throughput']:.2f} token/s")
return res
def _assert_throughput_improvement(
test_case,
test_name: str,
res_enabled: dict,
res_disabled: dict,
min_improvement_pct: float,
):
test_case.assertEqual(
WORLD_SIZE,
"8",
f"This test requires 8 GPUs to properly measure throughput improvement, got {WORLD_SIZE}",
)
enabled = res_enabled["total_throughput"]
disabled = res_disabled["total_throughput"]
improvement_pct = (enabled - disabled) / disabled * 100
print(f"\n=== {test_name} Throughput Comparison ===")
print(
f"Total: enabled={enabled:.2f}, disabled={disabled:.2f}, improvement={improvement_pct:.2f}%"
)
test_case.assertGreaterEqual(
improvement_pct,
min_improvement_pct,
f"{test_name}: Throughput improvement ({improvement_pct:.2f}%) < {min_improvement_pct}%",
)
class TestPrefillDelayerTokenUsageLowWatermark(CustomTestCase):
def test_1_with_low_watermark(self):
# The kv cache size here is deliberately small, thus we use smaller token usage
self._run(token_usage_low_watermark=0.5)
def test_2_without_low_watermark(self):
self._run(token_usage_low_watermark=None)
def _run(self, token_usage_low_watermark):
model = "Qwen/Qwen3-0.6B"
base_url = DEFAULT_URL_FOR_TEST
world_size = int(WORLD_SIZE)
process = _launch_server(
model=model,
base_url=base_url,
prefill_delayer=True,
other_args=["--max-total-tokens", "50000"],
# e.g. gen throughput is 370 tok/s on H200.
# Will need a different threshold on B200
max_delay_passes=3000,
token_usage_low_watermark=token_usage_low_watermark,
)
async def run_test():
client = openai.AsyncClient(base_url=f"{base_url}/v1", api_key="EMPTY")
long_prompt = "Hello " * 5000
async def send_blocking_request():
return await client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": long_prompt}],
max_tokens=10000,
extra_body={"data_parallel_rank": 0},
)
async def send_normal_request(dp_rank, req_idx):
start = time.time()
await client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": "Say hi"}],
max_tokens=10,
extra_body={"data_parallel_rank": dp_rank},
)
elapsed = time.time() - start
return dp_rank, req_idx, elapsed
asyncio.create_task(send_blocking_request())
await asyncio.sleep(3)
num_reqs_per_rank = 10
results = await asyncio.gather(
*[
send_normal_request(dp_rank, req_idx)
for dp_rank in range(1, world_size)
for req_idx in range(num_reqs_per_rank)
]
)
enabled = token_usage_low_watermark is not None
thresh = 5
for dp_rank, req_idx, elapsed in results:
print(f"DP rank {dp_rank} req {req_idx} completed in {elapsed:.2f}s")
self.assertTrue(
(elapsed < thresh) if enabled else (elapsed > thresh),
f"DP rank {dp_rank} req {req_idx}: elapsed={elapsed:.2f}s, thresh={thresh}, enabled={enabled}. "
f"Maybe you need a different `max_delay_passes` when using hardware other than H200.",
)
try:
asyncio.run(run_test())
metrics_text = _print_prefill_delayer_metrics(base_url, expect_metrics=True)
if token_usage_low_watermark is not None:
total = _sum_prometheus_metric_values(metrics_text, "token_watermark")
self.assertGreater(total, 0, "Expected token_watermark > 0")
print(f"total token_watermark: {total}")
finally:
kill_process_tree(process.pid)
class TestPrefillDelayerAccuracy(CustomTestCase):
def test_1_mgsm_en_has_prefill_delayer(self):
self._run_accuracy_test(prefill_delayer=True)
def test_2_mgsm_en_no_prefill_delayer(self):
self._run_accuracy_test(prefill_delayer=False)
def _run_accuracy_test(self, prefill_delayer: bool):
model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
base_url = DEFAULT_URL_FOR_TEST
process = _launch_server(
prefill_delayer=prefill_delayer,
model=model,
base_url=base_url,
other_args=[
# Not really needed, only to test support non-FCFS algorithms
"--schedule-policy",
"lpm",
# Use this to ensure prefill delayer will be run
"--max-total-tokens",
"4096",
],
)
try:
args = SimpleNamespace(
base_url=base_url,
model=model,
eval_name="mgsm_en",
num_examples=None,
num_threads=1024,
)
metrics = run_eval(args)
print(f"=== mgsm_en ({prefill_delayer=}) ===")
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.87)
finally:
kill_process_tree(process.pid)
def _launch_server(
*,
model,
base_url,
prefill_delayer: bool,
other_args,
max_delay_passes: int = 100,
token_usage_low_watermark: float = None,
):
os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
return popen_launch_server(
model,
base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
WORLD_SIZE,
"--enable-dp-attention",
"--dp",
WORLD_SIZE,
"--chunked-prefill-size",
"131072",
"--mem-fraction-static",
"0.6",
"--enable-metrics",
*(["--enable-prefill-delayer"] if prefill_delayer else []),
"--prefill-delayer-max-delay-passes",
str(max_delay_passes),
*(
[
"--prefill-delayer-token-usage-low-watermark",
str(token_usage_low_watermark),
]
if token_usage_low_watermark is not None
else []
),
*(other_args or []),
],
)
def _print_prefill_delayer_metrics(base_url: str, expect_metrics: bool) -> str:
metrics_response = requests.get(f"{base_url}/metrics")
assert metrics_response.status_code == 200
metrics_text = metrics_response.text
prefill_delayer_metrics = [
line for line in metrics_text.split("\n") if "prefill_delayer" in line
]
print("=== PrefillDelayer Metrics ===")
for line in prefill_delayer_metrics:
print(line)
if expect_metrics:
assert "sglang:prefill_delayer_wait_forward_passes" in metrics_text
assert "sglang:prefill_delayer_wait_seconds" in metrics_text
assert "sglang:prefill_delayer_outcomes_total" in metrics_text
return metrics_text
def _sum_prometheus_metric_values(metrics_text: str, label_value: str) -> int:
matches = re.findall(rf'{label_value}".*?\}} (\d+)', metrics_text)
return sum(int(m) for m in matches)
if __name__ == "__main__":
unittest.main()
test/srt/test_profile_v2.py
-102
View File
@@ -1,102 +0,0 @@
import os
import shutil
import tempfile
import unittest
from pathlib import Path
import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestStartProfile(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.output_dir = tempfile.mkdtemp()
envs.SGLANG_TORCH_PROFILER_DIR.set(cls.output_dir)
envs.SGLANG_PROFILE_V2.set(True)
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def setUp(self):
self._clear_profile_dir()
def test_profile_by_stage(self):
self._start_profile(
profile_by_stage=True,
num_steps=10,
)
self._post_request()
self._check_profile_output(pattern="*-prefill*", expect_existence=True)
self._check_profile_output(pattern="*-decode*", expect_existence=True)
def test_decode_only(self):
self._start_profile(
profile_by_stage=True,
profile_stages=["decode"],
num_steps=10,
)
self._post_request()
self._check_profile_output(pattern="*-prefill*", expect_existence=False) # NOTE
self._check_profile_output(pattern="*-decode*", expect_existence=True)
def _start_profile(self, **kwargs):
"""Start profiling with optional parameters."""
response = requests.post(
f"{DEFAULT_URL_FOR_TEST}/start_profile",
json=kwargs if kwargs else None,
)
self.assertEqual(response.status_code, 200)
def _post_request(self):
response = requests.post(
f"{DEFAULT_URL_FOR_TEST}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
},
},
)
self.assertEqual(response.status_code, 200)
def _clear_profile_dir(self):
if os.path.isdir(self.output_dir):
shutil.rmtree(self.output_dir)
def _check_profile_output(self, pattern: str, expect_existence: bool):
self.assertTrue(
os.path.isdir(self.output_dir), "Output directory does not exist."
)
self.assertEqual(
len(list(Path(self.output_dir).glob(pattern))) > 0,
expect_existence,
f"Does not find {pattern=} ({list(Path(self.output_dir).glob('**/*'))=})",
)
if __name__ == "__main__":
unittest.main()
test/srt/test_release_memory_occupation.py
-470
View File
@@ -1,470 +0,0 @@
"""Test memory release and resume operations for SGLang engine in hybrid RL training.
This test suite evaluates the SGLang engine's memory management capabilities, focusing
on releasing and resuming memory occupation for KV cache and model weights. It simulates
an RL workflow where the SGLang engine acts as a rollout engine for experience collection.
The process involves initializing the engine, sending a small number of requests to simulate
rollout, releasing memory to mimic offloading during RL training, resuming memory occupation,
updating weights with a trained HuggingFace model, and verifying the updated weights.
Detailed in our proposal (https://github.com/sgl-project/sglang/pull/7099), two test cases
are included:
1. Basic Release and Resume: Uses a lower mem_fraction_static (0.6) to control memory allocation
and avoid OOM errors carefully. This test simulates a scenario without multi-stage memory management,
ensuring the engine can release and resume memory occupation while maintaining functionality after
weight updates.
2. Multi-Stage Release and Resume: Employs a higher mem_fraction_static (0.85) to simulate higher
memory pressure, leveraging multi-stage memory management. It sequentially releases and resumes
KV cache and model weights, verifying memory deallocation and reallocation at each stage, and
ensuring correct weight updates and text generation.
3. Tensor Parallel Tests: Tests memory release and resume operations with different tensor parallel
configurations (tp=1, tp=2) to ensure proper memory management in distributed settings. For different
data parallel size, we test it in verl.
"""
import os
import time
import unittest
import torch
from transformers import AutoModelForCausalLM
import sglang as sgl
from sglang.srt.constants import (
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.test.test_utils import (
DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
CustomTestCase,
)
# (temporarily) set to true to observe memory usage in nvidia-smi more clearly
_DEBUG_EXTRA = False
def get_gpu_memory_gb():
return torch.cuda.device_memory_used() / 1024**3
class TestReleaseMemoryOccupation(CustomTestCase):
def _setup_engine(
self,
model_name,
mem_fraction_static=0.8,
tp_size=1,
ep_size=1,
enable_weights_cpu_backup=False,
):
"""Common setup for engine and HF model."""
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=mem_fraction_static,
tp_size=tp_size,
ep_size=ep_size,
enable_weights_cpu_backup=enable_weights_cpu_backup,
# disable_cuda_graph=True, # for debugging only
)
return engine
def _common_test_params(self):
"""Common test parameters."""
return {
"prompt": "Today is a sunny day and I like",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"expect_output_before_update_weights": " to spend it outdoors. I decided to",
"expect_output_after_update_weights": " to go for a walk. I like",
"prompt_moe": "The weather is nice today, and I want to",
"sampling_params_moe": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a",
"expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I",
"prompt_hybrid_mamba": "The weather is nice today, and I want to",
"sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
"expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
}
def _test_initial_generation(
self, engine, prompt, sampling_params, expect_output_before_update_weights
):
"""Test initial generation and memory allocation."""
print("generate (#1)")
outputs = engine.generate(prompt, sampling_params)["text"]
self.assertEqual(outputs, expect_output_before_update_weights)
if _DEBUG_EXTRA:
time.sleep(3)
def test_release_and_resume_occupation(self):
# Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
assert (
torch.cuda.device_count() >= 2
), "Need at least 2 GPUs for tensor parallel tests"
for tp_size in [1, 2]:
print(f"Testing tp_size={tp_size} for test_release_and_resume_occupation")
engine = self._setup_engine(
model_name=model_name, mem_fraction_static=0.6, tp_size=tp_size
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_release_and_resume_occupation_with_weights_cpu_backup(self):
# Test release and resume occupation with weights CPU backup
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
print("Testing test_release_and_resume_occupation_with_weights_cpu_backup")
engine = self._setup_engine(
model_name=model_name,
mem_fraction_static=0.6,
enable_weights_cpu_backup=True,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
print("generate post resume")
outputs = engine.generate(params["prompt"], params["sampling_params"])["text"]
self.assertEqual(outputs, params["expect_output_before_update_weights"])
engine.shutdown()
def test_multi_stage_release_and_resume(self):
# With multi-stage release and resume, we can set the memory fraction to 0.85 without concern of OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
for tp_size in [1, 2]:
if tp_size == 2 and torch.cuda.device_count() < 2:
continue
print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume")
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.85, # Higher memory pressure
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_release_kv_cache = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_kv_cache,
gpu_memory_usage_before_release,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_release_weights = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_after_release_kv_cache,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_release_cuda_graph = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_cuda_graph,
gpu_memory_usage_after_release_weights,
)
print(f"Release took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_release:.1f} → {gpu_memory_usage_after_release_kv_cache:.1f} → {gpu_memory_usage_after_release_weights:.1f} → {gpu_memory_usage_after_release_cuda_graph:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
gpu_memory_usage_before_resume = get_gpu_memory_gb()
# gpu_memory_usage_after_release_weights and gpu_memory_usage_before_resume should be close
self.assertAlmostEqual(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_before_resume,
delta=3.0,
)
print(f"Resume weights took {time.perf_counter() - t:.2f}s")
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_resume_cuda_graph = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_cuda_graph,
gpu_memory_usage_before_resume,
)
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_resume_weights = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_weights,
gpu_memory_usage_after_resume_cuda_graph,
)
# Update weights from a trained model to serving engine, and then destroy the trained model
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb()
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_kv_cache,
gpu_memory_usage_after_resume_weights,
)
print(f"Resume + update took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_resume:.1f} → {gpu_memory_usage_after_resume_cuda_graph:.1f} → {gpu_memory_usage_after_resume_weights:.1f} → {gpu_memory_usage_after_loaded_hf_model:.1f} → {gpu_memory_usage_after_resume_kv_cache:.1f} GB"
)
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_moe_model_release_and_resume(self):
# Test with MoE model
model_name = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
tp_size = ep_size = 2
print(
f"Testing tp_size={tp_size} and ep_size={ep_size} for test_moe_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.5,
tp_size=tp_size,
ep_size=ep_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_moe"],
params["sampling_params_moe"],
params["expect_output_before_update_weights_moe"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights_moe"])
engine.shutdown()
def test_hybrid_mamba_model_release_and_resume(self):
# Test with Hybrid Mamba model
model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
tp_size = 4
print(
f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_hybrid_mamba"],
params["sampling_params_hybrid_mamba"],
params["expect_output_before_update_weights_hybrid_mamba"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
engine.update_weights_from_disk(model_name)
# destroy the hf model
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(
params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"]
)["text"]
self.assertEqual(
outputs, params["expect_output_after_update_weights_hybrid_mamba"]
)
engine.shutdown()
if __name__ == "__main__":
unittest.main()