minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

bench_serving: add Zipfian shared-prefix sampling to generated-shared-prefix (#26378)

Browse Source
This commit is contained in:
Jimmy Shong
2026-05-28 14:39:46 -07:00
committed by GitHub
parent 97d129f8c6
commit f838adb7d4
4 changed files with 728 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/bench_serving.py
+71
View File
@@ -17,6 +17,7 @@ import asyncio
import copy
import importlib.util
import json
import math
import os
import random
import shutil
@@ -1930,6 +1931,44 @@ def run_benchmark(args_: argparse.Namespace):
)
def _finite_positive_float(value) -> float:
"""argparse type for a finite, strictly positive float."""
try:
parsed = float(value)
except (TypeError, ValueError) as exc:
raise argparse.ArgumentTypeError(
f"expected a finite float > 0, got {value!r}"
) from exc
if not math.isfinite(parsed) or parsed <= 0:
raise argparse.ArgumentTypeError(f"expected a finite float > 0, got {value!r}")
return parsed
def _validate_parsed_gsp_args(
parser: argparse.ArgumentParser, args: argparse.Namespace
) -> None:
"""Reject malformed GSP distribution/alpha combinations at parse time.
Invoked from the CLI entry point right after ``parser.parse_args()`` so
users see a clear argparse-style error before any server, model, or
tokenizer setup runs and masks the real cause with an unrelated network
failure.
"""
distribution = getattr(args, "gsp_group_distribution", None)
alpha = getattr(args, "gsp_zipf_alpha", None)
if distribution == "zipf" and alpha is None:
parser.error(
"--gsp-group-distribution=zipf requires --gsp-zipf-alpha "
"(a finite float > 0)"
)
if distribution == "uniform" and alpha is not None:
parser.error(
"--gsp-zipf-alpha is only meaningful with "
"--gsp-group-distribution=zipf; remove --gsp-zipf-alpha "
"or set --gsp-group-distribution=zipf"
)
class LoRAPathAction(argparse.Action):
def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, [])
@@ -2366,6 +2405,37 @@ if __name__ == "__main__":
action="store_true",
help="Keep requests in order without shuffling. By default, requests are shuffled randomly.",
)
group.add_argument(
"--gsp-group-distribution",
type=str,
choices=["uniform", "zipf"],
default="uniform",
help=(
"Prefix-group sampling distribution for generated-shared-prefix. "
"'uniform' (default) assigns each group an equal number of requests. "
"'zipf' samples each request's group by rank with "
"p(rank) = (1/rank**alpha) / sum_k(1/k**alpha); rank starts at 1 "
"and group index 0 is the hottest. Requires --gsp-zipf-alpha "
"(a finite float > 0) when set to 'zipf'. Total request count is "
"still num_groups * prompts_per_group, identical to uniform mode; "
"only the per-request group assignment changes. The on-disk "
"dataset cache uses a distinct key per (group_distribution, "
"zipf_alpha), so uniform-mode caches are never mixed with "
"zipf-mode caches and zipf runs with different alpha use "
"separate files."
),
)
group.add_argument(
"--gsp-zipf-alpha",
type=_finite_positive_float,
default=None,
help=(
"Zipf exponent alpha for --gsp-group-distribution=zipf, with "
"p(rank) = (1/rank**alpha) / sum_k(1/k**alpha) and rank starting "
"at 1. Must be a finite float strictly greater than 0; larger "
"values concentrate requests on lower-ranked (hotter) groups."
),
)
mooncake_group = parser.add_argument_group("mooncake dataset arguments")
mooncake_group.add_argument(
"--mooncake-slowdown-factor",
@@ -2413,4 +2483,5 @@ if __name__ == "__main__":
help="Custom HTTP headers in Key=Value format. Example: --header MyHeader=MY_VALUE MyAnotherHeader=myanothervalue",
)
args = parser.parse_args()
_validate_parsed_gsp_args(parser, args)
run_benchmark(args)
python/sglang/benchmark/datasets/generated_shared_prefix.py
+132 -35
View File
@@ -1,3 +1,4 @@
import math
import pickle
import random
import uuid
@@ -5,7 +6,7 @@ from argparse import Namespace
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import List
from typing import List, Optional
import numpy as np
from tqdm.asyncio import tqdm
@@ -19,6 +20,22 @@ from sglang.benchmark.datasets.common import (
)
def _zipf_group_probs(num_groups: int, alpha: float) -> np.ndarray:
"""Rank-based Zipf probability vector with rank starting at 1.
weight(rank) = 1 / rank ** alpha (rank in 1..num_groups)
probability(rank) = weight(rank) / sum_over_all_ranks(weight)
The returned array has length num_groups; element i corresponds to
group index i (rank i + 1), so group 0 is the hottest.
"""
if num_groups <= 0:
raise ValueError(f"num_groups must be > 0, got {num_groups}")
ranks = np.arange(1, num_groups + 1, dtype=np.float64)
weights = 1.0 / (ranks**alpha)
return weights / weights.sum()
@dataclass
class GeneratedSharedPrefixDataset(BaseDataset):
num_groups: int
@@ -32,10 +49,40 @@ class GeneratedSharedPrefixDataset(BaseDataset):
send_routing_key: bool
num_turns: int
ordered: bool
group_distribution: str = "uniform"
zipf_alpha: Optional[float] = None
@classmethod
def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset":
assert not getattr(args, "tokenize_prompt", False)
group_distribution = args.gsp_group_distribution
zipf_alpha = args.gsp_zipf_alpha
# Defensive validation for in-process callers that construct a
# Namespace by hand and bypass the argparse boundary in
# bench_serving.py. The CLI hook enforces the same rules first.
if group_distribution not in ("uniform", "zipf"):
raise ValueError(
f"--gsp-group-distribution must be 'uniform' or 'zipf', "
f"got {group_distribution!r}"
)
if group_distribution == "zipf":
if zipf_alpha is None:
raise ValueError(
"--gsp-group-distribution=zipf requires --gsp-zipf-alpha "
"(a finite float > 0)"
)
if not math.isfinite(zipf_alpha) or zipf_alpha <= 0:
raise ValueError(
f"--gsp-zipf-alpha must be a finite float > 0, got {zipf_alpha!r}"
)
elif zipf_alpha is not None:
raise ValueError(
"--gsp-zipf-alpha is only meaningful with "
"--gsp-group-distribution=zipf; remove --gsp-zipf-alpha "
"or set --gsp-group-distribution=zipf"
)
return cls(
num_groups=args.gsp_num_groups,
prompts_per_group=args.gsp_prompts_per_group,
@@ -48,6 +95,8 @@ class GeneratedSharedPrefixDataset(BaseDataset):
send_routing_key=getattr(args, "gsp_send_routing_key", False),
num_turns=getattr(args, "gsp_num_turns", 1),
ordered=getattr(args, "gsp_ordered", False),
group_distribution=group_distribution,
zipf_alpha=zipf_alpha,
)
def load(
@@ -66,6 +115,8 @@ class GeneratedSharedPrefixDataset(BaseDataset):
num_turns=self.num_turns,
fast_prepare=self.fast_prepare,
ordered=self.ordered,
group_distribution=self.group_distribution,
zipf_alpha=self.zipf_alpha,
)
@@ -77,13 +128,24 @@ def get_gen_prefix_cache_path(
question_len: int,
output_len: int,
tokenizer,
group_distribution: str = "uniform",
zipf_alpha: Optional[float] = None,
):
"""Create cache directory under ~/.cache/sglang/benchmark"""
"""Create cache directory under ~/.cache/sglang/benchmark.
The uniform-mode filename is preserved exactly as before so existing
on-disk caches remain valid. Non-default sampling modes get an extra
suffix encoding the parameters that affect the cached payload.
"""
cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"
suffix = ""
if group_distribution != "uniform":
suffix = f"_{group_distribution}_{zipf_alpha}"
cache_key = (
f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_"
f"{system_prompt_len}_{question_len}_{output_len}_"
f"{system_prompt_len}_{question_len}_{output_len}{suffix}_"
f"{tokenizer.__class__.__name__}.pkl"
)
return cache_dir / cache_key
@@ -102,8 +164,22 @@ def sample_generated_shared_prefix_requests(
num_turns: int = 1,
fast_prepare: bool = False,
ordered: bool = False,
group_distribution: str = "uniform",
zipf_alpha: Optional[float] = None,
) -> List[DatasetRow]:
"""Generate benchmark requests with shared system prompts using random tokens and caching."""
"""Generate benchmark requests with shared system prompts using random tokens and caching.
When group_distribution is "uniform" (default), each group receives exactly
prompts_per_group requests; behavior matches the legacy generator.
When group_distribution is "zipf", each request's group is sampled by rank
with probability 1/rank**zipf_alpha / sum_k(1/k**zipf_alpha); rank starts at
1 and group index 0 is the hottest. Sampling uses an isolated
numpy.random.default_rng(seed) so the shared question/system-prompt pool
stays byte-identical to uniform mode for the same seed and other args.
Zipf mode is cached on disk under a distinct key per (group_distribution,
zipf_alpha) value.
"""
cache_path = get_gen_prefix_cache_path(
seed,
num_groups,
@@ -112,18 +188,25 @@ def sample_generated_shared_prefix_requests(
question_len,
output_len,
tokenizer,
group_distribution=group_distribution,
zipf_alpha=zipf_alpha,
)
should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1
# range_ratio != 1 / num_turns > 1 perturb the payload but are not in the
# cache key; send_routing_key embeds a per-run uuid + timestamp that is
# meaningless to cache. Bypass for these pre-existing reasons only.
should_cache = range_ratio == 1 and not send_routing_key and num_turns == 1
# Try to load from cache first
if cache_path.exists() and should_cache:
if should_cache and cache_path.exists():
print(f"\nLoading cached generated input data from {cache_path}")
with open(cache_path, "rb") as f:
return pickle.load(f)
if not should_cache:
print(f"\nCache bypassed ({range_ratio=}, {send_routing_key=}, {num_turns=})")
print(
f"\nGenerating new input data... "
f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})"
f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=}, {group_distribution=}, {zipf_alpha=})"
)
run_random_str = uuid.uuid4().hex[:8]
@@ -150,12 +233,11 @@ def sample_generated_shared_prefix_requests(
).reshape(num_groups, prompts_per_group)
del system_prompt_len, question_len, output_len
# Generate system prompts for each group
system_prompts = [
gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups)
]
# Generate questions: shape (num_groups, prompts_per_group, num_turns)
# shape: (num_groups, prompts_per_group, num_turns)
questions = [
[
[
@@ -167,48 +249,64 @@ def sample_generated_shared_prefix_requests(
for g in range(num_groups)
]
# Combine system prompts with questions
# Per-slot group assignment. Uniform mode is the identity assignment
# [0,0,...,1,1,...,N-1,N-1]; zipf mode samples from the rank distribution
# using an isolated RNG so the module-level random / numpy.random state
# that compute_random_lens / gen_prompt rely on is never perturbed -- this
# keeps the system-prompt and question pool byte-identical to uniform mode
# for the same seed and other args.
total_slots = num_groups * prompts_per_group
if group_distribution == "uniform":
assignment = np.repeat(np.arange(num_groups), prompts_per_group)
else: # "zipf"
rng = np.random.default_rng(seed)
probs = _zipf_group_probs(num_groups, zipf_alpha)
assignment = rng.choice(num_groups, size=total_slots, replace=True, p=probs)
input_requests = []
total_input_tokens = 0
total_output_tokens = 0
for slot_idx, sampled_g in enumerate(
tqdm(assignment, desc="Generating shared-prefix prompts")
):
# src_(g,p) walks the question pool in uniform-enumeration order, so
# per-slot question text is reproducibly identical across modes.
src_g, src_p = divmod(slot_idx, prompts_per_group)
sampled_g = int(sampled_g)
for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
system_prompt = system_prompts[group_idx]
system_prompt = system_prompts[sampled_g]
routing_key = (
f"{run_random_str}_{run_start_timestamp}_{group_idx}"
f"{run_random_str}_{run_start_timestamp}_{sampled_g}"
if send_routing_key
else None
)
for prompt_idx in tqdm(
range(prompts_per_group), desc="Generating questions", leave=False
):
turn_questions = questions[group_idx][prompt_idx]
turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[
1:
]
full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts
prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0]))
output_len_val = int(output_lens[group_idx, prompt_idx])
turn_questions = questions[src_g][src_p]
turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[1:]
full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts
prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0]))
output_len_val = int(output_lens[src_g, src_p])
input_requests.append(
DatasetRow(
prompt=full_prompt,
prompt_len=prompt_len,
output_len=output_len_val,
routing_key=routing_key,
)
input_requests.append(
DatasetRow(
prompt=full_prompt,
prompt_len=prompt_len,
output_len=output_len_val,
routing_key=routing_key,
)
total_input_tokens += prompt_len
total_output_tokens += output_len_val
)
total_input_tokens += prompt_len
total_output_tokens += output_len_val
if not ordered:
random.shuffle(input_requests)
# Print statistics
print(f"\nGenerated shared prefix dataset statistics:")
print(f"Number of groups: {num_groups}")
print(f"Prompts per group: {prompts_per_group}")
print(f"Number of turns: {num_turns}")
print(f"Group distribution: {group_distribution}")
if group_distribution == "zipf":
print(f"Zipf alpha: {zipf_alpha}")
print(f"Total prompts: {len(input_requests)}")
if not fast_prepare:
print(f"Total input tokens: {total_input_tokens}")
@@ -221,7 +319,6 @@ def sample_generated_shared_prefix_requests(
f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n"
)
# Save to cache
if should_cache:
cache_path.parent.mkdir(parents=True, exist_ok=True)
print(f"Caching generated input data to {cache_path}")