bench_serving.py 54.8 KB
Newer Older
zhyncs's avatar
zhyncs committed
1
2
# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py
# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py
3

Ying Sheng's avatar
Ying Sheng committed
4
"""
5
Benchmark online serving with dynamic requests.
Ying Sheng's avatar
Ying Sheng committed
6
7

Usage:
8
python3 -m sglang.bench_serving --backend sglang --num-prompt 10
Ying Sheng's avatar
Ying Sheng committed
9

10
python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5
Ying Sheng's avatar
Ying Sheng committed
11
"""
zhyncs's avatar
zhyncs committed
12
13
14
15
16

import argparse
import asyncio
import json
import os
17
import pickle
zhyncs's avatar
zhyncs committed
18
19
20
21
22
23
import random
import resource
import sys
import time
import traceback
import warnings
24
from argparse import ArgumentParser
zhyncs's avatar
zhyncs committed
25
from dataclasses import dataclass, field
26
from datetime import datetime
27
from pathlib import Path
28
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
zhyncs's avatar
zhyncs committed
29
30
31
32
33
34
35
36
37
38
39
40
41

import aiohttp
import numpy as np
import requests
from tqdm.asyncio import tqdm
from transformers import (
    AutoTokenizer,
    PreTrainedTokenizer,
    PreTrainedTokenizerBase,
    PreTrainedTokenizerFast,
)

AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
42
ASSISTANT_SUFFIX = "Assistant:"
zhyncs's avatar
zhyncs committed
43

44
45
global args

zhyncs's avatar
zhyncs committed
46
47
48
49
50
51
52
53

@dataclass
class RequestFuncInput:
    prompt: str
    api_url: str
    prompt_len: int
    output_len: int
    model: str
54
    lora_name: str
55
    extra_request_body: Dict[str, Any]
zhyncs's avatar
zhyncs committed
56
57
58
59
60
61
62
63
64
65
66


@dataclass
class RequestFuncOutput:
    generated_text: str = ""
    success: bool = False
    latency: float = 0.0
    ttft: float = 0.0  # Time to first token
    itl: List[float] = field(default_factory=list)  # List of inter-token latencies
    prompt_len: int = 0
    error: str = ""
67
    output_len: int = 0
zhyncs's avatar
zhyncs committed
68
69
70
71
72
73


def remove_prefix(text: str, prefix: str) -> str:
    return text[len(prefix) :] if text.startswith(prefix) else text


74
75
76
77
def remove_suffix(text: str, suffix: str) -> str:
    return text[: -len(suffix)] if text.endswith(suffix) else text


78
79
80
81
82
83
84
85
def get_auth_headers() -> Dict[str, str]:
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        return {"Authorization": f"Bearer {api_key}"}
    else:
        return {}


86
# trt llm does not support ignore_eos
87
88
89
90
91
92
93
94
95
96
97
98
# https://github.com/triton-inference-server/tensorrtllm_backend/issues/505
async def async_request_trt_llm(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    assert api_url.endswith("generate_stream")

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        payload = {
            "accumulate_tokens": True,
            "text_input": request_func_input.prompt,
zhyncs's avatar
zhyncs committed
99
            "temperature": 0.000001,
100
101
102
            "top_p": 1.0,
            "max_tokens": request_func_input.output_len,
            "stream": True,
Ying Sheng's avatar
Ying Sheng committed
103
104
            "min_length": request_func_input.output_len,
            "end_id": 1048576,
105
            **request_func_input.extra_request_body,
106
        }
107
108
109
        if args.disable_ignore_eos:
            del payload["min_length"]
            del payload["end_id"]
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        ttft = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
        try:
            async with session.post(url=api_url, json=payload) as response:
                if response.status == 200:
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:")

                        data = json.loads(chunk)
                        output.generated_text += data["text_output"]
                        timestamp = time.perf_counter()
                        # First token
                        if ttft == 0.0:
Xu Song's avatar
Xu Song committed
131
                            ttft = timestamp - st
132
133
134
135
136
137
138
139
140
141
                            output.ttft = ttft

                        # Decoding phase
                        else:
                            output.itl.append(timestamp - most_recent_timestamp)

                        most_recent_timestamp = timestamp

                    output.latency = most_recent_timestamp - st
                    output.success = True
Ying Sheng's avatar
Ying Sheng committed
142
                    output.output_len = request_func_input.output_len
143
144
145
146
147
148
149
150
151
152
153
154
155
156

                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))

        if pbar:
            pbar.update(1)
        return output


zhyncs's avatar
zhyncs committed
157
158
159
160
161
162
163
164
165
166
# set ignore_eos True by default
async def async_request_openai_completions(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    assert api_url.endswith(
        "completions"
    ), "OpenAI Completions API URL must end with 'completions'."

Lianmin Zheng's avatar
Lianmin Zheng committed
167
168
    prompt = request_func_input.prompt

zhyncs's avatar
zhyncs committed
169
170
171
    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        payload = {
            "model": request_func_input.model,
Lianmin Zheng's avatar
Lianmin Zheng committed
172
            "prompt": prompt,
zhyncs's avatar
zhyncs committed
173
174
175
            "temperature": 0.0,
            "best_of": 1,
            "max_tokens": request_func_input.output_len,
176
            "stream": not args.disable_stream,
177
            "ignore_eos": not args.disable_ignore_eos,
178
            **request_func_input.extra_request_body,
zhyncs's avatar
zhyncs committed
179
        }
180
        headers = get_auth_headers()
zhyncs's avatar
zhyncs committed
181
182
183
184
185

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        generated_text = ""
186
        output_len = request_func_input.output_len
zhyncs's avatar
zhyncs committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
        ttft = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
        try:
            async with session.post(
                url=api_url, json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
201
                        latency = time.perf_counter() - st
zhyncs's avatar
zhyncs committed
202
                        if chunk == "[DONE]":
203
                            pass
zhyncs's avatar
zhyncs committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
                        else:
                            data = json.loads(chunk)

                            # NOTE: Some completion API might have a last
                            # usage summary response without a token so we
                            # want to check a token was generated
                            if data["choices"][0]["text"]:
                                timestamp = time.perf_counter()
                                # First token
                                if ttft == 0.0:
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft

                                # Decoding phase
218
219
                                else:
                                    output.itl.append(timestamp - most_recent_timestamp)
zhyncs's avatar
zhyncs committed
220
221
222

                                most_recent_timestamp = timestamp
                                generated_text += data["choices"][0]["text"]
Lzhang-hub's avatar
Lzhang-hub committed
223
                                output_len = (data.get("usage") or {}).get(
224
225
                                    "completion_tokens", output_len
                                )
zhyncs's avatar
zhyncs committed
226
227
228
229

                    output.generated_text = generated_text
                    output.success = True
                    output.latency = latency
230
                    output.output_len = output_len
zhyncs's avatar
zhyncs committed
231
232
233
234
235
236
237
238
239
240
241
242
243
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))

    if pbar:
        pbar.update(1)
    return output


244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
async def async_request_truss(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url

    prompt = request_func_input.prompt

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        payload = {
            "model": request_func_input.model,
            "prompt": prompt,
            "temperature": 0.0,
            "best_of": 1,
            "max_tokens": request_func_input.output_len,
            "stream": not args.disable_stream,
            "ignore_eos": not args.disable_ignore_eos,
            **request_func_input.extra_request_body,
        }
263
        headers = get_auth_headers()
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        generated_text = ""
        ttft = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
        try:
            async with session.post(
                url=api_url, json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
                        latency = time.perf_counter() - st
                        if chunk == "[DONE]":
                            pass
                        else:
                            data = json.loads(chunk)

                            # NOTE: Some completion API might have a last
                            # usage summary response without a token so we
                            # want to check a token was generated
                            if data["choices"][0]["delta"]["content"]:
                                timestamp = time.perf_counter()
                                # First token
                                if ttft == 0.0:
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft

                                # Decoding phase
                                else:
                                    output.itl.append(timestamp - most_recent_timestamp)

                                most_recent_timestamp = timestamp
                                generated_text += data["choices"][0]["delta"]["content"]

                    output.generated_text = generated_text
                    output.success = True
                    output.latency = latency
                    output.output_len = request_func_input.output_len
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))

    if pbar:
        pbar.update(1)
    return output


323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
async def async_request_sglang_generate(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    prompt = request_func_input.prompt

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        payload = {
            "text": prompt,
            "sampling_params": {
                "temperature": 0.0,
                "max_new_tokens": request_func_input.output_len,
                "ignore_eos": not args.disable_ignore_eos,
            },
            "stream": not args.disable_stream,
339
            "lora_path": request_func_input.lora_name,
340
341
            "return_logprob": args.return_logprob,
            "logprob_start_len": -1,
342
343
            **request_func_input.extra_request_body,
        }
344
        headers = get_auth_headers()
345
346
347
348
349

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        generated_text = ""
350
        output_len = request_func_input.output_len
351
352
353
        ttft = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
354
        last_output_len = 0
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        try:
            async with session.post(
                url=api_url, json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue
                        # print(chunk_bytes)

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
                        latency = time.perf_counter() - st
                        if chunk == "[DONE]":
                            pass
                        else:
                            data = json.loads(chunk)

                            # NOTE: Some completion API might have a last
                            # usage summary response without a token so we
                            # want to check a token was generated
                            if data["text"]:
                                timestamp = time.perf_counter()
378
379
380
                                generated_text = data["text"]
                                output_len = data["meta_info"]["completion_tokens"]

381
382
383
384
385
386
387
                                # First token
                                if ttft == 0.0:
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft

                                # Decoding phase
                                else:
388
389
390
391
392
393
394
                                    num_new_tokens = output_len - last_output_len
                                    if num_new_tokens == 0:
                                        continue
                                    adjust_itl = (
                                        timestamp - most_recent_timestamp
                                    ) / num_new_tokens
                                    output.itl.extend([adjust_itl] * num_new_tokens)
395
396

                                most_recent_timestamp = timestamp
Lianmin Zheng's avatar
Lianmin Zheng committed
397
                                last_output_len = output_len
398
399
400
401

                    output.generated_text = generated_text
                    output.success = True
                    output.latency = latency
402
                    output.output_len = output_len
403
404
405
406
407
408
409
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
410
            print(f"{output.error=}")
411
412
413
414
415
416

    if pbar:
        pbar.update(1)
    return output


417
async def async_request_gserver(
Lianmin Zheng's avatar
Lianmin Zheng committed
418
419
420
421
422
423
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    raise NotImplementedError()


424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
async def async_request_profile(api_url: str) -> RequestFuncOutput:
    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        output = RequestFuncOutput()
        try:
            async with session.post(url=api_url) as response:
                if response.status == 200:
                    output.success = True
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))

    return output


zhyncs's avatar
zhyncs committed
442
def get_model(pretrained_model_name_or_path: str) -> str:
443
    if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true":
zhyncs's avatar
zhyncs committed
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
        import huggingface_hub.constants
        from modelscope import snapshot_download

        model_path = snapshot_download(
            model_id=pretrained_model_name_or_path,
            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
            ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"],
        )

        return model_path
    return pretrained_model_name_or_path


def get_tokenizer(
    pretrained_model_name_or_path: str,
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
Lianmin Zheng's avatar
Lianmin Zheng committed
460
461
462
463
464
465
466
    if pretrained_model_name_or_path.endswith(
        ".json"
    ) or pretrained_model_name_or_path.endswith(".model"):
        from sglang.srt.hf_transformers_utils import get_tokenizer

        return get_tokenizer(pretrained_model_name_or_path)

zhyncs's avatar
zhyncs committed
467
468
469
470
471
472
473
474
475
    if pretrained_model_name_or_path is not None and not os.path.exists(
        pretrained_model_name_or_path
    ):
        pretrained_model_name_or_path = get_model(pretrained_model_name_or_path)
    return AutoTokenizer.from_pretrained(
        pretrained_model_name_or_path, trust_remote_code=True
    )


476
477
478
479
480
481
482
def get_dataset(args, tokenizer):
    if args.dataset_name == "sharegpt":
        input_requests = sample_sharegpt_requests(
            dataset_path=args.dataset_path,
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
            fixed_output_len=args.sharegpt_output_len,
483
            context_len=args.sharegpt_context_len,
484
            prompt_suffix=args.prompt_suffix,
485
            apply_chat_template=args.apply_chat_template,
486
487
488
489
490
491
492
493
494
495
496
497
        )
    elif args.dataset_name == "random":
        input_requests = sample_random_requests(
            input_len=args.random_input_len,
            output_len=args.random_output_len,
            num_prompts=args.num_prompts,
            range_ratio=args.random_range_ratio,
            tokenizer=tokenizer,
            dataset_path=args.dataset_path,
        )
    elif args.dataset_name == "generated-shared-prefix":
        input_requests = sample_generated_shared_prefix_requests(
498
499
500
501
502
            num_groups=args.gsp_num_groups,
            prompts_per_group=args.gsp_prompts_per_group,
            system_prompt_len=args.gsp_system_prompt_len,
            question_len=args.gsp_question_len,
            output_len=args.gsp_output_len,
503
            tokenizer=tokenizer,
504
            args=args,
505
506
507
508
509
510
        )
    else:
        raise ValueError(f"Unknown dataset: {args.dataset_name}")
    return input_requests


zhyncs's avatar
zhyncs committed
511
ASYNC_REQUEST_FUNCS = {
512
513
514
    "sglang": async_request_sglang_generate,
    "sglang-native": async_request_sglang_generate,
    "sglang-oai": async_request_openai_completions,
zhyncs's avatar
zhyncs committed
515
516
    "vllm": async_request_openai_completions,
    "lmdeploy": async_request_openai_completions,
517
    "trt": async_request_trt_llm,
518
    "gserver": async_request_gserver,
519
    "truss": async_request_truss,
zhyncs's avatar
zhyncs committed
520
521
522
523
524
525
526
527
}


@dataclass
class BenchmarkMetrics:
    completed: int
    total_input: int
    total_output: int
Ying Sheng's avatar
Ying Sheng committed
528
    total_output_retokenized: int
zhyncs's avatar
zhyncs committed
529
530
531
    request_throughput: float
    input_throughput: float
    output_throughput: float
Ying Sheng's avatar
Ying Sheng committed
532
    output_throughput_retokenized: float
533
534
    total_throughput: float
    total_throughput_retokenized: float
zhyncs's avatar
zhyncs committed
535
536
537
538
539
540
541
542
543
544
545
    mean_ttft_ms: float
    median_ttft_ms: float
    std_ttft_ms: float
    p99_ttft_ms: float
    mean_tpot_ms: float
    median_tpot_ms: float
    std_tpot_ms: float
    p99_tpot_ms: float
    mean_itl_ms: float
    median_itl_ms: float
    std_itl_ms: float
546
    p95_itl_ms: float
zhyncs's avatar
zhyncs committed
547
    p99_itl_ms: float
548
    max_itl_ms: float
zhyncs's avatar
zhyncs committed
549
550
    mean_e2e_latency_ms: float
    median_e2e_latency_ms: float
551
552
    std_e2e_latency_ms: float
    p99_e2e_latency_ms: float
553
    concurrency: float
zhyncs's avatar
zhyncs committed
554
555


Lianmin Zheng's avatar
Lianmin Zheng committed
556
SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
Lianmin Zheng's avatar
Lianmin Zheng committed
557
558


Lianmin Zheng's avatar
Lianmin Zheng committed
559
560
561
562
def download_and_cache_file(url: str, filename: Optional[str] = None):
    """Read and cache a file from a url."""
    if filename is None:
        filename = os.path.join("/tmp", url.split("/")[-1])
Lianmin Zheng's avatar
Lianmin Zheng committed
563

Lianmin Zheng's avatar
Lianmin Zheng committed
564
565
566
    # Check if the cache file already exists
    if os.path.exists(filename):
        return filename
Lianmin Zheng's avatar
Lianmin Zheng committed
567

Lianmin Zheng's avatar
Lianmin Zheng committed
568
    print(f"Downloading from {url} to {filename}")
Lianmin Zheng's avatar
Lianmin Zheng committed
569

Lianmin Zheng's avatar
Lianmin Zheng committed
570
571
572
    # Stream the response to show the progress bar
    response = requests.get(url, stream=True)
    response.raise_for_status()  # Check for request errors
Lianmin Zheng's avatar
Lianmin Zheng committed
573

Lianmin Zheng's avatar
Lianmin Zheng committed
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
    # Total size of the file in bytes
    total_size = int(response.headers.get("content-length", 0))
    chunk_size = 1024  # Download in chunks of 1KB

    # Use tqdm to display the progress bar
    with open(filename, "wb") as f, tqdm(
        desc=filename,
        total=total_size,
        unit="B",
        unit_scale=True,
        unit_divisor=1024,
    ) as bar:
        for chunk in response.iter_content(chunk_size=chunk_size):
            f.write(chunk)
            bar.update(len(chunk))

    return filename
Lianmin Zheng's avatar
Lianmin Zheng committed
591
592


zhyncs's avatar
zhyncs committed
593
594
595
596
597
def sample_sharegpt_requests(
    dataset_path: str,
    num_requests: int,
    tokenizer: PreTrainedTokenizerBase,
    fixed_output_len: Optional[int] = None,
598
    context_len: Optional[int] = None,
599
    prompt_suffix: Optional[str] = "",
600
    apply_chat_template=False,
zhyncs's avatar
zhyncs committed
601
602
603
604
) -> List[Tuple[str, int, int]]:
    if fixed_output_len is not None and fixed_output_len < 4:
        raise ValueError("output_len too small")

Lianmin Zheng's avatar
Lianmin Zheng committed
605
    # Download sharegpt if necessary
606
    if not os.path.isfile(dataset_path) and dataset_path == "":
Lianmin Zheng's avatar
Lianmin Zheng committed
607
        dataset_path = download_and_cache_file(SHAREGPT_URL)
zhyncs's avatar
zhyncs committed
608
609
610
611

    # Load the dataset.
    with open(dataset_path) as f:
        dataset = json.load(f)
612

zhyncs's avatar
zhyncs committed
613
    # Filter out the conversations with less than 2 turns.
614
615
616
617
618
    dataset = [
        data
        for data in dataset
        if len(data.get("conversations", data.get("conversation", []))) >= 2
    ]
zhyncs's avatar
zhyncs committed
619
620
    # Only keep the first two turns of each conversation.
    dataset = [
621
622
623
624
        (
            data.get("conversations", data.get("conversation", []))[0]["value"],
            data.get("conversations", data.get("conversation", []))[1]["value"],
        )
zhyncs's avatar
zhyncs committed
625
626
627
628
629
630
631
632
633
634
635
636
637
638
        for data in dataset
    ]

    # Shuffle the dataset.
    random.shuffle(dataset)

    # Filter out sequences that are too long or too short
    filtered_dataset: List[Tuple[str, int, int]] = []
    for i in range(len(dataset)):
        if len(filtered_dataset) == num_requests:
            break

        # Tokenize the prompts and completions.
        prompt = dataset[i][0]
639
        if prompt_suffix:
640
641
642
643
644
            prompt = (
                remove_suffix(prompt, ASSISTANT_SUFFIX)
                + prompt_suffix
                + ASSISTANT_SUFFIX
            )
645
646
647
648
649
650
651
652
653

        if apply_chat_template:
            prompt = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                add_generation_prompt=True,
                tokenize=False,
            )
            prompt = prompt.replace(tokenizer.bos_token, "")

Lianmin Zheng's avatar
Lianmin Zheng committed
654
        prompt_token_ids = tokenizer.encode(prompt)
zhyncs's avatar
zhyncs committed
655
        completion = dataset[i][1]
Lianmin Zheng's avatar
Lianmin Zheng committed
656
        completion_token_ids = tokenizer.encode(completion)
zhyncs's avatar
zhyncs committed
657
658
659
660
        prompt_len = len(prompt_token_ids)
        output_len = (
            len(completion_token_ids) if fixed_output_len is None else fixed_output_len
        )
661

662
        if prompt_len < 2 or output_len < 2:
zhyncs's avatar
zhyncs committed
663
664
            # Prune too short sequences.
            continue
665
666

        if context_len and prompt_len + output_len > context_len:
zhyncs's avatar
zhyncs committed
667
668
            # Prune too long sequences.
            continue
669

zhyncs's avatar
zhyncs committed
670
671
        filtered_dataset.append((prompt, prompt_len, output_len))

672
673
    print(f"#Input tokens: {np.sum([x[1] for x in filtered_dataset])}")
    print(f"#Output tokens: {np.sum([x[2] for x in filtered_dataset])}")
zhyncs's avatar
zhyncs committed
674
675
676
    return filtered_dataset


677
678
679
680
681
682
def sample_random_requests(
    input_len: int,
    output_len: int,
    num_prompts: int,
    range_ratio: float,
    tokenizer: PreTrainedTokenizerBase,
Lianmin Zheng's avatar
Lianmin Zheng committed
683
    dataset_path: str,
684
685
686
) -> List[Tuple[str, int, int]]:

    input_lens = np.random.randint(
Yineng Zhang's avatar
Yineng Zhang committed
687
        max(int(input_len * range_ratio), 1),
688
689
690
691
692
693
694
695
        input_len + 1,
        size=num_prompts,
    )
    output_lens = np.random.randint(
        int(output_len * range_ratio),
        output_len + 1,
        size=num_prompts,
    )
Lianmin Zheng's avatar
Lianmin Zheng committed
696
697
698
699
700

    if True:
        # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens

        # Download sharegpt if necessary
Lianmin Zheng's avatar
Lianmin Zheng committed
701
702
        if not os.path.isfile(dataset_path):
            dataset_path = download_and_cache_file(SHAREGPT_URL)
Lianmin Zheng's avatar
Lianmin Zheng committed
703
704
705
706
707

        # Load the dataset.
        with open(dataset_path) as f:
            dataset = json.load(f)
        # Filter out the conversations with less than 2 turns.
708
709
710
711
712
        dataset = [
            data
            for data in dataset
            if len(data.get("conversations", data.get("conversation", []))) >= 2
        ]
Lianmin Zheng's avatar
Lianmin Zheng committed
713
714
        # Only keep the first two turns of each conversation.
        dataset = [
715
716
717
718
            (
                data.get("conversations", data.get("conversation", []))[0]["value"],
                data.get("conversations", data.get("conversation", []))[1]["value"],
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
719
720
721
722
723
724
725
            for data in dataset
        ]
        # Shuffle the dataset.
        random.shuffle(dataset)

        # Filter out sequences that are too long or too short
        input_requests: List[Tuple[str, int, int]] = []
726
727
728
729
730
        for data in dataset:
            i = len(input_requests)
            if i == num_prompts:
                break

Lianmin Zheng's avatar
Lianmin Zheng committed
731
            # Tokenize the prompts and completions.
732
            prompt = data[0]
Lianmin Zheng's avatar
Lianmin Zheng committed
733
            prompt_token_ids = tokenizer.encode(prompt)
Lianmin Zheng's avatar
Lianmin Zheng committed
734
735
            prompt_len = len(prompt_token_ids)

736
737
738
739
            # Skip empty prompt
            if prompt_len == 0:
                continue

Yineng Zhang's avatar
Yineng Zhang committed
740
            if prompt_len > input_lens[i]:
Lianmin Zheng's avatar
Lianmin Zheng committed
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
                input_ids = prompt_token_ids[: input_lens[i]]
            else:
                ratio = (input_lens[i] + prompt_len - 1) // prompt_len
                input_ids = (prompt_token_ids * ratio)[: input_lens[i]]
            prompt = tokenizer.decode(input_ids)
            input_requests.append((prompt, int(input_lens[i]), int(output_lens[i])))
    else:
        # Sample token ids from random integers. This can cause some NaN issues.
        offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
        input_requests = []
        for i in range(num_prompts):
            prompt = tokenizer.decode(
                [
                    (offsets[i] + i + j) % tokenizer.vocab_size
                    for j in range(input_lens[i])
                ]
            )
            input_requests.append((prompt, int(input_lens[i]), int(output_lens[i])))
759
760
761
762
763
764

    print(f"#Input tokens: {np.sum(input_lens)}")
    print(f"#Output tokens: {np.sum(output_lens)}")
    return input_requests


765
766
767
768
769
770
771
def gen_prompt(tokenizer, token_num):
    """Generate a random prompt of specified token length using tokenizer vocabulary."""
    all_available_tokens = list(tokenizer.get_vocab().values())
    selected_tokens = random.choices(all_available_tokens, k=token_num)
    return tokenizer.decode(selected_tokens)


772
773
774
775
776
777
def get_gen_prefix_cache_path(args, tokenizer):
    """Create cache directory under ~/.cache/sglang/benchmark"""
    cache_dir = Path.home() / ".cache" / "sglang" / "benchmark"

    # Create a unique cache filename based on the generation parameters
    cache_key = (
778
779
        f"gen_shared_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_"
        f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_"
780
781
782
783
784
        f"{tokenizer.__class__.__name__}.pkl"
    )
    return cache_dir / cache_key


785
786
787
788
789
790
791
def sample_generated_shared_prefix_requests(
    num_groups: int,
    prompts_per_group: int,
    system_prompt_len: int,
    question_len: int,
    output_len: int,
    tokenizer: PreTrainedTokenizerBase,
792
    args: argparse.Namespace,
793
) -> List[Tuple[str, int, int]]:
794
795
796
797
798
799
800
    """Generate benchmark requests with shared system prompts using random tokens and caching."""
    cache_path = get_gen_prefix_cache_path(args, tokenizer)

    # Try to load from cache first
    if cache_path.exists():
        print(f"\nLoading cached generated input data from {cache_path}")
        with open(cache_path, "rb") as f:
801
802
            return pickle.load(f)

803
804
    print("\nGenerating new input data...")

805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
    # Generate system prompts for each group
    system_prompts = []
    for _ in range(num_groups):
        system_prompt = gen_prompt(tokenizer, system_prompt_len)
        system_prompts.append(system_prompt)

    # Generate questions
    questions = []
    for _ in range(num_groups * prompts_per_group):
        question = gen_prompt(tokenizer, question_len)
        questions.append(question)

    # Combine system prompts with questions
    input_requests = []
    total_input_tokens = 0
    total_output_tokens = 0

822
    for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
823
        system_prompt = system_prompts[group_idx]
824
825
826
        for prompt_idx in tqdm(
            range(prompts_per_group), desc="Generating questions", leave=False
        ):
827
828
829
830
831
832
833
834
            question = questions[group_idx * prompts_per_group + prompt_idx]
            full_prompt = f"{system_prompt}\n\n{question}"
            prompt_len = len(tokenizer.encode(full_prompt))

            input_requests.append((full_prompt, prompt_len, output_len))
            total_input_tokens += prompt_len
            total_output_tokens += output_len

835
836
837
838
    # Shuffle questions
    random.shuffle(input_requests)

    # Print statistics
839
840
841
842
843
844
845
846
847
848
849
850
    print(f"\nGenerated shared prefix dataset statistics:")
    print(f"Number of groups: {num_groups}")
    print(f"Prompts per group: {prompts_per_group}")
    print(f"Total prompts: {len(input_requests)}")
    print(f"Total input tokens: {total_input_tokens}")
    print(f"Total output tokens: {total_output_tokens}")
    print(
        f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
    )
    print(
        f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n"
    )
851
852
853
854
855
856

    # Save to cache
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    print(f"Caching generated input data to {cache_path}")
    with open(cache_path, "wb") as f:
        pickle.dump(input_requests, f)
857
858
859
860

    return input_requests


zhyncs's avatar
zhyncs committed
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
async def get_request(
    input_requests: List[Tuple[str, int, int]],
    request_rate: float,
) -> AsyncGenerator[Tuple[str, int, int], None]:
    input_requests = iter(input_requests)
    for request in input_requests:
        yield request

        if request_rate == float("inf"):
            # If the request rate is infinity, then we don't need to wait.
            continue

        # Sample the request interval from the exponential distribution.
        interval = np.random.exponential(1.0 / request_rate)
        # The next request will be sent after the interval.
        await asyncio.sleep(interval)


def calculate_metrics(
    input_requests: List[Tuple[str, int, int]],
    outputs: List[RequestFuncOutput],
    dur_s: float,
    tokenizer: PreTrainedTokenizerBase,
884
    backend: str,
zhyncs's avatar
zhyncs committed
885
) -> Tuple[BenchmarkMetrics, List[int]]:
Ying Sheng's avatar
Ying Sheng committed
886
887
    output_lens: List[int] = []
    retokenized_output_lens: List[int] = []
zhyncs's avatar
zhyncs committed
888
889
890
891
892
    total_input = 0
    completed = 0
    itls: List[float] = []
    tpots: List[float] = []
    ttfts: List[float] = []
zhyncs's avatar
zhyncs committed
893
    e2e_latencies: List[float] = []
zhyncs's avatar
zhyncs committed
894
895
    for i in range(len(outputs)):
        if outputs[i].success:
Ying Sheng's avatar
Ying Sheng committed
896
897
898
            output_len = outputs[i].output_len
            output_lens.append(output_len)
            retokenized_output_len = len(
Lianmin Zheng's avatar
Lianmin Zheng committed
899
                tokenizer.encode(outputs[i].generated_text, add_special_tokens=False)
Ying Sheng's avatar
Ying Sheng committed
900
901
            )
            retokenized_output_lens.append(retokenized_output_len)
zhyncs's avatar
zhyncs committed
902
903
904
905
906
            total_input += input_requests[i][1]
            if output_len > 1:
                tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1))
            itls += outputs[i].itl
            ttfts.append(outputs[i].ttft)
zhyncs's avatar
zhyncs committed
907
908
909

            e2e_latencies.append(outputs[i].latency)

zhyncs's avatar
zhyncs committed
910
911
            completed += 1
        else:
Ying Sheng's avatar
Ying Sheng committed
912
913
            output_lens.append(0)
            retokenized_output_lens.append(0)
zhyncs's avatar
zhyncs committed
914
915
916
917
918
919
920
921
922
923

    if completed == 0:
        warnings.warn(
            "All requests failed. This is likely due to a misconfiguration "
            "on the benchmark arguments.",
            stacklevel=2,
        )
    metrics = BenchmarkMetrics(
        completed=completed,
        total_input=total_input,
Ying Sheng's avatar
Ying Sheng committed
924
925
        total_output=sum(output_lens),
        total_output_retokenized=sum(retokenized_output_lens),
zhyncs's avatar
zhyncs committed
926
927
        request_throughput=completed / dur_s,
        input_throughput=total_input / dur_s,
Ying Sheng's avatar
Ying Sheng committed
928
929
        output_throughput=sum(output_lens) / dur_s,
        output_throughput_retokenized=sum(retokenized_output_lens) / dur_s,
930
931
932
        total_throughput=(total_input + sum(output_lens)) / dur_s,
        total_throughput_retokenized=(total_input + sum(retokenized_output_lens))
        / dur_s,
zhyncs's avatar
zhyncs committed
933
934
935
936
937
938
939
940
941
942
943
944
        mean_ttft_ms=np.mean(ttfts or 0)
        * 1000,  # ttfts is empty if streaming is not supported by backend
        median_ttft_ms=np.median(ttfts or 0) * 1000,
        std_ttft_ms=np.std(ttfts or 0) * 1000,
        p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000,
        mean_tpot_ms=np.mean(tpots or 0) * 1000,
        median_tpot_ms=np.median(tpots or 0) * 1000,
        std_tpot_ms=np.std(tpots or 0) * 1000,
        p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000,
        mean_itl_ms=np.mean(itls or 0) * 1000,
        median_itl_ms=np.median(itls or 0) * 1000,
        std_itl_ms=np.std(itls or 0) * 1000,
945
        p95_itl_ms=np.percentile(itls or 0, 95) * 1000,
zhyncs's avatar
zhyncs committed
946
        p99_itl_ms=np.percentile(itls or 0, 99) * 1000,
947
        max_itl_ms=np.max(itls or 0) * 1000,
zhyncs's avatar
zhyncs committed
948
949
        mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000,
        median_e2e_latency_ms=np.median(e2e_latencies) * 1000,
950
951
        std_e2e_latency_ms=np.std(e2e_latencies) * 1000,
        p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000,
952
        concurrency=np.sum(e2e_latencies) / dur_s,
zhyncs's avatar
zhyncs committed
953
954
    )

Ying Sheng's avatar
Ying Sheng committed
955
    return metrics, output_lens
zhyncs's avatar
zhyncs committed
956
957
958
959
960


async def benchmark(
    backend: str,
    api_url: str,
961
    base_url: str,
zhyncs's avatar
zhyncs committed
962
963
964
965
    model_id: str,
    tokenizer: PreTrainedTokenizerBase,
    input_requests: List[Tuple[str, int, int]],
    request_rate: float,
966
    max_concurrency: Optional[int],
zhyncs's avatar
zhyncs committed
967
    disable_tqdm: bool,
968
    lora_names: List[str],
969
    extra_request_body: Dict[str, Any],
970
    profile: bool,
971
    pd_seperated: bool = False,
zhyncs's avatar
zhyncs committed
972
973
974
975
976
977
):
    if backend in ASYNC_REQUEST_FUNCS:
        request_func = ASYNC_REQUEST_FUNCS[backend]
    else:
        raise ValueError(f"Unknown backend: {backend}")

978
    # Limit concurrency
979
980
981
982
983
984
985
986
987
    # From https://github.com/vllm-project/vllm/pull/9390
    semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None

    async def limited_request_func(request_func_input, pbar):
        if semaphore is None:
            return await request_func(request_func_input=request_func_input, pbar=pbar)
        async with semaphore:
            return await request_func(request_func_input=request_func_input, pbar=pbar)

988
    # Warmup
zhyncs's avatar
zhyncs committed
989
990
    print("Starting initial single prompt test run...")
    test_prompt, test_prompt_len, test_output_len = input_requests[0]
991
992
993
994
995
    if lora_names != None and len(lora_names) != 0:
        lora_name = lora_names[0]
    else:
        lora_name = None

zhyncs's avatar
zhyncs committed
996
997
998
999
1000
    test_input = RequestFuncInput(
        model=model_id,
        prompt=test_prompt,
        api_url=api_url,
        prompt_len=test_prompt_len,
1001
        output_len=min(test_output_len, 32),
1002
        lora_name=lora_name,
1003
        extra_request_body=extra_request_body,
zhyncs's avatar
zhyncs committed
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
    )
    test_output = await request_func(request_func_input=test_input)
    if not test_output.success:
        raise ValueError(
            "Initial test run failed - Please make sure benchmark arguments "
            f"are correctly specified. Error: {test_output.error}"
        )
    else:
        print("Initial test run completed. Starting main benchmark run...")

1014
1015
    # Flush cache
    if "sglang" in backend:
1016
        requests.post(base_url + "/flush_cache", headers=get_auth_headers())
1017
1018

    time.sleep(1.0)
1019

1020
    # Start profiler
1021
1022
1023
1024
1025
1026
1027
1028
    if profile:
        print("Starting profiler...")
        profile_output = await async_request_profile(
            api_url=base_url + "/start_profile"
        )
        if profile_output.success:
            print("Profiler started")

zhyncs's avatar
zhyncs committed
1029
1030
    pbar = None if disable_tqdm else tqdm(total=len(input_requests))

1031
    # Run all requests
zhyncs's avatar
zhyncs committed
1032
1033
1034
1035
    benchmark_start_time = time.perf_counter()
    tasks: List[asyncio.Task] = []
    async for request in get_request(input_requests, request_rate):
        prompt, prompt_len, output_len = request
1036
1037
1038
1039
1040
1041
        if lora_names != None and len(lora_names) != 0:
            idx = random.randint(0, len(lora_names) - 1)
            lora_name = lora_names[idx]
        else:
            lora_name = None

zhyncs's avatar
zhyncs committed
1042
1043
1044
1045
1046
1047
        request_func_input = RequestFuncInput(
            model=model_id,
            prompt=prompt,
            api_url=api_url,
            prompt_len=prompt_len,
            output_len=output_len,
1048
            lora_name=lora_name,
1049
            extra_request_body=extra_request_body,
zhyncs's avatar
zhyncs committed
1050
1051
1052
        )
        tasks.append(
            asyncio.create_task(
1053
                limited_request_func(request_func_input=request_func_input, pbar=pbar)
zhyncs's avatar
zhyncs committed
1054
1055
1056
1057
            )
        )
    outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks)

1058
    # Stop profiler
1059
1060
1061
1062
1063
1064
    if profile:
        print("Stopping profiler...")
        profile_output = await async_request_profile(api_url=base_url + "/stop_profile")
        if profile_output.success:
            print("Profiler stopped")

zhyncs's avatar
zhyncs committed
1065
1066
1067
    if pbar is not None:
        pbar.close()

1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
    if "sglang" in backend:
        server_info = requests.get(base_url + "/get_server_info")
        if pd_seperated:
            accept_length = server_info.json()["decode"][0].get(
                "avg_spec_accept_length", None
            )
        else:
            accept_length = server_info.json().get("avg_spec_accept_length", None)
    else:
        accept_length = None

1079
    # Compute metrics and print results
zhyncs's avatar
zhyncs committed
1080
    benchmark_duration = time.perf_counter() - benchmark_start_time
Ying Sheng's avatar
Ying Sheng committed
1081
    metrics, output_lens = calculate_metrics(
zhyncs's avatar
zhyncs committed
1082
1083
1084
1085
        input_requests=input_requests,
        outputs=outputs,
        dur_s=benchmark_duration,
        tokenizer=tokenizer,
1086
        backend=backend,
zhyncs's avatar
zhyncs committed
1087
1088
1089
    )

    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
1090
    print("{:<40} {:<10}".format("Backend:", backend))
zhyncs's avatar
zhyncs committed
1091
    print("{:<40} {:<10}".format("Traffic request rate:", request_rate))
1092
1093
1094
1095
1096
1097
    print(
        "{:<40} {:<10}".format(
            "Max reqeuest concurrency:",
            max_concurrency if max_concurrency else "not set",
        )
    )
zhyncs's avatar
zhyncs committed
1098
1099
1100
1101
    print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
    print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
    print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
Ying Sheng's avatar
Ying Sheng committed
1102
1103
1104
1105
1106
    print(
        "{:<40} {:<10}".format(
            "Total generated tokens (retokenized):", metrics.total_output_retokenized
        )
    )
zhyncs's avatar
zhyncs committed
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
    print(
        "{:<40} {:<10.2f}".format(
            "Request throughput (req/s):", metrics.request_throughput
        )
    )
    print(
        "{:<40} {:<10.2f}".format(
            "Input token throughput (tok/s):", metrics.input_throughput
        )
    )
    print(
        "{:<40} {:<10.2f}".format(
            "Output token throughput (tok/s):", metrics.output_throughput
        )
    )
1122
1123
1124
1125
1126
    print(
        "{:<40} {:<10.2f}".format(
            "Total token throughput (tok/s):", metrics.total_throughput
        )
    )
1127
    print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency))
1128
1129
    if accept_length:
        print("{:<40} {:<10.2f}".format("Accept length:", accept_length))
zhyncs's avatar
zhyncs committed
1130
1131
1132
1133
1134
1135
1136
1137
1138
    print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-"))
    print(
        "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms)
    )
    print(
        "{:<40} {:<10.2f}".format(
            "Median E2E Latency (ms):", metrics.median_e2e_latency_ms
        )
    )
zhyncs's avatar
zhyncs committed
1139
1140
1141
1142
    print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-"))
    print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms))
    print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms))
    print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms))
1143
    print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-"))
zhyncs's avatar
zhyncs committed
1144
1145
    print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms))
    print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms))
1146
    print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms))
zhyncs's avatar
zhyncs committed
1147
    print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms))
1148
    print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms))
zhyncs's avatar
zhyncs committed
1149
1150
    print("=" * 50)

zhyncs's avatar
zhyncs committed
1151
1152
1153
1154
1155
1156
    if (
        metrics.median_ttft_ms is not None
        and metrics.mean_itl_ms is not None
        and metrics.output_throughput is not None
    ):
        result = {
1157
            # Arguments
zhyncs's avatar
zhyncs committed
1158
1159
1160
            "backend": args.backend,
            "dataset_name": args.dataset_name,
            "request_rate": request_rate,
1161
            "max_concurrency": max_concurrency,
1162
1163
1164
1165
1166
1167
1168
            "sharegpt_output_len": args.sharegpt_output_len,
            "random_input_len": args.random_input_len,
            "random_output_len": args.random_output_len,
            "random_range_ratio": args.random_range_ratio,
            # Results
            "duration": benchmark_duration,
            "completed": metrics.completed,
1169
1170
1171
            "total_input_tokens": metrics.total_input,
            "total_output_tokens": metrics.total_output,
            "total_output_tokens_retokenized": metrics.total_output_retokenized,
1172
1173
1174
            "request_throughput": metrics.request_throughput,
            "input_throughput": metrics.input_throughput,
            "output_throughput": metrics.output_throughput,
1175
1176
            "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms,
            "median_e2e_latency_ms": metrics.median_e2e_latency_ms,
1177
1178
            "std_e2e_latency_ms": metrics.std_e2e_latency_ms,
            "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms,
1179
            "mean_ttft_ms": metrics.mean_ttft_ms,
1180
            "median_ttft_ms": metrics.median_ttft_ms,
1181
1182
1183
1184
1185
1186
            "std_ttft_ms": metrics.std_ttft_ms,
            "p99_ttft_ms": metrics.p99_ttft_ms,
            "mean_tpot_ms": metrics.mean_tpot_ms,
            "median_tpot_ms": metrics.median_tpot_ms,
            "std_tpot_ms": metrics.std_tpot_ms,
            "p99_tpot_ms": metrics.p99_tpot_ms,
1187
            "mean_itl_ms": metrics.mean_itl_ms,
1188
            "median_itl_ms": metrics.median_itl_ms,
1189
            "std_itl_ms": metrics.std_itl_ms,
1190
            "p95_itl_ms": metrics.p95_itl_ms,
1191
            "p99_itl_ms": metrics.p99_itl_ms,
1192
            "concurrency": metrics.concurrency,
1193
            "accept_length": accept_length,
zhyncs's avatar
zhyncs committed
1194
1195
1196
1197
        }
    else:
        print(f"Error running benchmark for request rate: {request_rate}")
        print("-" * 30)
1198

zhyncs's avatar
zhyncs committed
1199
1200
1201
1202
1203
1204
1205
    # Determine output file name
    if args.output_file:
        output_file_name = args.output_file
    else:
        now = datetime.now().strftime("%m%d")
        if args.dataset_name == "random":
            output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl"
1206
        else:
zhyncs's avatar
zhyncs committed
1207
            output_file_name = f"{args.backend}_{now}_{args.num_prompts}_sharegpt.jsonl"
1208

zhyncs's avatar
zhyncs committed
1209
1210
1211
    # Append results to a JSONL file
    with open(output_file_name, "a") as file:
        file.write(json.dumps(result) + "\n")
1212

1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
    result.update(
        {
            "input_lens": [output.prompt_len for output in outputs],
            "output_lens": output_lens,
            "ttfts": [output.ttft for output in outputs],
            "itls": [output.itl for output in outputs],
            "generated_texts": [output.generated_text for output in outputs],
            "errors": [output.error for output in outputs],
        }
    )
zhyncs's avatar
zhyncs committed
1223
1224
1225
    return result


1226
1227
1228
1229
1230
1231
1232
1233
1234
def check_chat_template(model_path):
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        return "chat_template" in tokenizer.init_kwargs
    except Exception as e:
        print(f"Fail to load tokenizer config with error={e}")
        return False


1235
1236
1237
1238
1239
1240
def set_global_args(args_: argparse.Namespace):
    """Set the global args."""
    global args
    args = args_


1241
1242
1243
1244
def run_benchmark(args_: argparse.Namespace):
    global args
    args = args_

1245
1246
1247
1248
    # Set default value for max_concurrency if not present
    if not hasattr(args, "max_concurrency"):
        args.max_concurrency = None

1249
1250
    print(f"benchmark_args={args}")

Lianmin Zheng's avatar
Lianmin Zheng committed
1251
    # Set global environments
1252
    set_ulimit()
zhyncs's avatar
zhyncs committed
1253
1254
1255
    random.seed(args.seed)
    np.random.seed(args.seed)

1256
1257
1258
1259
    extra_request_body = {}
    if args.extra_request_body:
        extra_request_body = json.loads(args.extra_request_body)

Lianmin Zheng's avatar
Lianmin Zheng committed
1260
    # Set url
zhyncs's avatar
zhyncs committed
1261
1262
1263
    if args.port is None:
        args.port = {
            "sglang": 30000,
1264
1265
            "sglang-native": 30000,
            "sglang-oai": 30000,
zhyncs's avatar
zhyncs committed
1266
1267
            "lmdeploy": 23333,
            "vllm": 8000,
1268
            "trt": 8000,
1269
            "gserver": 9988,
1270
            "truss": 8080,
zhyncs's avatar
zhyncs committed
1271
1272
1273
1274
1275
1276
1277
1278
        }.get(args.backend, 30000)

    model_url = (
        f"{args.base_url}/v1/models"
        if args.base_url
        else f"http://{args.host}:{args.port}/v1/models"
    )

1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
    if args.backend in ["sglang", "sglang-native"]:
        api_url = (
            f"{args.base_url}/generate"
            if args.base_url
            else f"http://{args.host}:{args.port}/generate"
        )
    elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]:
        api_url = (
            f"{args.base_url}/v1/completions"
            if args.base_url
            else f"http://{args.host}:{args.port}/v1/completions"
        )
    elif args.backend == "trt":
1292
1293
1294
1295
1296
1297
1298
1299
        api_url = (
            f"{args.base_url}/v2/models/ensemble/generate_stream"
            if args.base_url
            else f"http://{args.host}:{args.port}/v2/models/ensemble/generate_stream"
        )
        if args.model is None:
            print("Please provide a model using `--model` when using `trt` backend.")
            sys.exit(1)
1300
    elif args.backend == "gserver":
Lianmin Zheng's avatar
Lianmin Zheng committed
1301
1302
        api_url = args.base_url if args.base_url else f"{args.host}:{args.port}"
        args.model = args.model or "default"
1303
1304
1305
1306
1307
1308
    elif args.backend == "truss":
        api_url = (
            f"{args.base_url}/v1/models/model:predict"
            if args.base_url
            else f"http://{args.host}:{args.port}/v1/models/model:predict"
        )
1309
1310
1311
    base_url = (
        f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url
    )
1312

Lianmin Zheng's avatar
Lianmin Zheng committed
1313
    # Get model name
zhyncs's avatar
zhyncs committed
1314
    if args.model is None:
1315
1316
1317
1318
1319
        if args.backend == "truss":
            print(
                "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct"
            )
            sys.exit(1)
zhyncs's avatar
zhyncs committed
1320
        try:
1321
            response = requests.get(model_url, headers=get_auth_headers())
zhyncs's avatar
zhyncs committed
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
            model_list = response.json().get("data", [])
            args.model = model_list[0]["id"] if model_list else None
        except Exception as e:
            print(f"Failed to fetch model from {model_url}. Error: {e}")
            print(
                "Please specify the correct host and port using `--host` and `--port`."
            )
            sys.exit(1)

    if args.model is None:
        print("No model specified or found. Please provide a model using `--model`.")
        sys.exit(1)

1335
1336
1337
1338
1339
1340
    if not check_chat_template(args.model):
        print(
            "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n"
            "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n"
        )

zhyncs's avatar
zhyncs committed
1341
1342
    print(f"{args}\n")

Lianmin Zheng's avatar
Lianmin Zheng committed
1343
    # Read dataset
zhyncs's avatar
zhyncs committed
1344
1345
1346
1347
    backend = args.backend
    model_id = args.model
    tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
    tokenizer = get_tokenizer(tokenizer_id)
1348
    input_requests = get_dataset(args, tokenizer)
zhyncs's avatar
zhyncs committed
1349

1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
    return asyncio.run(
        benchmark(
            backend=backend,
            api_url=api_url,
            base_url=base_url,
            model_id=model_id,
            tokenizer=tokenizer,
            input_requests=input_requests,
            request_rate=args.request_rate,
            max_concurrency=args.max_concurrency,
            disable_tqdm=args.disable_tqdm,
1361
            lora_names=args.lora_name,
1362
1363
1364
            extra_request_body=extra_request_body,
            profile=args.profile,
            pd_seperated=args.pd_seperated,
Lianmin Zheng's avatar
Lianmin Zheng committed
1365
        )
1366
    )
zhyncs's avatar
zhyncs committed
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379


def set_ulimit(target_soft_limit=65535):
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
        except ValueError as e:
            print(f"Fail to set RLIMIT_NOFILE: {e}")


1380
1381
1382
1383
1384
1385
1386
class LoRAPathAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, [])
        for lora_name in values:
            getattr(namespace, self.dest).append(lora_name)


zhyncs's avatar
zhyncs committed
1387
if __name__ == "__main__":
1388
    parser = ArgumentParser(description="Benchmark the online serving throughput.")
zhyncs's avatar
zhyncs committed
1389
1390
1391
1392
    parser.add_argument(
        "--backend",
        type=str,
        choices=list(ASYNC_REQUEST_FUNCS.keys()),
1393
        default="sglang",
zhyncs's avatar
zhyncs committed
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
        help="Must specify a backend, depending on the LLM Inference Engine.",
    )
    parser.add_argument(
        "--base-url",
        type=str,
        default=None,
        help="Server or API base url if not using http host and port.",
    )
    parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0."
    )
    parser.add_argument(
        "--port",
        type=int,
        help="If not set, the default port is configured according to its default value for different LLM Inference Engines.",
    )
    parser.add_argument(
1411
1412
1413
        "--dataset-name",
        type=str,
        default="sharegpt",
1414
        choices=["sharegpt", "random", "generated-shared-prefix"],
1415
1416
1417
1418
        help="Name of the dataset to benchmark on.",
    )
    parser.add_argument(
        "--dataset-path", type=str, default="", help="Path to the dataset."
zhyncs's avatar
zhyncs committed
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
    )
    parser.add_argument(
        "--model",
        type=str,
        help="Name or path of the model. If not set, the default model will request /v1/models for conf.",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        help="Name or path of the tokenizer. If not set, using the model conf.",
    )
    parser.add_argument(
        "--num-prompts",
        type=int,
        default=1000,
        help="Number of prompts to process. Default is 1000.",
    )
    parser.add_argument(
        "--sharegpt-output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the output length from the ShareGPT dataset.",
    )
1442
1443
1444
1445
1446
1447
    parser.add_argument(
        "--sharegpt-context-len",
        type=int,
        default=None,
        help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.",
    )
1448
1449
1450
    parser.add_argument(
        "--random-input-len",
        type=int,
1451
        default=1024,
1452
1453
1454
1455
        help="Number of input tokens per request, used only for random dataset.",
    )
    parser.add_argument(
        "--random-output-len",
1456
        default=1024,
1457
1458
1459
1460
1461
1462
        type=int,
        help="Number of output tokens per request, used only for random dataset.",
    )
    parser.add_argument(
        "--random-range-ratio",
        type=float,
Yineng Zhang's avatar
Yineng Zhang committed
1463
        default=0.0,
1464
1465
1466
        help="Range of sampled ratio of input/output length, "
        "used only for random dataset.",
    )
zhyncs's avatar
zhyncs committed
1467
1468
1469
    parser.add_argument(
        "--request-rate",
        type=float,
1470
        default=float("inf"),
zhyncs's avatar
zhyncs committed
1471
        help="Number of requests per second. If this is inf, then all the requests are sent at time 0. "
min-xu-et's avatar
min-xu-et committed
1472
        "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.",
zhyncs's avatar
zhyncs committed
1473
    )
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
    parser.add_argument(
        "--max-concurrency",
        type=int,
        default=None,
        help="Maximum number of concurrent requests. This can be used "
        "to help simulate an environment where a higher level component "
        "is enforcing a maximum number of concurrent requests. While the "
        "--request-rate argument controls the rate at which requests are "
        "initiated, this argument will control how many are actually allowed "
        "to execute at a time. This means that when used in combination, the "
        "actual request rate may be lower than specified with --request-rate, "
        "if the server is not processing requests fast enough to keep up.",
    )
1487
    parser.add_argument("--output-file", type=str, help="Output JSONL file name.")
1488
1489
1490
1491
1492
    parser.add_argument(
        "--disable-tqdm",
        action="store_true",
        help="Specify to disable tqdm progress bar.",
    )
1493
1494
1495
1496
1497
    parser.add_argument(
        "--disable-stream",
        action="store_true",
        help="Disable streaming mode.",
    )
1498
    parser.add_argument(
1499
        "--return-logprob",
1500
        action="store_true",
1501
        help="Return logprob.",
1502
    )
1503
    parser.add_argument("--seed", type=int, default=1, help="The random seed.")
1504
    parser.add_argument(
1505
        "--disable-ignore-eos",
1506
        action="store_true",
1507
        help="Disable ignoring EOS.",
1508
    )
1509
1510
1511
1512
1513
1514
1515
    parser.add_argument(
        "--extra-request-body",
        metavar='{"key1": "value1", "key2": "value2"}',
        type=str,
        help="Append given JSON object to the request payload. You can use this to specify"
        "additional generate params like sampling params.",
    )
1516
1517
1518
1519
1520
    parser.add_argument(
        "--apply-chat-template",
        action="store_true",
        help="Apply chat template",
    )
1521
1522
1523
1524
1525
1526
1527
1528
1529
    parser.add_argument(
        "--profile",
        action="store_true",
        help="Use Torch Profiler. The endpoint must be launched with "
        "SGLANG_TORCH_PROFILER_DIR to enable profiler.",
    )
    parser.add_argument(
        "--lora-name",
        type=str,
1530
        nargs="*",
1531
        default=None,
1532
1533
        action=LoRAPathAction,
        help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...",
1534
    )
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
    parser.add_argument(
        "--prompt-suffix",
        type=str,
        default="",
        help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.",
    )
    parser.add_argument(
        "--pd-seperated",
        action="store_true",
        help="Benchmark PD disaggregation server",
    )
1546
1547
1548

    group = parser.add_argument_group("generated-shared-prefix dataset arguments")
    group.add_argument(
1549
        "--gsp-num-groups",
1550
1551
1552
1553
1554
        type=int,
        default=64,
        help="Number of system prompt groups for generated-shared-prefix dataset",
    )
    group.add_argument(
1555
        "--gsp-prompts-per-group",
1556
1557
1558
1559
1560
        type=int,
        default=16,
        help="Number of prompts per system prompt group for generated-shared-prefix dataset",
    )
    group.add_argument(
1561
        "--gsp-system-prompt-len",
1562
1563
1564
1565
1566
        type=int,
        default=2048,
        help="Target length in tokens for system prompts in generated-shared-prefix dataset",
    )
    group.add_argument(
1567
        "--gsp-question-len",
1568
1569
1570
1571
1572
        type=int,
        default=128,
        help="Target length in tokens for questions in generated-shared-prefix dataset",
    )
    group.add_argument(
1573
        "--gsp-output-len",
1574
1575
1576
1577
        type=int,
        default=256,
        help="Target length in tokens for outputs in generated-shared-prefix dataset",
    )
zhyncs's avatar
zhyncs committed
1578
    args = parser.parse_args()
1579
    run_benchmark(args)