@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 893b2affff
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 04:45 UTC

feat: add TxtSlicesDataset to allow sampling slices from txt file for benchmarking (#30156)

Signed-off-by: jdebache <jdebache@nvidia.com>

Julien Debache committed Apr 14, 2026 at 11:20 UTC 893b2affff548bf0f54978ca2ba4545c5b4d5fbe
6 files changed +870 -142
tests/benchmarks/test_sampling_params.py new
+258
@@ -0,0 +1,258 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import numpy as np
5 +import pytest
6 +
7 +from vllm.benchmarks.datasets.utils import get_sampling_params
8 +from vllm.tokenizers import TokenizerLike
9 +
10 +
11 +class _FakeTokenizer(TokenizerLike):
12 + """Minimal tokenizer implementing the TokenizerLike protocol
13 + for testing get_sampling_params."""
14 +
15 + def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
16 + self._vocab_size = vocab_size
17 + self._num_special_tokens = num_special_tokens
18 +
19 + # -- Properties required by TokenizerLike --
20 +
21 + @classmethod
22 + def from_pretrained(cls, path_or_repo_id, *a, **kw): # type: ignore[override]
23 + return cls()
24 +
25 + @property
26 + def vocab_size(self) -> int:
27 + return self._vocab_size
28 +
29 + @property
30 + def all_special_tokens(self) -> list[str]:
31 + return []
32 +
33 + @property
34 + def all_special_ids(self) -> list[int]:
35 + return []
36 +
37 + @property
38 + def bos_token_id(self) -> int:
39 + return 0
40 +
41 + @property
42 + def eos_token_id(self) -> int:
43 + return 1
44 +
45 + @property
46 + def pad_token_id(self) -> int:
47 + return 2
48 +
49 + @property
50 + def is_fast(self) -> bool:
51 + return False
52 +
53 + @property
54 + def max_token_id(self) -> int:
55 + return self._vocab_size - 1
56 +
57 + @property
58 + def max_chars_per_token(self) -> int:
59 + return 4
60 +
61 + @property
62 + def truncation_side(self) -> str:
63 + return "right"
64 +
65 + def num_special_tokens_to_add(self) -> int:
66 + return self._num_special_tokens
67 +
68 + def __call__(self, text, text_pair=None, **kw): # type: ignore[override]
69 + raise NotImplementedError
70 +
71 + def get_vocab(self) -> dict[str, int]:
72 + return {}
73 +
74 + def get_added_vocab(self) -> dict[str, int]:
75 + return {}
76 +
77 + def encode(self, text, **kw) -> list[int]: # type: ignore[override]
78 + raise NotImplementedError
79 +
80 + def apply_chat_template(self, messages, **kw): # type: ignore[override]
81 + raise NotImplementedError
82 +
83 + def convert_tokens_to_ids(self, tokens): # type: ignore[override]
84 + raise NotImplementedError
85 +
86 + def convert_tokens_to_string(self, tokens: list[str]) -> str:
87 + raise NotImplementedError
88 +
89 + def decode(self, ids, skip_special_tokens: bool = False) -> str: # type: ignore[override]
90 + raise NotImplementedError
91 +
92 + def convert_ids_to_tokens( # type: ignore[override]
93 + self, ids, skip_special_tokens: bool = False
94 + ) -> list[str]:
95 + raise NotImplementedError
96 +
97 +
98 +class TestGetSamplingParams:
99 + """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``."""
100 +
101 + # -- helpers --
102 +
103 + @staticmethod
104 + def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer:
105 + return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special)
106 +
107 + # -- return shape / dtype --
108 +
109 + def test_returns_three_arrays(self):
110 + rng = np.random.default_rng(0)
111 + result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok())
112 + assert len(result) == 3
113 + for arr in result:
114 + assert isinstance(arr, np.ndarray)
115 +
116 + @pytest.mark.parametrize("n", [1, 10, 100])
117 + def test_output_length_matches_num_requests(self, n: int):
118 + rng = np.random.default_rng(42)
119 + input_lens, output_lens, offsets = get_sampling_params(
120 + rng, n, 0.0, 64, 32, self._tok()
121 + )
122 + assert input_lens.shape == (n,)
123 + assert output_lens.shape == (n,)
124 + assert offsets.shape == (n,)
125 +
126 + # -- fixed lengths (range_ratio = 0) --
127 +
128 + def test_zero_range_ratio_gives_constant_lengths(self):
129 + rng = np.random.default_rng(7)
130 + input_lens, output_lens, _ = get_sampling_params(
131 + rng, 20, 0.0, 128, 64, self._tok()
132 + )
133 + assert np.all(input_lens == 128)
134 + assert np.all(output_lens == 64)
135 +
136 + def test_special_tokens_subtracted_from_input_only(self):
137 + rng = np.random.default_rng(7)
138 + input_lens, output_lens, _ = get_sampling_params(
139 + rng, 10, 0.0, 100, 50, self._tok(num_special=4)
140 + )
141 + # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96
142 + assert np.all(input_lens == 96)
143 + # special tokens are not subtracted from output length
144 + assert np.all(output_lens == 50)
145 +
146 + # -- range ratios --
147 +
148 + def test_input_range_bounds(self):
149 + rng = np.random.default_rng(0)
150 + ratio = 0.5
151 + base = 200
152 + input_lens, _, _ = get_sampling_params(
153 + rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok()
154 + )
155 + lo = int(np.floor(base * (1 - ratio)))
156 + hi = int(np.ceil(base * (1 + ratio)))
157 + assert np.all(input_lens >= lo)
158 + assert np.all(input_lens <= hi)
159 +
160 + def test_output_range_bounds(self):
161 + rng = np.random.default_rng(0)
162 + ratio = 0.3
163 + base = 100
164 + _, output_lens, _ = get_sampling_params(
165 + rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok()
166 + )
167 + lo = max(1, int(np.floor(base * (1 - ratio))))
168 + hi = int(np.ceil(base * (1 + ratio)))
169 + assert np.all(output_lens >= lo)
170 + assert np.all(output_lens <= hi)
171 +
172 + def test_output_low_clamped_to_one(self):
173 + """Even with a high ratio that would push output_low to 0,
174 + the function clamps it to 1."""
175 + rng = np.random.default_rng(0)
176 + # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
177 + _, output_lens, _ = get_sampling_params(
178 + rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok()
179 + )
180 + assert np.all(output_lens >= 1)
181 +
182 + # -- offsets bounded by vocab_size --
183 +
184 + @pytest.mark.parametrize("vocab", [100, 32000, 128256])
185 + def test_offsets_within_vocab(self, vocab: int):
186 + rng = np.random.default_rng(0)
187 + _, _, offsets = get_sampling_params(
188 + rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab)
189 + )
190 + assert np.all(offsets >= 0)
191 + assert np.all(offsets < vocab)
192 +
193 + # -- reproducibility --
194 +
195 + def test_same_seed_same_results(self):
196 + tok = self._tok()
197 + rr = {"input": 0.3, "output": 0.2}
198 + a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
199 + b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
200 + for arr_a, arr_b in zip(a, b):
201 + np.testing.assert_array_equal(arr_a, arr_b)
202 +
203 + def test_different_seed_different_results(self):
204 + tok = self._tok()
205 + rr = {"input": 0.3, "output": 0.2}
206 + a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok)
207 + b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok)
208 + # Extremely unlikely all three arrays match with different seeds
209 + assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b))
210 +
211 + # -- validation / error paths --
212 +
213 + @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
214 + def test_invalid_input_range_ratio(self, bad_ratio: float):
215 + rng = np.random.default_rng(0)
216 + with pytest.raises(ValueError, match="input_range_ratio"):
217 + get_sampling_params(
218 + rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok()
219 + )
220 +
221 + @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
222 + def test_invalid_output_range_ratio(self, bad_ratio: float):
223 + rng = np.random.default_rng(0)
224 + with pytest.raises(ValueError, match="output_range_ratio"):
225 + get_sampling_params(
226 + rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok()
227 + )
228 +
229 + def test_invalid_dict_missing_keys(self):
230 + rng = np.random.default_rng(0)
231 + with pytest.raises(ValueError, match="input.*output"):
232 + get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok())
233 +
234 + def test_input_len_zero_with_special_tokens(self):
235 + """input_len < num_special_tokens → real_input_len = 0, which is fine
236 + (range [0, 0])."""
237 + rng = np.random.default_rng(0)
238 + input_lens, _, _ = get_sampling_params(
239 + rng, 5, 0.0, 5, 50, self._tok(num_special=10)
240 + )
241 + # real_input_len = max(0, 5 - 10) = 0
242 + assert np.all(input_lens == 0)
243 +
244 + # -- edge cases --
245 +
246 + def test_single_request(self):
247 + rng = np.random.default_rng(0)
248 + i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok())
249 + assert i.shape == (1,)
250 + assert o.shape == (1,)
251 + assert off.shape == (1,)
252 +
253 + def test_large_num_requests(self):
254 + rng = np.random.default_rng(0)
255 + i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok())
256 + assert i.shape == (10_000,)
257 + assert o.shape == (10_000,)
258 + assert off.shape == (10_000,)
tests/benchmarks/test_txt_slices_dataset.py new
+68
@@ -0,0 +1,68 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import json
4 +from pathlib import Path
5 +
6 +import pytest
7 +from transformers import AutoTokenizer, PreTrainedTokenizerBase
8 +
9 +from vllm.benchmarks.datasets import CustomDataset
10 +from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
11 +
12 +
13 +@pytest.fixture(scope="session")
14 +def hf_tokenizer() -> PreTrainedTokenizerBase:
15 + # Use a small, commonly available tokenizer
16 + return AutoTokenizer.from_pretrained("gpt2")
17 +
18 +
19 +text_content = """
20 +Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
21 +incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
22 +exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
23 +Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
24 +nulla pariatur. Excepteur sint occaecat cupidatat non proident,
25 +sunt in culpa qui officia deserunt mollit anim id est laborum.
26 +"""
27 +
28 +
29 +@pytest.mark.benchmark
30 +def test_create_txt_slices_jsonl(
31 + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
32 +) -> None:
33 + """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
34 + txt_path = tmp_path / "input.txt"
35 + jsonl_path = tmp_path / "input.txt.jsonl"
36 +
37 + txt_path.write_text(text_content)
38 +
39 + create_txt_slices_jsonl(
40 + input_path=str(txt_path),
41 + output_path=str(jsonl_path),
42 + tokenizer_name="gpt2",
43 + num_prompts=10,
44 + input_len=10,
45 + output_len=10,
46 + )
47 +
48 + # Verify the JSONL file is valid and has the expected structure
49 + records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
50 +
51 + assert len(records) == 10
52 + for record in records:
53 + assert "prompt" in record
54 + assert "output_tokens" in record
55 + assert isinstance(record["prompt"], str)
56 + assert record["output_tokens"] == 10
57 +
58 + # Verify the JSONL file can be loaded by CustomDataset
59 + dataset = CustomDataset(dataset_path=str(jsonl_path))
60 + samples = dataset.sample(
61 + tokenizer=hf_tokenizer,
62 + num_requests=10,
63 + output_len=10,
64 + skip_chat_template=True,
65 + )
66 +
67 + assert len(samples) == 10
68 + assert all(sample.expected_output_len == 10 for sample in samples)
vllm/benchmarks/datasets/__init__.py new
+84
@@ -0,0 +1,84 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.benchmarks.datasets.datasets import (
5 + DEFAULT_NUM_PROMPTS,
6 + AIMODataset,
7 + ASRDataset,
8 + BenchmarkDataset,
9 + BlazeditDataset,
10 + BurstGPTDataset,
11 + ConversationDataset,
12 + CustomDataset,
13 + CustomMMDataset,
14 + HuggingFaceDataset,
15 + InstructCoderDataset,
16 + MLPerfDataset,
17 + MMStarDataset,
18 + MMVUDataset,
19 + MTBenchDataset,
20 + MultiModalConversationDataset,
21 + NextEditPredictionDataset,
22 + PrefixRepetitionRandomDataset,
23 + RandomDataset,
24 + RandomDatasetForReranking,
25 + RandomMultiModalDataset,
26 + SampleRequest,
27 + ShareGPTDataset,
28 + SonnetDataset,
29 + SpecBench,
30 + VisionArenaDataset,
31 + add_dataset_parser,
32 + add_random_dataset_base_args,
33 + add_random_multimodal_dataset_args,
34 + gen_prompt_decode_to_target_len,
35 + get_samples,
36 + is_valid_sequence,
37 + lora_path_on_disk,
38 + lora_tokenizer_cache,
39 + process_image,
40 + process_video,
41 + zeta_prompt,
42 +)
43 +from vllm.benchmarks.datasets.utils import RangeRatio
44 +
45 +__all__ = [
46 + "DEFAULT_NUM_PROMPTS",
47 + "AIMODataset",
48 + "ASRDataset",
49 + "BenchmarkDataset",
50 + "BlazeditDataset",
51 + "BurstGPTDataset",
52 + "ConversationDataset",
53 + "CustomDataset",
54 + "CustomMMDataset",
55 + "HuggingFaceDataset",
56 + "InstructCoderDataset",
57 + "MLPerfDataset",
58 + "MMStarDataset",
59 + "MMVUDataset",
60 + "MTBenchDataset",
61 + "MultiModalConversationDataset",
62 + "NextEditPredictionDataset",
63 + "PrefixRepetitionRandomDataset",
64 + "RandomDataset",
65 + "RandomDatasetForReranking",
66 + "RandomMultiModalDataset",
67 + "SampleRequest",
68 + "ShareGPTDataset",
69 + "SonnetDataset",
70 + "SpecBench",
71 + "VisionArenaDataset",
72 + "add_dataset_parser",
73 + "add_random_dataset_base_args",
74 + "add_random_multimodal_dataset_args",
75 + "gen_prompt_decode_to_target_len",
76 + "get_samples",
77 + "is_valid_sequence",
78 + "lora_path_on_disk",
79 + "lora_tokenizer_cache",
80 + "process_image",
81 + "process_video",
82 + "RangeRatio",
83 + "zeta_prompt",
84 +]
vllm/benchmarks/datasets/create_txt_slices_dataset.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Convert a plain-text file (local path or URL) into a JSONL dataset
5 +compatible with ``CustomDataset`` (``--dataset-name custom``), by
6 +randomly slicing the tokenized text into prompts.
7 +
8 +Each line of the output JSONL contains a ``prompt`` (decoded from a random
9 +slice of the tokenized source text) and an ``output_tokens`` count.
10 +
11 +Usage
12 +-----
13 +::
14 +
15 + python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\
16 + --input sonnet.txt \\
17 + --output sonnet_dataset.jsonl \\
18 + --tokenizer gpt2 \\
19 + --num-prompts 1000 \\
20 + --input-len 1024 \\
21 + --output-len 128
22 +
23 +The resulting JSONL file can then be used with the serving benchmark::
24 +
25 + python -m vllm.benchmarks.serve \\
26 + --dataset-name custom \\
27 + --dataset-path sonnet_dataset.jsonl \\
28 + ...
29 +"""
30 +
31 +from __future__ import annotations
32 +
33 +import argparse
34 +import json
35 +import logging
36 +import random
37 +import urllib.request
38 +
39 +import numpy as np
40 +from transformers import AutoTokenizer
41 +
42 +from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params
43 +
44 +logger = logging.getLogger(__name__)
45 +
46 +
47 +def load_text(path: str) -> str:
48 + """Load text from a local file or URL."""
49 + if path.startswith(("http://", "https://")):
50 + with urllib.request.urlopen(path) as response:
51 + return response.read().decode("utf-8")
52 + with open(path, encoding="utf-8") as f:
53 + return f.read()
54 +
55 +
56 +def create_txt_slices_jsonl(
57 + *,
58 + input_path: str,
59 + output_path: str,
60 + tokenizer_name: str,
61 + num_prompts: int,
62 + input_len: int,
63 + output_len: int,
64 + range_ratio: RangeRatio = 0.0,
65 + seed: int = 0,
66 + trust_remote_code: bool = False,
67 +) -> None:
68 + """Read *input_path*, slice it into prompts, and write JSONL to
69 + *output_path*."""
70 +
71 + tokenizer = AutoTokenizer.from_pretrained(
72 + tokenizer_name, trust_remote_code=trust_remote_code
73 + )
74 +
75 + text = load_text(input_path)
76 + if not text:
77 + raise ValueError("The text file is empty and cannot be sampled from.")
78 +
79 + token_ids = tokenizer(text, add_special_tokens=False).input_ids
80 + if not token_ids:
81 + raise ValueError("Tokenizing the text produced zero tokens; cannot sample.")
82 +
83 + rng_np = np.random.default_rng(seed)
84 + rng_py = random.Random(seed)
85 +
86 + input_lens, output_lens, _ = get_sampling_params(
87 + rng_np,
88 + num_prompts,
89 + range_ratio,
90 + input_len,
91 + output_len,
92 + tokenizer,
93 + )
94 +
95 + num_available_tokens = len(token_ids)
96 +
97 + records: list[dict[str, object]] = []
98 + for i in range(num_prompts):
99 + req_input_len = int(input_lens[i])
100 + req_output_len = int(output_lens[i])
101 +
102 + # Randomly select a start position and slice with cycling
103 + start_pos = rng_py.randint(0, num_available_tokens - 1)
104 + prompt_token_ids = [
105 + token_ids[(start_pos + j) % num_available_tokens]
106 + for j in range(req_input_len)
107 + ]
108 + prompt = tokenizer.decode(prompt_token_ids, skip_special_tokens=False)
109 +
110 + records.append({"prompt": prompt, "output_tokens": req_output_len})
111 +
112 + with open(output_path, "w", encoding="utf-8") as f:
113 + for record in records:
114 + f.write(json.dumps(record, ensure_ascii=False) + "\n")
115 +
116 + logger.info(
117 + "Wrote %d prompts to %s",
118 + len(records),
119 + output_path,
120 + )
121 +
122 +
123 +def main(argv: list[str] | None = None) -> None:
124 + parser = argparse.ArgumentParser(
125 + description="Convert a plain-text file into a JSONL dataset "
126 + "for CustomDataset (--dataset-name custom).",
127 + )
128 + parser.add_argument(
129 + "--input",
130 + required=True,
131 + help="Path or URL to the source text file.",
132 + )
133 + parser.add_argument(
134 + "--output",
135 + required=True,
136 + help="Path for the output JSONL file.",
137 + )
138 + parser.add_argument(
139 + "--tokenizer",
140 + required=True,
141 + help="HuggingFace tokenizer name or path.",
142 + )
143 + parser.add_argument(
144 + "--num-prompts",
145 + type=int,
146 + default=1000,
147 + help="Number of prompt samples to generate (default: 1000).",
148 + )
149 + parser.add_argument(
150 + "--input-len",
151 + type=int,
152 + default=1024,
153 + help="Target number of input tokens per prompt (default: 1024).",
154 + )
155 + parser.add_argument(
156 + "--output-len",
157 + type=int,
158 + default=128,
159 + help="Target number of output tokens per prompt (default: 128).",
160 + )
161 + parser.add_argument(
162 + "--range-ratio",
163 + type=str,
164 + default="0.0",
165 + help="Range ratio for input/output length sampling (default: 0.0). "
166 + "A single float applies to both ISL and OSL. "
167 + 'A JSON dict like \'{"input": 0.3, "output": 0.5}\' sets them '
168 + "independently. Values must be in [0, 1).",
169 + )
170 + parser.add_argument(
171 + "--seed",
172 + type=int,
173 + default=0,
174 + help="Random seed for reproducibility (default: 0).",
175 + )
176 + parser.add_argument(
177 + "--trust-remote-code",
178 + action="store_true",
179 + help="Trust remote code from HuggingFace.",
180 + )
181 +
182 + args = parser.parse_args(argv)
183 +
184 + logging.basicConfig(level=logging.INFO)
185 +
186 + # Parse --range-ratio: try float first, then JSON dict.
187 + range_ratio: RangeRatio
188 + try:
189 + range_ratio = float(args.range_ratio)
190 + except ValueError:
191 + import json as _json
192 +
193 + range_ratio = _json.loads(args.range_ratio)
194 +
195 + create_txt_slices_jsonl(
196 + input_path=args.input,
197 + output_path=args.output,
198 + tokenizer_name=args.tokenizer,
199 + num_prompts=args.num_prompts,
200 + input_len=args.input_len,
201 + output_len=args.output_len,
202 + range_ratio=range_ratio,
203 + seed=args.seed,
204 + trust_remote_code=args.trust_remote_code,
205 + )
206 +
207 +
208 +if __name__ == "__main__":
209 + main()
vllm/benchmarks/datasets/datasets.py renamed
+150 -142
@@ -22,8 +22,7 @@ import random
22 from abc import ABC, abstractmethod
23 from collections.abc import Callable, Iterator, Mapping
24 from contextlib import suppress
25 -from copy import deepcopy
26 -from dataclasses import dataclass
25 +from dataclasses import dataclass, replace
26 from functools import cache
27 from io import BytesIO
28 from tempfile import NamedTemporaryFile
@@ -35,6 +34,11 @@ from huggingface_hub import snapshot_download
34 from PIL import Image
35 from typing_extensions import deprecated
36
37 +from vllm.benchmarks.datasets.utils import (
38 + RangeRatio,
39 + _resolve_range_ratios,
40 + get_sampling_params,
41 +)
42 from vllm.inputs import MultiModalDataDict
43 from vllm.lora.request import LoRARequest
44 from vllm.lora.utils import get_adapter_absolute_path
@@ -60,10 +64,6 @@ logger = logging.getLogger(__name__)
64
65 DEFAULT_NUM_PROMPTS = 1000
66
63 -# -----------------------------------------------------------------------------
64 -# Data Classes
65 -# -----------------------------------------------------------------------------
66 -
67
68 @dataclass
69 class SampleRequest:
@@ -71,9 +71,9 @@ class SampleRequest:
71 Represents a single inference request for benchmarking.
72 """
73
74 - prompt: str | list[str]
74 + prompt: str | list[str] | list[dict]
75 prompt_len: int
76 - expected_output_len: int
76 + expected_output_len: int | None
77 multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
78 lora_request: LoRARequest | None = None
79 request_id: str | None = None
@@ -110,7 +110,7 @@ class BenchmarkDataset(ABC):
110 # default seed.
111 self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
112 self.disable_shuffle = disable_shuffle
113 - self.data = None
113 + self.data: Any | None = None
114
115 def apply_multimodal_chat_transformation(
116 self,
@@ -249,6 +249,7 @@ class BenchmarkDataset(ABC):
249 num_requests: int,
250 request_id_prefix: str = "",
251 no_oversample: bool = False,
252 + **kwargs,
253 ) -> list[SampleRequest]:
254 """
255 Abstract method to generate sample requests from the dataset.
@@ -296,8 +297,10 @@ class BenchmarkDataset(ABC):
297 needed = num_requests - len(requests)
298 additional = []
299 for i in range(needed):
299 - req = deepcopy(random.choice(requests))
300 - req.request_id = request_id_prefix + str(len(requests) + i)
300 + req = replace(
301 + random.choice(requests),
302 + request_id=request_id_prefix + str(len(requests) + i),
303 + )
304 additional.append(req)
305 requests.extend(additional)
306 logger.info("Oversampled requests to reach %d total samples.", num_requests)
@@ -533,7 +536,7 @@ class RandomDataset(BenchmarkDataset):
536 request_id_prefix: str = "",
537 no_oversample: bool = False,
538 prefix_len: int = DEFAULT_PREFIX_LEN,
536 - range_ratio: float = DEFAULT_RANGE_RATIO,
539 + range_ratio: RangeRatio = DEFAULT_RANGE_RATIO,
540 input_len: int = DEFAULT_INPUT_LEN,
541 output_len: int = DEFAULT_OUTPUT_LEN,
542 batchsize: int = 1,
@@ -542,24 +545,33 @@ class RandomDataset(BenchmarkDataset):
545 lora_assignment: str = "random",
546 **kwargs,
547 ) -> list[SampleRequest]:
545 - # validate total input tokens (prefix + sampled) is at least 1.
548 + resolved_input_rr, _ = _resolve_range_ratios(range_ratio)
549 +
550 num_special = int(tokenizer.num_special_tokens_to_add())
551 real_input_len = max(0, int(input_len) - num_special)
548 - min_sampled_input = math.floor(real_input_len * (1.0 - float(range_ratio)))
552 + min_sampled_input = math.floor(
553 + real_input_len * (1.0 - float(resolved_input_rr))
554 + )
555 min_total_input = int(prefix_len) + min_sampled_input
556 if min_total_input < 1:
557 raise ValueError(
558 "--random-input-len is too small: with tokenizer special "
553 - f"tokens {num_special} and --random-range-ratio {range_ratio}, "
559 + f"tokens {num_special} and "
560 + f"input range ratio {resolved_input_rr}, "
561 "the minimum possible total input tokens (prefix + sampled) is "
562 f"{min_total_input}. Increase --random-input-len and/or "
556 - "--random-prefix-len, or decrease --random-range-ratio so that "
557 - "prefix_len + floor(max(0, random_input_len - num_special)) "
558 - "* (1 - range_ratio) >= 1."
559 - )
560 -
561 - input_lens, output_lens, offsets = self.get_sampling_params(
562 - num_requests, range_ratio, input_len, output_len, tokenizer
563 + "--random-prefix-len, or decrease the input range ratio "
564 + "so that prefix_len + floor(max(0, random_input_len - "
565 + "num_special)) * (1 - input_range_ratio) >= 1."
566 + )
567 +
568 + input_lens, output_lens, offsets = get_sampling_params(
569 + self._rng,
570 + num_requests,
571 + range_ratio,
572 + input_len,
573 + output_len,
574 + tokenizer,
575 )
576
577 vocab_size = tokenizer.vocab_size
@@ -661,55 +673,6 @@ class RandomDataset(BenchmarkDataset):
673 )
674 return adjusted_tokens
675
664 - def get_sampling_params(
665 - self,
666 - num_requests: int,
667 - range_ratio: float,
668 - input_len: int,
669 - output_len: int,
670 - tokenizer: TokenizerLike,
671 - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
672 - """
673 - Get the sampling parameters for the dataset.
674 - """
675 - # Enforce range_ratio < 1
676 - if not (0.0 <= range_ratio < 1.0):
677 - raise ValueError("range_ratio must be in [0, 1).")
678 - num_special_tokens = int(tokenizer.num_special_tokens_to_add())
679 - real_input_len = max(0, int(input_len) - num_special_tokens)
680 - # Bounds use floor for low and ceil for high
681 - input_low = math.floor(real_input_len * (1 - range_ratio))
682 - input_high = math.ceil(real_input_len * (1 + range_ratio))
683 - output_low = math.floor(output_len * (1 - range_ratio))
684 - output_high = math.ceil(output_len * (1 + range_ratio))
685 - # Ensure the lower bound for output length is at least 1 to
686 - # prevent sampling 0 tokens.
687 - output_low = max(output_low, 1)
688 - output_high = max(output_high, 1)
689 -
690 - if input_low > input_high:
691 - raise ValueError(
692 - f"Invalid input sampling interval: low={input_low} > high={input_high}"
693 - )
694 - if output_low > output_high:
695 - raise ValueError(
696 - "Invalid output sampling interval: "
697 - f"low={output_low} > high={output_high}"
698 - )
699 -
700 - logger.info(
701 - "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
702 - input_low,
703 - input_high,
704 - output_low,
705 - output_high,
706 - )
707 -
708 - input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests)
709 - output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests)
710 - offsets = self._rng.integers(0, tokenizer.vocab_size, size=num_requests)
711 - return input_lens, output_lens, offsets
712 -
676 def generate_token_sequence(
677 self,
678 *,
@@ -776,8 +739,11 @@ class RandomDatasetForReranking(RandomDataset):
739 tokenizer: TokenizerLike,
740 num_requests: int,
741 request_id_prefix: str = "",
779 - range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
742 + no_oversample: bool = False,
743 + prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
744 + range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
745 input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
746 + output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
747 batchsize: int = 1,
748 is_reranker: bool = True,
749 **kwargs,
@@ -786,8 +752,13 @@ class RandomDatasetForReranking(RandomDataset):
752
753 query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len
754
789 - query_lens, _, query_offsets = self.get_sampling_params(
790 - 1, range_ratio, query_len_param, 0, tokenizer
755 + query_lens, _, query_offsets = get_sampling_params(
756 + self._rng,
757 + 1,
758 + range_ratio,
759 + query_len_param,
760 + 0,
761 + tokenizer,
762 )
763
764 query_len = int(query_lens[0])
@@ -800,8 +771,13 @@ class RandomDatasetForReranking(RandomDataset):
771 else:
772 doc_len_param = input_len - query_len - n_sep_tokens
773
803 - doc_lens, _, doc_offsets = self.get_sampling_params(
804 - num_requests, range_ratio, doc_len_param, 0, tokenizer
774 + doc_lens, _, doc_offsets = get_sampling_params(
775 + self._rng,
776 + num_requests,
777 + range_ratio,
778 + doc_len_param,
779 + 0,
780 + tokenizer,
781 )
782
783 vocab_size = tokenizer.vocab_size
@@ -1175,9 +1151,10 @@ class RandomMultiModalDataset(RandomDataset):
1151 request_id_prefix: str = "",
1152 no_oversample: bool = False,
1153 prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
1178 - range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO,
1154 + range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
1155 input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
1156 output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
1157 + batchsize: int = 1,
1158 limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT,
1159 base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST,
1160 num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,
@@ -1187,9 +1164,18 @@ class RandomMultiModalDataset(RandomDataset):
1164 enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT,
1165 **kwargs,
1166 ) -> list[SampleRequest]:
1190 - # Get the sampling parameters for the dataset
1191 - input_lens, output_lens, offsets = self.get_sampling_params(
1192 - num_requests, range_ratio, input_len, output_len, tokenizer
1167 + if batchsize != 1:
1168 + raise NotImplementedError(
1169 + "batchsize > 1 is not supported for RandomMultiModalDataset."
1170 + )
1171 +
1172 + input_lens, output_lens, offsets = get_sampling_params(
1173 + self._rng,
1174 + num_requests,
1175 + range_ratio,
1176 + input_len,
1177 + output_len,
1178 + tokenizer,
1179 )
1180
1181 (
@@ -1326,16 +1312,16 @@ class ShareGPTDataset(BenchmarkDataset):
1312 self,
1313 tokenizer: TokenizerLike,
1314 num_requests: int,
1315 + request_id_prefix: str = "",
1316 + no_oversample: bool = False,
1317 lora_path: str | None = None,
1318 max_loras: int | None = None,
1319 output_len: int | None = None,
1320 enable_multimodal_chat: bool = False,
1333 - request_id_prefix: str = "",
1334 - no_oversample: bool = False,
1321 lora_assignment: str = "random",
1322 **kwargs,
1337 - ) -> list:
1338 - samples: list = []
1323 + ) -> list[SampleRequest]:
1324 + samples: list[SampleRequest] = []
1325 ind = 0
1326 for entry in self.data:
1327 if len(samples) >= num_requests:
@@ -1449,8 +1435,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1435 type=str,
1436 default=None,
1437 action=_ValidateDatasetArgs,
1452 - help="Path to the sharegpt/sonnet dataset. "
1453 - "Or the huggingface dataset ID if using HF dataset.",
1438 + help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
1439 + "using HF dataset.",
1440 )
1441 parser.add_argument(
1442 "--no-oversample",
@@ -1648,12 +1634,12 @@ def add_random_dataset_base_args(
1634 )
1635 parser_or_group.add_argument(
1636 "--random-range-ratio",
1651 - type=float,
1652 - default=0.0,
1637 + type=str,
1638 + default="0.0",
1639 help="Range ratio for sampling input/output length, "
1654 - "used only for random sampling. Must be in the range [0, 1) to define "
1655 - "a symmetric sampling range"
1656 - "[length * (1 - range_ratio), length * (1 + range_ratio)].",
1640 + "used only for random sampling. A single float applies to both "
1641 + 'ISL and OSL. A JSON dict like \'{"input": 0.3, "output": 0.5}\' '
1642 + "sets them independently. Values must be in [0, 1).",
1643 )
1644 parser_or_group.add_argument(
1645 "--random-prefix-len",
@@ -1786,10 +1772,25 @@ def add_random_multimodal_dataset_args(
1772 )
1773
1774
1775 +def _parse_range_ratio(value: str) -> RangeRatio:
1776 + """Parse a ``--random-range-ratio`` CLI string.
1777 +
1778 + Accepts either a plain float (``"0.3"``) or a JSON dict
1779 + (``'{"input": 0.3, "output": 0.5}'``).
1780 + """
1781 + try:
1782 + return float(value)
1783 + except ValueError:
1784 + return json.loads(value)
1785 +
1786 +
1787 def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1788 if not hasattr(args, "request_id_prefix"):
1789 args.request_id_prefix = ""
1790
1791 + if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str):
1792 + args.random_range_ratio = _parse_range_ratio(args.random_range_ratio)
1793 +
1794 if args.dataset_name == "custom":
1795 dataset = CustomDataset(
1796 dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
@@ -2120,7 +2121,7 @@ class CustomDataset(BenchmarkDataset):
2121 # This will be the standardized format which load_data()
2122 # has to convert into depending on the filetype of dataset_path.
2123 # sample() will assume this standardized format of self.data
2123 - self.data = []
2124 + self.data: list[dict] = []
2125
2126 # Load the JSONL file
2127 if self.dataset_path.endswith(".jsonl"):
@@ -2149,15 +2150,15 @@ class CustomDataset(BenchmarkDataset):
2150 self,
2151 tokenizer: TokenizerLike,
2152 num_requests: int,
2153 + request_id_prefix: str = "",
2154 + no_oversample: bool = False,
2155 lora_path: str | None = None,
2156 max_loras: int | None = None,
2157 output_len: int | None = None,
2158 enable_multimodal_chat: bool = False,
2159 skip_chat_template: bool = False,
2157 - request_id_prefix: str = "",
2158 - no_oversample: bool = False,
2160 **kwargs,
2160 - ) -> list:
2161 + ) -> list[SampleRequest]:
2162 # load all data if needed
2163 self.num_available_samples = len(self.data)
2164 if num_requests <= 0:
@@ -2168,7 +2169,7 @@ class CustomDataset(BenchmarkDataset):
2169 num_requests,
2170 )
2171
2171 - sampled_requests = []
2172 + sampled_requests: list[SampleRequest] = []
2173 for i, item in enumerate(self.data):
2174 if len(sampled_requests) >= num_requests:
2175 break
@@ -2252,7 +2253,7 @@ class CustomMMDataset(CustomDataset):
2253 request_id_prefix: str = "",
2254 no_oversample: bool = False,
2255 **kwargs,
2255 - ) -> list:
2256 + ) -> list[SampleRequest]:
2257 # load all data if needed
2258 self.num_available_samples = len(self.data)
2259 if num_requests <= 0:
@@ -2340,9 +2341,13 @@ class SpecBench(CustomDataset):
2341 if not getattr(self, "disable_shuffle", False):
2342 random.shuffle(self.data)
2343
2343 - def sample(self, **kwargs) -> list:
2344 + def sample(
2345 + **kwargs,
2346 + ) -> list[SampleRequest]:
2347 # leverage CustomDataset sample
2345 - return super().sample(**kwargs)
2348 + return super().sample(
2349 + **kwargs,
2350 + )
2351
2352
2353 # -----------------------------------------------------------------------------
@@ -2381,14 +2386,14 @@ class SonnetDataset(BenchmarkDataset):
2386 self,
2387 tokenizer: TokenizerLike,
2388 num_requests: int,
2389 + request_id_prefix: str = "",
2390 + no_oversample: bool = False,
2391 prefix_len: int = DEFAULT_PREFIX_LEN,
2392 input_len: int = DEFAULT_INPUT_LEN,
2393 output_len: int = DEFAULT_OUTPUT_LEN,
2394 return_prompt_formatted: bool = False,
2388 - request_id_prefix: str = "",
2389 - no_oversample: bool = False,
2395 **kwargs,
2391 - ) -> list:
2396 + ) -> list[SampleRequest]:
2397 # Calculate average token length for a poem line.
2398 tokenized_lines = [tokenizer(line).input_ids for line in self.data]
2399 avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
@@ -2411,7 +2416,7 @@ class SonnetDataset(BenchmarkDataset):
2416 num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)
2417 prefix_lines = self.data[:num_prefix_lines]
2418
2414 - samples = []
2419 + samples: list[SampleRequest] = []
2420 ind = 0
2421 while len(samples) < num_requests:
2422 extra_lines = random.choices(
@@ -2482,11 +2487,11 @@ class BurstGPTDataset(BenchmarkDataset):
2487 self,
2488 tokenizer: TokenizerLike,
2489 num_requests: int,
2485 - max_loras: int | None = None,
2486 - lora_path: str | None = None,
2490 request_id_prefix: str = "",
2491 no_oversample: bool = False,
2492 lora_assignment: str = "random",
2493 + max_loras: int | None = None,
2494 + lora_path: str | None = None,
2495 **kwargs,
2496 ) -> list[SampleRequest]:
2497 samples = []
@@ -2574,15 +2579,15 @@ class ConversationDataset(HuggingFaceDataset):
2579 self,
2580 tokenizer: TokenizerLike,
2581 num_requests: int,
2577 - output_len: int | None = None,
2578 - enable_multimodal_chat: bool = False,
2582 request_id_prefix: str = "",
2583 no_oversample: bool = False,
2584 + output_len: int | None = None,
2585 + enable_multimodal_chat: bool = False,
2586 **kwargs,
2582 - ) -> list:
2587 + ) -> list[SampleRequest]:
2588 # Filter examples with at least 2 conversations
2589 filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2585 - sampled_requests = []
2590 + sampled_requests: list[SampleRequest] = []
2591 ind = 0
2592 dynamic_output = output_len is None
2593
@@ -2634,15 +2639,15 @@ class MultiModalConversationDataset(HuggingFaceDataset):
2639 self,
2640 tokenizer: TokenizerLike,
2641 num_requests: int,
2637 - output_len: int | None = None,
2638 - enable_multimodal_chat: bool = False,
2642 request_id_prefix: str = "",
2643 no_oversample: bool = False,
2644 + output_len: int | None = None,
2645 + enable_multimodal_chat: bool = False,
2646 **kwargs,
2642 - ) -> list:
2647 + ) -> list[SampleRequest]:
2648 # Filter examples with at least 2 conversations
2649 filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2645 - sampled_requests = []
2650 + sampled_requests: list[SampleRequest] = []
2651 ind = 0
2652 dynamic_output = output_len is None
2653
@@ -2703,12 +2708,12 @@ class VisionArenaDataset(HuggingFaceDataset):
2708 self,
2709 tokenizer: TokenizerLike,
2710 num_requests: int,
2706 - output_len: int | None = None,
2707 - enable_multimodal_chat: bool = False,
2711 request_id_prefix: str = "",
2712 no_oversample: bool = False,
2713 + output_len: int | None = None,
2714 + enable_multimodal_chat: bool = False,
2715 **kwargs,
2711 - ) -> list:
2716 + ) -> list[SampleRequest]:
2717 parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
2718 if parser_fn is None:
2719 raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2753,9 +2758,11 @@ class MMVUDataset(HuggingFaceDataset):
2758
2759 DEFAULT_OUTPUT_LEN = 128
2760 SUPPORTED_DATASET_PATHS = {
2756 - "yale-nlp/MMVU": lambda x: x["question"]
2757 - + " "
2758 - + (" ".join(f"{k}.{v}" for k, v in x["choices"].items())),
2761 + "yale-nlp/MMVU": lambda x: (
2762 + x["question"]
2763 + + " "
2764 + + (" ".join(f"{k}.{v}" for k, v in x["choices"].items()))
2765 + ),
2766 }
2767
2768 def __init__(self, **kwargs) -> None:
@@ -2770,12 +2777,12 @@ class MMVUDataset(HuggingFaceDataset):
2777 self,
2778 tokenizer: TokenizerLike,
2779 num_requests: int,
2773 - output_len: int | None = None,
2774 - enable_multimodal_chat: bool = False,
2780 request_id_prefix: str = "",
2781 no_oversample: bool = False,
2782 + output_len: int | None = None,
2783 + enable_multimodal_chat: bool = False,
2784 **kwargs,
2778 - ) -> list:
2785 + ) -> list[SampleRequest]:
2786 parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
2787 if parser_fn is None:
2788 raise ValueError(f"Unsupported dataset path: {self.hf_name}")
@@ -2838,15 +2845,15 @@ class InstructCoderDataset(HuggingFaceDataset):
2845 self,
2846 tokenizer: TokenizerLike,
2847 num_requests: int,
2848 + request_id_prefix: str = "",
2849 + no_oversample: bool = False,
2850 output_len: int | None = None,
2851 enable_multimodal_chat: bool = False,
2852 skip_chat_template: bool = False,
2844 - request_id_prefix: str = "",
2845 - no_oversample: bool = False,
2853 **kwargs,
2854 ) -> list[SampleRequest]:
2855 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2849 - sampled_requests = []
2856 + sampled_requests: list[SampleRequest] = []
2857 for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
2858 # apply template
2859 if not skip_chat_template:
@@ -2903,15 +2910,15 @@ class MTBenchDataset(HuggingFaceDataset):
2910 self,
2911 tokenizer: TokenizerLike,
2912 num_requests: int,
2913 + request_id_prefix: str = "",
2914 + no_oversample: bool = False,
2915 output_len: int | None = None,
2916 enable_multimodal_chat: bool = False,
2917 skip_chat_template: bool = False,
2909 - request_id_prefix: str = "",
2910 - no_oversample: bool = False,
2918 **kwargs,
2912 - ) -> list:
2919 + ) -> list[SampleRequest]:
2920 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2914 - sampled_requests = []
2921 + sampled_requests: list[SampleRequest] = []
2922
2923 for i, item in enumerate(self.data):
2924 if len(sampled_requests) >= num_requests:
@@ -2976,7 +2983,7 @@ class BlazeditDataset(HuggingFaceDataset):
2983 min_distance: float = 0.0,
2984 max_distance: float = 1.0,
2985 **kwargs,
2979 - ) -> list:
2986 + ) -> list[SampleRequest]:
2987 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2988 sampled_requests = []
2989
@@ -3050,12 +3057,12 @@ class AIMODataset(HuggingFaceDataset):
3057 self,
3058 tokenizer: TokenizerLike,
3059 num_requests: int,
3053 - output_len: int | None = None,
3060 request_id_prefix: str = "",
3061 no_oversample: bool = False,
3062 + output_len: int | None = None,
3063 **kwargs,
3057 - ) -> list:
3058 - sampled_requests = []
3064 + ) -> list[SampleRequest]:
3065 + sampled_requests: list[SampleRequest] = []
3066 ind = 0
3067 dynamic_output = output_len is None
3068
@@ -3228,18 +3235,18 @@ class ASRDataset(HuggingFaceDataset):
3235 self,
3236 tokenizer: TokenizerLike,
3237 num_requests: int,
3231 - output_len: int | None = None,
3238 request_id_prefix: str = "",
3239 no_oversample: bool = False,
3240 + output_len: int | None = None,
3241 **kwargs,
3235 - ) -> list:
3242 + ) -> list[SampleRequest]:
3243 output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
3244 if "openai" in getattr(tokenizer, "name_or_path", ""):
3245 prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
3246 else:
3247 prompt = ""
3248 prompt_len = len(tokenizer(prompt).input_ids)
3242 - sampled_requests = []
3249 + sampled_requests: list[SampleRequest] = []
3250 ind = 0
3251 skipped = 0
3252 asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")
@@ -3326,9 +3333,9 @@ class MLPerfDataset(HuggingFaceDataset):
3333 self,
3334 tokenizer: TokenizerLike,
3335 num_requests: int,
3329 - output_len: int | None = None,
3336 request_id_prefix: str = "",
3337 no_oversample: bool = False,
3338 + output_len: int | None = None,
3339 **kwargs,
3340 ) -> list[SampleRequest]:
3341 # Force dynamic output length based on reference completion.
@@ -3405,12 +3412,12 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
3412 self,
3413 tokenizer: TokenizerLike,
3414 num_requests: int,
3415 + request_id_prefix: str = "",
3416 + no_oversample: bool = False,
3417 prefix_len: int = DEFAULT_PREFIX_LEN,
3418 suffix_len: int = DEFAULT_SUFFIX_LEN,
3419 num_prefixes: int = DEFAULT_NUM_PREFIXES,
3420 output_len: int = DEFAULT_OUTPUT_LEN,
3412 - request_id_prefix: str = "",
3413 - no_oversample: bool = False,
3421 **kwargs,
3422 ) -> list[SampleRequest]:
3423 vocab_size = tokenizer.vocab_size
@@ -3421,7 +3428,7 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):
3428 f"to num_prefixes ({num_prefixes})"
3429 )
3430
3424 - def _generate_exact_length_tokens(target_length: int) -> list[int]:
3431 + def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
3432 """Generate tokens that decode and re-encode to exactly
3433 target_length."""
3434 # Generate random tokens
@@ -3491,10 +3498,10 @@ class MMStarDataset(HuggingFaceDataset):
3498 self,
3499 tokenizer: TokenizerLike,
3500 num_requests: int,
3494 - output_len: int | None = None,
3495 - enable_multimodal_chat: bool = False,
3501 request_id_prefix: str = "",
3502 no_oversample: bool = False,
3503 + output_len: int | None = None,
3504 + enable_multimodal_chat: bool = False,
3505 **kwargs,
3506 ) -> list[SampleRequest]:
3507 # If --hf-output-len is not set, use the default output length.
@@ -3516,6 +3523,7 @@ class MMStarDataset(HuggingFaceDataset):
3523 # if enable_multimodal_chat is False).
3524 prompt_len = len(tokenizer(question_text).input_ids)
3525
3526 + prompt: str | list[dict]
3527 if enable_multimodal_chat:
3528 # If multimodal content should be embedded in the chat message,
3529 # convert to [{"role":"user","content":[...]}]
vllm/benchmarks/datasets/utils.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Shared utilities for benchmark dataset sampling.
5 +"""
6 +
7 +import logging
8 +import math
9 +
10 +import numpy as np
11 +
12 +from vllm.tokenizers import TokenizerLike
13 +
14 +logger = logging.getLogger(__name__)
15 +
16 +# Type alias: a single float applies to both ISL and OSL; a dict allows
17 +# specifying them independently via ``{"input": …, "output": …}``.
18 +RangeRatio = float | dict[str, float]
19 +
20 +
21 +def _resolve_range_ratios(
22 + range_ratio: RangeRatio,
23 +) -> tuple[float, float]:
24 + """Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*.
25 +
26 + *range_ratio* is either a single float (used for both input and output)
27 + or a dict with ``"input"`` and ``"output"`` keys.
28 + """
29 + if isinstance(range_ratio, dict):
30 + try:
31 + return float(range_ratio["input"]), float(range_ratio["output"])
32 + except KeyError as exc:
33 + raise ValueError(
34 + "When range_ratio is a dict it must contain 'input' and "
35 + f"'output' keys, got: {sorted(range_ratio)}"
36 + ) from exc
37 + ratio = float(range_ratio)
38 + return ratio, ratio
39 +
40 +
41 +def get_sampling_params(
42 + rng: np.random.Generator,
43 + num_requests: int,
44 + range_ratio: RangeRatio,
45 + input_len: int,
46 + output_len: int,
47 + tokenizer: TokenizerLike,
48 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
49 + """
50 + Sample per-request input/output token lengths and vocab offsets.
51 +
52 + Lengths are drawn uniformly from integer ranges around the configured
53 + means, controlled by *range_ratio*. It may be a single ``float``
54 + (applied to both input and output) or a ``dict`` with ``"input"`` and
55 + ``"output"`` keys for independent control.
56 +
57 + Tokenizer special tokens are subtracted from ``input_len`` before
58 + computing the sampling interval.
59 +
60 + Returns:
61 + (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of
62 + shape ``(num_requests,)``.
63 + """
64 + input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio)
65 +
66 + if not (0.0 <= input_range_ratio < 1.0):
67 + raise ValueError("input_range_ratio must be in [0, 1).")
68 + if not (0.0 <= output_range_ratio < 1.0):
69 + raise ValueError("output_range_ratio must be in [0, 1).")
70 + num_special_tokens = int(tokenizer.num_special_tokens_to_add())
71 + real_input_len = max(0, int(input_len) - num_special_tokens)
72 + input_low = math.floor(real_input_len * (1 - input_range_ratio))
73 + input_high = math.ceil(real_input_len * (1 + input_range_ratio))
74 + output_low = math.floor(output_len * (1 - output_range_ratio))
75 + output_high = math.ceil(output_len * (1 + output_range_ratio))
76 + # Ensure the lower bound for output length is at least 1 to
77 + # prevent sampling 0 tokens.
78 + output_low = max(output_low, 1)
79 + output_high = max(output_high, 1)
80 +
81 + if input_low > input_high:
82 + raise ValueError(
83 + f"Invalid input sampling interval: low={input_low} > high={input_high}"
84 + )
85 + if output_low > output_high:
86 + raise ValueError(
87 + f"Invalid output sampling interval: low={output_low} > high={output_high}"
88 + )
89 +
90 + logger.info(
91 + "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
92 + input_low,
93 + input_high,
94 + output_low,
95 + output_high,
96 + )
97 +
98 + input_lens = rng.integers(input_low, input_high + 1, size=num_requests)
99 + output_lens = rng.integers(output_low, output_high + 1, size=num_requests)
100 + offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests)
101 + return input_lens, output_lens, offsets