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/benchmark/vllm/README.md b/benchmark/vllm/README.md index a4228bc9..189975f6 100644 --- a/benchmark/vllm/README.md +++ b/benchmark/vllm/README.md @@ -104,6 +104,7 @@ users can also directly run the vLLm benchmark scripts and change the benchmarki | pyt_vllm_gpt-oss-120b_w4a8 | [amd/gpt-oss120b-w-mxfp4-a-fp8](https://huggingface.co/amd/gpt-oss120b-w-mxfp4-a-fp8) | | pyt_vllm_kimi-k2.6 | [moonshotai/Kimi-K2.6](https://huggingface.co/moonshotai/Kimi-K2.6) | | pyt_vllm_kimi-k2.6_fp4 | [amd/Kimi-K2.6-MXFP4](https://huggingface.co/amd/Kimi-K2.6-MXFP4) | +| pyt_vllm_kimi-k3 | [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) | | pyt_vllm_llama-3.1-8b | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | | pyt_vllm_llama-3.1-8b_fp8 | [amd/Llama-3.1-8B-Instruct-FP8-KV](https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV) | | pyt_vllm_llama-3.1-405b | [meta-llama/Llama-3.1-405B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct) | @@ -129,6 +130,25 @@ users can also directly run the vLLm benchmark scripts and change the benchmarki | pyt_vllm_qwen3.5-397b-a17b | [Qwen/Qwen3.5-397B-A17B](https://huggingface.co/Qwen/Qwen3.5-397B-A17B) | | pyt_vllm_qwen3.5-397b-a17b_fp8 | [Qwen/Qwen3.5-397B-A17B-FP8](https://huggingface.co/Qwen/Qwen3.5-397B-A17B-FP8) | +>[!NOTE] +>`pyt_vllm_kimi-k3` is the one exception to the shared Docker image above. Kimi K3 requires +>vLLM >= 0.27.0, which is not yet in a tagged `vllm-openai-rocm` release, so it builds from +>the model-specific `vllm/vllm-openai-rocm:kimi-k3` image via +>[docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile](../../docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile). +>It needs an 8x MI350X/MI355X (gfx950) node — the ~1680 GB minimum footprint does not fit a +>single 8x MI300X node — and the checkpoint is ~1.56 TB, so make sure `HF_HUB_CACHE` has room. +>It is deliberately not tagged `vllm_default`; run it explicitly: +> +>```sh +>madengine run --tags pyt_vllm_kimi-k3 --keep-model-dir --live-output +>``` +> +>The config tracks the [MI355X recipe profile](https://recipes.vllm.ai/moonshotai/Kimi-K3?hardware=mi355x) +>for a text-only serving run, with two intentional deviations: MAD adds +>`--no-enable-prefix-caching` for benchmark hygiene (as it does for every model here), and the +>gsm8k accuracy stage is disabled because K3's always-on reasoning is returned inline over +>`/v1/completions` and exhausts the generation budget. + ### Standalone benchmarking ----------------------------- diff --git a/docker/pyt_atom.ubuntu.amd.Dockerfile b/docker/pyt_atom.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..a248512b --- /dev/null +++ b/docker/pyt_atom.ubuntu.amd.Dockerfile @@ -0,0 +1,39 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 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. +# +################################################################################# +ARG BASE_DOCKER=rocm/atom-dev:latest +FROM $BASE_DOCKER + +USER root +ENV WORKSPACE_DIR=/workspace +RUN mkdir -p $WORKSPACE_DIR +WORKDIR $WORKSPACE_DIR + +# record configuration for posterity +RUN pip3 list + +# Specify entrypoint to override upstream +ENTRYPOINT [""] 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/docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile b/docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..db33cd7b --- /dev/null +++ b/docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile @@ -0,0 +1,44 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 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 requires vLLM >= 0.27.0, which is not in a tagged vllm-openai-rocm +# release yet; the model-specific :kimi-k3 image is the only ROCm build with the +# KDA / Gated MLA / Stable LatentMoE support the checkpoint needs. Kept separate +# from docker/pyt_vllm so the ~35 other vLLM entries stay on the tagged release. +# Fold this back into docker/pyt_vllm once K3 lands in a vX.Y.Z ROCm image. +ARG BASE_DOCKER=vllm/vllm-openai-rocm:kimi-k3 +FROM $BASE_DOCKER + +USER root +ENV WORKSPACE_DIR=/workspace +RUN mkdir -p $WORKSPACE_DIR +WORKDIR $WORKSPACE_DIR + +# record configuration for posterity +RUN pip3 list + +# Specify entrypoint to override upstream +ENTRYPOINT [""] diff --git a/models.json b/models.json index 15c3cd5b..29e6904b 100644 --- a/models.json +++ b/models.json @@ -327,6 +327,47 @@ "args": "--model_repo amd/Kimi-K2.6-MXFP4 --config configs/default.yaml" }, + { + "name": "pyt_vllm_kimi-k3", + "url": "", + "dockerfile": "docker/pyt_vllm_kimi_k3", + "scripts": "scripts/vllm/run.sh", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_Kimi-K3.csv", + "tags": [ + "pyt", + "vllm", + "inference" + ], + "timeout": -1, + "skip_gpu_arch": "gfx942", + "args": + "--model_repo moonshotai/Kimi-K3 --config configs/default.yaml" + }, + { + "name": "pyt_atom_kimi-k3", + "url": "", + "dockerfile": "docker/pyt_atom", + "scripts": "scripts/atom/run.sh", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "multiple_results": "perf_Kimi-K3.csv", + "tags": [ + "pyt", + "atom", + "atom_default", + "inference" + ], + "timeout": -1, + "skip_gpu_arch": "gfx942, gfx906, gfx908, gfx90a, A100, H100, V100", + "args": + "--model_repo moonshotai/Kimi-K3 --config configs/default.yaml" + }, { "name": "pyt_vllm_llama-3.1-8b", "url": "", @@ -2194,6 +2235,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/atom/configs/accuracy.yaml b/scripts/atom/configs/accuracy.yaml new file mode 100644 index 00000000..b59d0bc9 --- /dev/null +++ b/scripts/atom/configs/accuracy.yaml @@ -0,0 +1,17 @@ +# ATOM Kimi-K3 accuracy benchmark (GSM8K 5-shot) +# Expected range: flexible-extract 0.9538–0.9591 + +- benchmark: accuracy + model: moonshotai/Kimi-K3 + tp: 8 + kv_cache_dtype: fp8 + num_fewshot: 5 + apply_chat_template: false + extra_args: + --trust-remote-code: true + --max-model-len: 16384 + --max-num-seqs: 64 + --max-num-batched-tokens: 16384 + --gpu-memory-utilization: 0.93 + --block-size: 128 + --no-enable_prefix_caching: true diff --git a/scripts/atom/configs/default.yaml b/scripts/atom/configs/default.yaml new file mode 100644 index 00000000..68680037 --- /dev/null +++ b/scripts/atom/configs/default.yaml @@ -0,0 +1,20 @@ +# ATOM Kimi-K3 serving benchmark configs +# 1k/4k ISL, 1k OSL, mc64/128/256 +# gfx950 (MI350X/MI355X) only — MXFP4 weights require gfx95x + +## Kimi-K3 MXFP4 TP8 +- benchmark: serving + model: moonshotai/Kimi-K3 + tp: 8 + inp: 1024 4096 + out: 1024 + kv_cache_dtype: fp8 + max_concurrency: 64 128 256 + extra_args: + --trust-remote-code: true + --max-model-len: 16384 + --max-num-seqs: 64 + --max-num-batched-tokens: 16384 + --gpu-memory-utilization: 0.93 + --block-size: 128 + --no-enable_prefix_caching: true diff --git a/scripts/atom/configs/perf.yaml b/scripts/atom/configs/perf.yaml new file mode 100644 index 00000000..c7f92dd4 --- /dev/null +++ b/scripts/atom/configs/perf.yaml @@ -0,0 +1,25 @@ +# ATOM Performance Recipes +# Source: https://github.com/ROCm/ATOM/pull/1718 (recipes/Kimi-K3.md) +# Merge strategy: User config wins; perf fills gaps (missing keys only) + +# Kimi-K3 MXFP4 (gfx950 / MI355) +- model: moonshotai/Kimi-K3 + extra_args: + --kv_cache_dtype: fp8 + --trust-remote-code: true + --max-model-len: 16384 + --max-num-seqs: 64 + --max-num-batched-tokens: 16384 + --gpu-memory-utilization: 0.93 + --block-size: 128 + --no-enable_prefix_caching: true + env: + ATOM_LOADER_USE_THREADPOOL: '1' + ATOM_LOADER_THREADPOOL_WORKERS: '16' + ATOM_SYNC_AFTER_LOAD: '1' + ATOM_DIST_TIMEOUT_SECONDS: '3600' + ATOM_USE_TRITON_GEMM: '1' + AITER_USE_GROUPED_GEMM: '0' + ATOM_USE_TRITON_MOE: '0' + AITER_FLYDSL_FORCE: '1' + AITER_FORCE_GFX1250: '0' diff --git a/scripts/atom/models.json b/scripts/atom/models.json new file mode 100644 index 00000000..8cab4590 --- /dev/null +++ b/scripts/atom/models.json @@ -0,0 +1,19 @@ +[ + { + "name": "atom-perf", + "url": "", + "dockerfile": "../../docker/pyt_atom", + "scripts": "run.sh", + "data": "huggingface", + "n_gpus": "-1", + "owner": "", + "training_precision": "", + "multiple_results": "perf_auto.csv", + "tags": [ + "atom-perf", + "inference" + ], + "timeout": 432000, + "args": "" + } +] diff --git a/scripts/atom/run.sh b/scripts/atom/run.sh new file mode 100755 index 00000000..199f872c --- /dev/null +++ b/scripts/atom/run.sh @@ -0,0 +1,98 @@ +#!/bin/bash +############################################################################### +# +# MIT License +# +# Copyright (c) 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. +# +################################################################################# +set -ex + +# Preliminary setup +export HF_HUB_CACHE="/myworkspace" +MAD_MODEL_NAME=$(echo $MAD_MODEL_NAME | tr "/" "_") + +PERF_ARGS="" +PROFILE=false +while [[ "$#" -gt 0 ]]; do + case $1 in + --model_repo) MODEL="$2"; shift ;; + --config) CONFIG_ARG="$2"; shift ;; + --benchmark) BENCHMARK_ARG="$2"; shift ;; + --output_csv) OUTPUT_CSV="$2"; shift ;; + --perf) PERF_ARGS="$PERF_ARGS --perf $2"; shift ;; + --perf-output) PERF_ARGS="$PERF_ARGS --perf-output $2"; shift ;; + --no-perf-merge) PERF_ARGS="$PERF_ARGS --no-perf-merge" ;; + --profile) PROFILE=true ;; + *) + echo "Unknown parameter passed: $1" >&2 + echo "Usage: run.sh --model_repo [--config ] [--benchmark ] [--output_csv ] [--perf ] [--perf-output ] [--no-perf-merge]" >&2 + exit 1 + ;; + esac + shift +done + +MODEL_NAME=$(basename $MODEL) + +# By default run all benchmarks in configs/default.yaml; accept either CLI or env variable overrides +if [[ -z "$BENCHMARK" ]]; then + BENCHMARK=${BENCHMARK_ARG:-"all"} +fi +if [[ -z "$CONFIG" ]]; then + CONFIG=${CONFIG_ARG:-"configs/default.yaml"} +fi + +OUTPUT_CSV="perf_${MODEL_NAME}.csv" +if [[ ! -z "$MAD_OUTPUT_CSV" ]]; then + OUTPUT_CSV="$MAD_OUTPUT_CSV" +fi + +# install lm-eval for accuracy testing +pip install -qqq lm-eval[api] hf-transfer + + +# install profiling dependencies (rocm-trace-lite) only when profiling +if $PROFILE; then + apt-get update || true + apt-get install -y g++ || true + apt-get install -y libsqlite3-dev || true + if [ ! -f /usr/include/sqlite3.h ]; then + SQV=$(apt-cache policy libsqlite3-0 | awk '/Candidate:/{print $2}') + ( cd /tmp && (apt-get download libsqlite3-0=$SQV libsqlite3-dev=$SQV || apt-get download libsqlite3-0 libsqlite3-dev) \ + && for d in libsqlite3-0_*.deb libsqlite3-dev_*.deb; do [ -f "$d" ] && dpkg -x "$d" /; done && ldconfig ) + fi + git clone https://github.com/amathews-amd/rocm-trace-lite.git + cd rocm-trace-lite + git checkout amathews-amd/sig_handle + sed -i 's/ _generate_perfetto(output, json_file)/ pass # perfetto JSON export disabled/' rocm_trace_lite/cmd_trace.py + sed -i 's/sys.exit(result.returncode if result is not None else 0)/sys.exit(getattr(result, "returncode", 0) if "result" in dir() else 0)/g' rocm_trace_lite/cmd_trace.py + make -j + make install + pip install -e . + cd .. +fi + +# Run benchmark; use -u to make python prints unbuffered +python3 -u run_atom.py --config $CONFIG --model $MODEL --benchmark $BENCHMARK $($PROFILE && echo --profile) $PERF_ARGS + +# move the output csv to parent directory +mv "perf_${MODEL_NAME}.csv" ../$OUTPUT_CSV diff --git a/scripts/atom/run_atom.py b/scripts/atom/run_atom.py new file mode 100644 index 00000000..c793dea7 --- /dev/null +++ b/scripts/atom/run_atom.py @@ -0,0 +1,857 @@ +################################################################################# +# +# MIT License +# +# Copyright (c) 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. +# +################################################################################# + +import os +import sys +import csv +import glob +import json +import yaml +import psutil +import shutil +import signal +import argparse +import itertools +import subprocess +import time +from typing import List, Dict + +SUPPORTED_LIST_ARGS = ['model', 'tp', 'inp', 'out', 'bs', 'num_prompts', 'max_concurrency'] +CSV_HEADER = [ + "hf_pipeline_tag", + "model", + "benchmark", + "tp", + "inp", + "out", + "kv_cache_dtype", + "num_prompts", + "max_concurrency", + "bs", + "cmd", + "performance", + "metric", + "unit", +] + + +# --------------------------------------------------------------------------- +# Perf overrides: auto-fill of a sibling perf.yaml into --config. +# +# perf.yaml entries are keyed by a space-separated `model` list and carry only +# `extra_args` and/or `env`. Before expansion, each base config entry is +# matched against perf entries by model-token intersection; the first matching +# perf entry's extra_args and env are merged into the base entry as a +# *gap-fill*: anything already present in the base is left untouched, and only +# keys missing from the base are added from perf. Other base fields +# (benchmark, tp, inp, out, ...) are not touched at all — perf is a perf- +# tuning sheet, not a run-shape sheet. +# --------------------------------------------------------------------------- + +def _tokens(model_field): + """Split a `model:` value into a set of repo tokens. + + The base `config.yaml` and `perf.yaml` both allow space-separated repos + on a single `model:` line; tokenize on whitespace so we can intersect. + """ + if model_field is None: + return set() + return {t.strip() for t in str(model_field).split() if t.strip()} + + +def _fill_missing(base, overrides): + """Return a copy of `base` with keys from `overrides` added only when missing. + + Base wins on key collisions: any key already present in `base` keeps its + value, regardless of what `overrides` says. Both `extra_args` and `env` + in this schema are flat scalar dicts, so a one-level operation is + sufficient. Inputs are not mutated. + """ + out = dict(base or {}) + for k, v in (overrides or {}).items(): + if k not in out: + out[k] = v + return out + + +def load_perf_entries(perf_path): + """Load perf.yaml into a normalized list of override entries. + + Returns an empty list when the file is missing or empty; each entry is + `{'models': set[str], 'extra_args': dict, 'env': dict}`. + """ + if not perf_path or not os.path.exists(perf_path): + return [] + with open(perf_path, 'r') as f: + raw = yaml.safe_load(f) or [] + entries = [] + for e in raw: + entries.append({ + 'models': _tokens(e.get('model')), + 'extra_args': e.get('extra_args', {}) or {}, + 'env': e.get('env', {}) or {}, + 'bench_serving': e.get('bench_serving', {}) or {}, + }) + return entries + + +def apply_perf_overrides(configs, perf_entries): + """Mutate `configs` in place, gap-filling matching perf overrides. + + For each base entry, the first perf entry whose model tokens intersect + the base entry's model tokens wins. Only keys missing from the base + entry's `extra_args` / `env` are added from perf — anything already set + in the base is preserved. Returns a per-entry summary list suitable for + logging, distinguishing keys that were added vs. skipped (already + present in base). + """ + summary = [] + for cfg in configs: + base_tokens = _tokens(cfg.get('model')) + matched = next((p for p in perf_entries if p['models'] & base_tokens), None) + if not matched: + continue + + base_extra = cfg.get('extra_args') or {} + base_env = cfg.get('env') or {} + cfg['extra_args'] = _fill_missing(base_extra, matched['extra_args']) + cfg['env'] = _fill_missing(base_env, matched['env']) + + # gap-fill per-model bench_serving knobs from perf.yaml (base wins) + if matched.get('bench_serving'): + cfg['bench_serving'] = _fill_missing(cfg.get('bench_serving') or {}, matched['bench_serving']) + + summary.append({ + 'model': cfg.get('model'), + 'matched_perf_models': sorted(matched['models']), + 'extra_args_added': sorted(k for k in matched['extra_args'] if k not in base_extra), + 'extra_args_skipped': sorted(k for k in matched['extra_args'] if k in base_extra), + 'env_added': sorted(k for k in matched['env'] if k not in base_env), + 'env_skipped': sorted(k for k in matched['env'] if k in base_env), + }) + return summary + + +def parse_args(): + parser = argparse.ArgumentParser(description='Run ATOM 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('--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('--bs', + type=str, + help='select batch 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, + ) + parser.add_argument('--perf', + type=str, + help='Path to perf.yaml (default: /perf.yaml)', + required=False, + default=None, + ) + parser.add_argument('--perf-output', + type=str, + help='Where to write the merged config ' + '(default: /perf_config.yaml). ' + 'The original --config file is never modified.', + required=False, + default=None, + ) + parser.add_argument('--no-perf-merge', + action='store_true', + help='Disable auto-merge of sibling perf.yaml into --config', + ) + parser.add_argument('--profile', + action='store_true', + help='Kernel-profiling run: serve under rtl trace + roctx markers and aggregate a kernel payload', + ) + args = parser.parse_args() + return args + +def dict_to_args(args_dict: Dict) -> str: + """Convert argument dictionary to command-line string""" + args_list = [] + for key, value in args_dict.items(): + if value is True: + args_list.append(key) + elif value is not False and value is not None: + args_list.append(f"{key} {value}") + return " ".join(args_list) + +def dict_to_env(env_dict: Dict) -> str: + """Convert environment dictionary to env var string""" + return " ".join(f"{k}={v}" for k, v in env_dict.items()) + +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] + + return filtered_configs + +def _wait_for_server(server, port=8000, timeout=5400, poll_interval=30): + """Wait until the server's /v1/models endpoint responds. + + Returns True once the endpoint is reachable. Returns False if the server + process exits before becoming ready (e.g. an OOM during initialization) or + if the timeout elapses -- instead of blocking on a dead endpoint for the + full timeout, which previously made a crashed server look like a hang. + """ + deadline = time.time() + timeout + url = f"http://localhost:{port}/v1/models" + while time.time() < deadline: + if server.poll() is not None: + print( + f"Server process exited with code {server.returncode} " + "before becoming ready", + flush=True, + ) + return False + probe = subprocess.run( + f"curl -s {url}", + shell=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if probe.returncode == 0: + return True + time.sleep(poll_interval) + print("Timed out waiting for server to start", flush=True) + return False + +def run_serving(model, config): + # num_prompts = multiplier * max_concurrency; multiplier from perf.yaml bench_serving (default 10), popped so it stays out of the CSV row + multiplier = int((config.pop("bench_serving", {}) or {}).get("num_prompts_multiplier") or 10) + if not config.get("num_prompts"): + config["num_prompts"] = str(multiplier * int(config["max_concurrency"])) + output_json = ( + f"{config['model']}_serving_{config['tp']}_{config['inp']}_{config['out']}_{config['num_prompts']}_{config['max_concurrency']}.json" + ) + server_cmd = ( + "python -m atom.entrypoints.openai_server " + f"--model {model} " + f"-tp {config['tp']} " + ) + + # Add kv_cache_dtype if specified + if 'kv_cache_dtype' in config and config['kv_cache_dtype']: + server_cmd += f"--kv_cache_dtype {config['kv_cache_dtype']} " + + # Get env and extra args from config (keep as dicts) + env_dict = config.pop('env', {}) + extra_args_dict = config.pop('extra_args', {}) + + # Convert dicts to command line strings + env_str = dict_to_env(env_dict) if env_dict else "" + extra_args_str = dict_to_args(extra_args_dict) if extra_args_dict else "" + + server_cmd = f"{env_str} {server_cmd} {extra_args_str}".strip() + config["cmd"] = server_cmd + + # start server + print(server_cmd) + server = subprocess.Popen(server_cmd, shell=True) + + try: + # wait for server to start; fail fast if the process dies (e.g. OOM + # during init) instead of blocking on a dead endpoint for 30 minutes + if not _wait_for_server(server): + print("Server failed to start") + return [config] + else: + print(f"Server at {server.pid} contacted successfully", flush=True) + + # run benchmark + bench_cmd = ( + "python -m atom.benchmarks.benchmark_serving " + f"--model {model} " + f"--backend vllm " + f"--base-url http://localhost:8000 " + f"--percentile-metrics ttft,tpot,itl,e2el " + f"--dataset-name random " + f"--ignore-eos " + f"--request-rate inf " + f"--random-range-ratio 0.8 " + f"--max-concurrency {config['max_concurrency']} " + f"--num-prompts {config['num_prompts']} " + f"--random-input-len {config['inp']} " + f"--random-output-len {config['out']} " + f"--save-result " + f"--result-dir ./ " + f"--result-filename {output_json}" + ) + + # Add --trust-remote-code if it was in extra_args (from recipe) + if '--trust-remote-code' in extra_args_dict: + bench_cmd += " --trust-remote-code" + bench_args = config.pop('bench_args', {}) + bench_args_str = "" + for k, v in bench_args.items(): + if isinstance(v, bool): + bench_args_str += f"{k} " + else: + bench_args_str += f"{k} {v} " + bench_cmd = f"{bench_cmd} {bench_args_str}".strip() + + config["cmd"] = f"{server_cmd};{bench_cmd}" + print(bench_cmd) + subprocess.run(bench_cmd, shell=True, check=True) + + # parse output json + results = [] + with open(output_json, "r", newline="", encoding="utf-8") as f: + output = json.load(f) + if "total_token_throughput" in output: + metrics = { + "throughput_tot": str(output["total_token_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"]), + "median_e2el": str(output["median_e2el_ms"]), + } + for metric, perf in metrics.items(): + if "throughput" in metric: + unit = "tok/sec" + else: + unit = "ms" + result = { + "performance": perf, + "metric": metric, + "unit": unit, + **config + } + results.append(result) + + return results + + 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 + +def run_serving_profile(model, config): + """Single-pass profiling for the ATOM serving benchmark (Component 5, kernel breakdown). + + Serves `atom.entrypoints.openai_server` under rocm-trace-lite (`rtl trace --mode full`) with roctx + prefill/decode/mixed phase markers (scripts/atom/profiling/), producing rtl_trace.db (GPU kernels + + per-pass roctx markers), then aggregates it into ./profile/kernel_summary_payload.json for the + client-perf-hub uploader. Mirrors the vLLM/SGLang producer's run_serving_profile. + """ + config.pop("bench_serving", None) # perf.yaml-only knob; keep out of CSV row + if not config.get("num_prompts"): + config["num_prompts"] = str(min(5, 10 * int(config["max_concurrency"]))) + output_json = ( + f"{config['model']}_serving_profile_{config['tp']}_{config['inp']}_{config['out']}_{config['num_prompts']}_{config['max_concurrency']}.json" + ) + profile_dir = config.get("profile_dir", f"profile_{config['model']}_{config['tp']}") + os.makedirs(profile_dir, exist_ok=True) + rtl_trace_path = os.path.join(profile_dir, "rtl_trace.db") + + server_cmd = ( + "python -m atom.entrypoints.openai_server " + f"--model {model} " + f"-tp {config['tp']} " + ) + if 'kv_cache_dtype' in config and config['kv_cache_dtype']: + server_cmd += f"--kv_cache_dtype {config['kv_cache_dtype']} " + + env_dict = config.pop('env', {}) + extra_args_dict = config.pop('extra_args', {}) + env_str = dict_to_env(env_dict) if env_dict else "" + extra_args_str = dict_to_args(extra_args_dict) if extra_args_dict else "" + + # roctx phase markers: PYTHONPATH puts sitecustomize.py on the path so the patch reaches ATOM's + # spawned engine-core/worker processes (spawn start method); ATOM_ROCTX_PHASE_MARKERS arms it. + profiling_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiling") + roctx_env = f"PYTHONPATH={profiling_dir}${{PYTHONPATH:+:$PYTHONPATH}} ATOM_ROCTX_PHASE_MARKERS=1" + rtl_env = f"{roctx_env} rtl trace --mode full -o {rtl_trace_path}" + server_cmd = f"{env_str} {rtl_env} {server_cmd} {extra_args_str}".strip() + config["cmd"] = server_cmd + + print("=" * 60) + print(" PROFILE: rocm-trace-lite (--mode full) + roctx phase markers") + print("=" * 60) + print(server_cmd) + # Redirect the server's stdout+stderr to a file rather than inheriting the console: ATOM's aiter + # backend prints one line per shape during CUDA-graph capture, which otherwise bloats the CI log + # past the Actions blob-store limit (making it an unreadable BlobNotFound). On failure we print the + # tail so the real traceback is retrievable; the full log is kept under profile_dir (--keep-model-dir). + server_log_path = os.path.join(profile_dir, "server_profile.log") + _server_log = open(server_log_path, "w") + + def _tail_server_log(n=200): + try: + _server_log.flush() + with open(server_log_path, "r", errors="replace") as _f: + _tail = _f.readlines()[-n:] + print(f"----- last {len(_tail)} lines of {server_log_path} (server stdout+stderr) -----") + print("".join(_tail)) + print("----- end server log tail -----", flush=True) + except Exception as _e: + print(f"(could not read server log {server_log_path}: {_e})") + + server = subprocess.Popen(server_cmd, shell=True, stdout=_server_log, stderr=subprocess.STDOUT) + try: + status = subprocess.run( + "timeout 5400 bash -c 'until curl -s http://localhost:8000/v1/models; do sleep 30; done' || exit 1", + shell=True, + ) + if status.returncode != 0: + print("Server failed to start for profiling run") + _tail_server_log() + return [config] + print(f"Server at {server.pid} ready for profiling", flush=True) + + bench_cmd = ( + "python -m atom.benchmarks.benchmark_serving " + f"--model {model} " + f"--backend vllm " + f"--base-url http://localhost:8000 " + f"--percentile-metrics ttft,tpot,itl,e2el " + f"--dataset-name random " + f"--ignore-eos " + f"--request-rate inf " + f"--random-range-ratio 0.8 " + f"--max-concurrency {config['max_concurrency']} " + f"--num-prompts {config['num_prompts']} " + f"--random-input-len {config['inp']} " + f"--random-output-len {config['out']} " + f"--save-result " + f"--result-dir ./ " + f"--result-filename {output_json}" + ) + if '--trust-remote-code' in extra_args_dict: + bench_cmd += " --trust-remote-code" + print(bench_cmd) + try: + subprocess.run(bench_cmd, shell=True, check=True) + except subprocess.CalledProcessError: + print("Benchmark failed for profiling run") + _tail_server_log() + raise + finally: + # SIGINT (not SIGKILL) so rtl's signal handler finalizes rtl_trace.db before exit. + try: + parent = psutil.Process(server.pid) + for child in parent.children(recursive=True): + child.send_signal(signal.SIGINT) + server.send_signal(signal.SIGINT) + server.communicate(timeout=180) + except Exception: + try: + server.send_signal(signal.SIGKILL) + server.communicate() + except Exception: + pass + del server + + if not os.path.exists(rtl_trace_path): + print(f"WARNING: RTL trace not found at {rtl_trace_path}") + return [config] + print(f"RTL trace: {rtl_trace_path}") + + os.makedirs("profile", exist_ok=True) + payload_path = os.path.join("profile", "kernel_summary_payload.json") + # Shared, framework-agnostic curated aggregator (single source of truth = the vLLM + # producer, which includes #380 pass-curation). The whole MAD-private repo is present + # at runtime (madengine MODEL_DIR), so reference it by repo-relative path rather than + # keeping a per-engine copy (which drifts and misses curation). + scripts_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + aggregator = os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiling", "kernel_summary_payload.py") + try: + subprocess.run( + ["python3", aggregator, rtl_trace_path, "-o", payload_path, "--full", + "--tp", str(config["tp"]), "--precision", str(config.get("dtype", "auto")), + "--isl", str(config["inp"]), "--osl", str(config["out"])], + check=True, + ) + print(f"Payload: {payload_path}") + except Exception as e: + print(f"WARNING: payload generation failed: {e}") + + result = {"performance": "N/A", "metric": "profile", "unit": "trace", **config} + result["cmd"] = server_cmd + return [result] + +def run_accuracy(model, config): + output_json = ( + f"{config['model']}_accuracy_{config['tp']}.json" + ) + server_cmd = ( + "python -m atom.entrypoints.openai_server " + f"--model {model} " + f"-tp {config['tp']} " + ) + + # Add kv_cache_dtype if specified + if 'kv_cache_dtype' in config and config['kv_cache_dtype']: + server_cmd += f"--kv_cache_dtype {config['kv_cache_dtype']} " + + # Get env and extra args from config (keep as dicts) + env_dict = config.pop('env', {}) + extra_args_dict = config.pop('extra_args', {}) + + # Convert dicts to command line strings + env_str = dict_to_env(env_dict) if env_dict else "" + extra_args_str = dict_to_args(extra_args_dict) if extra_args_dict else "" + + server_cmd = f"{env_str} {server_cmd} {extra_args_str}".strip() + config["cmd"] = server_cmd + + # start server + print(server_cmd) + server = subprocess.Popen(server_cmd, shell=True) + + try: + # wait for server to start; fail fast if the process dies (e.g. OOM + # during init) instead of blocking on a dead endpoint for 30 minutes + if not _wait_for_server(server): + print("Server failed to start") + return [config] + else: + print(f"Server at {server.pid} contacted successfully", flush=True) + + # run benchmark + num_concurrent = config.get("num_concurrent", 64) + max_gen_toks = config.get("max_gen_toks", 2048) + limit = config.get("limit") + num_fewshot = config.get("num_fewshot", 3) + apply_chat_template = config.get("apply_chat_template", True) + model_args = { + "model": model, + "max_gen_toks": max_gen_toks, + "num_concurrent": num_concurrent, + "max_retries": 3, + "base_url": "http://localhost:8000/v1/completions", + "tokenized_requests": False, + } + model_args = ",".join([f"{k}={v}" for k, v in model_args.items()]) + bench_cmd = ( + "lm_eval " + "--model local-completions " + f"--model_args {model_args} " + f"--tasks gsm8k " + f"--num_fewshot {num_fewshot} " + f"--output_path ./tmp " + ) + if limit is not None: + bench_cmd += f"--limit {limit} " + if apply_chat_template: + bench_cmd += "--apply_chat_template " + bench_args = config.pop('bench_args', {}) + bench_args_str = "" + for k, v in bench_args.items(): + if isinstance(v, bool): + bench_args_str += f"{k} " + else: + bench_args_str += f"{k} {v} " + bench_cmd = f"{bench_cmd} {bench_args_str}".strip() + config["cmd"] = f"{server_cmd};{bench_cmd}" + print(bench_cmd) + subprocess.run(bench_cmd, shell=True, check=True) + + # find output file and move into output_json + output_files = glob.glob("./tmp/*/*.json") + if len(output_files) == 0: + print("No output files found") + return [config] + elif len(output_files) > 1: + print(f"Multiple output files found: {output_files}") + return [config] + else: + output_file = output_files[0] + shutil.move(output_file, output_json) + shutil.rmtree("./tmp") + + # parse output json + results = [] + + with open(output_json, "r", newline="", encoding="utf-8") as f: + output = json.load(f) + if "results" in output and "gsm8k" in output["results"]: + gsm8k_results = output["results"]["gsm8k"] + + for metric_key in ("exact_match,flexible-extract", "exact_match,strict-match"): + if metric_key in gsm8k_results: + results.append({ + "performance": gsm8k_results[metric_key], + "metric": metric_key, + "unit": "percent", + **config + }) + + return results + + 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 + +def main(): + args = parse_args() + + # Load base config + with open(args.config, 'r') as f: + print(f"Loading configs from {args.config}") + configs = yaml.safe_load(f) or [] + + # Perf merge: auto-detect a sibling perf.yaml, deep-merge matching + # extra_args + env into each base entry in memory, and materialize the + # merged result to a sibling perf_config.yaml. The original --config file + # is never modified; the in-memory merged configs are what actually run. + if not args.no_perf_merge: + config_dir = os.path.dirname(os.path.abspath(args.config)) + perf_path = args.perf or os.path.join(config_dir, 'perf.yaml') + perf_output = args.perf_output or os.path.join(config_dir, 'perf_config.yaml') + + if os.path.abspath(perf_output) == os.path.abspath(args.config): + raise ValueError( + f"--perf-output must differ from --config to keep the source untouched; " + f"both resolved to {perf_output}" + ) + + perf_entries = load_perf_entries(perf_path) + if perf_entries: + summary = apply_perf_overrides(configs, perf_entries) + if summary: + print(f"Applied perf overrides from {perf_path} (base wins; perf fills gaps):") + for s in summary: + print( + f" model={s['model']} <- perf {s['matched_perf_models']}\n" + f" extra_args added: {s['extra_args_added']}\n" + f" extra_args skipped: {s['extra_args_skipped']} (already set in base)\n" + f" env added: {s['env_added']}\n" + f" env skipped: {s['env_skipped']} (already set in base)" + ) + with open(perf_output, 'w') as f: + yaml.safe_dump(configs, f, sort_keys=False, default_flow_style=False) + print( + f"Wrote merged config to {perf_output} " + f"(original {args.config} left untouched)" + ) + else: + print( + f"No perf entries matched any model in {args.config}; " + f"not writing {perf_output}" + ) + else: + print(f"No perf.yaml found at {perf_path} (or empty); skipping perf merge") + + # Expand and filter configs + 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, extrasaction="ignore") + 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 CHECK_LOCAL_DATA env var to control whether to check for local data + # By default, this is set to False and can be enabled through madengine additional_context + CHECK_LOCAL_DATA = os.environ.get('CHECK_LOCAL_DATA', 'false').lower() == 'true' + + # Use dataprovider if present for model weights + if CHECK_LOCAL_DATA and (MAD_DATAHOME := os.environ.get('MAD_DATAHOME')) and os.path.exists(os.path.join(MAD_DATAHOME, model)): + model = os.path.join(MAD_DATAHOME, model) + print("Found MAD_DATAHOME updating model path.") + elif config.get('extra_args', {}).get('--load_dummy', False): + print("Found --load_dummy in config, using dummy weights for benchmarking") + else: + # Explicitly download model before running benchmarks for easier debugging. + # NOTE: pass each exclude pattern with its own --exclude flag and avoid + # shell=True; otherwise the hf CLI treats the extra patterns as positional + # filenames and silently downloads 0 files. + download_command = [ + "hf", "download", model, + "--exclude", "original/*", + "--exclude", "*.tf", + "--exclude", "*.onnx", + "--exclude", "*.flax", + "--exclude", "*.rust", + ] + # Don't trust the CLI's exit code: the typer/click versions in this venv + # leak a clean Exit(0) up as an unhandled exception, so `hf` can return + # exit status 1 even after a fully successful download. Verify the + # download via huggingface_hub's cache lookup helper, which resolves the + # snapshot path through the hub's own cache layout/ref handling rather + # than us hard-coding refs/main + models--/snapshots/. + # + # We also capture hf's stdout/stderr instead of inheriting them, so the + # spurious typer/click traceback (plus the tqdm progress bars on stderr) + # don't land in the log on the happy path. On real failure we re-emit + # both streams for diagnosis before raising. + from huggingface_hub import try_to_load_from_cache + + result = subprocess.run( + download_command, check=False, capture_output=True, text=True + ) + cached_config = try_to_load_from_cache( + repo_id=model, filename="config.json" + ) + # try_to_load_from_cache returns either a str path, None (not cached), + # or a sentinel marking "known not to exist". Collapse all non-path + # outcomes (and stale entries that no longer exist on disk) into one + # failure branch. + if not ( + isinstance(cached_config, str) and os.path.exists(cached_config) + ): + sys.stdout.write(result.stdout or "") + sys.stderr.write(result.stderr or "") + raise RuntimeError( + f"hf download {model} failed (exit={result.returncode}); " + f"config.json for {model} not found in the HF cache" + ) + snapshot_path = os.path.dirname(cached_config) + # On the happy path, keep hf's stdout (the "✓ Downloaded" line and the + # cache path) but drop stderr, which on this venv contains the spurious + # typer/click traceback and the tqdm progress bars. + sys.stdout.write(result.stdout or "") + if result.returncode != 0: + print( + f"hf exited with code {result.returncode} but {snapshot_path} " + f"looks complete; treating as success (known typer/click Exit " + f"propagation issue in this venv)." + ) + print(f"Downloaded {model} to {snapshot_path}") + + # run benchmark + results = [] + benchmark = config["benchmark"] + profile = bool(config.pop('profile', False)) or getattr(args, 'profile', False) + if benchmark == "serving": + results = run_serving_profile(model, config) if profile else run_serving(model, config) + elif benchmark == "accuracy": + results = run_accuracy(model, config) + else: + raise ValueError(f"Unknown benchmark: {benchmark}") + + # Write results to csv + for result in results: + writer.writerow(result) + + +if __name__ == "__main__": + main() 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() diff --git a/scripts/vllm/configs/default.yaml b/scripts/vllm/configs/default.yaml index d3687efe..cd861b99 100644 --- a/scripts/vllm/configs/default.yaml +++ b/scripts/vllm/configs/default.yaml @@ -20,6 +20,42 @@ env: VLLM_ROCM_USE_AITER: 1 +## Kimi K3 is natively MXFP4 (QAT), so there is no separate amd/ FP4 repo. +## Flags and env mirror the gfx950 single-node TP8 profile from +## https://recipes.vllm.ai/moonshotai/Kimi-K3/hw/mi355x.json +## TP8 only: ~1680 GB minimum footprint fits an 8x MI355X node (2304 GB), not TP4. +- benchmark: serving + model: moonshotai/Kimi-K3 + tp: 8 + inp: 1024 + out: 1024 + dtype: auto + # capped at 128 to match --max-num-seqs; higher concurrency would just queue + max_concurrency: 1 8 32 128 + env: + VLLM_ROCM_USE_AITER: 1 + SAFETENSORS_FAST_GPU: 1 + # 1 selects the aiter a8w4 MoE path; 0 falls back to a16w4 + AITER_SITUV2_A8W4: 1 + AITER_BF16_FP8_MOE_BOUND: 0 + VLLM_USE_BREAKABLE_CUDAGRAPH: 0 + extra_args: + --moe-backend: auto + --load-format: auto + --gpu-memory-utilization: 0.95 + # MoonViT-V2 is only 401M params; TP on it is pure comm overhead + --mm-encoder-tp-mode: data + --max-num-seqs: 128 + --max-num-batched-tokens: 4096 + # K3 always thinks and returns reasoning_content + --reasoning-parser: kimi_k3 + # text-only benchmark; skips loading MoonViT and frees VRAM for KV cache + --language-model-only: true + bench_args: + # gsm8k runs over /v1/completions, where the always-on reasoning is inline and + # exhausts the hardcoded max_gen_toks=2048 budget, so the score is meaningless + --run_accuracy: False + - benchmark: serving model: deepseek-ai/DeepSeek-V3.2