minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
docs_new/docs/developer_guide/bench_serving.mdx
+22
View File
@@ -83,6 +83,8 @@ Generated Shared Prefix flags (for `generated-shared-prefix`):
- `--gsp-system-prompt-len`
- `--gsp-question-len`
- `--gsp-output-len`
- `--gsp-group-distribution {uniform,zipf}`: per-request prefix-group sampling distribution (default: `uniform`). With `zipf`, each request's group is sampled by rank with `p(rank) = (1/rank**alpha) / sum_k(1/k**alpha)`; rank starts at 1 and group index 0 is the hottest. 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.
- `--gsp-zipf-alpha FLOAT`: Zipf exponent for `--gsp-group-distribution=zipf`. Must be a finite float strictly greater than 0; larger values concentrate requests on lower-ranked (hotter) groups. Required when the distribution is `zipf`; must be omitted otherwise.
Image dataset flags (for `image`):
@@ -318,6 +320,26 @@ python3 -m sglang.bench_serving \
--num-prompts 1024
```
Zipfian / power-law prefix popularity (opt-in via `--gsp-group-distribution=zipf`):
```bash Command
python3 -m sglang.bench_serving \
--backend sglang \
--host 127.0.0.1 --port 30000 \
--model meta-llama/Llama-3.1-8B-Instruct \
--dataset-name generated-shared-prefix \
--gsp-num-groups 64 --gsp-prompts-per-group 16 \
--gsp-system-prompt-len 2048 --gsp-question-len 128 --gsp-output-len 256 \
--gsp-group-distribution zipf --gsp-zipf-alpha 1.2 \
--seed 42
```
`zipf` mode samples each request's prefix group from the rank-based distribution `p(rank) = (1/rank**alpha) / sum_k(1/k**alpha)` with rank starting at 1, so group index 0 is the hottest. The total request count stays `num_groups * prompts_per_group` — identical to `uniform` mode — and only the per-request group assignment changes. `alpha` must be a finite float strictly greater than 0; larger values concentrate requests on lower-ranked (hotter) groups.
The on-disk dataset cache at `~/.cache/sglang/benchmark/gen_shared_prefix_*.pkl` includes `group_distribution` and `zipf_alpha` in its key, so uniform-mode and zipf-mode runs (or two zipf runs with different alpha) never share a cache file. Uniform-mode filenames are unchanged from the legacy format, so existing caches remain valid.
This flag controls prefix-popularity shape only. It does not by itself reproduce any production trace or guarantee an observed cache-hit rate for a given engine.
6) Tokenized prompts (ids) for strict length control (sglang only):
```bash Command
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}")
test/registered/bench_fn/test_benchmark_datasets_api.py
+503 -1
View File
@@ -1,11 +1,17 @@
import asyncio
import json
import pickle
import random
import subprocess
import sys
import tempfile
import unittest
from collections import Counter
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from PIL import Image
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
@@ -16,6 +22,9 @@ from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset
from sglang.benchmark.datasets.common import DatasetRow
from sglang.benchmark.datasets.custom import sample_custom_requests
from sglang.benchmark.datasets.generated_shared_prefix import (
GeneratedSharedPrefixDataset,
_zipf_group_probs,
get_gen_prefix_cache_path,
sample_generated_shared_prefix_requests,
)
from sglang.benchmark.datasets.image import sample_image_requests
@@ -26,7 +35,7 @@ from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
register_cpu_ci(est_time=40, suite="base-a-test-cpu")
register_cpu_ci(est_time=7, suite="base-b-test-cpu")
@@ -133,6 +142,8 @@ def make_args(**overrides):
"gsp_send_routing_key": False,
"gsp_num_turns": 1,
"gsp_ordered": False,
"gsp_group_distribution": "uniform",
"gsp_zipf_alpha": None,
"seed": 1,
"mooncake_workload": "conversation",
"speed_bench_category": None,
@@ -148,8 +159,19 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
self.processor = DummyProcessor(self.tokenizer)
self.tmpdir = tempfile.TemporaryDirectory()
self.tmpdir_path = Path(self.tmpdir.name)
# Redirect ~ for the GSP on-disk cache to the per-test tempdir, so
# tests never read/write the real ~/.cache/sglang/benchmark. The Zipf
# tests in particular compare freshly generated rows against the
# uniform path, and a stale cache file from prior runs would silently
# short-circuit the uniform path and break that comparison.
self._home_patch = patch(
"sglang.benchmark.datasets.generated_shared_prefix.Path.home",
return_value=self.tmpdir_path,
)
self._home_patch.start()
def tearDown(self):
self._home_patch.stop()
self.tmpdir.cleanup()
def _write_sharegpt_json(self):
@@ -554,6 +576,486 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
with self.assertRaises(ValueError):
get_dataset(args, self.tokenizer, model_id="dummy-model")
# ------------------------------------------------------------------
# Generated-shared-prefix Zipf sampling
# ------------------------------------------------------------------
def _run_gsp(
self,
*,
mode="uniform",
alpha=None,
seed=42,
num_groups=4,
prompts_per_group=5,
num_turns=1,
send_routing_key=False,
ordered=True,
range_ratio=1.0,
system_prompt_len=4,
question_len=3,
output_len=2,
fast_prepare=True,
global_seed=None,
):
# GSP's own `seed` kwarg only feeds the cache filename; reproducibility
# of compute_random_lens / gen_prompt comes from seeding the module
# globals before calling. Tests must seed both random and numpy here.
seed_for_globals = global_seed if global_seed is not None else seed
random.seed(seed_for_globals)
np.random.seed(seed_for_globals)
return sample_generated_shared_prefix_requests(
num_groups=num_groups,
prompts_per_group=prompts_per_group,
system_prompt_len=system_prompt_len,
question_len=question_len,
output_len=output_len,
range_ratio=range_ratio,
tokenizer=self.tokenizer,
seed=seed,
send_routing_key=send_routing_key,
num_turns=num_turns,
fast_prepare=fast_prepare,
ordered=ordered,
group_distribution=mode,
zipf_alpha=alpha,
)
@staticmethod
def _row_fields(rows):
return [(r.prompt, r.prompt_len, r.output_len, r.routing_key) for r in rows]
def test_gsp_uniform_default_unchanged(self):
# Uniform mode returns the documented number of rows and is
# bit-reproducible under fixed seeding of the global RNGs.
rows_a = self._run_gsp(
mode="uniform", num_groups=3, prompts_per_group=4, seed=7
)
rows_b = self._run_gsp(
mode="uniform", num_groups=3, prompts_per_group=4, seed=7
)
self.assertEqual(len(rows_a), 3 * 4)
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
def test_gsp_uniform_cache_path_format_unchanged(self):
# The uniform-mode cache filename keeps its existing
# gen_shared_prefix_<seed>_<N>_<P>_<sysL>_<qL>_<outL>_<TokenizerCls>.pkl
# shape. The trailing class name is a transformers/tokenizers internal
# detail (TokenizersBackend / PreTrainedTokenizerFast depending on
# version), so we only pin the deterministic numeric portion.
path = get_gen_prefix_cache_path(
seed=7,
num_groups=3,
prompts_per_group=4,
system_prompt_len=16,
question_len=8,
output_len=4,
tokenizer=self.tokenizer,
)
self.assertTrue(path.name.startswith("gen_shared_prefix_7_3_4_16_8_4_"))
self.assertTrue(path.name.endswith(".pkl"))
self.assertEqual(path.parent, Path.home() / ".cache" / "sglang" / "benchmark")
def test_zipf_group_probs_helper(self):
# Rank-based probability vector: weight(rank) = 1 / rank ** alpha,
# normalized to sum to 1, with rank starting at 1.
probs_n3_a1 = _zipf_group_probs(3, 1.0)
expected_n3_a1 = np.array([6.0, 3.0, 2.0]) / 11.0
np.testing.assert_allclose(probs_n3_a1, expected_n3_a1, atol=1e-12)
self.assertAlmostEqual(float(probs_n3_a1.sum()), 1.0, places=12)
probs_n4_a15 = _zipf_group_probs(4, 1.5)
ranks = np.arange(1, 5, dtype=np.float64)
ref = 1.0 / ranks**1.5
ref = ref / ref.sum()
np.testing.assert_allclose(probs_n4_a15, ref, atol=1e-12)
# Three-decimal pin against a hand-computable reference.
np.testing.assert_allclose(
np.round(probs_n4_a15, 3),
np.array([0.598, 0.212, 0.115, 0.075]),
atol=1e-3,
)
def test_zipf_group_probs_not_lora_skewed_formula(self):
# The helper must NOT use the LoRA `skewed` alpha**-i exponential
# formula; for alpha=1.5, N=4 the two formulas differ noticeably.
actual = _zipf_group_probs(4, 1.5)
lora_weights = np.array([1.5**-i for i in range(4)], dtype=np.float64)
lora_probs = lora_weights / lora_weights.sum()
self.assertFalse(
np.allclose(actual, lora_probs, atol=1e-3),
"Zipf helper must use rank-based 1/rank**alpha, not LoRA alpha**-i",
)
def test_zipf_reproducible_with_seed(self):
# Same seed + same args -> identical rows, including order, under
# both the in-order and shuffled paths.
kwargs = dict(
mode="zipf", alpha=1.7, seed=11, num_groups=4, prompts_per_group=10
)
rows_a = self._run_gsp(**kwargs)
rows_b = self._run_gsp(**kwargs)
self.assertEqual(len(rows_a), 4 * 10)
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
# Also under the shuffled path.
rows_c = self._run_gsp(ordered=False, **kwargs)
rows_d = self._run_gsp(ordered=False, **kwargs)
self.assertEqual(self._row_fields(rows_c), self._row_fields(rows_d))
def test_zipf_different_seeds_differ(self):
# Different seeds -> at least one differing slot under Zipf sampling.
base = dict(mode="zipf", alpha=1.7, num_groups=4, prompts_per_group=10)
rows_a = self._run_gsp(seed=11, **base)
rows_b = self._run_gsp(seed=12, **base)
self.assertEqual(len(rows_a), len(rows_b))
self.assertNotEqual(self._row_fields(rows_a), self._row_fields(rows_b))
def test_zipf_does_not_perturb_global_random_state(self):
# The Zipf branch must consume zero draws from the global random /
# numpy.random state. Therefore the per-slot generated questions and
# system prompts under uniform and Zipf modes for the same args and
# the same global seed are byte-equal.
common = dict(
num_groups=4,
prompts_per_group=6,
system_prompt_len=4,
question_len=3,
output_len=2,
range_ratio=1.0,
seed=99,
ordered=True,
send_routing_key=False,
fast_prepare=True,
global_seed=99,
)
uniform_rows = self._run_gsp(mode="uniform", **common)
zipf_rows = self._run_gsp(mode="zipf", alpha=1.3, **common)
# Slot i in uniform mode pairs system_prompts[i // P] with
# questions[i // P][i % P], so the question substring after the
# delimiter is exactly the i-th question. Same construction is used by
# the Zipf branch (only the system prompt changes per slot), so the
# question substrings must match slot-by-slot under the same global
# seed.
delim = "\n\n"
def question_of(prompt):
return prompt.split(delim, 1)[1]
uniform_questions = [question_of(r.prompt) for r in uniform_rows]
zipf_questions = [question_of(r.prompt) for r in zipf_rows]
self.assertEqual(uniform_questions, zipf_questions)
# The set of system prompts (which the gen_prompt path generates) must
# also match between modes (set equality, since Zipf reuses prefixes).
def system_of(prompt):
return prompt.split(delim, 1)[0]
self.assertEqual(
set(system_of(r.prompt) for r in uniform_rows),
set(system_of(r.prompt) for r in zipf_rows),
)
def test_zipf_deterministic_per_group_counts(self):
# The per-group counts are deterministic and pinned for a known
# (num_groups, prompts_per_group, alpha, seed) tuple. Any drift in
# the Zipf sampling implementation will trip this assertion.
rows = self._run_gsp(
mode="zipf",
alpha=2.0,
seed=0,
num_groups=4,
prompts_per_group=25,
send_routing_key=True,
ordered=True,
)
self.assertEqual(len(rows), 4 * 25)
# routing_key format is "<uuid8>_<timestamp>_<group_idx>".
per_group = Counter(int(r.routing_key.rsplit("_", 1)[-1]) for r in rows)
# Pinned counts derived from the implementation for
# (N=4, P=25, alpha=2.0, seed=0) using numpy.random.default_rng(seed)
# and rng.choice over _zipf_group_probs(N, alpha).
self.assertEqual(
dict(per_group),
{0: 63, 1: 18, 2: 12, 3: 7},
)
# Independent skew sanity check: rank-1 (hottest) > rank-N (coldest).
self.assertGreater(per_group[0], per_group[3])
def test_zipf_uses_distinct_cache_from_uniform(self):
# The on-disk cache key includes group_distribution and zipf_alpha,
# so uniform mode, zipf alpha=1.0, and zipf alpha=2.0 each get their
# own file. Uniform mode never reads a zipf cache and vice versa.
from sglang.benchmark.datasets import generated_shared_prefix as gsp_mod
fake_home = self.tmpdir_path / "fakehome"
fake_home.mkdir()
common = dict(
num_groups=2,
prompts_per_group=3,
system_prompt_len=4,
question_len=3,
output_len=2,
range_ratio=1.0,
seed=5,
send_routing_key=False,
num_turns=1,
fast_prepare=True,
ordered=True,
)
with patch.object(gsp_mod.Path, "home", return_value=fake_home):
uniform_path = get_gen_prefix_cache_path(
seed=common["seed"],
num_groups=common["num_groups"],
prompts_per_group=common["prompts_per_group"],
system_prompt_len=common["system_prompt_len"],
question_len=common["question_len"],
output_len=common["output_len"],
tokenizer=self.tokenizer,
)
zipf_path_a = get_gen_prefix_cache_path(
seed=common["seed"],
num_groups=common["num_groups"],
prompts_per_group=common["prompts_per_group"],
system_prompt_len=common["system_prompt_len"],
question_len=common["question_len"],
output_len=common["output_len"],
tokenizer=self.tokenizer,
group_distribution="zipf",
zipf_alpha=1.5,
)
zipf_path_b = get_gen_prefix_cache_path(
seed=common["seed"],
num_groups=common["num_groups"],
prompts_per_group=common["prompts_per_group"],
system_prompt_len=common["system_prompt_len"],
question_len=common["question_len"],
output_len=common["output_len"],
tokenizer=self.tokenizer,
group_distribution="zipf",
zipf_alpha=2.0,
)
self.assertNotEqual(uniform_path, zipf_path_a)
self.assertNotEqual(zipf_path_a, zipf_path_b)
# Run each mode; each writes its own cache file.
self._run_gsp(mode="uniform", **common)
self._run_gsp(mode="zipf", alpha=1.5, **common)
self._run_gsp(mode="zipf", alpha=2.0, **common)
self.assertTrue(uniform_path.exists())
self.assertTrue(zipf_path_a.exists())
self.assertTrue(zipf_path_b.exists())
# Sentinel into the uniform cache: zipf must not read it.
sentinel = [DatasetRow(prompt="SENTINEL", prompt_len=1, output_len=1)]
with open(uniform_path, "wb") as f:
pickle.dump(sentinel, f)
zipf_rows = self._run_gsp(mode="zipf", alpha=1.5, **common)
self.assertNotEqual(zipf_rows, sentinel)
# Second zipf call with same args must load from cache (no
# regeneration). Mutate the zipf cache to a sentinel and confirm.
zipf_sentinel = [
DatasetRow(prompt="ZIPF_SENTINEL", prompt_len=1, output_len=1)
]
with open(zipf_path_a, "wb") as f:
pickle.dump(zipf_sentinel, f)
reloaded = self._run_gsp(mode="zipf", alpha=1.5, **common)
self.assertEqual(reloaded, zipf_sentinel)
def test_zipf_total_rows_and_unique_prompts(self):
# Total returned row count under Zipf equals num_groups *
# prompts_per_group (identical to uniform mode) and every prompt
# string is unique even when groups repeat.
rows = self._run_gsp(
mode="zipf",
alpha=2.5,
seed=3,
num_groups=4,
prompts_per_group=10,
send_routing_key=False,
)
self.assertEqual(len(rows), 4 * 10)
self.assertEqual(len({r.prompt for r in rows}), len(rows))
def test_zipf_ordered_preserves_generation_order(self):
# With ordered=True, output preserves the sampled order and matches
# an independently re-derived group sequence from default_rng(seed).
rows = self._run_gsp(
mode="zipf",
alpha=1.5,
seed=21,
num_groups=3,
prompts_per_group=8,
send_routing_key=True,
ordered=True,
)
observed_groups = [int(r.routing_key.rsplit("_", 1)[-1]) for r in rows]
# Independently reproduce the expected group sequence: an isolated
# default_rng(seed) over _zipf_group_probs(N, alpha) sampling
# N * P slots.
expected_rng = np.random.default_rng(21)
expected_probs = _zipf_group_probs(3, 1.5)
expected_groups = expected_rng.choice(
3, size=3 * 8, replace=True, p=expected_probs
).tolist()
self.assertEqual(observed_groups, expected_groups)
def test_zipf_shuffle_path_matches_uniform_shuffle(self):
# When ordered=False, both modes go through random.shuffle on a list
# of equal length, so the same global RNG seed yields the same
# permutation. Verified indirectly: two Zipf calls with the same
# global seed produce identical orderings.
kwargs = dict(
mode="zipf",
alpha=1.2,
seed=8,
num_groups=4,
prompts_per_group=6,
send_routing_key=False,
ordered=False,
global_seed=8,
)
rows_a = self._run_gsp(**kwargs)
rows_b = self._run_gsp(**kwargs)
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
# ------------------------------------------------------------------
# CLI / from_args validation
# ------------------------------------------------------------------
def test_from_args_rejects_invalid_distribution_and_alpha(self):
# Defensive validation in from_args protects in-process callers
# that build a Namespace by hand and bypass the argparse boundary
# in bench_serving.py. Covers: unknown distribution, zipf without
# alpha, uniform with alpha, and non-finite/non-positive alpha.
cases = [
{"gsp_group_distribution": "not-a-distribution", "gsp_zipf_alpha": None},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": None},
{"gsp_group_distribution": "uniform", "gsp_zipf_alpha": 1.0},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": 0.0},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": -0.5},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("nan")},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("inf")},
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("-inf")},
]
for case in cases:
args = make_args(dataset_name="generated-shared-prefix", **case)
with self.assertRaises(ValueError, msg=f"case={case}"):
GeneratedSharedPrefixDataset.from_args(args)
def test_bench_serving_help_and_invalid_choice_argparse(self):
# Subprocess-driven coverage of the live CLI: --help advertises both
# flags with the rank-based Zipf formula and the alpha constraint,
# and argparse rejects an unknown distribution choice.
help_res = subprocess.run(
[sys.executable, "-m", "sglang.bench_serving", "--help"],
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(help_res.returncode, 0, help_res.stderr)
out = help_res.stdout
# Both new flags appear.
self.assertIn("--gsp-group-distribution", out)
self.assertIn("--gsp-zipf-alpha", out)
# Rank-based Zipf formula and alpha constraint are documented.
self.assertIn("1/rank**alpha", out)
self.assertIn("rank starts at 1", out)
self.assertIn("finite float", out)
# Argparse rejects unknown distribution choice.
bad_choice_res = subprocess.run(
[
sys.executable,
"-m",
"sglang.bench_serving",
"--dataset-name",
"generated-shared-prefix",
"--gsp-group-distribution",
"invalid_name",
],
capture_output=True,
text=True,
timeout=90,
)
self.assertNotEqual(bad_choice_res.returncode, 0)
self.assertIn("invalid choice", (bad_choice_res.stderr + bad_choice_res.stdout))
def test_bench_serving_cli_rejects_zipf_without_alpha_before_server(self):
# Malformed CLI combinations (zipf with no alpha) must fail at
# argparse time so users see the GSP-flag error directly, not a
# downstream connection or model-fetch failure.
res = subprocess.run(
[
sys.executable,
"-m",
"sglang.bench_serving",
"--dataset-name",
"generated-shared-prefix",
"--gsp-group-distribution",
"zipf",
"--ready-check-timeout-sec",
"0",
],
capture_output=True,
text=True,
timeout=90,
)
# parser.error() exits with code 2 (argparse convention).
self.assertEqual(res.returncode, 2, res.stderr)
stderr = res.stderr + res.stdout
self.assertIn("--gsp-group-distribution", stderr)
self.assertIn("--gsp-zipf-alpha", stderr)
# The error must mention the GSP flags directly, not a network or
# model-discovery problem masquerading as the failure.
for forbidden in [
"HTTPConnectionPool",
"HTTPSConnectionPool",
"Connection refused",
"Failed to fetch model",
"Traceback",
]:
self.assertNotIn(forbidden, stderr)
def test_bench_serving_cli_rejects_uniform_with_alpha_before_server(self):
# The complementary malformation: uniform distribution with an
# explicit alpha value. Must also fail at argparse time.
res = subprocess.run(
[
sys.executable,
"-m",
"sglang.bench_serving",
"--dataset-name",
"generated-shared-prefix",
"--gsp-group-distribution",
"uniform",
"--gsp-zipf-alpha",
"1.0",
"--ready-check-timeout-sec",
"0",
],
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(res.returncode, 2, res.stderr)
stderr = res.stderr + res.stdout
self.assertIn("--gsp-group-distribution", stderr)
self.assertIn("--gsp-zipf-alpha", stderr)
for forbidden in [
"HTTPConnectionPool",
"HTTPSConnectionPool",
"Connection refused",
"Failed to fetch model",
"Traceback",
]:
self.assertNotIn(forbidden, stderr)
if __name__ == "__main__":
unittest.main()