554 lines
18 KiB
Python
554 lines
18 KiB
Python
"""
|
|
Benchmark the latency of running a single static batch.
|
|
This script does not launch a server and uses the low-level APIs.
|
|
It accepts arguments similar to those of launch_server.py.
|
|
|
|
# Usage (latency test)
|
|
## with dummy weights:
|
|
python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy
|
|
## sweep through multiple data points and store (append) the results in a jsonl file:
|
|
python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --result-filename out.jsonl
|
|
## do some changes, and store the results under a different run_name:
|
|
python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --result-filename out.jsonl --run-name after
|
|
## plot the results in series of lines:
|
|
python -m sglang.bench_latency --result-filename out.jsonl --graph-sql="select run_name, batch_size, prefill_throughput from results"
|
|
|
|
# Usage (correctness test):
|
|
python -m sglang.bench_latency --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct
|
|
|
|
## Reference output (of the correctness test above, can be gpu dependent):
|
|
input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]]
|
|
|
|
prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633],
|
|
[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633],
|
|
[ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]],
|
|
device='cuda:0')
|
|
|
|
prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141],
|
|
[-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781],
|
|
[-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]],
|
|
device='cuda:0')
|
|
|
|
========== Prompt 0 ==========
|
|
<s> The capital of France is Paris.
|
|
The capital of the United States is Washington, D.C.
|
|
|
|
|
|
========== Prompt 1 ==========
|
|
<s> The capital of the United Kindom is London.
|
|
The capital of the United Kingdom is London.
|
|
The capital of the
|
|
|
|
========== Prompt 2 ==========
|
|
<s> Today is a sunny day and I like to go for a walk in the park.
|
|
I'm going to the park
|
|
"""
|
|
|
|
import argparse
|
|
import dataclasses
|
|
import itertools
|
|
import json
|
|
import logging
|
|
import multiprocessing
|
|
import os
|
|
import sqlite3
|
|
import time
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from sglang.srt.configs.model_config import ModelConfig
|
|
from sglang.srt.hf_transformers_utils import get_tokenizer
|
|
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
|
from sglang.srt.sampling.sampling_params import SamplingParams
|
|
from sglang.srt.server import _set_envs_and_config
|
|
from sglang.srt.server_args import PortArgs, ServerArgs
|
|
from sglang.srt.utils import (
|
|
configure_logger,
|
|
kill_child_process,
|
|
suppress_other_loggers,
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class BenchArgs:
|
|
run_name: str = "before"
|
|
batch_size: Tuple[int] = (1,)
|
|
input_len: Tuple[int] = (1024,)
|
|
output_len: Tuple[int] = (16,)
|
|
result_filename: str = ""
|
|
correctness_test: bool = False
|
|
# This is only used for correctness test
|
|
cut_len: int = 4
|
|
# Plotting args
|
|
graph_sql: str = (
|
|
"select run_name, batch_size, prefill_throughput from results where run_name='before'"
|
|
)
|
|
graph_filename: str = "out.png"
|
|
|
|
@staticmethod
|
|
def add_cli_args(parser: argparse.ArgumentParser):
|
|
parser.add_argument("--run-name", type=str, default=BenchArgs.run_name)
|
|
parser.add_argument(
|
|
"--batch-size", type=int, nargs="+", default=BenchArgs.batch_size
|
|
)
|
|
parser.add_argument(
|
|
"--input-len", type=int, nargs="+", default=BenchArgs.input_len
|
|
)
|
|
parser.add_argument(
|
|
"--output-len", type=int, nargs="+", default=BenchArgs.output_len
|
|
)
|
|
parser.add_argument(
|
|
"--result-filename", type=str, default=BenchArgs.result_filename
|
|
)
|
|
parser.add_argument("--correctness-test", action="store_true")
|
|
parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len)
|
|
# graphing
|
|
parser.add_argument("--graph-sql", type=str, default=BenchArgs.graph_sql)
|
|
parser.add_argument(
|
|
"--graph-filename", type=str, default=BenchArgs.graph_filename
|
|
)
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, args: argparse.Namespace):
|
|
# use the default value's type to case the args into correct types.
|
|
attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
|
|
return cls(
|
|
**{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs}
|
|
)
|
|
|
|
|
|
def load_model(server_args, port_args, tp_rank):
|
|
suppress_other_loggers()
|
|
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
|
|
|
|
model_config = ModelConfig(
|
|
server_args.model_path,
|
|
server_args.trust_remote_code,
|
|
context_length=server_args.context_length,
|
|
model_override_args=json.loads(server_args.json_model_override_args),
|
|
)
|
|
model_runner = ModelRunner(
|
|
model_config=model_config,
|
|
mem_fraction_static=server_args.mem_fraction_static,
|
|
gpu_id=tp_rank,
|
|
tp_rank=tp_rank,
|
|
tp_size=server_args.tp_size,
|
|
nccl_port=port_args.nccl_port,
|
|
server_args=server_args,
|
|
)
|
|
rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
|
|
tokenizer = get_tokenizer(
|
|
server_args.tokenizer_path,
|
|
tokenizer_mode=server_args.tokenizer_mode,
|
|
trust_remote_code=server_args.trust_remote_code,
|
|
)
|
|
if server_args.tp_size > 1:
|
|
dist.barrier()
|
|
return model_runner, tokenizer
|
|
|
|
|
|
def prepare_inputs_for_correctness_test(bench_args, tokenizer):
|
|
prompts = [
|
|
"The capital of France is",
|
|
"The capital of the United Kindom is",
|
|
"Today is a sunny day and I like",
|
|
]
|
|
input_ids = [tokenizer.encode(p) for p in prompts]
|
|
sampling_params = SamplingParams(
|
|
temperature=0,
|
|
max_new_tokens=BenchArgs.output_len,
|
|
)
|
|
|
|
reqs = []
|
|
for i in range(len(prompts)):
|
|
assert len(input_ids[i]) > bench_args.cut_len
|
|
|
|
tmp_input_ids = input_ids[i][: bench_args.cut_len]
|
|
req = Req(
|
|
rid=i,
|
|
origin_input_text=prompts[i],
|
|
origin_input_ids=tmp_input_ids,
|
|
sampling_params=sampling_params,
|
|
)
|
|
req.prefix_indices = []
|
|
req.fill_ids = req.origin_input_ids
|
|
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
|
|
reqs.append(req)
|
|
|
|
return input_ids, reqs
|
|
|
|
|
|
def prepare_extend_inputs_for_correctness_test(
|
|
bench_args, input_ids, reqs, model_runner
|
|
):
|
|
for i in range(len(reqs)):
|
|
req = reqs[i]
|
|
req.fill_ids += input_ids[i][bench_args.cut_len :]
|
|
req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
|
|
i, : bench_args.cut_len
|
|
]
|
|
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
|
|
return reqs
|
|
|
|
|
|
def prepare_synthetic_inputs_for_latency_test(batch_size, input_len):
|
|
input_ids = np.ones((batch_size, input_len), dtype=np.int32)
|
|
sampling_params = SamplingParams(
|
|
temperature=0,
|
|
max_new_tokens=BenchArgs.output_len,
|
|
)
|
|
|
|
reqs = []
|
|
for i in range(len(input_ids)):
|
|
req = Req(
|
|
rid=i,
|
|
origin_input_text="",
|
|
origin_input_ids=list(input_ids[i]),
|
|
sampling_params=sampling_params,
|
|
)
|
|
req.prefix_indices = []
|
|
req.fill_ids = req.origin_input_ids
|
|
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
|
|
reqs.append(req)
|
|
|
|
return reqs
|
|
|
|
|
|
@torch.inference_mode()
|
|
def extend(reqs, model_runner):
|
|
batch = ScheduleBatch.init_new(
|
|
reqs=reqs,
|
|
req_to_token_pool=model_runner.req_to_token_pool,
|
|
token_to_kv_pool=model_runner.token_to_kv_pool,
|
|
tree_cache=None,
|
|
model_config=model_runner.model_config,
|
|
)
|
|
batch.prepare_for_extend()
|
|
model_worker_batch = batch.get_model_worker_batch()
|
|
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
|
|
logits_output = model_runner.forward(forward_batch)
|
|
next_token_ids = model_runner.sample(logits_output, forward_batch)
|
|
return next_token_ids, logits_output.next_token_logits, batch
|
|
|
|
|
|
@torch.inference_mode()
|
|
def decode(input_token_ids, batch, model_runner):
|
|
batch.output_ids = input_token_ids
|
|
batch.prepare_for_decode()
|
|
model_worker_batch = batch.get_model_worker_batch()
|
|
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
|
|
logits_output = model_runner.forward(forward_batch)
|
|
next_token_ids = model_runner.sample(logits_output, forward_batch)
|
|
return next_token_ids, logits_output.next_token_logits
|
|
|
|
|
|
def correctness_test(
|
|
server_args,
|
|
port_args,
|
|
bench_args,
|
|
tp_rank,
|
|
):
|
|
configure_logger(server_args, prefix=f" TP{tp_rank}")
|
|
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
|
|
|
|
# Load the model
|
|
model_runner, tokenizer = load_model(server_args, port_args, tp_rank)
|
|
|
|
# Prepare inputs
|
|
input_ids, reqs = prepare_inputs_for_correctness_test(bench_args, tokenizer)
|
|
rank_print(f"\n{input_ids=}\n")
|
|
|
|
if bench_args.cut_len > 0:
|
|
# Prefill
|
|
next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
|
|
rank_print(f"prefill logits (first half): {next_token_logits} \n")
|
|
|
|
# Prepare extend inputs
|
|
reqs = prepare_extend_inputs_for_correctness_test(
|
|
bench_args, input_ids, reqs, model_runner
|
|
)
|
|
|
|
# Extend
|
|
next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
|
|
rank_print(f"prefill logits (final): {next_token_logits} \n")
|
|
|
|
# Decode
|
|
output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))]
|
|
for _ in range(bench_args.output_len[0] - 1):
|
|
next_token_ids, _ = decode(next_token_ids, batch, model_runner)
|
|
next_token_ids_list = next_token_ids.tolist()
|
|
for i in range(len(reqs)):
|
|
output_ids[i].append(next_token_ids_list[i])
|
|
|
|
# Print
|
|
for i in range(len(reqs)):
|
|
rank_print(f"========== Prompt {i} ==========")
|
|
rank_print(tokenizer.decode(output_ids[i]), "\n")
|
|
|
|
|
|
def synchronize(device):
|
|
if device == "cuda":
|
|
torch.cuda.synchronize()
|
|
elif device == "xpu":
|
|
torch.xpu.synchronize()
|
|
|
|
|
|
def latency_test_run_once(
|
|
run_name, model_runner, rank_print, reqs, batch_size, input_len, output_len, device
|
|
):
|
|
max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len)
|
|
if batch_size > max_batch_size:
|
|
rank_print(
|
|
f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit"
|
|
)
|
|
return
|
|
|
|
# Clear the pools.
|
|
model_runner.req_to_token_pool.clear()
|
|
model_runner.token_to_kv_pool.clear()
|
|
|
|
measurement_results = {
|
|
"run_name": run_name,
|
|
"batch_size": batch_size,
|
|
"input_len": input_len,
|
|
"output_len": output_len,
|
|
}
|
|
|
|
tot_latency = 0
|
|
|
|
# Prefill
|
|
synchronize(device)
|
|
tic = time.time()
|
|
next_token_ids, _, batch = extend(reqs, model_runner)
|
|
synchronize(device)
|
|
prefill_latency = time.time() - tic
|
|
tot_latency += prefill_latency
|
|
throughput = input_len * batch_size / prefill_latency
|
|
rank_print(
|
|
f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s"
|
|
)
|
|
measurement_results["prefill_latency"] = prefill_latency
|
|
measurement_results["prefill_throughput"] = throughput
|
|
|
|
# Decode
|
|
decode_latencies = []
|
|
for i in range(output_len - 1):
|
|
synchronize(device)
|
|
tic = time.time()
|
|
next_token_ids, _ = decode(next_token_ids, batch, model_runner)
|
|
synchronize(device)
|
|
latency = time.time() - tic
|
|
tot_latency += latency
|
|
throughput = batch_size / latency
|
|
decode_latencies.append(latency)
|
|
if i < 5:
|
|
rank_print(
|
|
f"Decode. latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s"
|
|
)
|
|
|
|
# record decode timing from 2nd output
|
|
if output_len > 1:
|
|
med_decode_latency = np.median(decode_latencies)
|
|
med_decode_throughput = batch_size / med_decode_latency
|
|
rank_print(
|
|
f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s"
|
|
)
|
|
measurement_results["median_decode_latency"] = med_decode_latency
|
|
measurement_results["median_decode_throughput"] = med_decode_throughput
|
|
|
|
throughput = (input_len + output_len) * batch_size / tot_latency
|
|
rank_print(
|
|
f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s"
|
|
)
|
|
measurement_results["total_latency"] = tot_latency
|
|
measurement_results["total_throughput"] = throughput
|
|
return measurement_results
|
|
|
|
|
|
def latency_test(
|
|
server_args,
|
|
port_args,
|
|
bench_args,
|
|
tp_rank,
|
|
):
|
|
configure_logger(server_args, prefix=f" TP{tp_rank}")
|
|
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
|
|
|
|
# Load the model
|
|
model_runner, tokenizer = load_model(server_args, port_args, tp_rank)
|
|
|
|
# Prepare inputs for warm up
|
|
reqs = prepare_synthetic_inputs_for_latency_test(
|
|
bench_args.batch_size[0], bench_args.input_len[0]
|
|
)
|
|
|
|
# Warm up
|
|
rank_print("Warmup ...")
|
|
latency_test_run_once(
|
|
bench_args.run_name,
|
|
model_runner,
|
|
rank_print,
|
|
reqs,
|
|
bench_args.batch_size[0],
|
|
bench_args.input_len[0],
|
|
8, # shorter decoding to speed up the warmup
|
|
server_args.device,
|
|
)
|
|
rank_print("Benchmark ...")
|
|
|
|
# Run the sweep
|
|
result_list = []
|
|
for bs, il, ol in itertools.product(
|
|
bench_args.batch_size, bench_args.input_len, bench_args.output_len
|
|
):
|
|
reqs = prepare_synthetic_inputs_for_latency_test(bs, il)
|
|
ret = latency_test_run_once(
|
|
bench_args.run_name,
|
|
model_runner,
|
|
rank_print,
|
|
reqs,
|
|
bs,
|
|
il,
|
|
ol,
|
|
server_args.device,
|
|
)
|
|
if ret is not None:
|
|
result_list.append(ret)
|
|
|
|
# Write results in jsonlines format on rank 0.
|
|
if tp_rank == 0 and bench_args.result_filename:
|
|
import jsonlines
|
|
|
|
with jsonlines.open(bench_args.result_filename, "a") as f:
|
|
f.write_all(result_list)
|
|
|
|
|
|
def plot_latency_test(
|
|
server_args,
|
|
bench_args,
|
|
tp_rank,
|
|
):
|
|
assert tp_rank == 0
|
|
|
|
# read the jsonl file and put in sqlite
|
|
df = pd.read_json(bench_args.result_filename, lines=True)
|
|
conn = sqlite3.connect(":memory:")
|
|
cur = conn.cursor()
|
|
|
|
# get the columns and their types
|
|
column_names = list(df.iloc[0].keys())
|
|
type_dict = {
|
|
str: "TEXT",
|
|
np.int64: "INTEGER",
|
|
np.float64: "FLOAT",
|
|
}
|
|
column_types = [type_dict[type(i)] for i in list(df.iloc[0])]
|
|
|
|
# create the table
|
|
cur.execute(
|
|
f"""
|
|
CREATE TABLE IF NOT EXISTS results (
|
|
{", ".join([f"{name} {type}" for name, type in zip(column_names, column_types)])}
|
|
)
|
|
"""
|
|
)
|
|
conn.commit()
|
|
|
|
# write the results to DB
|
|
df.to_sql("results", conn, if_exists="replace", index=False)
|
|
conn.commit()
|
|
|
|
# read it back using sql
|
|
df = pd.read_sql_query(bench_args.graph_sql, conn)
|
|
conn.close()
|
|
|
|
# plot it and save to a file
|
|
import matplotlib.pyplot as plt
|
|
|
|
assert (
|
|
len(df.columns) == 3
|
|
), f"The sql should have fetched <series, x, y> columns, not {df.columns}"
|
|
for label in df[df.columns[0]].unique():
|
|
q = f"{df.columns[0]}=='{label}'"
|
|
series = df.query(q)
|
|
plt.plot(series[df.columns[1]], series[df.columns[2]], label=q, marker="o")
|
|
plt.xlabel(df.columns[1])
|
|
plt.ylabel(df.columns[2])
|
|
plt.legend()
|
|
plt.savefig(bench_args.graph_filename, dpi=300)
|
|
|
|
# if in kitty, just dump it to the terminal
|
|
if os.environ["TERM"] == "xterm-kitty":
|
|
os.system(
|
|
f"kitty icat --use-window-size 1,1,600,600 {bench_args.graph_filename}"
|
|
)
|
|
|
|
|
|
def main(server_args, bench_args):
|
|
_set_envs_and_config(server_args)
|
|
|
|
if server_args.model_path:
|
|
if bench_args.correctness_test:
|
|
work_func = correctness_test
|
|
else:
|
|
work_func = latency_test
|
|
elif os.path.isfile(bench_args.result_filename):
|
|
assert bench_args.graph_filename, "please provide a filename for the graph"
|
|
work_func = plot_latency_test
|
|
else:
|
|
raise ValueError(
|
|
"Provide --model-path for running the tests or "
|
|
"provide --result-filename for plotting the results"
|
|
)
|
|
|
|
port_args = PortArgs.init_new(server_args)
|
|
|
|
if server_args.tp_size == 1:
|
|
work_func(server_args, port_args, bench_args, 0)
|
|
else:
|
|
workers = []
|
|
for tp_rank in range(server_args.tp_size):
|
|
proc = multiprocessing.Process(
|
|
target=work_func,
|
|
args=(
|
|
server_args,
|
|
port_args,
|
|
bench_args,
|
|
tp_rank,
|
|
),
|
|
)
|
|
proc.start()
|
|
workers.append(proc)
|
|
|
|
for proc in workers:
|
|
proc.join()
|
|
|
|
proc.terminate()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
ServerArgs.add_cli_args(parser)
|
|
BenchArgs.add_cli_args(parser)
|
|
args = parser.parse_args()
|
|
server_args = ServerArgs.from_cli_args(args)
|
|
bench_args = BenchArgs.from_cli_args(args)
|
|
|
|
logging.basicConfig(
|
|
level=getattr(logging, server_args.log_level.upper()),
|
|
format="%(message)s",
|
|
)
|
|
|
|
try:
|
|
main(server_args, bench_args)
|
|
except Exception as e:
|
|
raise e
|
|
finally:
|
|
kill_child_process()
|