diff --git a/README.md b/README.md index 2a972a38..57c29271 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Below are blueprints of supported models along with their documentation. | [**xDiT diffusion inference**](benchmark/xdit/README.md) | Diffusion Transformer inference using xDiT | FLUX.1, FLUX.1 Kontext, FLUX.2, FLUX.2 Klein, HunyuanVideo, HunyuanVideo 1.5, LTX-2, Stable Diffusion 3.5, Wan 2.1, Wan 2.2, Z-Image Turbo | | [**JAX MaxText training**](benchmark/jax_maxtext/README.md) | Train LLMs on AMD Instinct GPUs using JAX MaxText | Llama 2 7B/70B, Llama 3/3.1 8B/70B, Llama 3.1 405B, Llama 3.3 70B, DeepSeek-V2-lite 16B, Mixtral-8x7B | | [**vLLM inference**](benchmark/vllm/README.md) | LLM Inference with vLLM on AMD Instinct GPUs | DeepSeek-R1, gpt-oss-20b/120b, Llama-2-70b, Llama-3.1-8b/405b, Llama-3.3-70b, Llama-4-Scout/Maverick, Mixtral-8x7b/8x22b, Phi-4, Qwen3-8b/32b/30b-a3b/235b-a22b | -| [**SGLang inference**](benchmark/sglang/README.md) | LLM Inference with SGLang on AMD Instinct GPUs | DeepSeek-R1-Distill-Qwen-32B | +| [**SGLang inference**](benchmark/sglang/README.md) | LLM Inference with SGLang on AMD Instinct GPUs | DeepSeek-R1-Distill-Qwen-32B, Kimi-K3 | | [**PyTorch training**](benchmark/pytorch_train/README.md) | Train LLMs on AMD Instinct GPUs using AMD's Primus | Llama 2/3/3.1/3.2/3.3/4, GPT-OSS 20B/120B, Qwen2/2.5/3, Flux, SDXL, DLRM, and others | | [**PyTorch inference**](benchmark/pytorch_inference/README.md) | Inference recipes for Multimodal, video and vision transformer models | Mochi video, Chai-1, CLIP (ViT-B-32), Wan2.1, Janus-Pro-7B, HunyuanVideo | | [**Megatron-LM training**](benchmark/megatron_lm/README.md) | Train LLMs on AMD Instinct GPUs using ROCm Megatron-LM | Llama 2 7B/70B, Llama 3/3.1 8B/70B, Llama 3.3 70B, DeepSeek-V2-lite, DeepSeek-V3, Mixtral 8x7B/8x22B, Qwen 2.5 7B/72B | diff --git a/benchmark/sglang/README.md b/benchmark/sglang/README.md index 0612327c..5467b4a8 100644 --- a/benchmark/sglang/README.md +++ b/benchmark/sglang/README.md @@ -80,6 +80,48 @@ users can also change the benchmarking parameters. Refer to the [Standalone benc | model_name | | --------------------------------------- | | pyt_sglang_deepseek-r1-distill-qwen-32b | +| pyt_sglang_kimi-k3 | +| pyt_sglang_kimi-k3_dspark | + +>[!NOTE] +>The two `pyt_sglang_kimi-k3*` entries are the exception to everything described above. They track the +>AMD day-0 recipes in [sgl-project/sglang#32548](https://github.com/sgl-project/sglang/issues/32548) +>(day-0 support: [#32541](https://github.com/sgl-project/sglang/pull/32541), see also the +>[SGLang K3 cookbook](https://docs.sglang.io/cookbook/autoregressive/Moonshotai/Kimi-K3)) and differ in +>four ways: +> +>- **Image.** They build from `lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727` via +> [docker/pyt_sglang_kimi_k3.ubuntu.amd.Dockerfile](../../docker/pyt_sglang_kimi_k3.ubuntu.amd.Dockerfile), +> not the shared `lmsysorg/sglang:v0.4.5-rocm630` above, which predates K3 support. +>- **Benchmark.** They measure *online serving* (`sglang serve` + `sglang.benchmark.serving`) through +> [scripts/sglang/run_sglang.py](../../scripts/sglang/run_sglang.py) and +> [scripts/sglang/configs/kimi_k3.yaml](../../scripts/sglang/configs/kimi_k3.yaml), rather than the +> offline latency/throughput path documented below. +>- **Hardware.** 8x MI350X/MI355X (gfx950) TP8 only, hence `skip_gpu_arch: gfx942`. The checkpoint is +> large, so make sure `HF_HUB_CACHE` has room; `pyt_sglang_kimi-k3_dspark` additionally pulls the +> [RadixArk/Kimi-K3-DSpark](https://huggingface.co/RadixArk/Kimi-K3-DSpark) draft checkpoint. +>- **Invocation.** They carry no sweep tag, so a tag run does not pull the checkpoint and hold 8 GPUs. +> Run them explicitly by name: +> +>```sh +>madengine run --tags pyt_sglang_kimi-k3 --keep-model-dir --live-output +>madengine run --tags pyt_sglang_kimi-k3_dspark --keep-model-dir --live-output +>``` +> +>The sweep is 8192-token input / 1024-token output at concurrency 2/4/8/16/32. The issue does not state +>its input and output lengths; they were recovered from its own tables, where +>`(E2EL - TTFT) / TPOT + 1` lands on ~1024 output tokens on every row and +>`concurrency x (inp + out) / E2EL` reproduces the reported total throughput only at 8192 input tokens. +>Keeping that shape is what makes MAD's numbers comparable to the issue's. +> +>The runner emits `--tp-size` where the issue writes `--tp`. There is no `--tp` server argument; +>`tp_size` has a single alias, `--tensor-parallel-size`, and `--tp` resolves only through argparse +>prefix matching, which a future `--tp*` option would silently break. + +>[!WARNING] +>The published performance tables were measured on **MI355X**. A commenter on the tracking issue reports +>much weaker results on **MI350X** with untuned AITER kernels. Both report as gfx950, so `skip_gpu_arch` +>cannot distinguish them — treat MI350X numbers from this recipe as unvalidated. ### Standalone benchmarking ----------------------------- diff --git a/docker/pyt_sglang_kimi_k3.ubuntu.amd.Dockerfile b/docker/pyt_sglang_kimi_k3.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..b8322b48 --- /dev/null +++ b/docker/pyt_sglang_kimi_k3.ubuntu.amd.Dockerfile @@ -0,0 +1,50 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# Kimi K3 day-0 support landed in sgl-project/sglang#32541 and is not in a +# tagged SGLang ROCm release yet; this model-specific image is the only ROCm +# build carrying the KDA / Stable LatentMoE / AITER A8W4 support the checkpoint +# needs. Kept separate from docker/pyt_sglang, which is still on the v0.4.5 +# rocm630 base that the existing SGLang entry is validated against. +# +# The tag is the day-0 image named in the AMD tracking issue +# https://github.com/sgl-project/sglang/issues/32548, which is what its MI355X +# performance tables were measured against. A newer rocm720-mi35x-k3-20260728 +# tag exists on Docker Hub but nothing published ties it to the recipe, so it is +# deliberately not adopted here. +# +# Fold this back into docker/pyt_sglang once K3 lands in a versioned ROCm image. +ARG BASE_DOCKER=lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 + +FROM $BASE_DOCKER + +USER root +ENV WORKSPACE_DIR=/workspace +RUN mkdir -p $WORKSPACE_DIR +WORKDIR $WORKSPACE_DIR + +# record configuration for posterity +RUN pip3 list diff --git a/models.json b/models.json index 15c3cd5b..af46396b 100644 --- a/models.json +++ b/models.json @@ -2194,6 +2194,46 @@ "args": "--model_repo deepseek-ai/DeepSeek-R1-Distill-Qwen-32B --test_option latency --num_gpu 8 --datatype bfloat16 --dataset random --batch_size 1,8,32 --lat_input_output_len '128:128;128:1024;1024:128;1024:1024'" }, + { + "name": "pyt_sglang_kimi-k3", + "url": "", + "dockerfile": "docker/pyt_sglang_kimi_k3", + "scripts": "scripts/sglang/run_kimi_k3.sh", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_Kimi-K3.csv", + "tags": [ + "pyt", + "sglang", + "inference" + ], + "timeout": -1, + "skip_gpu_arch": "gfx942", + "args": + "--model_repo moonshotai/Kimi-K3 --config configs/kimi_k3.yaml --variant nospec" + }, + { + "name": "pyt_sglang_kimi-k3_dspark", + "url": "", + "dockerfile": "docker/pyt_sglang_kimi_k3", + "scripts": "scripts/sglang/run_kimi_k3.sh", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_Kimi-K3.csv", + "tags": [ + "pyt", + "sglang", + "inference" + ], + "timeout": -1, + "skip_gpu_arch": "gfx942", + "args": + "--model_repo moonshotai/Kimi-K3 --config configs/kimi_k3.yaml --variant dspark" + }, { "name": "pyt_hy_video", "url": "", diff --git a/scripts/sglang/configs/kimi_k3.yaml b/scripts/sglang/configs/kimi_k3.yaml new file mode 100644 index 00000000..50de3dbd --- /dev/null +++ b/scripts/sglang/configs/kimi_k3.yaml @@ -0,0 +1,90 @@ +################################################################################# +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# +# Kimi K3 on MI355X, tracking the AMD day-0 recipes in +# https://github.com/sgl-project/sglang/issues/32548 (day-0 support: #32541). +# +# Two variants, selected with --variant because both serve the same model: +# nospec the non speculative-decoding recipe +# dspark the same server plus the DSpark draft checkpoint +# +# Sweep shape. The issue's tables do not state input/output lengths. They are +# recovered from the tables themselves: (E2EL - TTFT) / TPOT + 1 lands on ~1024 +# output tokens for every row, and concurrency * (inp + out) / E2EL reproduces +# the reported total throughput only at inp=8192 (e.g. concurrency 8: +# 8 * 9216 / 31.297 s = 2355.8 vs 2356.21 reported). MAD's usual 1024/1024 would +# produce numbers that cannot be compared against the tracking issue. +# +# max_concurrency mirrors the issue's rows exactly (2/4/8/16/32), so each variant +# expands into the same five measurements the tables report. +# +# --disable-radix-cache comes from the recipe itself, not from MAD benchmark +# hygiene, and so is spelled out in extra_args rather than forced by the runner. +# +# The runner emits --tp-size, not the --tp shown in the issue: ServerArgs +# declares tp_size with a single alias, --tensor-parallel-size, and there is no +# --tp option. "--tp 8" only parses through argparse prefix matching, which any +# future --tp* option would silently break. + +- benchmark: serving + variant: nospec + model: moonshotai/Kimi-K3 + tp: 8 + inp: 8192 + out: 1024 + dtype: bfloat16 + max_concurrency: 2 4 8 16 32 + env: &k3_env + SGLANG_USE_AITER: 1 + SGLANG_AITER_K3_OPT: 1 + AITER_FLYDSL_FORCE: 1 + # selects the AITER A8W4 MoE path for the natively-MXFP4 K3 weights + AITER_SITUV2_A8W4: 1 + extra_args: &k3_extra_args + --attention-backend: triton + --mem-fraction-static: 0.85 + --cuda-graph-max-bs-decode: 256 + --disable-radix-cache: true + # K3 always thinks; these split reasoning and tool calls out of the answer + --reasoning-parser: kimi_k3 + --tool-call-parser: kimi_k3 + +# DSpark reuses the anchors above so the shared recipe has one definition. It is +# reported separately in the issue because it also yields an accept-length +# column, which the runner records when the server reports one. +- benchmark: serving + variant: dspark + model: moonshotai/Kimi-K3 + tp: 8 + inp: 8192 + out: 1024 + dtype: bfloat16 + max_concurrency: 2 4 8 16 32 + env: *k3_env + extra_args: + <<: *k3_extra_args + --speculative-draft-model-path: RadixArk/Kimi-K3-DSpark + --speculative-algorithm: DSPARK diff --git a/scripts/sglang/run_kimi_k3.sh b/scripts/sglang/run_kimi_k3.sh new file mode 100755 index 00000000..4cd85790 --- /dev/null +++ b/scripts/sglang/run_kimi_k3.sh @@ -0,0 +1,64 @@ +#!/bin/bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# Entry point for the Kimi K3 SGLang serving benchmark. Separate from run.sh, +# which drives the offline bench_one_batch / bench_offline_throughput path and +# is gated to gfx94x; K3 is an online serving recipe on gfx950. +set -ex + +# Preliminary setup +if [[ -z "${HF_HUB_CACHE:-}" ]]; then + export HF_HUB_CACHE="/myworkspace" +fi +export HF_TOKEN=$MAD_SECRETS_HFTOKEN + +while [[ "$#" -gt 0 ]]; do + case $1 in + --model_repo) MODEL="$2"; shift ;; + --config) CONFIG_ARG="$2"; shift ;; + --variant) VARIANT_ARG="$2"; shift ;; + *) echo "Unknown parameter passed: $1"; exit 1 ;; + esac + shift +done + +# Accept either CLI or env variable overrides +if [[ -z "$CONFIG" ]]; then + CONFIG=${CONFIG_ARG:-"configs/kimi_k3.yaml"} +fi +if [[ -z "$VARIANT" ]]; then + VARIANT=${VARIANT_ARG:-"all"} +fi + +pip install -qqq hf-transfer + +# Run benchmark; use -u to make python prints unbuffered +python3 -u run_sglang.py --config $CONFIG --model $MODEL --variant $VARIANT + +# move the output csv to parent directory +MODEL_NAME=$(basename $MODEL) +OUTPUT_CSV="perf_${MODEL_NAME}.csv" +mv $OUTPUT_CSV ../ diff --git a/scripts/sglang/run_sglang.py b/scripts/sglang/run_sglang.py new file mode 100644 index 00000000..9e0bdaa7 --- /dev/null +++ b/scripts/sglang/run_sglang.py @@ -0,0 +1,355 @@ +################################################################################# +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +"""Config-driven online serving benchmark for SGLang. + +Structured after scripts/vllm/run_vllm.py, which solves the same problem for +vLLM, and emits the same perf CSV schema so MAD's multiple_results ingestion is +shared. The existing scripts/sglang/sglang_benchmark_report.sh stays as-is; it +drives the offline bench_one_batch / bench_offline_throughput path. +""" + +import os +import csv +import json +import yaml +import psutil +import signal +import argparse +import itertools +import subprocess +from typing import List, Dict + +SUPPORTED_LIST_ARGS = ['model', 'tp', 'inp', 'out', 'num_prompts', 'max_concurrency'] +CSV_HEADER = [ + "model", + "benchmark", + "variant", + "tp", + "inp", + "out", + "dtype", + "num_prompts", + "max_concurrency", + "cmd", + "performance", + "metric", + "unit", +] + +HOST = "127.0.0.1" +PORT = 30000 +# Kimi K3 is a ~1.5 TB checkpoint loaded over TP8; 30 minutes (what run_vllm.py +# allows) is not enough. Overridable for smaller models. +SERVER_START_TIMEOUT = int(os.environ.get("SGLANG_SERVER_START_TIMEOUT", 5400)) + + +def parse_args(): + parser = argparse.ArgumentParser(description='Run SGLang serving benchmark') + parser.add_argument('--config', + type=str, + help='config yaml file', + required=True, + ) + parser.add_argument('--model', + type=str, + help='select model from config', + required=False, + default=None, + ) + parser.add_argument('--variant', + type=str, + help='select variant from config', + required=False, + default=None, + ) + parser.add_argument('--benchmark', + type=str, + help='select benchmark from config', + required=False, + default=None, + ) + parser.add_argument('--tp', + type=str, + help='select tensor parallel size from config', + required=False, + default=None, + ) + parser.add_argument('--inp', + type=str, + help='select input size from config', + required=False, + default=None, + ) + parser.add_argument('--out', + type=str, + help='select output size from config', + required=False, + default=None, + ) + parser.add_argument('--num_prompts', + type=str, + help='select num prompts from config', + required=False, + default=None, + ) + parser.add_argument('--max_concurrency', + type=str, + help='select max concurrency from config', + required=False, + default=None, + ) + args = parser.parse_args() + return args + + +def expand_configs(args, configs: List[Dict]): + # Apply architecture specific overrides to config + cfgs = [] + arch = os.environ.get('MAD_SYSTEM_GPU_ARCHITECTURE', 'unknown') + for config in configs: + cfg = config.copy() + # pop all architecture specific overrides from config + arch_overrides = cfg.pop('arch_overrides', {}) + if arch_override := arch_overrides.get(arch, {}): + print(f"Detected {arch} architecture, applying override {arch_override} to config {cfg}") + cfg.update(arch_override) + cfgs.append(cfg) + + # Expand combinations from SUPPORTED_LIST_ARGS + print(f"Expanding configs for the following keys: {SUPPORTED_LIST_ARGS} into individual configs") + config_list = [] + for cfg in cfgs: + # split config into common args and list args + common_cfgs = {k: v for k, v in cfg.items() if k not in SUPPORTED_LIST_ARGS} + list_cfgs = {k: str(v).split(' ') for k, v in cfg.items() if k in SUPPORTED_LIST_ARGS} + # expand list args into one dict per combination + expanded_cfgs = [dict(zip(list_cfgs.keys(), x)) for x in itertools.product(*list_cfgs.values())] + for expanded_cfg in expanded_cfgs: + config_list.append({**common_cfgs, **expanded_cfg}) + + # filter config list according to command line args if specified + filtered_configs = config_list + for arg_name in SUPPORTED_LIST_ARGS: + if arg_val := getattr(args, arg_name): + print(f"Filtering configs by {arg_name}={arg_val}") + filtered_configs = [cfg for cfg in filtered_configs if cfg.get(arg_name, None) == arg_val] + + # filter configs by benchmark + if args.benchmark and args.benchmark != "all": + print(f"Filtering configs by benchmark={args.benchmark}") + filtered_configs = [cfg for cfg in filtered_configs if cfg["benchmark"] == args.benchmark] + + # filter configs by variant; variants share a model, so this is the only way + # to select between recipes such as K3 nospec and K3 dspark + if args.variant and args.variant != "all": + print(f"Filtering configs by variant={args.variant}") + filtered_configs = [cfg for cfg in filtered_configs if cfg.get("variant", None) == args.variant] + + return filtered_configs + + +def read_last_json_line(path: str): + """SGLang appends one JSON object per run to --output-file, so the result of + this run is the last non-empty line (vLLM writes a plain JSON document).""" + with open(path, "r", newline="", encoding="utf-8") as f: + lines = [line for line in f if line.strip()] + if not lines: + raise Exception(f"No benchmark results found in {path}") + return json.loads(lines[-1]) + + +def run_serving(model, config): + # by default use num_prompts = 10 * max_concurrency if not specified + if not config.get("num_prompts"): + config["num_prompts"] = str(10 * int(config["max_concurrency"])) + server_cmd = ( + "sglang serve " + f"--model-path {model} " + f"--dtype {config['dtype']} " + f"--tp-size {config['tp']} " + f"--trust-remote-code " + f"--host {HOST} " + f"--port {PORT} " + ) + # pop env and extra args from config + env = config.pop('env', "") + extra_args = config.pop('extra_args', "") + server_cmd = f"{env} {server_cmd} {extra_args}".strip() + config["cmd"] = server_cmd + + # start server + print(server_cmd, flush=True) + server = subprocess.Popen(server_cmd, shell=True) + results = [] + + try: + # wait for the server to become ready. /health only returns 200 once the + # server leaves the Starting state, whereas /v1/models answers earlier. + status = subprocess.run( + f"timeout {SERVER_START_TIMEOUT} bash -c " + f"'until curl -sf http://{HOST}:{PORT}/health; do sleep 30; done' || exit 1", + shell=True + ) + if status.returncode != 0: + raise Exception("Server failed to start") + else: + print(f"Server at {server.pid} contacted successfully", flush=True) + + # run serving benchmark + output_json = ( + f"{config['model']}_{config['variant']}_serving_{config['tp']}_{config['inp']}_" + f"{config['out']}_{config['num_prompts']}_{config['max_concurrency']}.jsonl" + ) + bench_cmd = ( + "python3 -m sglang.benchmark.serving " + f"--backend sglang " + f"--host {HOST} " + f"--port {PORT} " + f"--model {model} " + f"--dataset-name random " + f"--random-input-len {config['inp']} " + f"--random-output-len {config['out']} " + f"--random-range-ratio 1.0 " + f"--max-concurrency {config['max_concurrency']} " + f"--num-prompts {config['num_prompts']} " + f"--output-file {output_json}" + ) + config["cmd"] = f"{server_cmd};{bench_cmd}" + print(bench_cmd, flush=True) + subprocess.run(bench_cmd, shell=True, check=True) + + # parse output jsonl + output = read_last_json_line(output_json) + if "total_throughput" in output: + metrics = { + "throughput_tot": str(output["total_throughput"]), + "throughput_gen": str(output["output_throughput"]), + "median_ttft": str(output["median_ttft_ms"]), + "median_tpot": str(output["median_tpot_ms"]), + "median_itl": str(output["median_itl_ms"]), + # SGLang names this median_e2e_latency_ms, not vLLM's median_e2el_ms + "median_e2el": str(output["median_e2e_latency_ms"]), + } + # only reported under speculative decoding + if output.get("accept_length"): + metrics["accept_length"] = str(output["accept_length"]) + for metric, perf in metrics.items(): + if "throughput" in metric: + unit = "tok/sec" + elif metric == "accept_length": + unit = "tokens" + else: + unit = "ms" + result = { + "performance": perf, + "metric": metric, + "unit": unit, + **config + } + results.append(result) + + finally: + # kill server and children + parent = psutil.Process(server.pid) + for child in parent.children(recursive=True): + child.send_signal(signal.SIGINT) + server.send_signal(signal.SIGINT) + _ = server.communicate() + del server + + return results + + +def main(): + args = parse_args() + + # Load, expand and filter configs + with open(args.config, 'r') as f: + print(f"Loading configs from {args.config}") + configs = yaml.safe_load(f) + configs = expand_configs(args, configs) + print(f"Running configs: ", *configs, sep='\n') + + # Iterate over configs + for config in configs: + model = config['model'] + # Use model name for logging + config['model'] = os.path.basename(model) + + # Write header to csv + OUTPUT_CSV = "perf_" + os.path.basename(model) + ".csv" + header_write = 0 if os.path.exists(OUTPUT_CSV) else 1 + with open(OUTPUT_CSV, "a+", newline="") as outf: + writer = csv.DictWriter(outf, delimiter=",", fieldnames=CSV_HEADER) + if header_write: + writer.writeheader() + outf.flush() + + # Use huggingface token if present + if MAD_SECRETS_HFTOKEN := os.environ.get('MAD_SECRETS_HFTOKEN'): + os.environ['HF_TOKEN'] = MAD_SECRETS_HFTOKEN + else: + print("Warning: MAD_SECRETS_HFTOKEN is not set. If a gated model is used, please set MAD_SECRETS_HFTOKEN=") + # Use dataprovider if present for model weights + if MAD_DATAHOME := os.environ.get('MAD_DATAHOME'): + model = MAD_DATAHOME + else: + # Explicitly download model before running benchmarks for easier debugging + download_command=f"hf download {model} --exclude \"original/*\" \"*.tf\" \"*.onnx\" \"*.flax\" \"*.rust\"" + subprocess.run(download_command, shell=True, check=True) + # A speculative-decoding config needs its draft checkpoint too + draft = config.get("extra_args", {}).get("--speculative-draft-model-path") + if draft: + subprocess.run(f"hf download {draft}", shell=True, check=True) + + # concatenate env vars and extra args into the corresponding strings + env_vars = config.get("env", {}) + extra_args = config.get("extra_args", {}) + env_vars_str = " ".join(f"{k}={v}" for k, v in env_vars.items()) + extra_args_str = "" + for k, v in extra_args.items(): + if isinstance(v, bool): + extra_args_str += f" {k}" + else: + extra_args_str += f" {k} {v}" + config["env"] = env_vars_str + config["extra_args"] = extra_args_str + + # run benchmark + results = [] + benchmark = config["benchmark"] + if benchmark == "serving": + results = run_serving(model, config) + else: + raise ValueError(f"Unknown benchmark: {benchmark}") + + # Write results to csv + for result in results: + writer.writerow(result) + + +if __name__ == "__main__": + main()