chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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)