@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e56f44d9ec
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 15, 08:30 UTC

Support datasets in `vllm bench serve` and sync with benchmark_[serving,datasets].py (#18566)

Michael Goin committed May 27, 2025 at 19:59 UTC e56f44d9ec6f8f07f0d2c7936eea9bb2c0212bf2
3 files changed +692 -101
vllm/benchmarks/datasets.py
+175 -10
@@ -62,6 +62,7 @@ class SampleRequest:
62
63 class BenchmarkDataset(ABC):
64 DEFAULT_SEED = 0
65 + IS_MULTIMODAL = False
66
67 def __init__(
68 self,
@@ -316,13 +317,15 @@ class RandomDataset(BenchmarkDataset):
317 )
318
319 vocab_size = tokenizer.vocab_size
320 + num_special_tokens = tokenizer.num_special_tokens_to_add()
321 + real_input_len = input_len - num_special_tokens
322
323 prefix_token_ids = (np.random.randint(
324 0, vocab_size, size=prefix_len).tolist() if prefix_len > 0 else [])
325
326 # New sampling logic: [X * (1 - b), X * (1 + b)]
324 - input_low = int(input_len * (1 - range_ratio))
325 - input_high = int(input_len * (1 + range_ratio))
327 + input_low = int(real_input_len * (1 - range_ratio))
328 + input_high = int(real_input_len * (1 + range_ratio))
329 output_low = int(output_len * (1 - range_ratio))
330 output_high = int(output_len * (1 + range_ratio))
331
@@ -345,6 +348,17 @@ class RandomDataset(BenchmarkDataset):
348 vocab_size).tolist()
349 token_sequence = prefix_token_ids + inner_seq
350 prompt = tokenizer.decode(token_sequence)
351 + # After decoding the prompt we have to encode and decode it again.
352 + # This is done because in some cases N consecutive tokens
353 + # give a string tokenized into != N number of tokens.
354 + # For example for GPT2Tokenizer:
355 + # [6880, 6881] -> ['Ġcalls', 'here'] ->
356 + # [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
357 + # To avoid uncontrolled change of the prompt length,
358 + # the encoded sequence is truncated before being decode again.
359 + re_encoded_sequence = tokenizer.encode(
360 + prompt, add_special_tokens=False)[:input_lens[i]]
361 + prompt = tokenizer.decode(re_encoded_sequence)
362 total_input_len = prefix_len + int(input_lens[i])
363 requests.append(
364 SampleRequest(
@@ -637,6 +651,7 @@ class ConversationDataset(HuggingFaceDataset):
651 SUPPORTED_DATASET_PATHS = {
652 'lmms-lab/LLaVA-OneVision-Data', 'Aeala/ShareGPT_Vicuna_unfiltered'
653 }
654 + IS_MULTIMODAL = True
655
656 def sample(self,
657 tokenizer: PreTrainedTokenizerBase,
@@ -701,6 +716,7 @@ class VisionArenaDataset(HuggingFaceDataset):
716 "lmarena-ai/vision-arena-bench-v0.1":
717 lambda x: x["turns"][0][0]["content"]
718 }
719 + IS_MULTIMODAL = True
720
721 def sample(
722 self,
@@ -784,6 +800,64 @@ class InstructCoderDataset(HuggingFaceDataset):
800 return sampled_requests
801
802
803 +# -----------------------------------------------------------------------------
804 +# MT-Bench Dataset Implementation
805 +# -----------------------------------------------------------------------------
806 +
807 +
808 +class MTBenchDataset(HuggingFaceDataset):
809 + """
810 + MT-Bench Dataset.
811 + https://huggingface.co/datasets/philschmid/mt-bench
812 +
813 + We create a single turn dataset for MT-Bench.
814 + This is similar to Spec decoding benchmark setup in vLLM
815 + https://github.com/vllm-project/vllm/blob/9d98ab5ec/examples/offline_inference/eagle.py#L14-L18
816 + """ # noqa: E501
817 +
818 + DEFAULT_OUTPUT_LEN = 256 # avg len used in SD bench in vLLM
819 + SUPPORTED_DATASET_PATHS = {
820 + "philschmid/mt-bench",
821 + }
822 +
823 + def sample(
824 + self,
825 + tokenizer: PreTrainedTokenizerBase,
826 + num_requests: int,
827 + output_len: Optional[int] = None,
828 + enable_multimodal_chat: bool = False,
829 + **kwargs,
830 + ) -> list:
831 + output_len = (output_len
832 + if output_len is not None else self.DEFAULT_OUTPUT_LEN)
833 + sampled_requests = []
834 +
835 + for item in self.data:
836 + if len(sampled_requests) >= num_requests:
837 + break
838 + prompt = item["turns"][0]
839 +
840 + # apply template
841 + prompt = tokenizer.apply_chat_template(
842 + [{
843 + "role": "user",
844 + "content": prompt
845 + }],
846 + add_generation_prompt=True,
847 + tokenize=False,
848 + )
849 +
850 + prompt_len = len(tokenizer(prompt).input_ids)
851 + sampled_requests.append(
852 + SampleRequest(
853 + prompt=prompt,
854 + prompt_len=prompt_len,
855 + expected_output_len=output_len,
856 + ))
857 + self.maybe_oversample_requests(sampled_requests, num_requests)
858 + return sampled_requests
859 +
860 +
861 # -----------------------------------------------------------------------------
862 # AIMO Dataset Implementation
863 # -----------------------------------------------------------------------------
@@ -858,18 +932,18 @@ def _format_zeta_prompt(
932 sample: dict,
933 original_start_marker: str = "<|editable_region_start|>") -> dict:
934 """Format the zeta prompt for the Next Edit Prediction (NEP) dataset.
861 -
862 - This function formats examples from the NEP dataset
863 - into prompts and expected outputs. It could be
935 +
936 + This function formats examples from the NEP dataset
937 + into prompts and expected outputs. It could be
938 further extended to support more NEP datasets.
865 -
939 +
940 Args:
867 - sample: The dataset sample containing events,
941 + sample: The dataset sample containing events,
942 inputs, and outputs.
869 - original_start_marker: The marker indicating the
870 - start of the editable region. Defaults to
943 + original_start_marker: The marker indicating the
944 + start of the editable region. Defaults to
945 "<|editable_region_start|>".
872 -
946 +
947 Returns:
948 A dictionary with the formatted prompts and expected outputs.
949 """
@@ -919,3 +993,94 @@ class NextEditPredictionDataset(HuggingFaceDataset):
993 break
994 self.maybe_oversample_requests(samples, num_requests)
995 return samples
996 +
997 +
998 +# -----------------------------------------------------------------------------
999 +# ASR Dataset Implementation
1000 +# -----------------------------------------------------------------------------
1001 +
1002 +
1003 +class ASRDataset(HuggingFaceDataset):
1004 + """
1005 + Dataset class for processing a ASR dataset for transcription.
1006 + Tested on the following set:
1007 +
1008 + +----------------+----------------------------------------+--------------------------+-----------------------------+
1009 + | Dataset | Domain | Speaking Style | hf-subset |
1010 + +----------------+----------------------------------------+--------------------------+-----------------------------+
1011 + | TED-LIUM | TED talks | Oratory | release1, release2, release3|
1012 + | | | | release3-speaker-adaptation |
1013 + | VoxPopuli | European Parliament | Oratory | en, de, it, fr, ... |
1014 + | LibriSpeech | Audiobook | Narrated | "LIUM/tedlium" |
1015 + | GigaSpeech | Audiobook, podcast, YouTube | Narrated, spontaneous | xs, s, m, l, xl, dev, test |
1016 + | SPGISpeech | Financial meetings | Oratory, spontaneous | S, M, L, dev, test |
1017 + | AMI | Meetings | Spontaneous | ihm, sdm |
1018 + +----------------+----------------------------------------+--------------------------+-----------------------------+
1019 +
1020 + """ # noqa: E501
1021 +
1022 + SUPPORTED_DATASET_PATHS = {
1023 + "openslr/librispeech_asr",
1024 + "facebook/voxpopuli",
1025 + "LIUM/tedlium",
1026 + "edinburghcstr/ami",
1027 + "speechcolab/gigaspeech",
1028 + "kensho/spgispeech",
1029 + }
1030 +
1031 + DEFAULT_OUTPUT_LEN = 128
1032 + IS_MULTIMODAL = True
1033 +
1034 + # TODO Whisper-specific. Abstract interface when more models are supported.
1035 + TRANSCRIPTION_PREAMBLE = (
1036 + "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>")
1037 + skip_long_audios: bool = True
1038 +
1039 + def sample(
1040 + self,
1041 + tokenizer: PreTrainedTokenizerBase,
1042 + num_requests: int,
1043 + output_len: Optional[int] = None,
1044 + **kwargs,
1045 + ) -> list:
1046 + try:
1047 + import librosa
1048 + except ImportError as e:
1049 + raise ImportError(
1050 + "librosa is required for ASRDataset. Please install it "
1051 + "using `pip install librosa`.") from e
1052 +
1053 + output_len = (output_len
1054 + if output_len is not None else self.DEFAULT_OUTPUT_LEN)
1055 + prompt = ASRDataset.TRANSCRIPTION_PREAMBLE
1056 + prompt_len = len(tokenizer(prompt).input_ids)
1057 + sampled_requests = []
1058 + skipped = 0
1059 + for item in self.data:
1060 + if len(sampled_requests) >= num_requests:
1061 + break
1062 + audio = item["audio"]
1063 + y, sr = audio["array"], audio["sampling_rate"]
1064 + duration_s = librosa.get_duration(y=y, sr=sr)
1065 + # Whisper max supported duration
1066 + if self.skip_long_audios and duration_s > 30:
1067 + skipped += 1
1068 + continue
1069 +
1070 + mm_content = {"audio": (y, sr)}
1071 + sampled_requests.append(
1072 + SampleRequest(
1073 + prompt=prompt,
1074 + prompt_len=prompt_len,
1075 + expected_output_len=output_len,
1076 + multi_modal_data=mm_content,
1077 + ))
1078 + if skipped:
1079 + logger.warning(
1080 + "%d samples discarded from dataset due to"
1081 + " their length being greater than"
1082 + " what Whisper supports.",
1083 + skipped,
1084 + )
1085 + self.maybe_oversample_requests(sampled_requests, num_requests)
1086 + return sampled_requests
vllm/benchmarks/endpoint_request_func.py
+223 -3
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """The request function for API endpoints."""
3
4 +import io
5 import json
6 import os
7 import sys
@@ -24,11 +25,11 @@ class RequestFuncInput:
25 output_len: int
26 model: str
27 model_name: Optional[str] = None
27 - best_of: int = 1
28 logprobs: Optional[int] = None
29 extra_body: Optional[dict] = None
30 multi_modal_content: Optional[dict] = None
31 ignore_eos: bool = False
32 + language: Optional[str] = None
33
34
35 @dataclass
@@ -71,7 +72,7 @@ async def async_request_openai_completions(
72 if request_func_input.model_name else request_func_input.model,
73 "prompt": request_func_input.prompt,
74 "temperature": 0.0,
74 - "best_of": request_func_input.best_of,
75 + "repetition_penalty": 1.0,
76 "max_tokens": request_func_input.output_len,
77 "logprobs": request_func_input.logprobs,
78 "stream": True,
@@ -154,7 +155,226 @@ async def async_request_openai_completions(
155 return output
156
157
158 +async def async_request_openai_chat_completions(
159 + request_func_input: RequestFuncInput,
160 + pbar: Optional[tqdm] = None,
161 +) -> RequestFuncOutput:
162 + api_url = request_func_input.api_url
163 + assert api_url.endswith(("chat/completions", "profile")), (
164 + "OpenAI Chat Completions API URL must end with 'chat/completions'.")
165 +
166 + async with aiohttp.ClientSession(trust_env=True,
167 + timeout=AIOHTTP_TIMEOUT) as session:
168 + content = [{"type": "text", "text": request_func_input.prompt}]
169 + if request_func_input.multi_modal_content:
170 + content.append(request_func_input.multi_modal_content)
171 + payload = {
172 + "model":
173 + request_func_input.model_name
174 + if request_func_input.model_name else request_func_input.model,
175 + "messages": [
176 + {
177 + "role": "user",
178 + "content": content
179 + },
180 + ],
181 + "temperature":
182 + 0.0,
183 + "max_completion_tokens":
184 + request_func_input.output_len,
185 + "stream":
186 + True,
187 + "stream_options": {
188 + "include_usage": True,
189 + },
190 + }
191 + if request_func_input.ignore_eos:
192 + payload["ignore_eos"] = request_func_input.ignore_eos
193 + if request_func_input.extra_body:
194 + payload.update(request_func_input.extra_body)
195 + headers = {
196 + "Content-Type": "application/json",
197 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
198 + }
199 +
200 + output = RequestFuncOutput()
201 + output.prompt_len = request_func_input.prompt_len
202 +
203 + generated_text = ""
204 + ttft = 0.0
205 + st = time.perf_counter()
206 + most_recent_timestamp = st
207 + try:
208 + async with session.post(url=api_url, json=payload,
209 + headers=headers) as response:
210 + if response.status == 200:
211 + async for chunk_bytes in response.content:
212 + chunk_bytes = chunk_bytes.strip()
213 + if not chunk_bytes:
214 + continue
215 +
216 + chunk = chunk_bytes.decode("utf-8").removeprefix(
217 + "data: ")
218 + if chunk != "[DONE]":
219 + timestamp = time.perf_counter()
220 + data = json.loads(chunk)
221 +
222 + if choices := data.get("choices"):
223 + content = choices[0]["delta"].get("content")
224 + # First token
225 + if ttft == 0.0:
226 + ttft = timestamp - st
227 + output.ttft = ttft
228 +
229 + # Decoding phase
230 + else:
231 + output.itl.append(timestamp -
232 + most_recent_timestamp)
233 +
234 + generated_text += content or ""
235 + elif usage := data.get("usage"):
236 + output.output_tokens = usage.get(
237 + "completion_tokens")
238 +
239 + most_recent_timestamp = timestamp
240 +
241 + output.generated_text = generated_text
242 + output.success = True
243 + output.latency = most_recent_timestamp - st
244 + else:
245 + output.error = response.reason or ""
246 + output.success = False
247 + except Exception:
248 + output.success = False
249 + exc_info = sys.exc_info()
250 + output.error = "".join(traceback.format_exception(*exc_info))
251 +
252 + if pbar:
253 + pbar.update(1)
254 + return output
255 +
256 +
257 +async def async_request_openai_audio(
258 + request_func_input: RequestFuncInput,
259 + pbar: Optional[tqdm] = None,
260 +) -> RequestFuncOutput:
261 + # Lazy import without PlaceholderModule to avoid vllm dep.
262 + import soundfile
263 +
264 + api_url = request_func_input.api_url
265 + assert api_url.endswith(("transcriptions", "translations")), (
266 + "OpenAI Chat Completions API URL must end with 'transcriptions' ")
267 + "or `translations`."
268 +
269 + async with aiohttp.ClientSession(trust_env=True,
270 + timeout=AIOHTTP_TIMEOUT) as session:
271 + content = [{"type": "text", "text": request_func_input.prompt}]
272 + payload = {
273 + "model":
274 + request_func_input.model_name
275 + if request_func_input.model_name else request_func_input.model,
276 + "temperature":
277 + 0.0,
278 + "max_completion_tokens":
279 + request_func_input.output_len,
280 + "stream":
281 + True,
282 + "language":
283 + "en",
284 + # Flattened due to multipart/form-data
285 + "stream_include_usage":
286 + True,
287 + "stream_continuous_usage_stats":
288 + True,
289 + }
290 + if request_func_input.extra_body:
291 + payload.update(request_func_input.extra_body)
292 + headers = {
293 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
294 + }
295 +
296 + # Send audio file
297 + def to_bytes(y, sr):
298 + buffer = io.BytesIO()
299 + soundfile.write(buffer, y, sr, format="WAV")
300 + buffer.seek(0)
301 + return buffer
302 +
303 + with to_bytes(*request_func_input.multi_modal_content["audio"]) as f:
304 + form = aiohttp.FormData()
305 + form.add_field("file", f, content_type="audio/wav")
306 + for key, value in payload.items():
307 + form.add_field(key, str(value))
308 +
309 + output = RequestFuncOutput()
310 + output.prompt_len = request_func_input.prompt_len
311 +
312 + generated_text = ""
313 + ttft = 0.0
314 + st = time.perf_counter()
315 + most_recent_timestamp = st
316 + try:
317 + async with session.post(url=api_url,
318 + data=form,
319 + headers=headers) as response:
320 + if response.status == 200:
321 + async for chunk_bytes in response.content:
322 + chunk_bytes = chunk_bytes.strip()
323 + if not chunk_bytes:
324 + continue
325 +
326 + chunk = chunk_bytes.decode("utf-8").removeprefix(
327 + "data: ")
328 + if chunk != "[DONE]":
329 + timestamp = time.perf_counter()
330 + data = json.loads(chunk)
331 +
332 + if choices := data.get("choices"):
333 + content = choices[0]["delta"].get(
334 + "content")
335 + # First token
336 + if ttft == 0.0:
337 + ttft = timestamp - st
338 + output.ttft = ttft
339 +
340 + # Decoding phase
341 + else:
342 + output.itl.append(
343 + timestamp - most_recent_timestamp)
344 +
345 + generated_text += content or ""
346 + elif usage := data.get("usage"):
347 + output.output_tokens = usage.get(
348 + "completion_tokens")
349 +
350 + most_recent_timestamp = timestamp
351 +
352 + output.generated_text = generated_text
353 + output.success = True
354 + output.latency = most_recent_timestamp - st
355 + else:
356 + output.error = response.reason or ""
357 + output.success = False
358 + except Exception:
359 + output.success = False
360 + exc_info = sys.exc_info()
361 + output.error = "".join(traceback.format_exception(*exc_info))
362 +
363 + if pbar:
364 + pbar.update(1)
365 + return output
366 +
367 +
368 # TODO: Add more request functions for different API protocols.
369 ASYNC_REQUEST_FUNCS = {
159 - "openai-comp": async_request_openai_completions,
370 + "vllm": async_request_openai_completions,
371 + "openai": async_request_openai_completions,
372 + "openai-chat": async_request_openai_chat_completions,
373 + "openai-audio": async_request_openai_audio,
374 }
375 +
376 +OPENAI_COMPATIBLE_BACKENDS = [
377 + k for k, v in ASYNC_REQUEST_FUNCS.items()
378 + if v in (async_request_openai_completions,
379 + async_request_openai_chat_completions)
380 +]
vllm/benchmarks/serve.py
+294 -88
@@ -7,7 +7,7 @@ to launch the vLLM OpenAI API server:
7
8 On the client side, run:
9 vllm bench serve \
10 - --endpoint-type <endpoint_type. Default 'openi-comp'> \
10 + --endpoint-type <endpoint_type. Default 'openai'> \
11 --label <benchmark result label. Default using endpoint_type> \
12 --model <your_model> \
13 --dataset-name <dataset_name. Default 'random'> \
@@ -22,7 +22,7 @@ import os
22 import random
23 import time
24 import warnings
25 -from collections.abc import AsyncGenerator
25 +from collections.abc import AsyncGenerator, Iterable
26 from dataclasses import dataclass
27 from datetime import datetime
28 from typing import Any, Optional
@@ -31,7 +31,14 @@ import numpy as np
31 from tqdm.asyncio import tqdm
32 from transformers import PreTrainedTokenizerBase
33
34 +from vllm.benchmarks.datasets import (AIMODataset, ASRDataset, BurstGPTDataset,
35 + ConversationDataset, HuggingFaceDataset,
36 + InstructCoderDataset, MTBenchDataset,
37 + NextEditPredictionDataset, RandomDataset,
38 + SampleRequest, ShareGPTDataset,
39 + SonnetDataset, VisionArenaDataset)
40 from vllm.benchmarks.endpoint_request_func import (ASYNC_REQUEST_FUNCS,
41 + OPENAI_COMPATIBLE_BACKENDS,
42 RequestFuncInput,
43 RequestFuncOutput)
44 from vllm.benchmarks.utils import (convert_to_pytorch_benchmark_format,
@@ -71,53 +78,18 @@ class BenchmarkMetrics:
78 percentiles_e2el_ms: list[tuple[float, float]]
79
80
74 -def sample_random_requests(
75 - prefix_len: int,
76 - input_len: int,
77 - output_len: int,
78 - num_prompts: int,
79 - range_ratio: float,
80 - tokenizer: PreTrainedTokenizerBase,
81 -) -> list[tuple[str, int, int]]:
82 - prefix_token_ids = np.random.randint(0,
83 - tokenizer.vocab_size,
84 - size=prefix_len).tolist()
85 -
86 - input_lens = np.random.randint(
87 - int(input_len * range_ratio),
88 - input_len + 1,
89 - size=num_prompts,
90 - )
91 - output_lens = np.random.randint(
92 - int(output_len * range_ratio),
93 - output_len + 1,
94 - size=num_prompts,
95 - )
96 - offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
97 - input_requests = []
98 - for i in range(num_prompts):
99 - prompt = tokenizer.decode(prefix_token_ids +
100 - [(offsets[i] + i + j) % tokenizer.vocab_size
101 - for j in range(input_lens[i])])
102 -
103 - input_requests.append((prompt, int(prefix_len + input_lens[i]),
104 - int(output_lens[i]), None))
105 -
106 - return input_requests
107 -
108 -
81 async def get_request(
110 - input_requests: list[tuple[str, int, int]],
82 + input_requests: list[SampleRequest],
83 request_rate: float,
84 burstiness: float = 1.0,
113 -) -> AsyncGenerator[tuple[str, int, int], None]:
85 +) -> AsyncGenerator[SampleRequest, None]:
86 """
87 Asynchronously generates requests at a specified rate
88 with OPTIONAL burstiness.
89
90 Args:
91 input_requests:
120 - A list of input requests, each represented as a tuple.
92 + A list of input requests, each represented as a SampleRequest.
93 request_rate:
94 The rate at which requests are generated (requests/s).
95 burstiness (optional):
@@ -129,7 +101,7 @@ async def get_request(
101 in more bursty requests, while a higher burstiness value
102 (burstiness > 1) results in a more uniform arrival of requests.
103 """
132 - input_requests = iter(input_requests)
104 + input_requests: Iterable[SampleRequest] = iter(input_requests)
105
106 # Calculate scale parameter theta to maintain the desired request_rate.
107 assert burstiness > 0, (
@@ -151,7 +123,7 @@ async def get_request(
123
124
125 def calculate_metrics(
154 - input_requests: list[tuple[str, int, int]],
126 + input_requests: list[SampleRequest],
127 outputs: list[RequestFuncOutput],
128 dur_s: float,
129 tokenizer: PreTrainedTokenizerBase,
@@ -184,7 +156,7 @@ def calculate_metrics(
156 if outputs[i].success:
157 output_len = outputs[i].output_tokens
158
187 - if output_len is None:
159 + if not output_len:
160 # We use the tokenizer to count the number of output tokens
161 # for some serving backends instead of looking at
162 # len(outputs[i].itl) since multiple output tokens may be
@@ -194,7 +166,7 @@ def calculate_metrics(
166 tokenizer(outputs[i].generated_text,
167 add_special_tokens=False).input_ids)
168 actual_output_lens.append(output_len)
197 - total_input += input_requests[i][1]
169 + total_input += input_requests[i].prompt_len
170 tpot = 0
171 if output_len > 1:
172 latency_minus_ttft = outputs[i].latency - outputs[i].ttft
@@ -277,19 +249,19 @@ async def benchmark(
249 model_id: str,
250 model_name: str,
251 tokenizer: PreTrainedTokenizerBase,
280 - input_requests: list[tuple[str, int, int]],
252 + input_requests: list[SampleRequest],
253 logprobs: Optional[int],
282 - best_of: int,
254 request_rate: float,
255 burstiness: float,
256 disable_tqdm: bool,
257 profile: bool,
258 selected_percentile_metrics: list[str],
288 - selected_percentiles: list[str],
259 + selected_percentiles: list[float],
260 ignore_eos: bool,
261 goodput_config_dict: dict[str, float],
262 max_concurrency: Optional[int],
292 - lora_modules: Optional[list[str]],
263 + lora_modules: Optional[Iterable[str]],
264 + extra_body: Optional[dict],
265 ):
266 if endpoint_type in ASYNC_REQUEST_FUNCS:
267 request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -298,11 +270,13 @@ async def benchmark(
270
271 print("Starting initial single prompt test run...")
272 test_prompt, test_prompt_len, test_output_len, test_mm_content = (
301 - input_requests[0])
302 - if endpoint_type != "openai-chat" and test_mm_content is not None:
303 - # multi-modal benchmark is only available on OpenAI Chat endpoint.
304 - raise ValueError("Multi-modal content is only supported on "
305 - "'openai-chat' endpoint_type.")
273 + input_requests[0].prompt,
274 + input_requests[0].prompt_len,
275 + input_requests[0].expected_output_len,
276 + input_requests[0].multi_modal_data,
277 + )
278 +
279 + assert test_mm_content is None or isinstance(test_mm_content, dict)
280 test_input = RequestFuncInput(
281 model=model_id,
282 model_name=model_name,
@@ -311,9 +285,9 @@ async def benchmark(
285 prompt_len=test_prompt_len,
286 output_len=test_output_len,
287 logprobs=logprobs,
314 - best_of=best_of,
288 multi_modal_content=test_mm_content,
289 ignore_eos=ignore_eos,
290 + extra_body=extra_body,
291 )
292
293 test_output = await request_func(request_func_input=test_input)
@@ -338,9 +312,9 @@ async def benchmark(
312 prompt_len=test_prompt_len,
313 output_len=test_output_len,
314 logprobs=logprobs,
341 - best_of=best_of,
315 multi_modal_content=test_mm_content,
343 - ignore_eos=ignore_eos)
316 + ignore_eos=ignore_eos,
317 + extra_body=extra_body)
318 profile_output = await request_func(request_func_input=profile_input)
319 if profile_output.success:
320 print("Profiler started")
@@ -374,7 +348,12 @@ async def benchmark(
348 benchmark_start_time = time.perf_counter()
349 tasks: list[asyncio.Task] = []
350 async for request in get_request(input_requests, request_rate, burstiness):
377 - prompt, prompt_len, output_len, mm_content = request
351 + prompt, prompt_len, output_len, mm_content = (
352 + request.prompt,
353 + request.prompt_len,
354 + request.expected_output_len,
355 + request.multi_modal_data,
356 + )
357 req_model_id, req_model_name = model_id, model_name
358 if lora_modules:
359 req_lora_module = next(lora_modules)
@@ -387,9 +366,9 @@ async def benchmark(
366 prompt_len=prompt_len,
367 output_len=output_len,
368 logprobs=logprobs,
390 - best_of=best_of,
369 multi_modal_content=mm_content,
392 - ignore_eos=ignore_eos)
370 + ignore_eos=ignore_eos,
371 + extra_body=extra_body)
372 tasks.append(
373 asyncio.create_task(
374 limited_request_func(request_func_input=request_func_input,
@@ -405,7 +384,6 @@ async def benchmark(
384 prompt_len=test_prompt_len,
385 output_len=test_output_len,
386 logprobs=logprobs,
408 - best_of=best_of,
387 )
388 profile_output = await request_func(request_func_input=profile_input)
389 if profile_output.success:
@@ -567,7 +545,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
545 parser.add_argument(
546 "--endpoint-type",
547 type=str,
570 - default="openai-comp",
548 + default="openai",
549 choices=list(ASYNC_REQUEST_FUNCS.keys()),
550 )
551 parser.add_argument(
@@ -596,9 +574,16 @@ def add_cli_args(parser: argparse.ArgumentParser):
574 "--dataset-name",
575 type=str,
576 default="random",
599 - choices=["random"],
577 + choices=["sharegpt", "burstgpt", "sonnet", "random", "hf"],
578 help="Name of the dataset to benchmark on.",
579 )
580 + parser.add_argument(
581 + "--dataset-path",
582 + type=str,
583 + default=None,
584 + help="Path to the sharegpt/sonnet dataset. "
585 + "Or the huggingface dataset ID if using HF dataset.",
586 + )
587 parser.add_argument(
588 "--max-concurrency",
589 type=int,
@@ -624,13 +609,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
609 help=
610 "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501
611 )
627 - parser.add_argument(
628 - "--best-of",
629 - type=int,
630 - default=1,
631 - help="Generates `best_of` sequences per prompt and "
632 - "returns the best one.",
633 - )
612 parser.add_argument("--use-beam-search", action="store_true")
613 parser.add_argument(
614 "--num-prompts",
@@ -691,6 +669,17 @@ def add_cli_args(parser: argparse.ArgumentParser):
669 action="store_true",
670 help="Specify to save benchmark results to a json file",
671 )
672 + parser.add_argument(
673 + "--save-detailed",
674 + action="store_true",
675 + help="When saving the results, whether to include per request "
676 + "information such as response, error, ttfs, tpots, etc.",
677 + )
678 + parser.add_argument(
679 + "--append-result",
680 + action="store_true",
681 + help="Append the benchmark result to the existing json file.",
682 + )
683 parser.add_argument(
684 "--metadata",
685 metavar="KEY=VALUE",
@@ -733,6 +722,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
722 default="99",
723 help="Comma-separated list of percentiles for selected metrics. "
724 "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
725 + "Default value is \"99\"."
726 "Use \"--percentile-metrics\" to select metrics.",
727 )
728 parser.add_argument(
@@ -745,7 +735,41 @@ def add_cli_args(parser: argparse.ArgumentParser):
735 "separated by spaces. Allowed request level metric names are "
736 "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of "
737 "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "
748 - "and the blog: https://hao-ai-lab.github.io/blogs/distserve")
738 + "and the blog: https://hao-ai-lab.github.io/blogs/distserve",
739 + )
740 +
741 + # group for dataset specific arguments
742 + sonnet_group = parser.add_argument_group("sonnet dataset options")
743 + sonnet_group.add_argument(
744 + "--sonnet-input-len",
745 + type=int,
746 + default=550,
747 + help=
748 + "Number of input tokens per request, used only for sonnet dataset.",
749 + )
750 + sonnet_group.add_argument(
751 + "--sonnet-output-len",
752 + type=int,
753 + default=150,
754 + help=
755 + "Number of output tokens per request, used only for sonnet dataset.",
756 + )
757 + sonnet_group.add_argument(
758 + "--sonnet-prefix-len",
759 + type=int,
760 + default=200,
761 + help=
762 + "Number of prefix tokens per request, used only for sonnet dataset.",
763 + )
764 +
765 + sharegpt_group = parser.add_argument_group("sharegpt dataset options")
766 + sharegpt_group.add_argument(
767 + "--sharegpt-output-len",
768 + type=int,
769 + default=None,
770 + help="Output length for each request. Overrides the output length "
771 + "from the ShareGPT dataset.",
772 + )
773
774 random_group = parser.add_argument_group("random dataset options")
775 random_group.add_argument(
@@ -765,9 +789,11 @@ def add_cli_args(parser: argparse.ArgumentParser):
789 random_group.add_argument(
790 "--random-range-ratio",
791 type=float,
768 - default=1.0,
769 - help="Range of sampled ratio of input/output length, "
770 - "used only for random sampling.",
792 + default=0.0,
793 + help="Range ratio for sampling input/output length, "
794 + "used only for random sampling. Must be in the range [0, 1) to define "
795 + "a symmetric sampling range"
796 + "[length * (1 - range_ratio), length * (1 + range_ratio)].",
797 )
798 random_group.add_argument(
799 "--random-prefix-len",
@@ -778,6 +804,54 @@ def add_cli_args(parser: argparse.ArgumentParser):
804 " request is [random-prefix-len, "
805 " random-prefix-len + random-prefix-len * random-range-ratio).")
806
807 + hf_group = parser.add_argument_group("hf dataset options")
808 + hf_group.add_argument("--hf-subset",
809 + type=str,
810 + default=None,
811 + help="Subset of the HF dataset.")
812 + hf_group.add_argument("--hf-split",
813 + type=str,
814 + default=None,
815 + help="Split of the HF dataset.")
816 + hf_group.add_argument(
817 + "--hf-output-len",
818 + type=int,
819 + default=None,
820 + help="Output length for each request. Overrides the output lengths "
821 + "from the sampled HF dataset.",
822 + )
823 +
824 + sampling_group = parser.add_argument_group("sampling parameters")
825 + sampling_group.add_argument(
826 + "--top-p",
827 + type=float,
828 + default=None,
829 + help="Top-p sampling parameter. Only has effect on "
830 + "openai-compatible backends.",
831 + )
832 + sampling_group.add_argument(
833 + "--top-k",
834 + type=int,
835 + default=None,
836 + help="Top-k sampling parameter. Only has effect on "
837 + "openai-compatible backends.",
838 + )
839 + sampling_group.add_argument(
840 + "--min-p",
841 + type=float,
842 + default=None,
843 + help="Min-p sampling parameter. Only has effect on "
844 + "openai-compatible backends.",
845 + )
846 + sampling_group.add_argument(
847 + "--temperature",
848 + type=float,
849 + default=None,
850 + help="Temperature sampling parameter. Only has effect on "
851 + "openai-compatible backends. If not specified, default to greedy "
852 + "decoding (i.e. temperature==0.0).",
853 + )
854 +
855 parser.add_argument(
856 '--tokenizer-mode',
857 type=str,
@@ -826,27 +900,142 @@ def main(args: argparse.Namespace):
900 tokenizer = get_tokenizer(tokenizer_id,
901 tokenizer_mode=tokenizer_mode,
902 trust_remote_code=args.trust_remote_code)
829 - # TODO: This should be refactored to use the benchmark_dataset.py
830 - # in later PRs.
903 +
904 if args.dataset_name is None:
905 raise ValueError(
906 "Please specify '--dataset-name' and the corresponding "
907 "'--dataset-path' if required.")
835 - elif args.dataset_name == "random":
836 - input_requests = sample_random_requests(
837 - prefix_len=args.random_prefix_len,
838 - input_len=args.random_input_len,
839 - output_len=args.random_output_len,
840 - num_prompts=args.num_prompts,
841 - range_ratio=args.random_range_ratio,
908 +
909 + if args.dataset_name == "sonnet":
910 + dataset = SonnetDataset(dataset_path=args.dataset_path)
911 + # For the "sonnet" dataset, formatting depends on the backend.
912 + if args.backend == "openai-chat":
913 + input_requests = dataset.sample(
914 + num_requests=args.num_prompts,
915 + input_len=args.sonnet_input_len,
916 + output_len=args.sonnet_output_len,
917 + prefix_len=args.sonnet_prefix_len,
918 + tokenizer=tokenizer,
919 + return_prompt_formatted=False,
920 + )
921 + else:
922 + assert tokenizer.chat_template or tokenizer.default_chat_template, (
923 + "Tokenizer/model must have chat template for sonnet dataset.")
924 + input_requests = dataset.sample(
925 + num_requests=args.num_prompts,
926 + input_len=args.sonnet_input_len,
927 + output_len=args.sonnet_output_len,
928 + prefix_len=args.sonnet_prefix_len,
929 + tokenizer=tokenizer,
930 + return_prompt_formatted=True,
931 + )
932 +
933 + elif args.dataset_name == "hf":
934 + # all following datasets are implemented from the
935 + # HuggingFaceDataset base class
936 + if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
937 + dataset_class = VisionArenaDataset
938 + args.hf_split = "train"
939 + args.hf_subset = None
940 + elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:
941 + dataset_class = InstructCoderDataset
942 + args.hf_split = "train"
943 + elif args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS:
944 + dataset_class = MTBenchDataset
945 + args.hf_split = "train"
946 + elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
947 + dataset_class = ConversationDataset
948 + args.hf_split = "train"
949 + elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
950 + dataset_class = AIMODataset
951 + args.hf_split = "train"
952 + elif args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS: # noqa: E501
953 + dataset_class = NextEditPredictionDataset
954 + args.hf_split = "train"
955 + elif args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS:
956 + dataset_class = ASRDataset
957 + args.hf_split = "train"
958 + else:
959 + supported_datasets = set([
960 + dataset_name for cls in HuggingFaceDataset.__subclasses__()
961 + for dataset_name in cls.SUPPORTED_DATASET_PATHS
962 + ])
963 + raise ValueError(
964 + f"Unsupported dataset path: {args.dataset_path}. "
965 + "Huggingface dataset only supports dataset_path"
966 + f" from one of following: {supported_datasets}. "
967 + "Please consider contributing if you would "
968 + "like to add support for additional dataset formats.")
969 +
970 + if dataset_class.IS_MULTIMODAL and endpoint_type not in [
971 + "openai-chat",
972 + "openai-audio",
973 + ]:
974 + # multi-modal benchmark is only available on OpenAI Chat backend.
975 + raise ValueError(
976 + "Multi-modal content is only supported on 'openai-chat' and "
977 + "'openai-audio' backend.")
978 + input_requests = dataset_class(
979 + dataset_path=args.dataset_path,
980 + dataset_subset=args.hf_subset,
981 + dataset_split=args.hf_split,
982 + random_seed=args.seed,
983 + ).sample(
984 + num_requests=args.num_prompts,
985 tokenizer=tokenizer,
986 + output_len=args.hf_output_len,
987 )
988
989 else:
846 - raise ValueError(f"Unknown dataset: {args.dataset_name}")
847 -
990 + # For datasets that follow a similar structure, use a mapping.
991 + dataset_mapping = {
992 + "sharegpt":
993 + lambda: ShareGPTDataset(random_seed=args.seed,
994 + dataset_path=args.dataset_path).sample(
995 + tokenizer=tokenizer,
996 + num_requests=args.num_prompts,
997 + output_len=args.sharegpt_output_len,
998 + ),
999 + "burstgpt":
1000 + lambda: BurstGPTDataset(random_seed=args.seed,
1001 + dataset_path=args.dataset_path).
1002 + sample(tokenizer=tokenizer, num_requests=args.num_prompts),
1003 + "random":
1004 + lambda: RandomDataset(dataset_path=args.dataset_path).sample(
1005 + tokenizer=tokenizer,
1006 + num_requests=args.num_prompts,
1007 + prefix_len=args.random_prefix_len,
1008 + input_len=args.random_input_len,
1009 + output_len=args.random_output_len,
1010 + range_ratio=args.random_range_ratio,
1011 + ),
1012 + }
1013 +
1014 + try:
1015 + input_requests = dataset_mapping[args.dataset_name]()
1016 + except KeyError as err:
1017 + raise ValueError(f"Unknown dataset: {args.dataset_name}") from err
1018 goodput_config_dict = check_goodput_args(args)
1019
1020 + # Collect the sampling parameters.
1021 + sampling_params = {
1022 + k: v
1023 + for k, v in {
1024 + "top_p": args.top_p,
1025 + "top_k": args.top_k,
1026 + "min_p": args.min_p,
1027 + "temperature": args.temperature,
1028 + }.items() if v is not None
1029 + }
1030 +
1031 + # Sampling parameters are only supported by openai-compatible backend.
1032 + if sampling_params and args.backend not in OPENAI_COMPATIBLE_BACKENDS:
1033 + raise ValueError("Sampling parameters are only supported by "
1034 + "openai-compatible backends.")
1035 +
1036 + if "temperature" not in sampling_params:
1037 + sampling_params["temperature"] = 0.0 # Default to greedy decoding.
1038 +
1039 # Avoid GC processing "static" data - reduce pause times.
1040 gc.collect()
1041 gc.freeze()
@@ -861,7 +1050,6 @@ def main(args: argparse.Namespace):
1050 tokenizer=tokenizer,
1051 input_requests=input_requests,
1052 logprobs=args.logprobs,
864 - best_of=args.best_of,
1053 request_rate=args.request_rate,
1054 burstiness=args.burstiness,
1055 disable_tqdm=args.disable_tqdm,
@@ -874,10 +1062,11 @@ def main(args: argparse.Namespace):
1062 goodput_config_dict=goodput_config_dict,
1063 max_concurrency=args.max_concurrency,
1064 lora_modules=args.lora_modules,
1065 + extra_body=sampling_params,
1066 ))
1067
1068 # Save config and results to json
880 - if args.save_result:
1069 + if args.save_result or args.append_result:
1070 result_json: dict[str, Any] = {}
1071
1072 # Setup
@@ -887,7 +1076,6 @@ def main(args: argparse.Namespace):
1076 result_json["label"] = label
1077 result_json["model_id"] = model_id
1078 result_json["tokenizer_id"] = tokenizer_id
890 - result_json["best_of"] = args.best_of
1079 result_json["num_prompts"] = args.num_prompts
1080
1081 # Metadata
@@ -910,6 +1098,19 @@ def main(args: argparse.Namespace):
1098 # Merge with benchmark result
1099 result_json = {**result_json, **benchmark_result}
1100
1101 + if not args.save_detailed:
1102 + # Remove fields with too many data points
1103 + for field in [
1104 + "input_lens",
1105 + "output_lens",
1106 + "ttfts",
1107 + "itls",
1108 + "generated_texts",
1109 + "errors",
1110 + ]:
1111 + if field in result_json:
1112 + del result_json[field]
1113 +
1114 # Save to file
1115 base_model_id = model_id.split("/")[-1]
1116 max_concurrency_str = (f"-concurrency{args.max_concurrency}"
@@ -920,6 +1121,11 @@ def main(args: argparse.Namespace):
1121 file_name = args.result_filename
1122 if args.result_dir:
1123 file_name = os.path.join(args.result_dir, file_name)
923 - with open(file_name, "w", encoding='utf-8') as outfile:
1124 + with open(file_name,
1125 + mode="a+" if args.append_result else "w",
1126 + encoding="utf-8") as outfile:
1127 + # Append a newline.
1128 + if args.append_result and outfile.tell() != 0:
1129 + outfile.write("\n")
1130 json.dump(result_json, outfile)
1131 save_to_pytorch_benchmark_format(args, result_json, file_name)