diff --git a/python/sglang/srt/debug_utils/schedule_simulator/__init__.py b/python/sglang/srt/debug_utils/schedule_simulator/__init__.py index fabd41320..af7e995da 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/__init__.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/__init__.py @@ -1,4 +1,5 @@ from sglang.srt.debug_utils.schedule_simulator.data_source import ( + generate_gsp_requests, generate_random_requests, load_from_request_logger, ) @@ -40,6 +41,7 @@ __all__ = [ "AttentionBalancednessRecorder", "load_from_request_logger", "generate_random_requests", + "generate_gsp_requests", "create_arg_parser", "main", ] diff --git a/python/sglang/srt/debug_utils/schedule_simulator/data_source/__init__.py b/python/sglang/srt/debug_utils/schedule_simulator/data_source/__init__.py index ae1092090..116f4e4b5 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/data_source/__init__.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/data_source/__init__.py @@ -2,7 +2,12 @@ from sglang.srt.debug_utils.schedule_simulator.data_source.data_loader import ( load_from_request_logger, ) from sglang.srt.debug_utils.schedule_simulator.data_source.data_synthesis import ( + generate_gsp_requests, generate_random_requests, ) -__all__ = ["load_from_request_logger", "generate_random_requests"] +__all__ = [ + "load_from_request_logger", + "generate_random_requests", + "generate_gsp_requests", +] diff --git a/python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py b/python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py index fb2d3ccd1..5ce19d053 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py @@ -26,6 +26,51 @@ def generate_random_requests( ) ) + print( + f"Generated {len(requests)} random requests " + f"(input_len={input_len}, output_len={output_len}, range_ratio={range_ratio})" + ) + return requests + + +def generate_gsp_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + range_ratio: float = 1.0, + seed: Optional[int] = None, +) -> List[SimRequest]: + if seed is not None: + random.seed(seed) + + requests = [] + idx = 0 + for group_idx in range(num_groups): + group_id = f"g{group_idx}" + prefix_len = _random_len(system_prompt_len, range_ratio) + for _ in range(prompts_per_group): + q_len = _random_len(question_len, range_ratio) + osl = _random_len(output_len, range_ratio) + requests.append( + SimRequest( + request_id=f"gsp{idx}", + input_len=prefix_len + q_len, + output_len=osl, + group_id=group_id, + prefix_len=prefix_len, + ) + ) + idx += 1 + + random.shuffle(requests) + print( + f"Generated {len(requests)} GSP requests " + f"({num_groups} groups x {prompts_per_group} prompts, " + f"system_prompt_len={system_prompt_len}, question_len={question_len}, " + f"output_len={output_len})" + ) return requests diff --git a/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py b/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py index 4ad2cbb1d..b09064c68 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py @@ -9,6 +9,7 @@ from sglang.srt.debug_utils.schedule_simulator.data_source.data_loader import ( load_from_request_logger, ) from sglang.srt.debug_utils.schedule_simulator.data_source.data_synthesis import ( + generate_gsp_requests, generate_random_requests, ) from sglang.srt.debug_utils.schedule_simulator.metrics import ( @@ -36,13 +37,29 @@ def create_arg_parser() -> argparse.ArgumentParser: data_group.add_argument( "--synthetic", action="store_true", help="Use synthetic data generation" ) + data_group.add_argument( + "--synth-gsp", + action="store_true", + help="Use generated-shared-prefix (GSP) data generation", + ) - parser.add_argument("--synth-num-requests", type=int, default=1000) - parser.add_argument("--synth-input-len", type=int, default=1024) - parser.add_argument("--synth-output-len", type=int, default=256) - parser.add_argument("--synth-range-ratio", type=float, default=1.0) + # Shared synthetic arguments parser.add_argument("--synth-seed", type=int, default=None) + # Random dataset arguments (aligned with bench_serving.py --random-* options) + parser.add_argument("--synth-random-num-requests", type=int, default=1000) + parser.add_argument("--synth-random-input-len", type=int, default=1024) + parser.add_argument("--synth-random-output-len", type=int, default=256) + parser.add_argument("--synth-random-range-ratio", type=float, default=0.0) + + # GSP dataset arguments (aligned with bench_serving.py --gsp-* options) + parser.add_argument("--synth-gsp-num-groups", type=int, default=64) + parser.add_argument("--synth-gsp-prompts-per-group", type=int, default=16) + parser.add_argument("--synth-gsp-system-prompt-len", type=int, default=2048) + parser.add_argument("--synth-gsp-question-len", type=int, default=128) + parser.add_argument("--synth-gsp-output-len", type=int, default=256) + parser.add_argument("--synth-gsp-range-ratio", type=float, default=1.0) + parser.add_argument("--num-gpus", type=int, default=8) parser.add_argument( "--router", type=str, choices=["random", "round_robin"], default="round_robin" @@ -59,19 +76,23 @@ def _load_requests(args: argparse.Namespace) -> List[SimRequest]: if args.input: requests = load_from_request_logger(args.input) print(f"Loaded {len(requests)} requests from {args.input}") - else: - requests = generate_random_requests( - num_requests=args.synth_num_requests, - input_len=args.synth_input_len, - output_len=args.synth_output_len, - range_ratio=args.synth_range_ratio, + elif args.synth_gsp: + requests = generate_gsp_requests( + num_groups=args.synth_gsp_num_groups, + prompts_per_group=args.synth_gsp_prompts_per_group, + system_prompt_len=args.synth_gsp_system_prompt_len, + question_len=args.synth_gsp_question_len, + output_len=args.synth_gsp_output_len, + range_ratio=args.synth_gsp_range_ratio, seed=args.synth_seed, ) - print( - f"Generated {len(requests)} synthetic requests " - f"(synth_input_len={args.synth_input_len}, " - f"synth_output_len={args.synth_output_len}, " - f"synth_range_ratio={args.synth_range_ratio})" + else: + requests = generate_random_requests( + num_requests=args.synth_random_num_requests, + input_len=args.synth_random_input_len, + output_len=args.synth_random_output_len, + range_ratio=args.synth_random_range_ratio, + seed=args.synth_seed, ) return requests diff --git a/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py b/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py index b6f2b64d1..059d4751a 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import List +from typing import List, Optional from sglang.srt.debug_utils.schedule_simulator.request import SimRequest @@ -25,8 +25,15 @@ class GPUState: def batch_size(self) -> int: return len(self.running_requests) - def total_seq_len(self) -> int: - return sum(req.seq_len() for req in self.running_requests) + def total_seq_len(self, extra_reqs: Optional[List[SimRequest]] = None) -> int: + seen_groups = set() + total = 0 + for req in self.running_requests + (extra_reqs or []): + is_shared = req.group_id is not None and req.group_id in seen_groups + total += req.seq_len() - (req.prefix_len if is_shared else 0) + if req.group_id is not None: + seen_groups.add(req.group_id) + return total def is_valid(self) -> bool: return self.total_seq_len() <= self.max_total_tokens diff --git a/python/sglang/srt/debug_utils/schedule_simulator/request.py b/python/sglang/srt/debug_utils/schedule_simulator/request.py index 8cdcad53d..6c0ec1356 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/request.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/request.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import Optional @dataclass @@ -7,6 +8,8 @@ class SimRequest: input_len: int output_len: int decoded_tokens: int = 0 + group_id: Optional[str] = None + prefix_len: int = 0 def seq_len(self) -> int: return self.input_len + self.decoded_tokens diff --git a/python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py b/python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py index 4970d355b..2eb20714e 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py @@ -12,5 +12,5 @@ class FIFOScheduler(SchedulerPolicy): gpu_state.evict_request(gpu_state.running_requests[-1]) for req in list(gpu_state.pending_requests): - if gpu_state.total_seq_len() + req.seq_len() <= gpu_state.max_total_tokens: + if gpu_state.total_seq_len(extra_reqs=[req]) <= gpu_state.max_total_tokens: gpu_state.start_request(req) diff --git a/test/registered/debug_utils/test_schedule_simulator.py b/test/registered/debug_utils/test_schedule_simulator.py index 0a9e5b04c..ce5c1084a 100644 --- a/test/registered/debug_utils/test_schedule_simulator.py +++ b/test/registered/debug_utils/test_schedule_simulator.py @@ -15,6 +15,7 @@ from sglang.srt.debug_utils.schedule_simulator import ( SimulationResult, Simulator, StepRecord, + generate_gsp_requests, generate_random_requests, load_from_request_logger, ) @@ -67,6 +68,76 @@ class TestGPUState(CustomTestCase): ] self.assertEqual(gpu.total_seq_len(), 100 + 210) + def test_total_seq_len_shared_prefix(self): + gpu = GPUState(gpu_id=0, max_total_tokens=10000) + gpu.running_requests = [ + SimRequest( + request_id="r1", + input_len=150, + output_len=50, + group_id="g0", + prefix_len=100, + ), + SimRequest( + request_id="r2", + input_len=150, + output_len=50, + group_id="g0", + prefix_len=100, + ), + ] + self.assertEqual(gpu.total_seq_len(), 150 + 50) + + def test_total_seq_len_shared_prefix_with_decoded(self): + gpu = GPUState(gpu_id=0, max_total_tokens=10000) + gpu.running_requests = [ + SimRequest( + request_id="r1", + input_len=150, + output_len=50, + decoded_tokens=10, + group_id="g0", + prefix_len=100, + ), + SimRequest( + request_id="r2", + input_len=150, + output_len=50, + decoded_tokens=5, + group_id="g0", + prefix_len=100, + ), + ] + self.assertEqual(gpu.total_seq_len(), 160 + 55) + + def test_total_seq_len_multiple_groups(self): + gpu = GPUState(gpu_id=0, max_total_tokens=10000) + gpu.running_requests = [ + SimRequest( + request_id="r1", + input_len=150, + output_len=50, + group_id="g0", + prefix_len=100, + ), + SimRequest( + request_id="r2", + input_len=150, + output_len=50, + group_id="g0", + prefix_len=100, + ), + SimRequest( + request_id="r3", + input_len=200, + output_len=50, + group_id="g1", + prefix_len=150, + ), + SimRequest(request_id="r4", input_len=80, output_len=20), + ] + self.assertEqual(gpu.total_seq_len(), 150 + 50 + 200 + 80) + class TestRouters(CustomTestCase): def test_round_robin(self): @@ -224,6 +295,69 @@ class TestDataSynthesis(CustomTestCase): for a, b in zip(r1, r2): self.assertEqual(a.input_len, b.input_len) + def test_generate_gsp_basic(self): + requests = generate_gsp_requests( + num_groups=4, + prompts_per_group=3, + system_prompt_len=100, + question_len=50, + output_len=25, + seed=42, + ) + self.assertEqual(len(requests), 12) + for req in requests: + self.assertIsNotNone(req.group_id) + self.assertEqual(req.prefix_len, 100) + self.assertEqual(req.input_len, 150) + self.assertEqual(req.output_len, 25) + + def test_generate_gsp_group_assignment(self): + requests = generate_gsp_requests( + num_groups=3, + prompts_per_group=2, + system_prompt_len=100, + question_len=50, + output_len=25, + seed=42, + ) + group_counts = {} + for req in requests: + group_counts[req.group_id] = group_counts.get(req.group_id, 0) + 1 + self.assertEqual(len(group_counts), 3) + for count in group_counts.values(): + self.assertEqual(count, 2) + + def test_generate_gsp_with_range_ratio(self): + requests = generate_gsp_requests( + num_groups=4, + prompts_per_group=5, + system_prompt_len=100, + question_len=50, + output_len=25, + range_ratio=0.5, + seed=42, + ) + for req in requests: + self.assertGreaterEqual(req.prefix_len, 50) + self.assertLessEqual(req.prefix_len, 100) + self.assertGreaterEqual(req.input_len - req.prefix_len, 25) + self.assertLessEqual(req.input_len - req.prefix_len, 50) + + def test_generate_gsp_shuffled(self): + requests = generate_gsp_requests( + num_groups=4, + prompts_per_group=10, + system_prompt_len=100, + question_len=50, + output_len=25, + seed=42, + ) + group_ids = [req.group_id for req in requests] + is_sorted = all( + group_ids[i] <= group_ids[i + 1] for i in range(len(group_ids) - 1) + ) + self.assertFalse(is_sorted) + class TestSimulator(CustomTestCase): def test_basic_run(self): @@ -377,26 +511,26 @@ class TestCLI(CustomTestCase): def test_cli_synthetic(self): result = self._run_cli( "--synthetic", - "--synth-num-requests", + "--synth-random-num-requests", "100", - "--synth-input-len", + "--synth-random-input-len", "512", - "--synth-output-len", + "--synth-random-output-len", "128", - "--synth-range-ratio", + "--synth-random-range-ratio", "0.5", "--num-gpus", "4", ) self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") - self.assertIn("Generated 100 synthetic requests", result.stdout) + self.assertIn("Generated 100 random requests", result.stdout) def test_cli_log_level(self): result = self._run_cli( "--synthetic", - "--synth-num-requests", + "--synth-random-num-requests", "10", - "--synth-output-len", + "--synth-random-output-len", "5", "--num-gpus", "2", @@ -410,12 +544,14 @@ class TestCLI(CustomTestCase): # 4 requests, input_len=10, output_len=2, 2 GPUs, all fit in memory result = self._run_cli( "--synthetic", - "--synth-num-requests", + "--synth-random-num-requests", "4", - "--synth-input-len", + "--synth-random-input-len", "10", - "--synth-output-len", + "--synth-random-output-len", "2", + "--synth-random-range-ratio", + "1.0", "--synth-seed", "42", "--num-gpus", @@ -438,12 +574,14 @@ class TestCLI(CustomTestCase): def test_e2e_queuing_due_to_token_limit(self): result = self._run_cli( "--synthetic", - "--synth-num-requests", + "--synth-random-num-requests", "4", - "--synth-input-len", + "--synth-random-input-len", "100", - "--synth-output-len", + "--synth-random-output-len", "3", + "--synth-random-range-ratio", + "1.0", "--synth-seed", "42", "--num-gpus", @@ -468,12 +606,14 @@ step=5 | GPU0[R=0:- Q=0:-]""", def test_e2e_retraction_due_to_token_growth(self): result = self._run_cli( "--synthetic", - "--synth-num-requests", + "--synth-random-num-requests", "2", - "--synth-input-len", + "--synth-random-input-len", "50", - "--synth-output-len", + "--synth-random-output-len", "10", + "--synth-random-range-ratio", + "1.0", "--synth-seed", "42", "--num-gpus", @@ -495,6 +635,58 @@ step=10 | GPU0[R=1:syn1 Q=0:-] step=13 | GPU0[R=0:- Q=0:-]""", ) + def test_cli_gsp_basic(self): + result = self._run_cli( + "--synth-gsp", + "--synth-gsp-num-groups", + "4", + "--synth-gsp-prompts-per-group", + "8", + "--synth-gsp-system-prompt-len", + "100", + "--synth-gsp-question-len", + "50", + "--synth-gsp-output-len", + "10", + "--synth-seed", + "42", + "--num-gpus", + "2", + ) + self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") + self.assertIn("Generated 32 GSP requests", result.stdout) + self.assertIn("4 groups x 8 prompts", result.stdout) + + def test_e2e_gsp_shared_prefix_enables_batching(self): + for has_long_prefix in [True, False]: + prefix_len, question_len = (50, 10) if has_long_prefix else (10, 50) + result = self._run_cli( + "--synth-gsp", + "--synth-gsp-num-groups", + "1", + "--synth-gsp-prompts-per-group", + "2", + "--synth-gsp-system-prompt-len", + str(prefix_len), + "--synth-gsp-question-len", + str(question_len), + "--synth-gsp-output-len", + "2", + "--synth-seed", + "42", + "--num-gpus", + "1", + "--max-total-tokens", + "80", + "--log-level", + "2", + ) + self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") + if has_long_prefix: + self.assertIn("R=2:", result.stdout) + else: + self.assertNotIn("R=2:", result.stdout) + if __name__ == "__main__": unittest.main()