[Auto Sync] Improve profilers and simplify bench_one_batch_server.py (#13866)

This commit is contained in:
Lianmin Zheng
2025-11-25 12:13:31 -08:00
committed by GitHub
parent e99ca6ac74
commit 1ab6ce0e62
12 changed files with 485 additions and 522 deletions

View File

@@ -16,7 +16,6 @@ import argparse
import dataclasses
import itertools
import json
import logging
import multiprocessing
import os
import random
@@ -38,132 +37,13 @@ from sglang.profiler import run_profile
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_blackwell, kill_process_tree
from sglang.test.nightly_bench_utils import save_results_as_pydantic_models
from sglang.test.test_utils import is_in_ci, write_github_step_summary
logger = logging.getLogger(__name__)
class ProfileLinks(BaseModel):
"""Pydantic model for profile trace links."""
extend: Optional[str] = None
decode: Optional[str] = None
class BenchmarkResult(BaseModel):
"""Pydantic model for benchmark results table data, for a single isl and osl"""
model_path: str
run_name: str
batch_size: int
input_len: int
output_len: int
latency: float
ttft: float
input_throughput: float
output_throughput: float
overall_throughput: float
last_gen_throughput: float
acc_length: Optional[float] = None
profile_links: Optional[ProfileLinks] = None
@staticmethod
def help_str() -> str:
return f"""
Note: To view the traces through perfetto-ui, please:
1. open with Google Chrome
2. allow popup
"""
def to_markdown_row(
self, trace_dir, base_url: str = "", relay_base: str = ""
) -> str:
"""Convert this benchmark result to a markdown table row."""
# Calculate costs (assuming H100 pricing for now)
hourly_cost_per_gpu = 2 # $2/hour for one H100
hourly_cost = hourly_cost_per_gpu * 1 # Assuming tp_size = 1 for simplicity
input_util = 0.7
accept_length = (
round(self.acc_length, 2) if self.acc_length is not None else "n/a"
)
itl = 1 / (self.output_throughput / self.batch_size) * 1000
input_cost = 1e6 / (self.input_throughput * input_util) / 3600 * hourly_cost
output_cost = 1e6 / self.output_throughput / 3600 * hourly_cost
def get_perfetto_relay_link_from_trace_file(trace_file: str):
import os
from urllib.parse import quote
rel_path = os.path.relpath(trace_file, trace_dir)
raw_file_link = f"{base_url}/{rel_path}"
relay_link = (
f"{relay_base}?src={quote(raw_file_link, safe='')}"
if relay_base and quote
else raw_file_link
)
return relay_link
# Handle profile links
profile_link = "NA | NA"
if self.profile_links:
if self.profile_links.extend or self.profile_links.decode:
# Create a combined link or use the first available one
trace_files = [self.profile_links.extend, self.profile_links.decode]
if any(trace_file is None for trace_file in trace_files):
logger.error("Some trace files are None", f"{trace_files=}")
trace_files_relay_links = [
(
f"[trace]({get_perfetto_relay_link_from_trace_file(trace_file)})"
if trace_file
else "N/A"
)
for trace_file in trace_files
]
profile_link = " | ".join(trace_files_relay_links)
# Build the row
return f"| {self.batch_size} | {self.input_len} | {self.latency:.2f} | {self.input_throughput:.2f} | {self.output_throughput:.2f} | {accept_length} | {itl:.2f} | {input_cost:.2f} | {output_cost:.2f} | {profile_link} |\n"
def generate_markdown_report(trace_dir, results: List["BenchmarkResult"]) -> str:
"""Generate a markdown report from a list of BenchmarkResult object from a single run."""
import os
# Build model header with run_name if it's not "default"
model_header = results[0].model_path
if results[0].run_name and results[0].run_name != "default":
model_header += f" ({results[0].run_name})"
# Include GPU config in model header if available
gpu_config = os.getenv("GPU_CONFIG", "")
if gpu_config:
model_header += f" [{gpu_config}]"
summary = f"### {model_header}\n"
# summary += (
# f"Input lens: {result.input_len}. Output lens: {result.output_len}.\n"
# )
summary += "| batch size | input len | latency (s) | input throughput (tok/s) | output throughput (tok/s) | acc length | ITL (ms) | input cost ($/1M) | output cost ($/1M) | profile (extend) | profile (decode)|\n"
summary += "| ---------- | --------- | ----------- | ------------------------- | ------------------------- | ---------- | -------- | ----------------- | ------------------ | --------------- | -------------- |\n"
# all results should share the same isl & osl
for result in results:
base_url = os.getenv("TRACE_BASE_URL", "").rstrip("/")
relay_base = os.getenv(
"PERFETTO_RELAY_URL",
"",
).rstrip("/")
summary += result.to_markdown_row(trace_dir, base_url, relay_base)
return summary
@dataclasses.dataclass
class BenchArgs:
run_name: str = "default"
seed: int = 42
batch_size: Tuple[int] = (1,)
input_len: Tuple[int] = (1024,)
output_len: Tuple[int] = (16,)
@@ -171,24 +51,25 @@ class BenchArgs:
return_logprob: bool = False
client_stream_interval: int = 1
input_len_step_percentage: float = 0.0
result_filename: str = "result.jsonl"
base_url: str = ""
skip_warmup: bool = False
show_report: bool = False
profile: bool = False
profile_steps: int = 3
profile_steps: int = 5
profile_by_stage: bool = False
profile_filename_prefix: str = None
append_to_github_summary: bool = True
profile_prefix: Optional[str] = None
profile_output_dir: Optional[str] = None
dataset_path: str = ""
parallel_batch: bool = False
dataset_name: str = "random"
output_path: Optional[str] = None
parallel_batch: bool = False
result_filename: str = "result.jsonl"
pydantic_result_filename: Optional[str] = None
append_to_github_summary: bool = True
seed: int = 42
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
parser.add_argument("--run-name", type=str, default=BenchArgs.run_name)
parser.add_argument("--seed", type=int, default=BenchArgs.seed)
parser.add_argument(
"--batch-size", type=int, nargs="+", default=BenchArgs.batch_size
)
@@ -199,13 +80,6 @@ class BenchArgs:
"--output-len", type=int, nargs="+", default=BenchArgs.output_len
)
parser.add_argument("--temperature", type=float, default=BenchArgs.temperature)
parser.add_argument(
"--dataset-name",
type=str,
default=BenchArgs.dataset_name,
choices=["mmmu", "random"],
help="Name of the dataset to benchmark on.",
)
parser.add_argument("--return-logprob", action="store_true")
parser.add_argument(
"--client-stream-interval",
@@ -217,9 +91,6 @@ class BenchArgs:
type=float,
default=BenchArgs.input_len_step_percentage,
)
parser.add_argument(
"--result-filename", type=str, default=BenchArgs.result_filename
)
parser.add_argument("--base-url", type=str, default=BenchArgs.base_url)
parser.add_argument("--skip-warmup", action="store_true")
parser.add_argument("--show-report", action="store_true")
@@ -228,17 +99,41 @@ class BenchArgs:
"--profile-steps", type=int, default=BenchArgs.profile_steps
)
parser.add_argument("--profile-by-stage", action="store_true")
parser.add_argument(
"--profile-prefix",
type=str,
default=BenchArgs.profile_prefix,
)
parser.add_argument(
"--profile-output-dir",
type=str,
default=BenchArgs.profile_output_dir,
)
parser.add_argument(
"--dataset-path",
type=str,
default=BenchArgs.dataset_path,
help="Path to the dataset.",
)
parser.add_argument(
"--dataset-name",
type=str,
default=BenchArgs.dataset_name,
choices=["mmmu", "random"],
help="Name of the dataset to benchmark on.",
)
parser.add_argument("--parallel-batch", action="store_true")
parser.add_argument(
"--profile-filename-prefix",
"--result-filename",
type=str,
default=BenchArgs.profile_filename_prefix,
default=BenchArgs.result_filename,
help="Store the results line by line in the JSON Line format to this file.",
)
parser.add_argument(
"--pydantic-result-filename",
type=str,
default=BenchArgs.pydantic_result_filename,
help="Store the results as pydantic models in the JSON format to this file.",
)
parser.add_argument(
"--no-append-to-github-summary",
@@ -246,25 +141,44 @@ class BenchArgs:
dest="append_to_github_summary",
help="Disable appending the output of this run to github ci summary",
)
parser.add_argument(
"--output-path",
type=str,
default=BenchArgs.output_path,
help="Path to save benchmark results as JSON format. If not specified, results will only be saved to result-filename.",
)
parser.add_argument("--seed", type=int, default=BenchArgs.seed)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
# use the default value's type to cast the args into correct types.
attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
kwargs = {}
for attr, attr_type in attrs:
val = getattr(args, attr)
if attr_type is type(None):
kwargs[attr] = val
else:
kwargs[attr] = attr_type(val)
return cls(**kwargs)
attrs = [attr.name for attr in dataclasses.fields(cls)]
return cls(**{attr: getattr(args, attr) for attr in attrs})
class BenchOneCaseResult(BaseModel):
run_name: str
batch_size: int
input_len: int
output_len: int
latency: float
input_throughput: float
output_throughput: float
overall_throughput: float
last_ttft: float
last_gen_throughput: float
acc_length: float
profile_link: Optional[str] = None
def dump_to_jsonl(self, result_filename: str):
with open(result_filename, "a") as fout:
res = {
"run_name": self.run_name,
"batch_size": self.batch_size,
"input_len": self.input_len,
"output_len": self.output_len,
"latency": round(self.latency, 4),
"input_throughput": round(self.input_throughput, 2),
"output_throughput": round(self.output_throughput, 2),
"overall_throughput": round(self.overall_throughput, 2),
"last_ttft": round(self.last_ttft, 4),
"last_gen_throughput": round(self.last_gen_throughput, 2),
"acc_length": round(self.acc_length, 2),
}
fout.write(json.dumps(res) + "\n")
def launch_server_internal(server_args):
@@ -309,15 +223,18 @@ def run_one_case(
run_name: str,
result_filename: str,
tokenizer: PreTrainedTokenizer | AutoProcessor,
dataset_name="",
profile: bool = False,
profile_steps: int = 3,
profile_steps: int = BenchArgs.profile_steps,
profile_by_stage: bool = False,
profile_filename_prefix: str = None,
dataset_path: str = "",
profile_prefix: Optional[str] = BenchArgs.profile_prefix,
profile_output_dir: Optional[str] = BenchArgs.profile_output_dir,
dataset_name: str = BenchArgs.dataset_name,
dataset_path: str = BenchArgs.dataset_path,
parallel_batch: bool = False,
):
requests.post(url + "/flush_cache")
# Load input token ids
# TODO: reuse bench_serving.get_dataset ?
if dataset_name == "mmmu":
input_requests = sample_mmmu_requests(
@@ -338,6 +255,7 @@ def run_one_case(
return_text=False,
)
# Load sampling parameters
use_structured_outputs = False
if use_structured_outputs:
texts = []
@@ -351,23 +269,6 @@ def run_one_case(
else:
json_schema = None
profile_link = None
if profile:
output_dir, profile_name = None, None
if profile_filename_prefix:
output_dir = os.path.dirname(profile_filename_prefix)
profile_name = os.path.basename(profile_filename_prefix)
profile_link: str = run_profile(
url,
profile_steps,
["CPU", "GPU"],
output_dir,
profile_name,
profile_by_stage,
)
tic = time.perf_counter()
payload = {
"sampling_params": {
"temperature": temperature,
@@ -394,14 +295,28 @@ def run_one_case(
payload["input_ids"] = input_ids
# Turn on profiler
profile_link = None
if profile:
profile_link: str = run_profile(
url=url,
num_steps=profile_steps,
activities=["CPU", "GPU"],
output_dir=profile_output_dir,
profile_by_stage=profile_by_stage,
profile_prefix=profile_prefix,
)
# Run the request
tic = time.perf_counter()
response = requests.post(
url + "/generate",
json=payload,
stream=True,
)
# The TTFT of the last request in the batch
ttft = 0.0
# Get the TTFT of the last request in the batch
last_ttft = 0.0
for chunk in response.iter_lines(decode_unicode=False):
chunk = chunk.decode("utf-8")
if chunk and chunk.startswith("data:"):
@@ -416,206 +331,116 @@ def run_one_case(
or data["meta_info"]["finish_reason"]["type"] == "length"
)
if data["meta_info"]["completion_tokens"] == 1:
ttft = time.perf_counter() - tic
last_ttft = time.perf_counter() - tic
# Compute metrics
latency = time.perf_counter() - tic
input_throughput = batch_size * input_len / ttft
output_throughput = batch_size * output_len / (latency - ttft)
input_throughput = batch_size * input_len / last_ttft
output_throughput = batch_size * output_len / (latency - last_ttft)
overall_throughput = batch_size * (input_len + output_len) / latency
server_info = requests.get(url + "/get_server_info").json()
acc_length = server_info["internal_states"][0].get("avg_spec_accept_length", None)
last_gen_throughput = server_info["internal_states"][0]["last_gen_throughput"]
internal_state = server_info.get("internal_states", [{}])
last_gen_throughput = internal_state[0].get("last_gen_throughput", None) or -1
acc_length = internal_state[0].get("avg_spec_accept_length", None) or -1
# Print results
print(f"batch size: {batch_size}")
print(f"input_len: {input_len}")
print(f"output_len: {output_len}")
print(f"latency: {latency:.2f} s")
print(f"ttft: {ttft:.2f} s")
print(f"last generation throughput: {last_gen_throughput:.2f} tok/s")
print(f"input throughput: {input_throughput:.2f} tok/s")
if output_len != 1:
print(f"output throughput: {output_throughput:.2f} tok/s")
print(f"last_ttft: {last_ttft:.2f} s")
print(f"last generation throughput: {last_gen_throughput:.2f} tok/s")
if acc_length > 0:
print(f"acc_length: {acc_length:.2f} ")
# Dump results
result = BenchOneCaseResult(
run_name=run_name,
batch_size=batch_size,
input_len=input_len,
output_len=output_len,
latency=latency,
input_throughput=input_throughput,
output_throughput=output_throughput,
overall_throughput=overall_throughput,
last_ttft=last_ttft,
last_gen_throughput=last_gen_throughput,
acc_length=acc_length,
profile_link=profile_link,
)
# Save and return the results
if result_filename:
with open(result_filename, "a") as fout:
res = {
"run_name": run_name,
"batch_size": batch_size,
"input_len": input_len,
"output_len": output_len,
"latency": round(latency, 4),
"output_throughput": round(output_throughput, 2),
"overall_throughput": round(overall_throughput, 2),
"last_gen_throughput": round(last_gen_throughput, 2),
}
fout.write(json.dumps(res) + "\n")
result.dump_to_jsonl(result_filename)
return (
batch_size,
latency,
ttft,
input_throughput,
output_throughput,
overall_throughput,
last_gen_throughput,
acc_length,
profile_link,
)
return result
def save_results_as_json(result: List[Tuple], bench_args: BenchArgs, model: str):
"""Save benchmark results as JSON using Pydantic models."""
json_results = []
# Generate all parameter combinations to match with results
param_combinations = list(
itertools.product(
bench_args.batch_size, bench_args.input_len, bench_args.output_len
def should_skip_due_to_token_capacity(
batch_size, input_len, output_len, skip_token_capacity_threshold
):
if batch_size * (input_len + output_len) > skip_token_capacity_threshold:
print(
"=" * 8
+ f"Skip benchmark {batch_size=} * ({input_len=} + {output_len=}) = {batch_size * (input_len + output_len)} > {skip_token_capacity_threshold=} due to kv cache limit."
+ "=" * 8
)
)
for i, (
batch_size,
latency,
ttft,
input_throughput,
output_throughput,
overall_throughput,
last_gen_throughput,
acc_length,
profile_link,
) in enumerate(result):
# Get the corresponding parameters for this result
bs, input_len, output_len = param_combinations[i]
# Parse profile links if available
profile_links = None
if profile_link:
profile_links = parse_profile_links(
profile_link, batch_size, input_len, output_len
)
benchmark_result = BenchmarkResult(
model_path=model,
run_name=bench_args.run_name,
batch_size=batch_size,
input_len=input_len,
output_len=output_len,
latency=latency,
ttft=ttft,
input_throughput=input_throughput,
output_throughput=output_throughput,
overall_throughput=overall_throughput,
last_gen_throughput=last_gen_throughput,
acc_length=acc_length,
profile_links=profile_links,
)
json_results.append(benchmark_result.model_dump())
# Save to JSON file
with open(bench_args.output_path, "w", encoding="utf-8") as f:
json.dump(json_results, f, indent=2, ensure_ascii=False)
print(f"Results saved as JSON to {bench_args.output_path}")
def parse_profile_links(
profile_dir: str, batch_size: int, input_len: int, output_len: int
) -> Optional[ProfileLinks]:
"""Parse profile directory to extract extend and decode trace file links."""
if not profile_dir or not os.path.exists(profile_dir):
return None
extend_link = None
decode_link = None
# Look for extend/prefill trace files
for file in os.listdir(profile_dir):
if file.endswith(".trace.json.gz") or file.endswith(".trace.json"):
if "extend" in file.lower() or "prefill" in file.lower():
extend_link = os.path.join(profile_dir, file)
elif "decode" in file.lower():
decode_link = os.path.join(profile_dir, file)
# If no specific extend/decode files found, try to find files with batch/input/output info
if not extend_link or not decode_link:
for file in os.listdir(profile_dir):
if file.endswith(".trace.json.gz") or file.endswith(".trace.json"):
if f"_batch{batch_size}_input{input_len}_output{output_len}_" in file:
if "prefill" in file.lower() or "extend" in file.lower():
extend_link = os.path.join(profile_dir, file)
elif "decode" in file.lower():
decode_link = os.path.join(profile_dir, file)
if extend_link or decode_link:
return ProfileLinks(extend=extend_link, decode=decode_link)
return None
return True
return False
def get_report_summary(
result: List[Tuple], server_args: ServerArgs, bench_args: BenchArgs
results: List[BenchOneCaseResult], bench_args: BenchArgs, server_args: ServerArgs
):
import tabulate
summary = (
f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}.\n"
)
summary += "| batch size | input len | latency (s) | input throughput (tok/s) | output throughput (tok/s) | acc length | ITL (ms) | input cost ($/1M) | output cost ($/1M) |"
headers = [
"batch size",
"latency (s)",
"input throughput (tok/s)",
"output throughput (tok/s)",
"acc length",
"ITL (ms)",
"input cost ($/1M)",
"output cost ($/1M)",
]
if bench_args.profile:
headers.append("profile")
rows = []
summary += " profile |"
for (
batch_size,
latency,
ttft,
input_throughput,
output_throughput,
_,
_,
acc_length,
trace_link,
) in result:
if is_blackwell():
hourly_cost_per_gpu = 4 # $4/hour for one B200
else:
hourly_cost_per_gpu = 2 # $2/hour for one H100
summary += "\n"
summary += "| ---------- | --------- | ----------- | ------------------------- | ------------------------- | ---------- | -------- | ----------------- | ------------------ |"
if bench_args.profile:
summary += "-------------|"
summary += "\n"
if is_blackwell():
hourly_cost_per_gpu = 4 # $4/hour for one B200
else:
hourly_cost_per_gpu = 2 # $2/hour for one H100
input_util = 0.7
# sort result by input_len
results.sort(key=lambda x: x.input_len)
for res in results:
hourly_cost = hourly_cost_per_gpu * server_args.tp_size
input_util = 0.7
accept_length = round(acc_length, 2) if acc_length is not None else "n/a"
itl = 1 / (output_throughput / batch_size) * 1000
input_cost = 1e6 / (input_throughput * input_util) / 3600 * hourly_cost
output_cost = 1e6 / output_throughput / 3600 * hourly_cost
row = [
batch_size,
latency,
input_throughput,
output_throughput,
accept_length,
itl,
input_cost,
output_cost,
]
if trace_link:
row.append(f"[Profile]({trace_link})")
rows.append(row)
accept_length = round(res.acc_length, 2) if res.acc_length > 0 else "n/a"
line = (
f"| {res.batch_size} | "
f"{res.input_len} | "
f"{res.latency:.2f} | "
f"{res.input_throughput:.2f} | "
f"{res.output_throughput:.2f} | "
f"{accept_length} | "
f"{1 / (res.output_throughput/res.batch_size) * 1000:.2f} | "
f"{1e6 / (res.input_throughput * input_util) / 3600 * hourly_cost:.2f} | "
f"{1e6 / res.output_throughput / 3600 * hourly_cost:.2f} |"
)
if bench_args.profile:
if res.profile_link:
line += f" [Profile]({res.profile_link}) |"
else:
line += f" n/a |"
line += "\n"
summary += line
summary += tabulate.tabulate(
rows, headers=headers, tablefmt="github", floatfmt=".2f"
)
return summary
@@ -625,47 +450,59 @@ def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
else:
proc, base_url = launch_server_process(server_args)
# Get tokenizer
server_info = requests.get(base_url + "/get_server_info").json()
if "tokenizer_path" in server_info:
tokenizer_path = server_info["tokenizer_path"]
elif "prefill" in server_info:
tokenizer_path = server_info["prefill"][0]["tokenizer_path"]
if bench_args.dataset_name == "mmmu":
# mmmu implies this is a MLLM
tokenizer = get_processor(tokenizer_path)
else:
tokenizer = get_tokenizer(tokenizer_path)
# warmup
# Get token capacity
internal_state = server_info.get("internal_states", [{}])
skip_token_capacity_threshold = (
internal_state[0].get("memory_usage", {}).get("token_capacity", 1000000000)
)
# Warmup
if not bench_args.skip_warmup:
print("=" * 8 + " Warmup Begin " + "=" * 8)
run_one_case(
base_url,
batch_size=16,
input_len=1024,
output_len=16,
temperature=bench_args.temperature,
return_logprob=bench_args.return_logprob,
stream_interval=bench_args.client_stream_interval,
input_len_step_percentage=bench_args.input_len_step_percentage,
dataset_name=bench_args.dataset_name,
run_name="",
result_filename="",
tokenizer=tokenizer,
dataset_path=bench_args.dataset_path,
parallel_batch=bench_args.parallel_batch,
)
print(f"Warmup with batch_size={bench_args.batch_size}")
for bs in bench_args.batch_size:
run_one_case(
base_url,
batch_size=bs,
input_len=1024,
output_len=16,
temperature=bench_args.temperature,
return_logprob=bench_args.return_logprob,
stream_interval=bench_args.client_stream_interval,
input_len_step_percentage=bench_args.input_len_step_percentage,
run_name="",
result_filename="",
tokenizer=tokenizer,
dataset_name=bench_args.dataset_name,
dataset_path=bench_args.dataset_path,
parallel_batch=bench_args.parallel_batch,
)
print("=" * 8 + " Warmup End " + "=" * 8 + "\n")
# benchmark
result = []
bench_result = []
results = []
profile_results = []
try:
# Benchmark all cases
for bs, il, ol in itertools.product(
bench_args.batch_size, bench_args.input_len, bench_args.output_len
):
result.append(
if should_skip_due_to_token_capacity(
bs, il, ol, skip_token_capacity_threshold
):
continue
results.append(
run_one_case(
base_url,
bs,
@@ -676,45 +513,54 @@ def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
stream_interval=bench_args.client_stream_interval,
input_len_step_percentage=bench_args.input_len_step_percentage,
run_name=bench_args.run_name,
dataset_name=bench_args.dataset_name,
result_filename=bench_args.result_filename,
tokenizer=tokenizer,
dataset_name=bench_args.dataset_name,
dataset_path=bench_args.dataset_path,
parallel_batch=bench_args.parallel_batch,
profile_filename_prefix=bench_args.profile_filename_prefix,
)
)
# Profile all cases
if bench_args.profile:
try:
for bs, il, ol in itertools.product(
bench_args.batch_size, bench_args.input_len, bench_args.output_len
):
bench_result.append(
(
run_one_case(
base_url,
bs,
il,
ol,
temperature=bench_args.temperature,
return_logprob=bench_args.return_logprob,
stream_interval=bench_args.client_stream_interval,
input_len_step_percentage=bench_args.input_len_step_percentage,
run_name=bench_args.run_name,
result_filename=bench_args.result_filename,
tokenizer=tokenizer,
dataset_name=bench_args.dataset_name,
profile=bench_args.profile,
profile_steps=bench_args.profile_steps,
profile_by_stage=bench_args.profile_by_stage,
dataset_path=bench_args.dataset_path,
parallel_batch=bench_args.parallel_batch,
profile_filename_prefix=bench_args.profile_filename_prefix,
)[-1],
if should_skip_due_to_token_capacity(
bs, il, ol, skip_token_capacity_threshold
):
continue
profile_prefix = (
bench_args.profile_prefix or ""
) + f"bs-{bs}-il-{il}"
profile_results.append(
run_one_case(
base_url,
bs,
il,
ol,
temperature=bench_args.temperature,
return_logprob=bench_args.return_logprob,
stream_interval=bench_args.client_stream_interval,
input_len_step_percentage=bench_args.input_len_step_percentage,
run_name=bench_args.run_name,
result_filename=bench_args.result_filename,
tokenizer=tokenizer,
dataset_name=bench_args.dataset_name,
dataset_path=bench_args.dataset_path,
parallel_batch=bench_args.parallel_batch,
profile=bench_args.profile,
profile_steps=bench_args.profile_steps,
profile_by_stage=bench_args.profile_by_stage,
profile_prefix=profile_prefix,
profile_output_dir=bench_args.profile_output_dir,
)
)
result = [t1[:-1] + t2 for t1, t2 in zip(result, bench_result)]
# Replace the profile link
for res, profile_res in zip(results, profile_results):
res.profile_link = profile_res.profile_link
except Exception as e:
print(f"Error profiling, there will be no profile trace dump: {e}")
finally:
@@ -723,20 +569,26 @@ def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
print(f"\nResults are saved to {bench_args.result_filename}")
# Save results as JSON if output_path is specified
if bench_args.output_path:
save_results_as_json(result, bench_args, model=server_args.model_path)
if not bench_args.show_report:
return
summary = get_report_summary(result, server_args, bench_args)
# Print summary
summary = get_report_summary(results, bench_args, server_args)
print(summary)
if is_in_ci() and bench_args.append_to_github_summary:
write_github_step_summary(summary)
# Save results as pydantic models in the JSON format
if bench_args.pydantic_result_filename:
save_results_as_pydantic_models(
results,
pydantic_result_filename=bench_args.pydantic_result_filename,
model_path=server_args.model_path,
)
def main():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
BenchArgs.add_cli_args(parser)
@@ -749,7 +601,3 @@ def main():
bench_args = BenchArgs.from_cli_args(args)
run_benchmark(server_args, bench_args)
if __name__ == "__main__":
main()

View File

@@ -18,26 +18,19 @@ import requests
PROFILER_DIR = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
def _run_profile(
def run_profile(
url: Optional[str],
num_steps: int,
activities: List[str],
output_dir: Optional[str] = None,
profile_name: Optional[str] = None,
profile_by_stage: bool = False,
merge_profiles: bool = False,
profile_prefix: Optional[str] = None,
) -> str:
if output_dir is None:
output_dir = PROFILER_DIR
output_dir = os.path.normpath(output_dir)
output_dir = os.path.abspath(output_dir)
output_dir = Path(output_dir)
# Add "profile_name/timestamp" to the path.
if profile_name:
output_dir = output_dir / profile_name
output_dir = output_dir / str(time.time())
output_dir = Path(os.path.abspath(os.path.normpath(output_dir))) / str(time.time())
output_dir.mkdir(exist_ok=True, parents=True)
print(f"Dump profiling traces to {output_dir}")
@@ -62,6 +55,7 @@ def _run_profile(
"activities": activities,
"profile_by_stage": profile_by_stage,
"merge_profiles": merge_profiles,
"profile_prefix": profile_prefix,
}
response = requests.post(url=url + "/start_profile", json=json_data)
@@ -71,28 +65,6 @@ def _run_profile(
return trace_link
def run_profile(
url: Optional[str],
num_steps: int,
activities: List[str],
output_dir: Optional[str] = None,
profile_name: Optional[str] = None,
profile_by_stage: bool = False,
merge_profiles: bool = False,
):
# step based profile will self terminate on num_steps constraints
link = _run_profile(
url,
num_steps,
activities,
output_dir,
profile_name,
profile_by_stage,
merge_profiles,
)
return link
if __name__ == "__main__":
parser = ArgumentParser(description="Benchmark the online serving throughput.")
parser.add_argument(
@@ -107,12 +79,6 @@ if __name__ == "__main__":
default=None,
help="Profile directory to dump profile traces.",
)
parser.add_argument(
"--profile-name",
type=str,
default=None,
help="The name of this profile run.",
)
parser.add_argument(
"--num-steps",
type=int,
@@ -126,6 +92,11 @@ if __name__ == "__main__":
default=False,
help="Whether to profile prefill and decode separately",
)
parser.add_argument(
"--profile-prefix",
type=str,
help="The prefix of this profiler file.",
)
parser.add_argument(
"--cpu",
action=argparse.BooleanOptionalAction,
@@ -152,7 +123,7 @@ if __name__ == "__main__":
action=argparse.BooleanOptionalAction,
type=bool,
default=False,
help="Whether to use rpd profiler (https://github.com/ROCm/rocmProfileData)",
help="Whether to use ROCM rpd profiler (https://github.com/ROCm/rocmProfileData)",
)
parser.add_argument(
"--merge-profiles",
@@ -172,12 +143,13 @@ if __name__ == "__main__":
activities.append("MEM")
if args.rpd:
activities.append("RPD")
run_profile(
args.url,
args.num_steps,
activities,
args.output_dir,
args.profile_name,
args.profile_by_stage,
args.merge_profiles,
url=args.url,
num_steps=args.num_steps,
activities=activities,
output_dir=args.output_dir,
profile_by_stage=args.profile_by_stage,
profile_prefix=args.profile_prefix,
merge_profiles=args.merge_profiles,
)

View File

@@ -1545,7 +1545,6 @@ def _execute_server_warmup(
try:
warmup_timeout = envs.SGLANG_WARMUP_TIMEOUT.get()
if server_args.disaggregation_mode == "null":
logger.info(f"Start of co-locate warmup ...")
res = requests.post(
url + request_name,
json=json_data,

View File

@@ -1321,17 +1321,23 @@ class SetInternalStateReqOutput(BaseReq):
class ProfileReqInput(BaseReq):
# The output directory
output_dir: Optional[str] = None
# Specify the steps to start the profiling
start_step: Optional[int] = None
# If set, it profile as many as this number of steps.
# If it is set, profiling is automatically stopped after this step, and
# the caller doesn't need to run stop_profile.
start_step: Optional[int] = None
num_steps: Optional[int] = None
# The activities to record. The choices are ["CPU", "GPU", "MEM", "RPD"]
activities: Optional[List[str]] = None
# Whether profile by stages (e.g., prefill and decode) separately
profile_by_stage: bool = False
# Whether to record source information (file and line number) for the ops.
with_stack: Optional[bool] = None
# Whether to save information about operators input shapes.
record_shapes: Optional[bool] = None
# Merge profiles from all ranks into a single trace
merge_profiles: bool = False
# The prefix of the profile filenames
profile_prefix: Optional[str] = None
@@ -1351,7 +1357,6 @@ class ProfileReq(BaseReq):
with_stack: Optional[bool] = None
record_shapes: Optional[bool] = None
profile_id: Optional[str] = None
# Merge profiles from all ranks into a single trace
merge_profiles: bool = False
profile_prefix: Optional[str] = None

View File

@@ -31,18 +31,22 @@ class SchedulerProfilerMixin:
self.torch_profiler_output_dir: Optional[Path] = None
self.profiler_activities: Optional[List[str]] = None
self.profile_id: Optional[str] = None
self.profiler_start_forward_ct: Optional[int] = None
self.profiler_target_forward_ct: Optional[int] = None
self.profiler_target_prefill_ct: Optional[int] = None
self.profiler_target_decode_ct: Optional[int] = None
self.profiler_prefill_ct: Optional[int] = None
self.profiler_decode_ct: Optional[int] = None
self.profiler_target_prefill_ct: Optional[int] = None
self.profiler_target_decode_ct: Optional[int] = None
self.profile_by_stage: bool = False
self.profile_steps: Optional[int] = None
self.profile_in_progress: bool = False
self.rpd_profiler = None
self.merge_profiles = False
# For ROCM
self.rpd_profiler = None
def init_profile(
self,
output_dir: Optional[str],
@@ -81,12 +85,11 @@ class SchedulerProfilerMixin:
self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1)
if num_steps:
self.profile_steps = num_steps
if self.profile_by_stage:
self.profiler_target_prefill_ct = num_steps
self.profiler_target_decode_ct = num_steps
self.profiler_prefill_ct = 0
self.profiler_decode_ct = 0
self.profiler_target_prefill_ct = num_steps
self.profiler_target_decode_ct = num_steps
elif start_step:
self.profiler_target_forward_ct = (
self.profiler_start_forward_ct + num_steps
@@ -119,7 +122,7 @@ class SchedulerProfilerMixin:
activity_map[a] for a in activities if a in activity_map
]
if "RPD" in activities:
if "RPD" in activities: # for ROCM
from rpdTracerControl import rpdTracerControl
rpdTracerControl.skipCreate()
@@ -217,6 +220,11 @@ class SchedulerProfilerMixin:
self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True)
if self.profile_prefix:
stage_prefix = self.profile_prefix + "-"
else:
stage_prefix = ""
stage_suffix = f"-{stage.name}" if stage else ""
logger.info("Stop profiling" + stage_suffix + "...")
if self.torch_profiler is not None:
@@ -233,7 +241,12 @@ class SchedulerProfilerMixin:
if getattr(self, "moe_ep_size", 1) > 1:
filename_parts.append(f"EP-{getattr(self, 'moe_ep_rank', 0)}")
filename = "-".join(filename_parts) + stage_suffix + ".trace.json.gz"
filename = (
stage_prefix
+ "-".join(filename_parts)
+ stage_suffix
+ ".trace.json.gz"
)
self.torch_profiler.export_chrome_trace(
os.path.join(self.torch_profiler_output_dir, filename)

View File

@@ -433,19 +433,19 @@ class TokenizerManager(TokenizerCommunicatorMixin):
self.auto_create_handle_loop()
obj.normalize_batch_and_arguments()
external_trace_header = None
if request:
if "trace_context" in request.headers:
trace_set_remote_propagate_context(request.headers["trace_context"])
else:
external_trace_header = extract_trace_headers(request.headers)
if self.enable_trace:
external_trace_header = None
if request:
if "trace_context" in request.headers:
trace_set_remote_propagate_context(request.headers["trace_context"])
else:
external_trace_header = extract_trace_headers(request.headers)
self._trace_request_start(obj, created_time, external_trace_header)
if self.server_args.tokenizer_worker_num > 1:
self._attach_multi_http_worker_info(obj)
if self.enable_trace:
self._trace_request_start(obj, created_time, external_trace_header)
if self.log_requests:
max_length, skip_names, _ = self.log_request_metadata
logger.info(

View File

@@ -90,27 +90,6 @@ class ForwardMode(IntEnum):
# Split Prefill for PD multiplexing
SPLIT_PREFILL = auto()
def __str__(self):
if self == ForwardMode.EXTEND:
return "EXTEND"
elif self == ForwardMode.DECODE:
return "DECODE"
elif self == ForwardMode.MIXED:
return "MIXED"
elif self == ForwardMode.IDLE:
return "IDLE"
elif self == ForwardMode.TARGET_VERIFY:
return "TARGET_VERIFY"
elif self == ForwardMode.DRAFT_EXTEND:
return "DRAFT_EXTEND"
elif self == ForwardMode.DRAFT_EXTEND_V2:
return "DRAFT_EXTEND_V2"
elif self == ForwardMode.PREBUILT:
return "PREBUILT"
elif self == ForwardMode.SPLIT_PREFILL:
return "SPLIT_PREFILL"
return "UNKNOWN"
def is_prefill(self):
return self.is_extend()

View File

@@ -3663,11 +3663,6 @@ def reserve_rope_cache_for_long_sequences(
"""Pre-expand RoPE cache for long sequences and speculative decoding."""
from sglang.srt.environ import envs
if logger is None:
import logging
logger = logging.getLogger(__name__)
SAFETY_FACTOR = envs.SGLANG_SPEC_EXPANSION_SAFETY_FACTOR.value
MARGIN = envs.SGLANG_ROPE_CACHE_SAFETY_MARGIN.value
ALIGN = envs.SGLANG_ROPE_CACHE_ALIGN.value
@@ -3689,23 +3684,13 @@ def reserve_rope_cache_for_long_sequences(
# 3) Align to reduce reallocation frequency
reserve = (reserve + ALIGN - 1) // ALIGN * ALIGN
logger.info(
f"RoPE cache reserve={reserve} (cap={base_ctx}, steps={steps}, draft={draft}, k={SAFETY_FACTOR}, margin={MARGIN})"
)
# Recursively expand all RoPE layers
def reserve_rope_cache_recursive(module):
for child in module.children():
if hasattr(child, "_ensure_cos_sin_cache_length") and hasattr(
child, "cos_sin_cache"
):
old_len = child.cos_sin_cache.shape[0]
child._ensure_cos_sin_cache_length(reserve - 1)
new_len = child.cos_sin_cache.shape[0]
if new_len > old_len:
logger.info(
f"Expanded RoPE cache from {old_len} to {new_len} positions"
)
else:
reserve_rope_cache_recursive(child)

View File

@@ -0,0 +1,156 @@
import json
import logging
import os
from typing import List, Optional
from pydantic import BaseModel
logger = logging.getLogger(__name__)
# TODO:
# There is huge redundancy between BenchmarkResult and BenchOneCaseResult, and redundancy between to_markdown_row, generate_markdown_report, get_report_summary.
# We should refactor them to reduce the code duplication.
# 1. Delete the BenchmarkResult use BenchOneCaseResult directly.
# 2. Merge all related markdown rendering functions into BenchOneCaseResult
class BenchmarkResult(BaseModel):
"""Pydantic model for benchmark results table data, for a single isl and osl"""
model_path: str
run_name: str
batch_size: int
input_len: int
output_len: int
latency: float
input_throughput: float
output_throughput: float
overall_throughput: float
last_ttft: float
last_gen_throughput: float
acc_length: Optional[float] = None
profile_link_extend: Optional[str] = None
profile_link_decode: Optional[str] = None
@staticmethod
def help_str() -> str:
return f"""
Note: To view the traces through perfetto-ui, please:
1. open with Google Chrome
2. allow popup
"""
def to_markdown_row(
self, trace_dir, base_url: str = "", relay_base: str = ""
) -> str:
"""Convert this benchmark result to a markdown table row."""
hourly_cost_per_gpu = 2 # $2/hour for one H100
hourly_cost = hourly_cost_per_gpu * 1 # Assuming tp_size = 1 for simplicity
input_util = 0.7
accept_length = round(self.acc_length, 2) if self.acc_length > 0 else "n/a"
itl = 1 / (self.output_throughput / self.batch_size) * 1000
input_cost = 1e6 / (self.input_throughput * input_util) / 3600 * hourly_cost
output_cost = 1e6 / self.output_throughput / 3600 * hourly_cost
def get_perfetto_relay_link_from_trace_file(trace_file: str):
from urllib.parse import quote
rel_path = os.path.relpath(trace_file, trace_dir)
raw_file_link = f"{base_url}/{rel_path}"
relay_link = (
f"{relay_base}?src={quote(raw_file_link, safe='')}"
if relay_base
else raw_file_link
)
return relay_link
# Handle profile links
profile_link = "NA | NA"
if self.profile_link_extend or self.profile_link_decode:
# Create a combined link or use the first available one
trace_files = [self.profile_link_extend, self.profile_link_decode]
if any(trace_file is None for trace_file in trace_files):
logger.error("Some trace files are None", f"{trace_files=}")
trace_files_relay_links = [
(
f"[trace]({get_perfetto_relay_link_from_trace_file(trace_file)})"
if trace_file
else "N/A"
)
for trace_file in trace_files
]
profile_link = " | ".join(trace_files_relay_links)
# Build the row
return f"| {self.batch_size} | {self.input_len} | {self.latency:.2f} | {self.input_throughput:.2f} | {self.output_throughput:.2f} | {accept_length} | {itl:.2f} | {input_cost:.2f} | {output_cost:.2f} | {profile_link} |\n"
def generate_markdown_report(trace_dir, results: List[BenchmarkResult]) -> str:
"""Generate a markdown report from a list of BenchmarkResult object from a single run."""
# Build model header with run_name if it's not "default"
model_header = results[0].model_path
if results[0].run_name and results[0].run_name != "default":
model_header += f" ({results[0].run_name})"
# Include GPU config in model header if available
gpu_config = os.getenv("GPU_CONFIG", "")
if gpu_config:
model_header += f" [{gpu_config}]"
summary = f"### {model_header}\n"
summary += "| batch size | input len | latency (s) | input throughput (tok/s) | output throughput (tok/s) | acc length | ITL (ms) | input cost ($/1M) | output cost ($/1M) | profile (extend) | profile (decode)|\n"
summary += "| ---------- | --------- | ----------- | ------------------------- | ------------------------- | ---------- | -------- | ----------------- | ------------------ | ---------------- | --------------- |\n"
# all results should share the same isl & osl
for result in results:
base_url = os.getenv("TRACE_BASE_URL", "").rstrip("/")
relay_base = os.getenv(
"PERFETTO_RELAY_URL",
"",
).rstrip("/")
summary += result.to_markdown_row(trace_dir, base_url, relay_base)
return summary
def save_results_as_pydantic_models(
results: List, pydantic_result_filename: str, model_path: str
):
"""Save benchmark results as JSON using Pydantic models."""
json_results = []
for res in results:
profile_link_extend = None
profile_link_decode = None
if res.profile_link:
for file in os.listdir(res.profile_link):
if file.endswith(".trace.json.gz") or file.endswith(".trace.json"):
if "extend" in file.lower() or "prefill" in file.lower():
profile_link_extend = os.path.join(res.profile_link, file)
elif "decode" in file.lower():
profile_link_decode = os.path.join(res.profile_link, file)
benchmark_result = BenchmarkResult(
model_path=model_path,
run_name=res.run_name,
batch_size=res.batch_size,
input_len=res.input_len,
output_len=res.output_len,
latency=res.latency,
input_throughput=res.input_throughput,
output_throughput=res.output_throughput,
overall_throughput=res.overall_throughput,
last_gen_throughput=res.last_gen_throughput,
last_ttft=res.last_ttft,
acc_length=res.acc_length,
profile_link_extend=profile_link_extend,
profile_link_decode=profile_link_decode,
)
json_results.append(benchmark_result.model_dump())
with open(pydantic_result_filename, "w", encoding="utf-8") as f:
json.dump(json_results, f, indent=2, ensure_ascii=False)

View File

@@ -10,6 +10,7 @@ python3 -m sglang.test.send_one --profile --profile-by-stage
import argparse
import dataclasses
import json
from typing import Optional
import requests
import tabulate
@@ -38,7 +39,7 @@ class BenchArgs:
profile: bool = False
profile_steps: int = 3
profile_by_stage: bool = False
profile_name_prefix: str = ""
profile_prefix: Optional[str] = None
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
@@ -51,9 +52,6 @@ class BenchArgs:
default=BenchArgs.different_prompts,
)
parser.add_argument("--temperature", type=float, default=BenchArgs.temperature)
parser.add_argument(
"--profile-name-prefix", type=str, default=BenchArgs.profile_name_prefix
)
parser.add_argument(
"--max-new-tokens", type=int, default=BenchArgs.max_new_tokens
)
@@ -74,6 +72,9 @@ class BenchArgs:
"--profile-steps", type=int, default=BenchArgs.profile_steps
)
parser.add_argument("--profile-by-stage", action="store_true")
parser.add_argument(
"--profile-prefix", type=str, default=BenchArgs.profile_prefix
)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
@@ -84,6 +85,7 @@ class BenchArgs:
def send_one_prompt(args: BenchArgs):
base_url = f"http://{args.host}:{args.port}"
# Construct the input
if args.image:
args.prompt = (
"Human: Describe this image in a very short sentence.\n\nAssistant:"
@@ -140,15 +142,14 @@ def send_one_prompt(args: BenchArgs):
if args.profile:
print(f"Running profiler with {args.profile_steps} steps...")
run_profile(
base_url,
args.profile_steps,
["CPU", "GPU"],
None,
None,
args.profile_by_stage,
args.profile_name_prefix,
url=base_url,
num_steps=args.profile_steps,
activities=["CPU", "GPU"],
profile_by_stage=args.profile_by_stage,
profile_prefix=args.profile_prefix,
)
# Send the request
response = requests.post(
f"{base_url}/generate",
json=json_data,
@@ -176,6 +177,7 @@ def send_one_prompt(args: BenchArgs):
print(ret)
return 0, 0
# Print results
if "spec_verify_ct" in ret["meta_info"] and ret["meta_info"]["spec_verify_ct"] > 0:
acc_length = (
ret["meta_info"]["completion_tokens"] / ret["meta_info"]["spec_verify_ct"]

View File

@@ -151,7 +151,10 @@ def send_single(
if profile:
run_profile(
base_url, profile_steps, ["CPU", "GPU"], None, None, profile_by_stage
url=base_url,
num_steps=profile_steps,
activities=["CPU", "GPU"],
profile_by_stage=profile_by_stage,
)
response = requests.post(