sglangv0.5.2 & support Qwen3-Next-80B-A3B-Instruct

This commit is contained in:
maxiao1
2025-09-13 17:00:20 +08:00
commit 118f1fc726
2037 changed files with 515371 additions and 0 deletions

View File

@@ -0,0 +1,91 @@
## Run synthetic multi-turn benchmark
```
# SGLang server with radix cache disabled
python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --disable-radix-cache
# SGLang server with radix cache on and first-come-first-serve policy
python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --schedule-policy fcfs
# The default SGLang server with radix cache on and long-prefix-match policy
python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000
# SGLang server with hierarchical radix cache enabled
python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --enable-hierarchical-cache
```
```
python bench_multiturn.py --model-path Qwen/Qwen2.5-14B-Instruct
```
Note: The performance gain of hierarchical caching depends on the ratio of reusable tokens to GPU memory capacity. The more tokens to be reused, the larger the model, and the more constrained the GPU memory size, the greater the benefit one can expect from hierarchical caching.
# Benchmark with more datasets
## Download Dataset
```bash
./download.sh {sharegpt|ultragpt|loogle|nextqa|all}
```
This script will automatically download the required dataset to the current working directory
## Multiturn Benchmark
### Supported Datasets
- sharegpt
- ultrachat
- loogle
### Example Usage:
```bash
python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \
--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \
--port 8001 --enable-multiturn --disable-shuffle
```
This uses `mistralai/Mistral-7B-Instruct-v0.3` model with `sglang` as backend. The dataset
is `longdep_qa.json`. We send `10 conversations` with `10 req/s` to port 8001. We enable
multiturn chat without shuffling the order of conversations (i.e. following the original
order in the dataset file).
### Note:
The requests of multiple conversations are sent in a round robin fashion.
For example, if we have 3 conversations A, B, C whose rounds are `[2, 3, 4]` correspondingly,
multiturn chat will send the requests to the backend in the following order: `[A1, B1, C1, A2, B2, C2, B3, C3, C4]`
This has implications on the cache reuse patterns: the cache reuse distance is the largest
under this request pattern (which means a prefix-aware local scheduler in the backend can
yield the most benefit compared to a FIFO scheduler)
## Shared Prefix Benchmark
### Supported Datasets
- loogle
### Example Usage:
```bash
python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \
--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \
--port 8001 --enable-shared-prefix --disable-shuffle
```
### Note:
Shared Prefix benchmark sends the questions for the same prompt together. For example,
if we have 3 shared prefix A, B, C, which have [2, 3, 4] questions correspondingly,
the shared prefix benchmark will send the requests to the
backend in the following order: `[A+Q1, A+Q2, B+Q1, B+Q2, B+Q3, C+Q1, C+Q2, C+Q3]`.
## Multi Modality Benchmark (WIP)
### Supported Datasets:
- nextqa
### Example Usage:
```bash
Server:
python3 -m sglang.launch_server --model-path lmms-lab/LLaVA-NeXT-Video-7B --tp 2 --dp 1 --port 8001 \
--host 0.0.0.0 --mem-fraction-static 0.9 --tokenizer-path llava-hf/llava-1.5-7b-hf \
--json-model-override-args "{\"architectures\": [\"LlavaVidForCausalLM\"], \"model_type\":\"llava\", \"mm_spatial_pool_stride\":2}"
Client:
python3 bench_serving.py --model lmms-lab/LLaVA-NeXT-Video-7B --backend sglang --dataset-path \
NExTVideo --dataset-name nextqa --request-rate 10 --num-prompts 1 --disable-shuffle --port 8001 \ --enable-multiturn --max-frames 16 --tokenizer llava-hf/llava-1.5-7b-hf --fixed-output-len 2048
```
Note: for the server args, `tokenizer-path`, overriding architecture are necessary.
## Supported Backend
- sglang (oai)
- vllm (oai)
- lmdeploy (oai)

View File

@@ -0,0 +1,101 @@
import json
import queue
import time
import requests
from bench_multiturn import (
ReadyQueue,
WorkloadGenerator,
gen_payload,
log_to_jsonl_file,
parse_args,
)
from tqdm.asyncio import tqdm
from sglang.bench_serving import get_tokenizer
class ContextWorkloadGenerator(WorkloadGenerator):
def __init__(self, args):
# Construct the base URL for requests
self.baseurl = f"http://{args.host}:{args.port}/"
self.url = self.baseurl + "generate"
self.tokenizer = get_tokenizer(args.model_path)
self.distribution = args.distribution
self.request_rate = args.request_rate
self.start_time = None
self.finished_time = None
self.sent_requests = 0
self.completed_requests = 0
self.dataset = json.load(open(args.dataset_path))
num_requests = min(args.num_clients, len(self.dataset["queries"]))
init_requests = []
for i in range(num_requests):
context_id = self.dataset["queries"][i]["context"]
init_requests.append(
(
i,
gen_payload(
self.dataset["contexts"][context_id]
+ self.dataset["queries"][i]["question"],
len(
self.tokenizer(
self.dataset["queries"][i]["reference_answer"]
)["input_ids"]
),
),
)
)
self.ready_queue = ReadyQueue(init_requests=init_requests)
self.response_queue = queue.Queue()
self.pbar = tqdm(total=num_requests)
self.performance_metrics = {
"ttft": [],
"latency": [],
"itl": [],
"prompt_len": [],
"cached_tokens": [],
"generated_len": [],
}
self.max_parallel = args.max_parallel
self.logfile = args.log_file
def response_handler(self):
while True:
try:
client_id, response = self.response_queue.get(
timeout=10
) # Block until response is available
if not response.success:
raise ValueError(f"Request failed with error: {response.error}")
self.performance_metrics["ttft"].append(response.ttft)
self.performance_metrics["itl"].extend(response.itl)
self.performance_metrics["latency"].append(response.latency)
self.performance_metrics["prompt_len"].append(response.prompt_len)
self.performance_metrics["cached_tokens"].append(response.cached_tokens)
self.performance_metrics["generated_len"].append(response.generated_len)
self.completed_requests += 1
except queue.Empty:
if self.pbar.n == self.pbar.total:
break
if __name__ == "__main__":
args = parse_args()
args.num_rounds = 1
args.max_parallel = 24
flush_cache_url = f"http://{args.host}:{args.port}/flush_cache"
for request_rate in [24, 16, 12, 8, 4, 2, 1]:
args.request_rate = request_rate
requests.post(flush_cache_url)
time.sleep(1)
performance_data = ContextWorkloadGenerator(args).run()
log_to_jsonl_file(performance_data, args.log_file, args.tag)

View File

@@ -0,0 +1,567 @@
import argparse
import asyncio
import json
import logging
import os
import queue
import random
import threading
import time
from dataclasses import dataclass
from functools import wraps
import aiohttp
from sglang.bench_serving import (
RequestFuncOutput,
get_tokenizer,
remove_prefix,
sample_random_requests,
)
# Set up logger
logger = logging.getLogger(__name__)
# Set up JSONL file for debug logging
debug_log_file = None
# Create a lock for thread-safe debug log writing
debug_log_lock = threading.Lock()
def write_debug_log(data):
global debug_log_file
"""Write debug information to a JSONL file"""
if debug_log_file is None:
return
# Acquire lock for thread-safe writing
with debug_log_lock:
# Write as JSONL (JSON Line format)
debug_log_file.write(json.dumps(data) + "\n")
debug_log_file.flush()
def parse_args():
parser = argparse.ArgumentParser(
description="Script to benchmark concurrent requests to a server."
)
parser.add_argument(
"--model-path",
type=str,
default="/data/models/Qwen3-0.6B",
help="model path compatible with Hugging Face Transformers",
)
parser.add_argument(
"--dataset-path",
type=str,
default="/data/models/ShareGPT_V3_unfiltered_cleaned_split/ShareGPT_V3_unfiltered_cleaned_split.json",
help="local dataset to sample tokens from",
)
parser.add_argument(
"--host",
type=str,
default="localhost",
help="Server hostname or IP (default: localhost)",
)
parser.add_argument(
"--port",
type=int,
default=30000,
help="Server port (default: 30000)",
)
parser.add_argument(
"--duration",
type=int,
default=600,
help="Duration to run the benchmark in seconds (default: 300 seconds)",
)
parser.add_argument(
"--log-level",
type=str,
default="info",
choices=["debug", "info"],
help="Set the logging level (default: info)",
)
parser.add_argument(
"--debug-log-file",
type=str,
default="debug.log.jsonl",
help="File to write debug logs in JSONL format",
)
return parser.parse_args()
def load_config():
config_path = os.getenv("CONFIG_PATH")
if not config_path:
raise ValueError("Environment variable 'CONFIG_PATH' is not set.")
with open(config_path, "r") as f:
config = json.load(f)
required_keys = [
"num_rounds",
"num_clients",
"round_ratios",
"mean_new_tokens_per_round",
"mean_return_tokens_per_round",
"mean_inter_round_interval",
]
for key in required_keys:
if key not in config:
raise KeyError(f"Missing required configuration key: {key}")
num_rounds = config["num_rounds"]
assert len(config["round_ratios"]) == num_rounds
assert len(config["mean_new_tokens_per_round"]) == num_rounds
assert len(config["mean_return_tokens_per_round"]) == num_rounds
assert len(config["mean_inter_round_interval"]) == num_rounds
print(config)
return config
@dataclass
class UserData:
user_id: int
current_round: int
total_rounds: int
prompt: str
return_tokens: int
start: int
def synchronized():
def _decorator(func):
@wraps(func)
def wrapper(self, *args, **kwargs):
with self.lock:
return func(self, *args, **kwargs)
return wrapper
return _decorator
class UserGenerator:
def __init__(self, config, model_path, dataset_path):
self.tokenizer_path = model_path
self.tokenizer = get_tokenizer(self.tokenizer_path)
self.dataset_path = dataset_path
self.user_id = 0
self.lock = threading.Lock()
self.num_rounds = config["num_rounds"]
self.cumulative_ratios = [
sum(config["round_ratios"][: i + 1])
for i in range(len(config["round_ratios"]))
]
self.mean_new_tokens_per_round = config["mean_new_tokens_per_round"]
self.mean_return_tokens_per_round = config["mean_return_tokens_per_round"]
self.mean_inter_round_interval = config["mean_inter_round_interval"]
self.sigma = 100
self.range_ratio = 0.8
assert self.range_ratio <= 1
self.candidate_inputs = [
[
r
for r in sample_random_requests(
input_len=(
self.mean_new_tokens_per_round[i] * (2 - self.range_ratio)
),
output_len=(
self.mean_return_tokens_per_round[i] * (2 - self.range_ratio)
),
num_prompts=config["num_clients"],
range_ratio=self.range_ratio / (2 - self.range_ratio),
tokenizer=self.tokenizer,
dataset_path=self.dataset_path,
random_sample=False,
)
]
for i in range(self.num_rounds)
]
self.multiturn_queue = []
self.user_stats = [0 for _ in range(self.num_rounds)]
self.input_stats = [[0, 0] for _ in range(self.num_rounds)]
self.output_stats = [[0, 0] for _ in range(self.num_rounds)]
def gen(self):
user_id = self.user_id
self.user_id += 1
rand_ratio = random.randint(0, self.cumulative_ratios[-1])
i = len(self.cumulative_ratios)
for idx, cumulative_ratio in enumerate(self.cumulative_ratios):
if rand_ratio >= cumulative_ratio:
continue
else:
i = idx + 1
break
total_rounds = i
current_round = 0
candidate_input = random.sample(self.candidate_inputs[current_round], 1)[0]
self.input_stats[0][0] += candidate_input.prompt_len
self.input_stats[0][1] += 1
prompt = f"{user_id} " + candidate_input.prompt
return_tokens = int(
random.gauss(self.mean_return_tokens_per_round[current_round], self.sigma)
)
if return_tokens <= 0:
return_tokens = self.mean_return_tokens_per_round[current_round]
start = 0
user_data = UserData(
user_id, current_round, total_rounds, prompt, return_tokens, start
)
self.user_stats[total_rounds - 1] += 1
return user_data
@synchronized()
def push(self, user_data, generated_text, len_itl):
self.output_stats[user_data.current_round][0] += len_itl + 1
self.output_stats[user_data.current_round][1] += 1
user_data.current_round += 1
if user_data.current_round >= user_data.total_rounds:
return
candidate_input = random.sample(
self.candidate_inputs[user_data.current_round], 1
)[0]
self.input_stats[user_data.current_round][0] += candidate_input.prompt_len
self.input_stats[user_data.current_round][1] += 1
user_data.prompt += generated_text + candidate_input.prompt
user_data.return_tokens = int(
random.gauss(
self.mean_return_tokens_per_round[user_data.current_round], self.sigma
)
)
if user_data.return_tokens <= 0:
user_data.return_tokens = self.mean_return_tokens_per_round[
user_data.current_round
]
interval = random.gauss(
self.mean_inter_round_interval[user_data.current_round], self.sigma
)
if interval <= 0:
interval = self.mean_inter_round_interval[user_data.current_round]
user_data.start = time.perf_counter() + interval
if len(self.multiturn_queue) == 0:
self.multiturn_queue.append(user_data)
else:
i = len(self.multiturn_queue)
for idx, d in enumerate(self.multiturn_queue):
if user_data.start < d.start:
i = idx
break
self.multiturn_queue.insert(idx, user_data)
@synchronized()
def pop(self):
if (
len(self.multiturn_queue)
and time.perf_counter() > self.multiturn_queue[0].start
):
return self.multiturn_queue.pop(0)
return self.gen()
def gen_payload(prompt, output_len):
payload = {
"text": prompt,
"sampling_params": {
"temperature": 0.0,
"max_new_tokens": output_len,
"ignore_eos": True,
},
"stream": True,
"stream_options": {"include_usage": True},
"lora_path": "",
"return_logprob": False,
"logprob_start_len": -1,
}
return payload
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)
async def async_request_sglang_generate(
user_data,
url,
atomic_counter,
):
"""
Sends a streaming request to the server. Gathers text token-by-token.
"""
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
headers = {}
generated_text = ""
ttft = 0.0
st = time.perf_counter()
most_recent_timestamp = st
output = RequestFuncOutput()
payload = gen_payload(user_data.prompt, user_data.return_tokens)
write_debug_log({"timestamp": st, "user_data": user_data.__dict__})
try:
async with session.post(url=url, json=payload, headers=headers) as response:
if response.status == 200:
prompt_tokens = 0
cached_tokens = 0
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)
if data.get("text"):
timestamp = time.perf_counter()
# First token
if ttft == 0.0:
ttft = time.perf_counter() - st
output.ttft = ttft
prompt_tokens = (data.get("meta_info") or {}).get(
"prompt_tokens", 0
)
cached_tokens = (data.get("meta_info") or {}).get(
"cached_tokens", 0
)
# Decoding phase
else:
output.itl.append(timestamp - most_recent_timestamp)
most_recent_timestamp = timestamp
generated_text = data["text"]
output.generated_text = generated_text
output.success = True
output.latency = latency
output.prompt_len = prompt_tokens
output.cached_tokens = cached_tokens
else:
output.error = response.reason or ""
output.success = False
except Exception as e:
output.success = False
output.error = str(e)
print(f"Request failed: {e}")
atomic_counter.increment(1)
return output
class AtomicCounter:
def __init__(self, initial_value=0):
self._value = initial_value
self.lock = threading.Lock()
@synchronized()
def increment(self, amount=1):
self._value += amount
@synchronized()
def get(self):
return self._value
class WorkloadGenerator:
def __init__(self, args):
config = load_config()
user_generator = UserGenerator(
config,
args.model_path,
args.dataset_path,
)
self.url = f"http://{args.host}:{args.port}/generate"
self.tokenizer = user_generator.tokenizer
self.start_time = None
self.finished_time = None
self.duration = args.duration
self.done = False
self.sent_requests = 0
self.completed_requests = 0
self.user_generator = user_generator
self.response_queue = queue.Queue()
self.performance_metrics = {
"ttft": [],
"latency": [],
"prompt_len": [],
"cached_tokens": [],
}
self.max_parallel = config["num_clients"]
self.atomic_counter = AtomicCounter()
async def handle_request(self, user_data):
try:
response = await async_request_sglang_generate(
user_data, self.url, self.atomic_counter
)
self.response_queue.put((user_data, response))
except Exception as e:
print(f"Request failed: {e}")
self.completed_requests += 1
def request_sender(self):
async def request_loop():
while True:
if self.sent_requests - self.completed_requests < self.max_parallel:
new_request = self.user_generator.pop()
if new_request:
asyncio.create_task(self.handle_request(new_request))
self.sent_requests += 1
else:
await asyncio.sleep(0.05)
continue
if time.perf_counter() - self.start_time > self.duration:
self.done = True
break
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(request_loop())
loop.close()
def response_handler(self):
while True:
try:
user_data, response = self.response_queue.get(timeout=10)
logger.info(
f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%"
)
if not response.success:
raise ValueError(f"Request failed with error: {response.error}")
self.user_generator.push(
user_data, response.generated_text, len(response.itl)
)
self.performance_metrics["ttft"].append(response.ttft)
self.performance_metrics["latency"].append(response.latency)
self.performance_metrics["prompt_len"].append(response.prompt_len)
self.performance_metrics["cached_tokens"].append(response.cached_tokens)
self.completed_requests += 1
self.finished_time = time.perf_counter()
except queue.Empty:
if self.done:
break
except ValueError as e:
print(f"Error processing response for client {user_data}: {e}")
continue
def run(self):
request_thread = threading.Thread(target=self.request_sender, daemon=True)
response_thread = threading.Thread(target=self.response_handler, daemon=True)
self.start_time = time.perf_counter()
request_thread.start()
response_thread.start()
request_thread.join()
response_thread.join()
performance_data = {
"summary": {
"total_requests": len(self.performance_metrics["ttft"]),
"average_ttft": sum(self.performance_metrics["ttft"])
/ len(self.performance_metrics["ttft"]),
"p90_ttft": sorted(self.performance_metrics["ttft"])[
int(0.9 * len(self.performance_metrics["ttft"]))
],
"median_ttft": sorted(self.performance_metrics["ttft"])[
len(self.performance_metrics["ttft"]) // 2
],
"average_latency": sum(self.performance_metrics["latency"])
/ len(self.performance_metrics["latency"]),
"p90_latency": sorted(self.performance_metrics["latency"])[
int(0.9 * len(self.performance_metrics["latency"]))
],
"median_latency": sorted(self.performance_metrics["latency"])[
len(self.performance_metrics["latency"]) // 2
],
"throughput": self.atomic_counter.get()
/ (self.finished_time - self.start_time),
"cache_hit_rate": (
0
if sum(self.performance_metrics["prompt_len"]) == 0
else sum(self.performance_metrics["cached_tokens"])
/ sum(self.performance_metrics["prompt_len"])
),
},
}
print("All requests completed")
print("Performance metrics summary:")
print(f" Total requests: {performance_data['summary']['total_requests']}")
print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}")
print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}")
print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}")
print(
f" Average latency: {performance_data['summary']['average_latency']:.2f}"
)
print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}")
print(f" Median latency: {performance_data['summary']['median_latency']:.2f}")
print(
f" Throughput: {performance_data['summary']['throughput']:.2f} requests per second"
)
print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}")
user_stats = self.user_generator.user_stats
input_stats = self.user_generator.input_stats
output_stats = self.user_generator.output_stats
print(f"round_ratios: {user_stats}")
print(
f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}"
)
print(
f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}"
)
return performance_data
def main():
global debug_log_file
args = parse_args()
if args.log_level == "debug":
logging.basicConfig(level=logging.DEBUG)
logger.info("use log_level debug")
# Initialize debug log file
debug_log_file = open(args.debug_log_file, "w")
else:
logging.basicConfig(level=logging.INFO)
logger.info("use log_level info")
performance_data = WorkloadGenerator(args).run()
# Close debug log file if it was opened
if debug_log_file:
debug_log_file.close()
if __name__ == "__main__":
main()

42
benchmark/hicache/bench_mix.sh Executable file
View File

@@ -0,0 +1,42 @@
#!/bin/bash
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib
rm -rf nohup.out && \
nohup python3 -m sglang.launch_server \
--attention-backend triton \
--model-path /code/models/Qwen3-32B/ \
--log-level info \
--tp 4 --mem-frac 0.25 \
--host 0.0.0.0 --port 33301 \
--enable-metrics --enable-cache-report \
--page-size 64 \
--enable-hierarchical-cache \
--hicache-ratio 2.5 --hicache-size 0 \
--hicache-io-backend kernel \
--hicache-mem-layout layer_first \
--hicache-write-policy write_through \
&
##################################################
export CONFIG_PATH=/tmp/bench_mix_config.json
# num_clients: Maximum number of concurrent client requests to be simulated
# round_ratios: Distribution of requests across rounds. Given sum(round_ratios) total requests,
# round_ratios[i] denotes the number of requests that will execute for (i+1) rounds
echo '{
"num_rounds": 10,
"num_clients": 60,
"round_ratios": [50, 25, 15, 15, 10, 10, 9, 8, 7, 6],
"mean_new_tokens_per_round": [1000, 400, 350, 300, 280, 260, 240, 220, 210, 200],
"mean_return_tokens_per_round": [100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
"mean_inter_round_interval": [30, 30, 30, 30, 30, 30, 30, 30, 30, 30]
}' > ${CONFIG_PATH}
rm -rf bench_mix.out && \
nohup python3 /sgl-workspace/sglang/benchmark/hicache/bench_mix.py \
--model-path /code/models/Qwen3-32B/ \
--dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \
--port 33301 \
--duration 600 \
> bench_mix.out &

View File

@@ -0,0 +1,517 @@
import argparse
import asyncio
import json
import queue
import random
import threading
import time
from datetime import datetime
from typing import Optional
import aiohttp
import numpy as np
import requests
from tqdm.asyncio import tqdm
from sglang.bench_serving import (
RequestFuncOutput,
get_tokenizer,
remove_prefix,
sample_random_requests,
)
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)
def parse_args():
parser = argparse.ArgumentParser(
description="Script to benchmark concurrent requests to a server."
)
parser.add_argument(
"--num-clients",
type=int,
default=256,
help="Number of concurrent clients",
)
parser.add_argument(
"--max-parallel",
type=int,
default=128,
help="Maximum number of parallel requests",
)
parser.add_argument(
"--request-length",
type=int,
default=512,
help="Length of each new request",
)
parser.add_argument(
"--output-length",
type=int,
default=64,
help="Length of each output",
)
parser.add_argument(
"--num-rounds",
type=int,
default=5,
help="Number of rounds per client",
)
parser.add_argument(
"--distribution",
type=str,
default="poisson",
choices=["poisson", "uniform"],
help="Distribution type for request intervals (poisson or uniform)",
)
parser.add_argument(
"--request-rate",
type=float,
default=1.0,
help="Average number of requests per second",
)
parser.add_argument(
"--host",
type=str,
default="localhost",
help="Server hostname or IP (default: localhost)",
)
parser.add_argument(
"--port",
type=int,
default=30000,
help="Server port (default: 30000)",
)
parser.add_argument(
"--model-path",
type=str,
default="meta-llama/Llama-3.1-8B-Instruct",
help="model path compatible with Hugging Face Transformers",
)
parser.add_argument(
"--dataset-path",
type=str,
default="",
help="local dataset to sample tokens from",
)
parser.add_argument(
"--log-file",
type=str,
default="performance_metrics.jsonl",
help="File to log performance metrics",
)
parser.add_argument(
"--disable-auto-run",
action="store_true",
help="If set, disable automatically testing with a range of request rates.",
)
parser.add_argument(
"--disable-random-sample",
action="store_true",
help="If set, disable random sampling of requests from the ShareGPT dataset.",
)
parser.add_argument(
"--sub-question-input-length",
type=int,
default=0,
help="Length of the sub question input for each request, if set 0 use request_length",
)
parser.add_argument(
"--ready-queue-policy",
type=str,
default="random",
help="Policy for popping requests from the ready queue (random or fifo)",
)
parser.add_argument(
"--tag",
type=str,
default="",
help="Tag of a certain run in the log file",
)
parser.add_argument("--seed", type=int, default=1, help="The random seed.")
parser.add_argument(
"--lora-path",
type=str,
default="",
help="String of LoRA path. Currently we only support benchmarking on a single LoRA adaptor.",
)
return parser.parse_args()
async def async_request_sglang_generate(
payload,
url,
pbar: Optional[tqdm] = None,
):
"""
Sends a streaming request to the server. Gathers text token-by-token.
"""
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
headers = {}
generated_text = ""
ttft = 0.0
st = time.perf_counter()
most_recent_timestamp = st
output = RequestFuncOutput()
try:
async with session.post(url=url, json=payload, headers=headers) as response:
if response.status == 200:
prompt_tokens = 0
cached_tokens = 0
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)
if data["text"]:
timestamp = time.perf_counter()
# First token
if ttft == 0.0:
ttft = time.perf_counter() - st
output.ttft = ttft
prompt_tokens = (data.get("meta_info") or {}).get(
"prompt_tokens", 0
)
cached_tokens = (data.get("meta_info") or {}).get(
"cached_tokens", 0
)
# Decoding phase
else:
output.itl.append(timestamp - most_recent_timestamp)
most_recent_timestamp = timestamp
generated_text = data["text"]
output.generated_text = generated_text
output.success = True
output.latency = latency
output.prompt_len = prompt_tokens
output.cached_tokens = cached_tokens
output.generated_len = len(output.itl) + 1
else:
output.error = response.reason or ""
output.success = False
except Exception as e:
output.success = False
output.error = str(e)
print(f"Request failed: {e}")
if pbar:
pbar.update(1)
return output
def gen_payload(prompt, output_len, lora_path=""):
payload = {
"text": prompt,
"sampling_params": {
"temperature": 0.0,
"max_new_tokens": output_len,
"ignore_eos": True,
},
"stream": True,
"stream_options": {"include_usage": True},
"lora_path": lora_path,
"return_logprob": False,
"logprob_start_len": -1,
}
return payload
def log_to_jsonl_file(data, file_path="performance_metrics.jsonl", tag=""):
"""Append the data with a timestamp and tag to the specified JSONL file."""
timestamped_data = {"timestamp": datetime.now().isoformat(), "tag": tag, **data}
try:
with open(file_path, "a") as file:
file.write(
json.dumps(timestamped_data) + "\n"
) # Write as a single line in JSONL format
except IOError as e:
print(f"Error writing to JSONL file: {e}")
class ReadyQueue:
"""
Thread-safe queue that can pop requests in different orders based on given policy.
"""
def __init__(self, init_requests=None, policy="random"):
self.lock = threading.Lock()
self.requests = init_requests or []
self.policy = policy
def append(self, item):
with self.lock:
self.requests.append(item)
def pop(self):
with self.lock:
if not self.requests:
return None
if self.policy == "random":
index = random.randrange(len(self.requests))
return self.requests.pop(index)
elif self.policy == "fifo":
return self.requests.pop(0)
else:
# todo, varying thinking time of clients
raise ValueError(f"{self.policy} not implemented")
class WorkloadGenerator:
def __init__(self, args):
# Construct the base URL for requests
self.url = f"http://{args.host}:{args.port}/generate"
self.tokenizer = get_tokenizer(args.model_path)
self.distribution = args.distribution
self.request_rate = args.request_rate
self.start_time = None
self.finished_time = None
self.sent_requests = 0
self.completed_requests = 0
self.candidate_inputs = sample_random_requests(
input_len=args.request_length,
output_len=args.output_length,
num_prompts=args.num_clients,
range_ratio=1.0,
tokenizer=self.tokenizer,
dataset_path=args.dataset_path,
random_sample=not args.disable_random_sample,
)
self.candidate_inputs = [i.prompt for i in self.candidate_inputs]
if args.sub_question_input_length != 0:
sub_question_input_length = args.sub_question_input_length
else:
sub_question_input_length = args.request_length
self.sub_question_inputs = sample_random_requests(
input_len=sub_question_input_length,
output_len=args.output_length,
num_prompts=args.num_clients * max(args.num_rounds - 1, 1),
range_ratio=1.0,
tokenizer=self.tokenizer,
dataset_path=args.dataset_path,
random_sample=not args.disable_random_sample,
)
init_requests = [
(
i,
gen_payload(
self.candidate_inputs[i], args.output_length, args.lora_path
),
)
for i in range(args.num_clients)
]
self.client_records = {
i: {"round": 0, "history": init_requests[i][1]["text"]}
for i in range(args.num_clients)
}
self.ready_queue = ReadyQueue(
init_requests=init_requests, policy=args.ready_queue_policy
)
self.candidate_inputs = self.candidate_inputs[args.num_clients :]
self.response_queue = queue.Queue()
self.pbar = tqdm(total=args.num_clients * args.num_rounds)
self.performance_metrics = {
"ttft": [],
"latency": [],
"prompt_len": [],
"cached_tokens": [],
"generated_len": [],
}
self.num_rounds = args.num_rounds
self.max_parallel = args.max_parallel
self.output_length = args.output_length
async def handle_request(self, item):
try:
client_id, payload = item
response = await async_request_sglang_generate(payload, self.url, self.pbar)
if self.pbar.n == self.pbar.total:
self.finished_time = time.perf_counter()
self.response_queue.put((client_id, response))
except Exception as e:
print(f"Request failed: {e}")
def request_sender(self):
async def request_loop():
while True:
if self.sent_requests - self.completed_requests < self.max_parallel:
new_request = self.ready_queue.pop()
if new_request:
asyncio.create_task(self.handle_request(new_request))
self.sent_requests += 1
else:
await asyncio.sleep(0.05)
continue
if self.pbar.n == self.pbar.total:
break
# Calculate Poisson-distributed wait time
if self.distribution == "poisson":
sleep_time = random.expovariate(self.request_rate)
elif self.distribution == "uniform":
avg_interval = (
1.0 / self.request_rate if self.request_rate > 0 else 1.0
)
sleep_time = random.uniform(0, 2 * avg_interval)
else:
raise ValueError("Invalid distribution type")
await asyncio.sleep(sleep_time) # Wait before sending the next request
# Create and run the event loop for asynchronous requests
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(request_loop())
loop.close()
def response_handler(self):
while True:
try:
client_id, response = self.response_queue.get(
timeout=10
) # Block until response is available
if not response.success:
raise ValueError(f"Request failed with error: {response.error}")
self.client_records[client_id]["history"] += response.generated_text
self.client_records[client_id]["round"] += 1
self.performance_metrics["ttft"].append(response.ttft)
self.performance_metrics["latency"].append(response.latency)
self.performance_metrics["prompt_len"].append(response.prompt_len)
self.performance_metrics["cached_tokens"].append(response.cached_tokens)
self.performance_metrics["generated_len"].append(response.generated_len)
self.completed_requests += 1
if self.client_records[client_id]["round"] < self.num_rounds:
# append new request to client's history
self.client_records[client_id][
"history"
] += self.sub_question_inputs.pop().prompt
self.ready_queue.append(
(
client_id,
gen_payload(
self.client_records[client_id]["history"],
self.output_length,
args.lora_path,
),
)
)
except queue.Empty:
if self.pbar.n == self.pbar.total:
break
except ValueError as e:
print(f"Error processing response for client {client_id}: {e}")
continue
def run(self):
request_thread = threading.Thread(target=self.request_sender, daemon=True)
response_thread = threading.Thread(target=self.response_handler, daemon=True)
self.start_time = time.perf_counter()
request_thread.start()
response_thread.start()
request_thread.join()
response_thread.join()
self.pbar.close()
duration = self.finished_time - self.start_time
performance_data = {
"summary": {
"total_requests": len(self.performance_metrics["ttft"]),
"request_rate": self.request_rate,
"average_ttft": sum(self.performance_metrics["ttft"])
/ len(self.performance_metrics["ttft"]),
"p90_ttft": sorted(self.performance_metrics["ttft"])[
int(0.9 * len(self.performance_metrics["ttft"]))
],
"median_ttft": sorted(self.performance_metrics["ttft"])[
len(self.performance_metrics["ttft"]) // 2
],
"average_latency": sum(self.performance_metrics["latency"])
/ len(self.performance_metrics["latency"]),
"p90_latency": sorted(self.performance_metrics["latency"])[
int(0.9 * len(self.performance_metrics["latency"]))
],
"median_latency": sorted(self.performance_metrics["latency"])[
len(self.performance_metrics["latency"]) // 2
],
"input_token_throughput": sum(self.performance_metrics["prompt_len"])
/ duration,
"output_token_throughput": sum(
self.performance_metrics["generated_len"]
)
/ duration,
"throughput": self.pbar.total / duration,
"cache_hit_rate": (
0
if sum(self.performance_metrics["prompt_len"]) == 0
else sum(self.performance_metrics["cached_tokens"])
/ sum(self.performance_metrics["prompt_len"])
),
},
}
print("All requests completed")
print("Performance metrics summary:")
print(
f" Total requests: {performance_data['summary']['total_requests']} at {performance_data['summary']['request_rate']} requests per second"
)
print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}")
print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}")
print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}")
print(
f" Average latency: {performance_data['summary']['average_latency']:.2f}"
)
print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}")
print(f" Median latency: {performance_data['summary']['median_latency']:.2f}")
print(
f" Input token throughput: {performance_data['summary']['input_token_throughput']:.2f} tokens per second"
)
print(
f" Output token throughput: {performance_data['summary']['output_token_throughput']:.2f} tokens per second"
)
print(
f" Request Throughput: {performance_data['summary']['throughput']:.2f} requests per second"
)
print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}")
return performance_data
if __name__ == "__main__":
args = parse_args()
flush_cache_url = f"http://{args.host}:{args.port}/flush_cache"
random.seed(args.seed)
np.random.seed(args.seed)
if args.disable_auto_run:
print("Running with specified request rate...")
request_rates = [args.request_rate]
else:
print("Auto-running with different request rates...")
request_rates = [16, 14, 12, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
for rate in request_rates:
args.request_rate = rate
requests.post(flush_cache_url)
time.sleep(1)
performance_data = WorkloadGenerator(args).run()
log_to_jsonl_file(performance_data, args.log_file, tag=args.tag)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,590 @@
import json
import os
import pickle
import random
from pathlib import Path
from typing import List, Optional, Tuple, Union
import numpy as np
from nextqa import NExTQALoader
# from nextqa.video import , VideoPrompt
from tqdm.asyncio import tqdm
from transformers import PreTrainedTokenizerBase
SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
from sglang.bench_serving import (
download_and_cache_file,
gen_prompt,
get_gen_prefix_cache_path,
)
from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path
from sglang.srt.entrypoints.openai.protocol import ChatCompletionMessageContentPart
from sglang.utils import encode_video_base64
# type of content fields, can be only prompts or with images/videos
MsgContent = Union[str, List[ChatCompletionMessageContentPart]]
# A list of all the conversations. Each conversation is a list of
# tuples. If multiturn is not enabled, the length of list is 1,
# containing only the first Q&A pair.
# For the shared prefix workload (synthetic, loogle, nextqa), it
# is a list of conversations sharing the same prefix (synthetic,
# doc, video)
SampleOutput = List[List[Tuple[MsgContent, int, int]]]
def common_filter_chat(
num_requests: int,
new_dataset: List,
tokenizer: PreTrainedTokenizerBase,
min_prompt_len: Optional[int],
min_output_len: Optional[int],
max_prompt_len: Optional[int],
max_output_len: Optional[int],
fixed_output_len: Optional[int],
) -> SampleOutput:
# Filter out sequences that are too long or too short
filtered_dataset: SampleOutput = []
l = 0
input_tokens = 0
output_tokens = 0
while l < num_requests:
for i in range(len(new_dataset)):
if l == num_requests:
break
processed = []
for j in new_dataset[i]:
# Tokenize the prompts and completions.
prompt = j[0]
prompt_token_ids = tokenizer.encode(prompt)
prompt_len = len(prompt_token_ids)
completion = j[1]
completion_token_ids = tokenizer.encode(completion)
output_len = (
len(completion_token_ids)
if fixed_output_len is None
else fixed_output_len
)
if (
min_prompt_len is not None
and prompt_len < min_prompt_len
or min_output_len is not None
and output_len < min_output_len
or max_prompt_len is not None
and prompt_len > max_prompt_len
or max_output_len is not None
and output_len > max_output_len
):
# Prune too short sequences.
continue
input_tokens += prompt_len
output_tokens += output_len
processed.append((prompt, prompt_len, output_len))
if len(processed) != 0:
filtered_dataset.append(processed)
l += 1
print(f"#Input tokens: {input_tokens}")
print(f"#Output tokens: {output_tokens}")
return filtered_dataset
def sample_sharegpt_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
disable_shuffle: bool = False,
enable_multiturn: bool = True,
fixed_output_len: Optional[int] = None,
) -> SampleOutput:
if fixed_output_len is not None and fixed_output_len < 4:
raise ValueError("output_len too small")
# Download sharegpt if necessary
if not os.path.isfile(dataset_path):
dataset_path = download_and_cache_file(SHAREGPT_URL)
# Load the dataset.
with open(dataset_path) as f:
dataset = json.load(f)
# Filter out the conversations with less than 2 turns.
dataset = [data for data in dataset if len(data["conversations"]) >= 2]
# Keep one conversation in one list
new_dataset = []
for data in dataset:
if len(data["conversations"]) % 2 != 0:
continue
if data["conversations"][0]["from"] != "human":
continue
chat = []
total_len = 2
if enable_multiturn:
total_len = len(data["conversations"])
for i in range(0, total_len, 2):
# One user One Assistant
chat.append(
(
data["conversations"][i]["value"],
data["conversations"][i + 1]["value"],
)
)
new_dataset.append(chat)
if not disable_shuffle:
# Shuffle the dataset.
random.shuffle(new_dataset)
# Filter out sequences that are too long or too short
filtered_dataset: SampleOutput = common_filter_chat(
num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len
)
return filtered_dataset
def sample_ultrachat_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
disable_shuffle: bool = False,
enable_multiturn: bool = True,
fixed_output_len: Optional[int] = None,
) -> SampleOutput:
if fixed_output_len is not None and fixed_output_len < 4:
raise ValueError("output_len too small")
# Load the dataset
dataset = []
with open(dataset_path) as f:
while True:
line = f.readline()
if not line:
break
dataset.append(json.loads(line))
# Filter out the conversations with less than 2 turns.
dataset = [data for data in dataset if len(data["data"]) >= 2]
# Keep one conversation in one list
new_dataset = []
for data in dataset:
if len(data["data"]) % 2 != 0:
continue
chat = []
total_len = 2
if enable_multiturn:
total_len = len(data["data"])
for i in range(0, total_len, 2):
# One user One Assistant
chat.append((data["data"][i], data["data"][i + 1]))
new_dataset.append(chat)
# Shuffle the dataset.
if not disable_shuffle:
random.shuffle(new_dataset)
# Filter out sequences that are too long or too short
filtered_dataset: SampleOutput = common_filter_chat(
num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len
)
return filtered_dataset
def sample_loogle_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
disable_shuffle: bool = False,
enable_multiturn: bool = True,
enable_shared_prefix: bool = False,
fixed_output_len: Optional[int] = None,
) -> SampleOutput:
if fixed_output_len is not None and fixed_output_len < 4:
raise ValueError("output_len too small")
# Load the dataset
dataset = []
with open(dataset_path) as f:
while True:
line = f.readline()
if not line:
break
dataset.append(json.loads(line))
# Keep one conversation in one list
new_dataset = []
# TODO: Add shared prefix support for loogle
# NOTE: Now we preprocess it only for chat
for data in dataset:
chat = []
if (
"qa_pairs" not in data
or data["qa_pairs"] == "none"
or len(data["qa_pairs"]) == 0
):
# If Q is none (for summarization),
# We add a question for summarization
# And keep the summary up to 1024 words
chat.append(
(
"Input: "
+ data["input"]
+ " Question: "
+ "Please summarize the input",
data["input"][:1024],
)
)
new_dataset.append(chat)
else:
qa_pairs = eval(data["qa_pairs"])
for i, qa in enumerate(qa_pairs):
if i == 0 or enable_shared_prefix:
# Combine input with the first Q
chat.append(
("Input: " + data["input"] + " Question: " + qa["Q"], qa["A"])
)
elif enable_multiturn:
chat.append((qa["Q"], qa["A"]))
new_dataset.append(chat)
# Shuffle the dataset.
if not disable_shuffle:
random.shuffle(new_dataset)
# Filter out sequences that are too long or too short
filtered_dataset: SampleOutput = common_filter_chat(
num_requests, new_dataset, tokenizer, 4, None, None, None, fixed_output_len
)
return filtered_dataset
def sample_nextqa_requests(
dataset_path: str,
num_requests: int,
tokenizer: PreTrainedTokenizerBase,
max_frames: int, # Specific for video
model_path: str,
disable_shuffle: bool = False,
enable_multiturn: bool = True, # No multiturn support for now
backend: str = "sglang-oai",
chat_template_name: Optional[str] = None,
fixed_output_len: Optional[int] = None,
) -> SampleOutput:
"""
Example of messages:
message = {
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": base64_data}},
{"type": "text", "text": video.prompt},
],
}
"""
if fixed_output_len is None:
fixed_output_len = 4096
# TODO: Check for multiturn
dataset = NExTQALoader(video_dir=dataset_path, max_frames=max_frames)
new_dataset = []
for v in dataset:
new_dataset.append(v)
if not disable_shuffle:
random.shuffle(new_dataset)
# TODO: prompt len can get from server side
filtered_dataset = []
l = 0
while l < num_requests:
for i in range(len(new_dataset)):
if l == num_requests:
break
video = new_dataset[i]
# text prompt
prompt = video.prompt
# NOTE: Chat Template is a must for video benchmark because we have to
# add special image token for later expansion
if backend == "sglang" or backend == "sglang-native":
if "chat_template" in tokenizer.init_kwargs:
chat_template = get_chat_template(tokenizer.get_chat_template())
elif chat_template_name is not None:
chat_template = get_chat_template(chat_template_name)
else:
chat_template = get_chat_template_by_model_path(model_path)
prompt = chat_template.image_token + prompt
prompt_token_ids = tokenizer(prompt).input_ids
prompt_len = len(prompt_token_ids)
output_len = fixed_output_len # max output len, not real output len
# video input
base64_data = encode_video_base64(video.path, video.num_frames)
# NOTE: This will be replaced by the expanded length from the server
prompt_len += video.num_frames
# add to content
content = [
{"type": "image_url", "image_url": {"url": base64_data}},
{"type": "text", "text": prompt},
]
filtered_dataset.append([(content, prompt_len, output_len)])
l += 1
return filtered_dataset
def sample_random_requests(
input_len: int,
output_len: int,
num_prompts: int,
range_ratio: float,
tokenizer: PreTrainedTokenizerBase,
dataset_path: str,
disable_shuffle: bool = False,
) -> SampleOutput:
input_lens = np.random.randint(
max(int(input_len * range_ratio), 1),
input_len + 1,
size=num_prompts,
)
output_lens = np.random.randint(
int(output_len * range_ratio),
output_len + 1,
size=num_prompts,
)
if True:
# Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens
# Download sharegpt if necessary
if not os.path.isfile(dataset_path):
dataset_path = download_and_cache_file(SHAREGPT_URL)
# Load the dataset.
with open(dataset_path) as f:
dataset = json.load(f)
# Filter out the conversations with less than 2 turns.
dataset = [data for data in dataset if len(data["conversations"]) >= 2]
# Only keep the first two turns of each conversation.
dataset = [
(data["conversations"][0]["value"], data["conversations"][1]["value"])
for data in dataset
]
if not disable_shuffle:
# Shuffle the dataset.
random.shuffle(dataset)
# Filter out sequences that are too long or too short
input_requests: SampleOutput = []
for data in dataset:
i = len(input_requests)
if i == num_prompts:
break
# Tokenize the prompts and completions.
prompt = data[0]
prompt_token_ids = tokenizer.encode(prompt)
prompt_len = len(prompt_token_ids)
# Skip empty prompt
if prompt_len == 0:
continue
if prompt_len > input_lens[i]:
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]))])
print(f"#Input tokens: {np.sum(input_lens)}")
print(f"#Output tokens: {np.sum(output_lens)}")
return input_requests
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)
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 = (
f"gsp_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_"
f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_"
f"{tokenizer.__class__.__name__}.pkl"
)
return cache_dir / cache_key
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,
args,
disable_shuffle: bool = False,
) -> SampleOutput:
"""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:
return pickle.load(f)
print("\nGenerating new input data...")
# 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
for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
system_prompt = system_prompts[group_idx]
input_requests.append([])
for prompt_idx in tqdm(
range(prompts_per_group), desc="Generating questions", leave=False
):
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[-1].append((full_prompt, prompt_len, output_len))
total_input_tokens += prompt_len
total_output_tokens += output_len
if not disable_shuffle:
# Shuffle questions
random.shuffle(input_requests)
# Print statistics
print(f"\nGenerated shared prefix dataset statistics:")
print(f"Number of groups: {num_groups}")
print(f"Prompts per group: {prompts_per_group}")
print(f"Total prompts: {len(input_requests) * prompts_per_group}")
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"
)
# 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)
return input_requests
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,
disable_shuffle=args.disable_shuffle,
enable_multiturn=args.enable_multiturn,
fixed_output_len=args.fixed_output_len,
)
elif args.dataset_name == "ultrachat":
input_requests = sample_ultrachat_requests(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
tokenizer=tokenizer,
disable_shuffle=args.disable_shuffle,
enable_multiturn=args.enable_multiturn,
fixed_output_len=args.fixed_output_len,
)
elif args.dataset_name == "loogle":
input_requests = sample_loogle_requests(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
tokenizer=tokenizer,
disable_shuffle=args.disable_shuffle,
enable_multiturn=args.enable_multiturn,
enable_shared_prefix=args.enable_shared_prefix,
fixed_output_len=args.fixed_output_len,
)
elif args.dataset_name == "nextqa":
input_requests = sample_nextqa_requests(
dataset_path=args.dataset_path,
num_requests=args.num_prompts,
tokenizer=tokenizer,
max_frames=args.max_frames,
model_path=args.model,
disable_shuffle=args.disable_shuffle,
enable_multiturn=args.enable_multiturn,
backend=args.backend,
chat_template_name=args.chat_template,
fixed_output_len=args.fixed_output_len,
)
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(
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,
args=args,
tokenizer=tokenizer,
)
else:
raise ValueError(f"Unknown dataset: {args.dataset_name}")
return input_requests

66
benchmark/hicache/download.sh Executable file
View File

@@ -0,0 +1,66 @@
#!/usr/bin/bash
# The usage function
usage() {
echo "Usage: $0 {sharegpt|ultragpt|loogle|nextqa|all}"
exit 1
}
# The download function
download() {
case "$1" in
sharegpt)
echo $1
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
;;
ultragpt)
echo $1
# Questions about the world
wget https://cloud.tsinghua.edu.cn/seafhttp/files/be1d7b87-22ca-449e-a6a7-c61d1ea7e010/ultrachat_release_230407.json
# Writing and Creation
wget https://cloud.tsinghua.edu.cn/seafhttp/files/61742d2a-25e2-4d08-b2b9-15f47ae50ace/ultrachat_material_release_230417.json
wget https://cloud.tsinghua.edu.cn/seafhttp/files/f71f6aa6-d346-4b16-85b7-8502efa3d608/ultrachat_material_release_230412.json
# External materials
wget https://cloud.tsinghua.edu.cn/seafhttp/files/42d22e28-e899-4975-a70f-5eda163e265d/ultrachat_existent_material_release_230420.json.gz
gunzip ultrachat_existent_material_release_230420.json.gz
;;
loogle)
echo $1
git lfs install
git clone git@hf.co:datasets/bigainlco/LooGLE
unzip LooGLE/data.zip
;;
nextqa)
echo $1
git lfs install
git clone https://huggingface.co/datasets/lmms-lab/NExTQA
unzip NExTQA/videos.zip
;;
*)
usage
exit 1
;;
esac
}
# Arg check
if [ "$#" -ne 1 ]; then
usage
fi
# Invoke
case "$1" in
sharegpt|ultragpt|loogle|nextqa)
download "$1"
;;
all)
download sharegpt
download ultragpt
download loogle
download nextqa
;;
*)
usage
;;
esac

159
benchmark/hicache/nextqa.py Normal file
View File

@@ -0,0 +1,159 @@
import os
import sys
from typing import List
import av
from datasets import load_dataset
def find_video_files(video_dir) -> List[str]:
if os.path.isfile(video_dir):
return [video_dir]
video_files = []
for root, dirs, files in os.walk(video_dir):
for file in files:
if file.endswith((".mp4", ".avi", ".mov")):
video_files.append(os.path.join(root, file))
# if file is dir
elif os.path.isdir(file):
video_files.extend(find_video_files(file))
return video_files
def video_frames(video_path, max_frames) -> int:
container = av.open(video_path)
total_frames = container.streams.video[0].frames
return min(total_frames, max_frames)
class Video:
def __init__(self, video_path, num_frames):
self.path = video_path
self.num_frames = num_frames
def __str__(self):
return f"Video({self.path}, {self.num_frames})"
def __iter__(self):
return iter((self.path, self.num_frames))
class VideoPrompt(Video):
def __init__(self, video_path, num_frames, prompt):
super().__init__(video_path, num_frames)
self.prompt = prompt
def __str__(self):
return f"VideoPrompt({self.path}, {self.num_frames}, {self.prompt})"
def __iter__(self):
return iter((self.path, self.num_frames, self.prompt))
class VideoLoader:
pass
class VideoFileLoader(VideoLoader):
"""
Load all the videos in a directory
"""
def __init__(self, video_dir, batch_size=1, max_frames=sys.maxsize):
super().__init__()
self.video_dir = video_dir
self.video_files = find_video_files(video_dir)
self.batch_size = batch_size
self.max_frames = max_frames
print(f"batch_size: {batch_size}, max_frames: {max_frames}")
def __iter__(self): # (file, number of frames)
if self.batch_size == 1:
for video_file in self.video_files:
yield Video(video_file, video_frames(video_file, self.max_frames))
else:
batch = []
for video_file in self.video_files:
video = Video(video_file, video_frames(video_file, self.max_frames))
batch.append(video)
if len(batch) == self.batch_size:
yield batch
batch = []
class NExTQALoader(VideoLoader):
"""
Load vdideos and prompts from NExT dataset
set: train, test or validation
"""
def __init__(
self, video_dir, batch_size=1, max_frames=sys.maxsize, dset="test", task="OE"
):
"""
task: 'MV' or 'OE'
"""
super().__init__()
self.task = task
print(f"Loading the {dset} data of {task} from lmms-lab/NExTQA")
self.ds = load_dataset("lmms-lab/NExTQA", task)
self.ds = self.ds[dset]
# self.n = ds.num_rows
self.video_dir = video_dir
self.video_files = find_video_files(video_dir)
self.video_to_path = dict()
for video_file in self.video_files:
video_id = video_file.split("/")[-1].split(".")[0]
self.video_to_path[video_id] = video_file
self.batch_size = batch_size
self.max_frames = max_frames
def get_video_prompt(self, entry, max_frames) -> VideoPrompt:
# Get video
video_id = entry["video"]
video_path = self.video_to_path[video_id]
assert os.path.exists(video_path), f"Video not found: {video_path}"
num_frames = min(entry["frame_count"], max_frames)
video = Video(video_path, num_frames)
prompt = entry["question"] + "?"
if self.task == "MC": # add choices
prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}'
return VideoPrompt(video_path, num_frames, prompt)
def __iter__(self):
if self.batch_size == 1:
for entry in self.ds:
yield self.get_video_prompt(entry, self.max_frames)
else:
batch = []
for entry in self.ds:
video = self.get_video_prompt(entry, self.max_frames)
batch.append(video)
if len(batch) == self.batch_size:
yield batch
batch = []
# main
if __name__ == "__main__":
video_dir = "./videos"
# video_loader = VideoFileLoader(video_dir, batch_size=16)
# for batch in video_loader:
# print(f"Number of videos in batch: {len(batch)}")
# for video_file, num_frames in batch:
# print(f"Video: {video_file} number of frames: {num_frames}")
video_loader = NExTQALoader(video_dir, batch_size=16, dset="test", task="OE")
for batch in video_loader:
print(f"Number of videos in batch: {len(batch)}")
for video_file, num_frames, prompt in batch:
print(
f"Video: {video_file} number of frames: {num_frames}, prompt: {prompt}"
)
# break
# for video_file, prompt in batch:
# print(f"Video: {video_file} prompt: {prompt}")
# break