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

Support shared prefix (gsp) in schedule simulator (#16353)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 08:03:16 +08:00
committed by GitHub
parent c88aaf22c7
commit ef5ac6f01e
8 changed files with 311 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/debug_utils/schedule_simulator/__init__.py
View File

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

7
python/sglang/srt/debug_utils/schedule_simulator/data_source/__init__.py
View File

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

45
python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py
View File

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

51
python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py
View File

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

13
python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py
View File

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

3
python/sglang/srt/debug_utils/schedule_simulator/request.py
View File

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

2
python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py
View File

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

224
test/registered/debug_utils/test_schedule_simulator.py
View File

@@ -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()