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

Support multiple engines for router simulation in schedule simulator (#16368)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 11:46:05 +08:00
committed by GitHub
parent e797f0c570
commit d7aa0ce72f
7 changed files with 72 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/test_schedule_simulator.py
+44 -47
View File
@@ -144,42 +144,37 @@ class TestGPUState(CustomTestCase):
class TestRouters(CustomTestCase):
def test_round_robin(self):
router = RoundRobinRouter()
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
router = RoundRobinRouter(num_gpus=4)
req = SimRequest(request_id="r1", input_len=100, output_len=50)
results = [router.route(req, gpu_states) for _ in range(8)]
results = [router.route(req) for _ in range(8)]
self.assertEqual(results, [0, 1, 2, 3, 0, 1, 2, 3])
def test_random_router(self):
router = RandomRouter()
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
router = RandomRouter(num_gpus=4)
req = SimRequest(request_id="r1", input_len=100, output_len=50)
results = [router.route(req, gpu_states) for _ in range(100)]
results = [router.route(req) for _ in range(100)]
self.assertTrue(all(0 <= r < 4 for r in results))
def test_sticky_router_same_group_same_gpu(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
reqs = [
SimRequest(request_id=f"r{i}", input_len=100, output_len=50, group_id="g0")
for i in range(10)
]
results = [router.route(req, gpu_states) for req in reqs]
results = [router.route(req) for req in reqs]
self.assertEqual(len(set(results)), 1)
def test_sticky_router_no_group_fallback(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
reqs = [
SimRequest(request_id=f"r{i}", input_len=100, output_len=50)
for i in range(100)
]
results = [router.route(req, gpu_states) for req in reqs]
results = [router.route(req) for req in reqs]
self.assertTrue(all(0 <= r < 4 for r in results))
def test_sticky_router_multiple_groups(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
for group_id in ["g0", "g1", "g2"]:
reqs = [
SimRequest(
@@ -190,7 +185,7 @@ class TestRouters(CustomTestCase):
)
for i in range(5)
]
results = [router.route(req, gpu_states) for req in reqs]
results = [router.route(req) for req in reqs]
self.assertEqual(len(set(results)), 1)
@@ -405,8 +400,8 @@ class TestSimulator(CustomTestCase):
for i in range(10)
]
sim = Simulator(
num_gpus=2,
router=RoundRobinRouter(),
num_gpus_per_engine=2,
router=RoundRobinRouter(num_gpus=2),
scheduler=FIFOScheduler(),
recorders=[
BatchSizeBalancednessRecorder(),
@@ -424,8 +419,8 @@ class TestSimulator(CustomTestCase):
SimRequest(request_id=f"r{i}", input_len=10, output_len=3) for i in range(4)
]
sim = Simulator(
num_gpus=2,
router=RoundRobinRouter(),
num_gpus_per_engine=2,
router=RoundRobinRouter(num_gpus=2),
scheduler=FIFOScheduler(),
max_total_tokens=10000,
)
@@ -436,7 +431,9 @@ class TestSimulator(CustomTestCase):
def test_empty_requests(self):
sim = Simulator(
num_gpus=2, router=RoundRobinRouter(), scheduler=FIFOScheduler()
num_gpus_per_engine=2,
router=RoundRobinRouter(num_gpus=2),
scheduler=FIFOScheduler(),
)
result = sim.run([])
self.assertEqual(result.summary, {})
@@ -447,8 +444,8 @@ class TestSimulator(CustomTestCase):
SimRequest(request_id=f"r{i}", input_len=10, output_len=3) for i in range(4)
]
sim = Simulator(
num_gpus=2,
router=RoundRobinRouter(),
num_gpus_per_engine=2,
router=RoundRobinRouter(num_gpus=2),
scheduler=FIFOScheduler(),
max_total_tokens=10000,
)
@@ -460,23 +457,18 @@ class TestSimulator(CustomTestCase):
self.assertEqual(len([r for r in result.step_records if r.step == 0]), 2)
def test_preemption_due_to_token_growth(self):
# 2 requests on 1 GPU, each input_len=50, output_len=10
# max_total_tokens=110, so initially both can run (100 tokens)
# After 5 decode steps, total = 50+5 + 50+5 = 110, still ok
# After 6 decode steps, total = 50+6 + 50+6 = 112 > 110, need preempt
requests = [
SimRequest(request_id="r0", input_len=50, output_len=10),
SimRequest(request_id="r1", input_len=50, output_len=10),
]
sim = Simulator(
num_gpus=1,
router=RoundRobinRouter(),
num_gpus_per_engine=1,
router=RoundRobinRouter(num_gpus=1),
scheduler=FIFOScheduler(),
max_total_tokens=110,
)
result = sim.run(requests)
# Check that preemption happened at some point
found_preemption = False
for record in result.step_records:
if record.running_count == 1 and record.pending_count == 1:
@@ -523,7 +515,7 @@ class TestCLI(CustomTestCase):
output_file = f.name
result = self._run_cli(
"--input", input_file, "--num-gpus", "2", "--output", output_file
"--input", input_file, "--num-gpus-per-engine", "2", "--output", output_file
)
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
self.assertIn("Loaded 2 requests", result.stdout)
@@ -562,7 +554,7 @@ class TestCLI(CustomTestCase):
"2",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"2",
"--router",
"sticky",
@@ -586,7 +578,7 @@ class TestCLI(CustomTestCase):
"128",
"--synth-random-range-ratio",
"0.5",
"--num-gpus",
"--num-gpus-per-engine",
"4",
)
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
@@ -599,7 +591,7 @@ class TestCLI(CustomTestCase):
"10",
"--synth-random-output-len",
"5",
"--num-gpus",
"--num-gpus-per-engine",
"2",
"--log-level",
"1",
@@ -608,7 +600,6 @@ class TestCLI(CustomTestCase):
self.assertIn("step=", result.stdout)
def test_e2e_simple_no_queuing(self):
# 4 requests, input_len=10, output_len=2, 2 GPUs, all fit in memory
result = self._run_cli(
"--synthetic",
"--synth-random-num-requests",
@@ -621,7 +612,7 @@ class TestCLI(CustomTestCase):
"1.0",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"2",
"--max-total-tokens",
"10000",
@@ -651,7 +642,7 @@ class TestCLI(CustomTestCase):
"1.0",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"1",
"--max-total-tokens",
"210",
@@ -683,7 +674,7 @@ step=5 | GPU0[R=0:- Q=0:-]""",
"1.0",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"1",
"--max-total-tokens",
"110",
@@ -717,7 +708,7 @@ step=13 | GPU0[R=0:- Q=0:-]""",
"10",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"2",
)
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
@@ -741,7 +732,7 @@ step=13 | GPU0[R=0:- Q=0:-]""",
"2",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"1",
"--max-total-tokens",
"80",
@@ -769,21 +760,25 @@ class TestLargerScale(CustomTestCase):
result = self._run_main(
"--synthetic",
"--synth-random-num-requests",
"2000",
"500000",
"--synth-random-input-len",
"32000",
"--synth-random-output-len",
"2000",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"8",
"--num-engines",
"250",
"--router",
"random",
"--max-total-tokens",
"2000000",
"--stop-criteria",
"exist_no_pending",
"--max-steps",
"1500",
)
self._assert_in_range(
result.summary["attention_compute_balancedness_mean"], 0.95, 1.0, "attn"
@@ -797,9 +792,9 @@ class TestLargerScale(CustomTestCase):
return self._run_main(
"--synth-gsp",
"--synth-gsp-num-groups",
"200",
"50000",
"--synth-gsp-prompts-per-group",
"20",
"100",
"--synth-gsp-system-prompt-len",
"31000",
"--synth-gsp-question-len",
@@ -808,8 +803,10 @@ class TestLargerScale(CustomTestCase):
"8000",
"--synth-seed",
"42",
"--num-gpus",
"--num-gpus-per-engine",
"8",
"--num-engines",
"250",
"--router",
router,
"--max-total-tokens",
@@ -823,22 +820,22 @@ class TestLargerScale(CustomTestCase):
def test_gsp_workload_random_policy(self):
result = self._run_gsp_workload("random")
self._assert_in_range(
result.summary["attention_compute_balancedness_mean"], 0.88, 0.98, "attn"
result.summary["attention_compute_balancedness_mean"], 0.90, 0.97, "attn"
)
self._assert_in_range(
result.summary["batch_size_balancedness_mean"], 0.88, 0.98, "bs"
result.summary["batch_size_balancedness_mean"], 0.90, 0.97, "bs"
)
self._assert_in_range(result.summary["avg_batch_size"], 24, 27, "avg_bs")
self._assert_in_range(result.summary["avg_batch_size"], 14, 17, "avg_bs")
def test_gsp_workload_sticky_policy(self):
result = self._run_gsp_workload("sticky")
self._assert_in_range(
result.summary["attention_compute_balancedness_mean"], 0.63, 0.70, "attn"
result.summary["attention_compute_balancedness_mean"], 0.64, 0.71, "attn"
)
self._assert_in_range(
result.summary["batch_size_balancedness_mean"], 0.63, 0.71, "bs"
result.summary["batch_size_balancedness_mean"], 0.64, 0.71, "bs"
)
self._assert_in_range(result.summary["avg_batch_size"], 33, 37, "avg_bs")
self._assert_in_range(result.summary["avg_batch_size"], 31, 36, "avg_bs")
if __name__ == "__main__":