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

docs_new/docs/developer_guide/bench_serving.mdx

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

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

@@ -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 should_cache and cache_path.exists():
-    if cache_path.exists() and should_cache:
         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[sampled_g]
-        system_prompt = system_prompts[group_idx]
         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(
+        turn_questions = questions[src_g][src_p]
-            range(prompts_per_group), desc="Generating questions", leave=False
+        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
-            turn_questions = questions[group_idx][prompt_idx]
+        prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0]))
-            turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[
+        output_len_val = int(output_lens[src_g, src_p])
-                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])
 
-            input_requests.append(
+        input_requests.append(
-                DatasetRow(
+            DatasetRow(
-                    prompt=full_prompt,
+                prompt=full_prompt,
-                    prompt_len=prompt_len,
+                prompt_len=prompt_len,
-                    output_len=output_len_val,
+                output_len=output_len_val,
-                    routing_key=routing_key,
+                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

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