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

[AMD CI] Migrate and Add More Testcases (#17116)

Browse Source 
Co-authored-by: yctseng0211 <yctseng@amd.com>
This commit is contained in:
Bingxu Chen
2026-01-19 08:07:39 -08:00
committed by GitHub
co-authored by yctseng0211
parent 20b0523eca
commit 2ea02f0642
19 changed files with 311 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py
+1 -1
View File
@@ -19,7 +19,7 @@ from sglang.test.test_utils import (
register_amd_ci(est_time=3600, suite="stage-c-test-large-8-gpu-amd-mi35x")
DEEPSEEK_R1_MODEL_PATH = "amd/DeepSeek-R1-MXFP4-Preview"
SERVER_LAUNCH_TIMEOUT = 1200
SERVER_LAUNCH_TIMEOUT = 1800
class TestDeepseekR1MXFP4(CustomTestCase):
test/registered/amd/test_deepseek_v3_basic.py
+84
View File
@@ -0,0 +1,84 @@
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_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_amd_ci,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_amd_ci(est_time=952, suite="stage-c-test-large-8-gpu-amd")
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/registered/amd/test_deepseek_v3_mtp.py
+116
View File
@@ -0,0 +1,116 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_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_amd_ci,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_amd_ci(est_time=980, suite="stage-c-test-large-8-gpu-amd")
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)
if is_in_amd_ci():
self.assertGreater(avg_spec_accept_length, 2.8)
else:
self.assertGreater(avg_spec_accept_length, 2.9)
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"
)
if is_in_amd_ci():
self.assertGreater(acc_length, 2.8)
else:
self.assertGreater(acc_length, 2.9)
if is_in_amd_ci():
self.assertGreater(speed, 15)
else:
self.assertGreater(speed, 130)
if __name__ == "__main__":
unittest.main()
test/registered/attention/test_wave_attention_kernels.py
+1 -1
View File
@@ -24,7 +24,7 @@ from sglang.srt.layers.attention.wave_ops.prefill_attention import (
from sglang.test.ci.ci_register import register_amd_ci
# Wave attention kernel unit tests (AMD only - requires wave_lang)
register_amd_ci(est_time=60, suite="stage-a-test-1")
register_amd_ci(est_time=60, suite="stage-a-test-1-amd")
class TestWaveAttention(unittest.TestCase):
test/registered/core/test_deterministic.py
+5 -1
View File
@@ -9,15 +9,18 @@ test into unit tests so that's easily reproducible in CI.
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_deterministic_utils import (
COMMON_SERVER_ARGS,
TestDeterministicBase,
)
from sglang.test.test_utils import is_in_amd_ci
register_cuda_ci(est_time=278, suite="stage-b-test-large-1-gpu")
register_amd_ci(est_time=278, suite="stage-b-test-small-1-gpu-amd")
@unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.")
class TestFlashinferDeterministic(TestDeterministicBase):
# Test with flashinfer attention backend
@classmethod
@@ -32,6 +35,7 @@ class TestFlashinferDeterministic(TestDeterministicBase):
return args
@unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.")
class TestFa3Deterministic(TestDeterministicBase):
# Test with fa3 attention backend
@classmethod
test/registered/hicache/test_hicache_storage_3fs_backend.py
+2 -1
View File
@@ -10,10 +10,11 @@ import unittest
from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=200, suite="stage-b-test-large-2-gpu")
register_amd_ci(est_time=300, suite="stage-b-test-large-2-gpu")
class HiCacheStorage3FSBackendBaseMixin(HiCacheStorageBaseMixin):
test/registered/hicache/test_hicache_storage_file_backend.py
+2 -1
View File
@@ -18,7 +18,7 @@ import requests
from sglang.bench_serving import get_tokenizer
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
@@ -31,6 +31,7 @@ from sglang.test.test_utils import (
)
register_cuda_ci(est_time=200, suite="stage-b-test-large-2-gpu")
register_amd_ci(est_time=526, suite="stage-b-test-large-2-gpu-amd")
class HiCacheStorageBaseMixin:
test/registered/openai_server/basic/test_serving_rerank.py
+1 -1
View File
@@ -7,7 +7,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
# Keep consistent with other openai_server/basic unit tests.
register_cuda_ci(est_time=10, suite="stage-b-test-large-1-gpu")
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu")
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
try:
from sglang.srt.entrypoints.openai.serving_rerank import (
test/registered/quant/test_awq_dequant.py
+1 -1
View File
@@ -19,7 +19,7 @@ from sglang.srt.layers.quantization.awq_triton import (
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=2, suite="stage-a-test-1")
register_amd_ci(est_time=2, suite="stage-a-test-1-amd")
device = "cuda"
test/registered/quant/test_fused_rms_fp8_group_quant.py
+1 -1
View File
@@ -7,7 +7,7 @@ import torch.nn.functional as F
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=10, suite="stage-a-test-1")
register_amd_ci(est_time=10, suite="stage-a-test-1-amd")
def _fp8_available() -> bool:
test/registered/quant/test_int4fp8_moe.py
+58
View File
@@ -0,0 +1,58 @@
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_amd_ci(est_time=313, suite="stage-b-test-small-1-gpu-amd")
class TestMixtralAccuracy(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"2",
"--mem-fraction-static",
"0.9",
"--context-length",
"38768",
"--quantization",
"quark_int4fp8_moe",
# The default aiter attention backend raises segmentation faults and other errors - as quark_int4fp8_moe is not related to attention, let's just use triton here.
"--attention-backend",
"triton",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=45 * 60,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1400,
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.assertGreater(metrics["accuracy"], 0.56)
test/registered/rl/test_lora_load_from_tensor.py
+1 -1
View File
@@ -11,7 +11,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=90, suite="stage-b-test-small-1-gpu")
register_amd_ci(est_time=90, suite="stage-b-test-small-1-gpu")
register_amd_ci(est_time=90, suite="stage-b-test-small-1-gpu-amd")
MODEL_PATH = "Qwen/Qwen3-0.6B"
LORA_REPO = "charent/self_cognition_Alice"
test/registered/rotary/test_rope_rocm.py
+123
View File
@@ -0,0 +1,123 @@
import unittest
import torch
from sglang.srt.layers.rotary_embedding import RotaryEmbedding
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=3, suite="stage-b-test-small-1-gpu-amd")
torch.manual_seed(0)
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_CASES = [
(64, 64, 32, 8000, True, torch.bfloat16, "cuda", 32, 32, 1, 1),
(256, 128, 4096, 10000, True, torch.bfloat16, "cuda", 2, 512, 4, 2),
(512, 128, 311, 10000, True, torch.bfloat16, "cuda", 3, 39, 4, 2),
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8),
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4),
(512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2),
]
@unittest.skipIf(_use_aiter, reason="SGLANG_USE_AITER=1 will not use vllm path.")
class TestRotaryEmbeddingNative(CustomTestCase):
# Compare RotaryEmbedding.forward_hip() to forward_native().
def _run_case(
self,
head_size: int,
rotary_dim: int,
max_pos: int,
base: int,
is_neox: bool,
dtype: torch.dtype,
device: str,
batch_size: int,
seq_len: int,
num_q: int,
num_kv: int,
) -> None:
rope_ref = RotaryEmbedding(
head_size, rotary_dim, max_pos, base, is_neox, dtype
).to(device)
rope_hip = RotaryEmbedding(
head_size, rotary_dim, max_pos, base, is_neox, dtype
).to(device)
pos_ids = torch.arange(seq_len, device=device).repeat(batch_size)
query = torch.randn(
batch_size * seq_len, num_q * head_size, dtype=dtype, device=device
)
key = torch.randn(
batch_size * seq_len, num_kv * head_size, dtype=dtype, device=device
)
q_ref, k_ref = rope_ref.forward_native(pos_ids, query.clone(), key.clone())
q_hip, k_hip = rope_hip.forward_hip(pos_ids, query.clone(), key.clone())
torch.testing.assert_close(q_ref, q_hip, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(k_ref, k_hip, atol=1e-2, rtol=1e-2)
def test_all_cases(self) -> None:
"""Drive over the full parameter matrix using subTest()."""
for case in _CASES:
with self.subTest(case=case):
self._run_case(*case)
@unittest.skipIf(not _use_aiter, reason="Requires AMD GPU plus SGLANG_USE_AITER=1")
class TestRotaryEmbeddingAITer(CustomTestCase):
# NOTE: Slightly relaxed tolerance (2e-2 vs 1e-2) for AITER RoPE kernel.
# Minor precision differences under investigation.
# See: https://github.com/sgl-project/sglang/pull/15318
@staticmethod
def _run_case_aiter(
head_size: int,
rotary_dim: int,
max_pos: int,
base: int,
is_neox: bool,
dtype: torch.dtype,
device: str,
batch_size: int,
seq_len: int,
num_q: int,
num_kv: int,
) -> None:
from aiter.rotary_embedding import RotaryEmbedding as AiterRotaryEmbedding
rope_ref = AiterRotaryEmbedding(
head_size, rotary_dim, max_pos, base, is_neox, dtype
).to(device)
rope_hip = AiterRotaryEmbedding(
head_size, rotary_dim, max_pos, base, is_neox, dtype
).to(device)
pos_ids = torch.arange(seq_len, device=device).repeat(batch_size)
query = torch.randn(
batch_size * seq_len, num_q * head_size, dtype=dtype, device=device
)
key = torch.randn(
batch_size * seq_len, num_kv * head_size, dtype=dtype, device=device
)
q_ref, k_ref = rope_ref.forward_native(pos_ids, query.clone(), key.clone())
q_hip, k_hip = rope_hip.forward_hip(pos_ids, query.clone(), key.clone())
torch.testing.assert_close(q_ref, q_hip, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(k_ref, k_hip, atol=2e-2, rtol=2e-2)
def test_all_cases(self) -> None:
for case in _CASES:
with self.subTest(case=case):
self._run_case_aiter(*case)
if __name__ == "__main__":
unittest.main()
test/registered/test_srt_backend.py
+1 -1
View File
@@ -20,7 +20,7 @@ from sglang.test.test_programs import (
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
register_cuda_ci(est_time=80, suite="stage-a-test-1")
register_amd_ci(est_time=120, suite="stage-a-test-1")
register_amd_ci(est_time=120, suite="stage-a-test-1-amd")
class TestSRTBackend(CustomTestCase):
test/registered/utils/test_bench_typebaseddispatcher.py
+264
View File
@@ -0,0 +1,264 @@
import timeit
from typing import Any, Callable, List, Tuple, Type
from sglang.test.ci.ci_register import register_amd_ci
from sglang.utils import TypeBasedDispatcher
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
class TypeBasedDispatcherList:
def __init__(self, mapping: List[Tuple[Type, Callable]]):
self._mapping = mapping
self._fallback_fn = None
def add_fallback_fn(self, fallback_fn: Callable):
self._fallback_fn = fallback_fn
def __iadd__(self, other: "TypeBasedDispatcher"):
self._mapping.extend(other._mapping)
return self
def __call__(self, obj: Any):
for ty, fn in self._mapping:
if isinstance(obj, ty):
return fn(obj)
if self._fallback_fn is not None:
return self._fallback_fn(obj)
raise ValueError(f"Invalid object: {obj}")
def create_test_mapping(num_types=30):
types = [type(f"RequestType{i}", (), {}) for i in range(num_types)]
def create_handler(i):
def handler(req):
return f"handler{i}"
return handler
handlers = [create_handler(i) for i in range(num_types)]
return list(zip(types, handlers))
def test_inheritance():
print("\n" + "=" * 60)
print("test for inheritance")
print("=" * 60)
class BaseRequest:
pass
def base_handler(req):
return "base_handler"
class DerivedRequest(BaseRequest):
pass
mapping = [(BaseRequest, base_handler)]
dict_dispatcher = TypeBasedDispatcher(mapping)
derived_obj = DerivedRequest()
expected = "base_handler"
# This test will fail with the current implementation, but pass with the suggested MRO-based fix
result_dict = dict_dispatcher(derived_obj)
assert result_dict == expected, f"Expected '{expected}', but got '{result_dict}'"
print("Pass: dict dispatcher handles inheritance.")
def benchmark_with_inheritance():
"""Performance test with inheritance scenarios"""
print("\nBenchmarking with inheritance scenarios...")
# Create type hierarchy with inheritance relationships
class BaseType:
pass
class ChildType1(BaseType):
pass
class ChildType2(BaseType):
pass
class GrandChildType(ChildType1):
pass
class UnrelatedType:
pass
def base_handler(obj):
return "handled"
mapping = [(BaseType, base_handler)]
dispatcher = TypeBasedDispatcher(mapping)
test_cases = [
BaseType(),
ChildType1(),
ChildType2(),
GrandChildType(),
UnrelatedType(),
]
# Test first call (includes MRO lookup)
first_call_times = []
for case in test_cases:
if not isinstance(case, UnrelatedType):
time_taken = timeit.timeit(lambda: dispatcher(case), number=1000)
first_call_times.append(time_taken)
# Test subsequent calls (using cache)
cached_call_times = []
for case in test_cases:
if not isinstance(case, UnrelatedType):
time_taken = timeit.timeit(lambda: dispatcher(case), number=1000)
cached_call_times.append(time_taken)
print(
f"First call (with MRO lookup): {sum(first_call_times)/len(first_call_times):.6f}s avg"
)
print(f"Cached call: {sum(cached_call_times)/len(cached_call_times):.6f}s avg")
print(f"Caching improvement: {sum(first_call_times)/sum(cached_call_times):.2f}x")
def benchmark_dispatchers():
mapping = create_test_mapping(30)
list_dispatcher = TypeBasedDispatcherList(mapping)
dist_dispatcher = TypeBasedDispatcher(mapping)
test_cases = []
for _, (ty, _) in enumerate(mapping):
test_cases.append(ty())
test_scenarios = [
("the first", [test_cases[0]] * 1000),
("the middle", [test_cases[len(test_cases) // 2]] * 1000),
("the last", [test_cases[-1]] * 1000),
("the random", test_cases * 1000),
]
print("=" * 60)
print("TypeBasedDispatcher benchmark test")
print("=" * 60)
for scenario_name, cases in test_scenarios:
print(f"\ntest scenario: {scenario_name}")
print(f"\ntest numbers: {len(cases)}")
list_time = timeit.timeit(
lambda: [list_dispatcher(case) for case in cases], number=10
)
dict_time = timeit.timeit(
lambda: [dist_dispatcher(case) for case in cases], number=10
)
print(f"for list: {list_time:.4f} s")
print(f"for dict: {dict_time:.4f} s")
print(f"improvement: {list_time/dict_time:.2f} x")
print(f"time reduce: {(1-dict_time/list_time) * 100:.1f} %")
def test_memory_usage():
import sys
mapping = create_test_mapping(30)
list_dispatcher = TypeBasedDispatcherList(mapping)
dict_dispatcher = TypeBasedDispatcher(mapping)
print("\n" + "=" * 60)
print("compare memory used:")
print("=" * 60)
list_size = sys.getsizeof(list_dispatcher._mapping)
dict_size = sys.getsizeof(dict_dispatcher._mapping)
print(f"memory used by list version: {list_size} bytes")
print(f"memory used by dict version: {dict_size} bytes")
print(f"compare memory used by the two version: {dict_size - list_size} bytes")
def test_edge_case():
"""test for edge case"""
print("\n" + "=" * 60)
print("test for edge case")
print("=" * 60)
mapping = create_test_mapping(30)
list_dispatcher = TypeBasedDispatcherList(mapping)
dict_dispatcher = TypeBasedDispatcher(mapping)
test_obj = mapping[0][0]()
result1 = list_dispatcher(test_obj)
result2 = dict_dispatcher(test_obj)
assert result1 == result2
print("Pass for normal test")
class UnkownType:
pass
try:
list_dispatcher(UnkownType())
print("exception was thrown from list version as expected")
except ValueError:
print("exception thrown from list version was processed...")
try:
dict_dispatcher(UnkownType())
print("exception was thrown from dict version as expected")
except ValueError:
print("exception thrown from dict version was processed...")
def simulate_real_workload():
"""simulate real workload"""
print("\n" + "=" * 60)
print("simulate real workload")
print("=" * 60)
mapping = create_test_mapping(30)
request_distribution = {
0: 0.2,
5: 0.3,
10: 0.1,
15: 0.15,
}
list_dispatcher = TypeBasedDispatcherList(mapping)
dict_dispatcher = TypeBasedDispatcher(mapping)
test_requests = []
for idx, prob in request_distribution.items():
count = int(1000 * prob)
test_requests.extend([mapping[idx][0]()] * count)
remaining = 1000 - len(test_requests)
for i in range(remaining):
test_requests.append(mapping[i % len(mapping)][0]())
list_time = timeit.timeit(
lambda: [list_dispatcher(req) for req in test_requests], number=100
)
dict_time = timeit.timeit(
lambda: [dict_dispatcher(req) for req in test_requests], number=100
)
print(f"list version: {list_time:.4f} s")
print(f"dict version: {dict_time:.4f} s")
print(f"improvement: {list_time/dict_time:.2f} x")
if __name__ == "__main__":
benchmark_dispatchers()
test_memory_usage()
test_edge_case()
simulate_real_workload()
test_inheritance()
benchmark_with_inheritance()
test/registered/utils/test_type_based_dispatcher.py
+225
View File
@@ -0,0 +1,225 @@
# tests/benchmarks/test_type_dispatcher_e2e.py
"""
E2E test for TypeBasedDispatcher optimization.
Tests real-world scenarios with actual request types.
"""
import timeit
import unittest
from sglang.srt.managers.io_struct import SamplingParams
from sglang.test.ci.ci_register import register_amd_ci
from sglang.utils import TypeBasedDispatcher
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
class TestTypeBasedDispatcher(unittest.TestCase):
"""Unit tests for TypeBasedDispatcher e2e performance."""
def test_type_dispatcher_e2e_performance(self):
"""End-to-end performance test with real request types"""
print("E2E Performance Test for TypeBasedDispatcher")
print("=" * 50)
from sglang.srt.managers.io_struct import (
AbortReq,
BatchTokenizedEmbeddingReqInput,
BatchTokenizedGenerateReqInput,
ClearHiCacheReqInput,
CloseSessionReqInput,
DestroyWeightsUpdateGroupReqInput,
ExpertDistributionReq,
FlushCacheReqInput,
FreezeGCReq,
GetInternalStateReq,
GetLoadReqInput,
GetWeightsByNameReqInput,
InitWeightsSendGroupForRemoteInstanceReqInput,
InitWeightsUpdateGroupReqInput,
LoadLoRAAdapterReqInput,
OpenSessionReqInput,
ProfileReq,
ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput,
RpcReqInput,
SendWeightsToRemoteInstanceReqInput,
SetInternalStateReq,
SlowDownReqInput,
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
UnloadLoRAAdapterReqInput,
UpdateWeightFromDiskReqInput,
UpdateWeightsFromIPCReqInput,
UpdateWeightsFromTensorReqInput,
)
mapping = [
(TokenizedGenerateReqInput, lambda req: "generate_handled"),
(TokenizedEmbeddingReqInput, lambda req: "embedding_handled"),
(BatchTokenizedGenerateReqInput, lambda req: "batch_generate_handled"),
(
BatchTokenizedEmbeddingReqInput,
lambda req: "batch_generate_embedding_handled",
),
(FlushCacheReqInput, lambda req: "flush_cache_handled"),
(ClearHiCacheReqInput, lambda req: "clear_hicache_handled"),
(AbortReq, lambda req: "abort_handled"),
(OpenSessionReqInput, lambda req: "open_session_handled"),
(CloseSessionReqInput, lambda req: "close_session_handled"),
(
UpdateWeightFromDiskReqInput,
lambda req: "update_weights_from_disk_handled",
),
(
InitWeightsUpdateGroupReqInput,
lambda req: "init_weights_update_group_handled",
),
(
DestroyWeightsUpdateGroupReqInput,
lambda req: "destroy_weights_update_group_handled",
),
(
InitWeightsSendGroupForRemoteInstanceReqInput,
lambda req: "init_weights_send_group_for_remote_instance_handled",
),
(
SendWeightsToRemoteInstanceReqInput,
lambda req: "send_weights_to_remote_instance_handled",
),
(
UpdateWeightsFromTensorReqInput,
lambda req: "update_weights_from_tensor_handled",
),
(
UpdateWeightsFromIPCReqInput,
lambda req: "update_weights_from_ipc_handled",
),
(GetWeightsByNameReqInput, lambda req: "get_weights_by_name_handled"),
(
ReleaseMemoryOccupationReqInput,
lambda req: "release_memory_occupation_handled",
),
(
ResumeMemoryOccupationReqInput,
lambda req: "resume_memory_occupation_handled",
),
(SlowDownReqInput, lambda req: "slow_down_handled"),
(ProfileReq, lambda req: "profile_handled"),
(FreezeGCReq, lambda req: "freeze_gc_handled"),
(GetInternalStateReq, lambda req: "get_internal_state_handled"),
(SetInternalStateReq, lambda req: "set_internal_state_handled"),
(RpcReqInput, lambda req: "rpc_request_handled"),
(ExpertDistributionReq, lambda req: "expert_distribution_handled"),
(LoadLoRAAdapterReqInput, lambda req: "load_lora_adapter_handled"),
(UnloadLoRAAdapterReqInput, lambda req: "unload_lora_adapter_handled"),
(GetLoadReqInput, lambda req: "get_load_handled"),
]
# Create requests that conforms to the real distribution
test_requests = []
test_requests.append(
TokenizedGenerateReqInput(
input_text="",
input_ids=[1, 2],
mm_inputs=dict(),
sampling_params=SamplingParams(),
return_logprob=False,
logprob_start_len=0,
top_logprobs_num=0,
token_ids_logprob=[1, 2],
stream=False,
)
)
test_requests.append(
TokenizedEmbeddingReqInput(
input_text="",
input_ids=[1, 2],
image_inputs=dict(),
token_type_ids=[1, 2],
sampling_params=SamplingParams(),
)
)
test_requests.append(
BatchTokenizedGenerateReqInput(
batch=[
TokenizedGenerateReqInput(
input_text="",
input_ids=[1, 2],
mm_inputs=dict(),
sampling_params=SamplingParams(),
return_logprob=False,
logprob_start_len=0,
top_logprobs_num=0,
token_ids_logprob=[1, 2],
stream=False,
)
]
)
)
test_requests.append(
BatchTokenizedEmbeddingReqInput(
batch=[
TokenizedEmbeddingReqInput(
input_text="",
input_ids=[1, 2],
image_inputs=dict(),
token_type_ids=[1, 2],
sampling_params=SamplingParams(),
)
]
)
)
test_requests.append(FlushCacheReqInput())
test_requests.append(ClearHiCacheReqInput())
test_requests.append(AbortReq())
test_requests.append(OpenSessionReqInput(capacity_of_str_len=0))
test_requests.append(CloseSessionReqInput(session_id=""))
test_requests.append(UpdateWeightFromDiskReqInput(model_path=""))
test_requests.append(
InitWeightsUpdateGroupReqInput(
master_address="",
master_port=0,
rank_offset=0,
world_size=0,
group_name="",
)
)
test_requests.append(DestroyWeightsUpdateGroupReqInput())
test_requests.append(
InitWeightsSendGroupForRemoteInstanceReqInput(
master_address="", ports="", group_name="", world_size=0, group_rank=0
)
)
test_requests.append(
SendWeightsToRemoteInstanceReqInput(master_address="", ports="")
)
test_requests.append(
UpdateWeightsFromTensorReqInput(serialized_named_tensors=[])
)
test_requests.append(GetWeightsByNameReqInput(name=""))
test_requests.append(ReleaseMemoryOccupationReqInput())
test_requests.append(RpcReqInput(method=""))
test_requests.append(GetLoadReqInput())
dispatcher = TypeBasedDispatcher(mapping)
# test
time_taken = timeit.timeit(
lambda: [dispatcher(req) for req in test_requests],
number=100, # Average of 100 runs
)
print(f"Total requests: {len(test_requests)}")
print(f"Time taken: {time_taken:.4f}s")
print(f"Requests per second: {len(test_requests) * 100 / time_taken:.0f}")
return time_taken
if __name__ == "__main__":
unittest.main()