diff --git a/README.md b/README.md index 0e6efcf..821b82c 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # sskj — 多平台大模型推理性能基准测试项目 +**更新(2026-08-18 12:24:19 CST)**:新增 Kimi-K3 / 601-604 / SGLang TP32×EP32 的真实 Prefill MoE backend 验收实验;固定比较 Marlin 与 FlashInfer MXFP4 在 16K→1、C=8/16、Chunk=8K/16K 下的 TTFT、E2E 与 Input TPS,并提供四节点一致候选镜像构建、原始证据和自动汇总,详见 `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md`。 + **更新(2026-08-17 11:56:47 CST)**:完成 Kimi-K3 / RTX PRO 6000D / SGLang FlashInfer MXFP4 MoE Phase 2 correctness harness;SwiGLU 控制矩阵 8/8 通过,覆盖 edge routing、EP global ID、NaN/Inf、确定性、CUDA Graph 和 Kimi 真尺寸,SiTU 目标矩阵 8/8 稳定复现同一 C++ 缺口,详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md`。 > 历史更新见 `git log`。项目目的与工作流见下方。 @@ -47,6 +49,7 @@ | `experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/` | RTX 6000D + vLLM | | `experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/` | RTX 6000D + SGLang | | `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32/` | RTX 6000D + SGLang,Kimi-K3(TP32×EP32,部署手册见 docs/KIMI_K3_DEPLOY.md) | +| `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/` | RTX 6000D + SGLang,Kimi-K3 真实 Prefill 的 Marlin/FlashInfer MXFP4 与 Chunk 8K/16K 对照 | | `experiments/pro6000/kimi3_pro6000_pd_rdma/` | RTX 6000D + SGLang,Kimi-K3 PD 分离(MoonCake RDMA,8 节点,见 deploy_pd.sh + docs/KIMI_K3_DEPLOY.md 附录 B) | 每个目录内:`run_bench.sh` 跑固定并发矩阵;`run_adaptive_concurrency.sh` 从 C=1 指数倍增搜饱和点;`run_adaptive_concurrency_add16.sh` 从 C=16 线性 +16 步进、带 TTFT SLO 停止与回退(当前主力用法,见 `experiments/ADAPTIVE_CONCURRENCY_USAGE.md`)。 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/Dockerfile b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/Dockerfile new file mode 100644 index 0000000..17fbeb2 --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/Dockerfile @@ -0,0 +1,34 @@ +ARG BASE_IMAGE=lmsysorg/sglang:kimi-k3-fiv617situ-warm +FROM ${BASE_IMAGE} + +ENV FLASHINFER_DISABLE_JIT=1 \ + FLASHINFER_DISABLE_VERSION_CHECK=1 \ + PYTHONUNBUFFERED=1 + +RUN rm -rf \ + /usr/local/lib/python3.12/dist-packages/flashinfer \ + /usr/local/lib/python3.12/dist-packages/flashinfer_python-*.dist-info + +COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/flashinfer_python.whl +RUN python3 -m pip install --no-deps --force-reinstall /tmp/flashinfer_python.whl && \ + rm -f /tmp/flashinfer_python.whl + +# The wheel contains Python/C++ sources. These three libraries are the exact +# SM120 AOT artifacts that passed the Phase 3/4 correctness gate. +COPY aot/ /usr/local/lib/python3.12/dist-packages/flashinfer/data/aot/ + +# Use the complete tested SGLang Python package, not a pair of loose runtime +# mounts. This preserves API compatibility between the two changed modules and +# their surrounding Kimi/MoE integration. +COPY sglang/ /sgl-workspace/sglang/python/sglang/ + +# Existing Kimi image compatibility fix: tcgen05 attention-residual is valid on +# SM100/SM103, not consumer Blackwell SM120. Bake the fallback into the image. +RUN python3 -c "from pathlib import Path; p=Path('/sgl-workspace/sglang/python/sglang/srt/layers/attn_residual.py'); s=p.read_text(); old=' _FAST_SUPPORTED = major >= 10'; new=' _FAST_SUPPORTED = major in (10, 11)'; assert old in s; p.write_text(s.replace(old,new))" && \ + python3 -m compileall -q /sgl-workspace/sglang/python/sglang && \ + python3 -c "from flashinfer.fused_moe import core; assert core.CUTLASS_FUSED_MOE_SUPPORTS_SITU is True" + +LABEL ai.meta-stone.purpose="Kimi-K3 SM120 FlashInfer MXFP4 Phase 5 acceptance" \ + ai.meta-stone.sglang.patch="fb929bb" \ + ai.meta-stone.flashinfer.patch="7d1a4ee" + diff --git a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md new file mode 100644 index 0000000..cd5e012 --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md @@ -0,0 +1,70 @@ +# Kimi-K3 / Pro6000D / SGLang TP32 EP32 MoE Backend Prefill + +This experiment is the real-serving acceptance gate for Kimi-K3 FlashInfer +MXFP4 MoE on NVIDIA SM120. It does not modify the existing baseline experiment. + +## Fixed matrix + +| Dimension | Values | +|---|---| +| Nodes | 601-604, 8 GPUs per node | +| Parallelism | TP32, EP32, DP1 | +| MoE runner | `marlin`, `flashinfer_mxfp4` | +| Input / output | 16K / 1 token | +| Concurrency | 8, 16 | +| Chunked prefill | 8K, 16K | +| Requests | 40 per repeat | +| Repeats | 3 | + +The eight configurations are identical except for MoE runner and chunk size. +Radix cache is disabled. `OSL=1` makes TTFT, E2E latency and input TPS the +primary metrics; TPOT is not meaningful for this matrix. + +## Unique entrypoint + +Run on 601 only. The sudo password is supplied at runtime and is never written +to Git or result files. + +```bash +cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill + +read -rsp 'sudo password: ' SUDO_PASSWORD; echo +export SUDO_PASSWORD +export RUN_ID=kimi3-moe-prefill-$(date +%Y%m%d-%H%M%S) + +bash run_moe_backend_prefill.sh build-image +bash run_moe_backend_prefill.sh run +``` + +For an unattended run, use the same entrypoint under tmux: + +```bash +umask 077 +printf '%s\n' "$SUDO_PASSWORD" > /tmp/kimi3-phase5-sudo +unset SUDO_PASSWORD + +tmux new-session -d -s kimi3-moe-prefill \ + "cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill && \ + SUDO_PASSWORD_FILE=/tmp/kimi3-phase5-sudo RUN_ID=\"${RUN_ID}\" \ + bash run_moe_backend_prefill.sh all 2>&1 | tee /data/hzy/${RUN_ID}.log" +``` + +Delete `/tmp/kimi3-phase5-sudo` after the run. The file is mode `0600` and is +never copied to result directories. + +## Outputs + +```text +results// + build/ per-node image build and smoke evidence + service/ exact docker commands and all four server logs + raw/ one benchmark JSONL per repeat + bench/ human-readable benchmark logs + gpu/ per-node GPU snapshots + results.csv flat repeat-level metrics + summary.json median metrics and FlashInfer-vs-Marlin deltas + orchestrator.log lifecycle and progress +``` + +The run aborts on a failed service start, any failed request, a missing result, +or an image capability smoke failure. Containers are removed on exit. diff --git a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/config.env b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/config.env new file mode 100644 index 0000000..7f067cf --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/config.env @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Kimi-K3 real-serving MoE backend acceptance on 601-604. + +EXPERIMENT="kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill" +MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}" +SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-kimi-k3}" +PORT="${PORT:-30000}" + +HEAD_HOST="${HEAD_HOST:-174.1.60.1}" +NODE_SSH_USER="${NODE_SSH_USER:-user}" +NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4) +DIST_PORT="${DIST_PORT:-20000}" + +BASE_IMAGE="${BASE_IMAGE:-lmsysorg/sglang:kimi-k3-fiv617situ-warm}" +CANDIDATE_IMAGE="${CANDIDATE_IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}" +FLASHINFER_TARGET="${FLASHINFER_TARGET:-/tmp/hzy-fi-situ-phase3-v3}" +FLASHINFER_WHEEL="${FLASHINFER_WHEEL:-/data/hzy/artifacts/flashinfer-kimi-sm120-situ-phase3-v3/flashinfer_python-0.6.18-py3-none-any.whl}" +SGLANG_SOURCE="${SGLANG_SOURCE:-/data/hzy/src/sglang-kimi-sm120/python/sglang}" + +DATASET_PATH="${DATASET_PATH:-${ROOT_DIR}/dataset/ShareGPT_V3_unfiltered_cleaned_split.json}" +BENCH_MODULE="${BENCH_MODULE:-sglang.benchmark.serving}" +RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/results}" + +INPUT_LEN="${INPUT_LEN:-16384}" +OUTPUT_LEN="${OUTPUT_LEN:-1}" +NUM_PROMPTS="${NUM_PROMPTS:-40}" +REPEATS="${REPEATS:-3}" +WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}" +HEALTH_WAIT_S="${HEALTH_WAIT_S:-2400}" + +BACKENDS=(marlin flashinfer_mxfp4) +CHUNKED_PREFILL_SIZES=(8192 16384) +CONCURRENCIES=(8 16) + diff --git a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/run_moe_backend_prefill.sh b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/run_moe_backend_prefill.sh new file mode 100755 index 0000000..e22792f --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/run_moe_backend_prefill.sh @@ -0,0 +1,429 @@ +#!/usr/bin/env bash +# Compare Marlin and FlashInfer MXFP4 on real four-node Kimi-K3 prefill. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" + +ACTION="${1:-all}" +RUN_ID="${RUN_ID:-kimi3-moe-prefill-$(date '+%Y%m%d-%H%M%S')}" +RESULT_ROOT="${RESULT_BASE}/${RUN_ID}" +BUILD_CONTEXT="/tmp/${EXPERIMENT}_image_context" +CONTAINER_PREFIX="${EXPERIMENT}" +SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10) + +mkdir -p "${RESULT_ROOT}"/{build,service,raw,bench,gpu} +log_init "${RESULT_ROOT}/orchestrator.log" + +usage() { + cat <<'EOF' +Usage: SUDO_PASSWORD=... bash run_moe_backend_prefill.sh ACTION + +Actions: + build-image Build the same candidate image on 601-604. + run Run the fixed 8-case matrix with an existing image. + all Build the image, run the matrix, and summarize it (default). + summarize Rebuild CSV summaries from an existing RUN_ID. + stop Remove this experiment's containers from all four nodes. +EOF +} + +require_password() { + if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then + [[ -r "$SUDO_PASSWORD_FILE" ]] || { + echo "ERROR: cannot read SUDO_PASSWORD_FILE=$SUDO_PASSWORD_FILE" >&2 + exit 2 + } + IFS= read -r SUDO_PASSWORD <"$SUDO_PASSWORD_FILE" + fi + if [[ -z "${SUDO_PASSWORD:-}" ]]; then + echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE at runtime" >&2 + exit 2 + fi +} + +is_head() { + [[ "$1" == "$HEAD_HOST" ]] +} + +run_host() { + local host="$1" + shift + if is_head "$host"; then + "$@" + else + ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" "$@" + fi +} + +sudo_host() { + local host="$1" + shift + require_password + if is_head "$host"; then + printf '%s\n' "$SUDO_PASSWORD" | sudo -S -p '' -- "$@" + return + fi + local remote_cmd + printf -v remote_cmd '%q ' "$@" + printf '%s\n' "$SUDO_PASSWORD" | \ + ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \ + "sudo -S -p '' -- ${remote_cmd}" +} + +check_inputs() { + local path + for path in "$MODEL_PATH" "$DATASET_PATH" "$FLASHINFER_TARGET/flashinfer" \ + "$FLASHINFER_WHEEL" "$SGLANG_SOURCE"; do + [[ -e "$path" ]] || { echo "ERROR: required path missing: $path" >&2; exit 2; } + done + [[ -d "$FLASHINFER_TARGET/flashinfer/data/aot/fused_moe_120" ]] || { + echo "ERROR: fused_moe_120 AOT artifact is missing" >&2 + exit 2 + } +} + +prepare_build_context() { + log "preparing candidate image context at ${BUILD_CONTEXT}" + rm -rf "$BUILD_CONTEXT" + mkdir -p "$BUILD_CONTEXT/aot" "$BUILD_CONTEXT/sglang" + cp "$SCRIPT_DIR/Dockerfile" "$BUILD_CONTEXT/Dockerfile" + cp "$FLASHINFER_WHEEL" "$BUILD_CONTEXT/flashinfer_python-0.6.18-py3-none-any.whl" + cp -a "$FLASHINFER_TARGET/flashinfer/data/aot/." "$BUILD_CONTEXT/aot/" + cp -a "$SGLANG_SOURCE/." "$BUILD_CONTEXT/sglang/" + find "$BUILD_CONTEXT" -type d -name __pycache__ -prune -exec rm -rf {} + + find "$BUILD_CONTEXT" -type f -name '*.pyc' -delete + du -sh "$BUILD_CONTEXT" | tee "$RESULT_ROOT/build/context_size.txt" +} + +stage_context() { + local host="$1" + if is_head "$host"; then + return + fi + log "staging image context on ${host}" + ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \ + "rm -rf '$BUILD_CONTEXT' && mkdir -p '$BUILD_CONTEXT'" + tar -C "$BUILD_CONTEXT" -cf - . | \ + ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \ + "tar -C '$BUILD_CONTEXT' -xf -" +} + +build_one_node() { + local host="$1" + local log_path="$RESULT_ROOT/build/${host}.log" + log "building ${CANDIDATE_IMAGE} on ${host}" + sudo_host "$host" docker build \ + --build-arg "BASE_IMAGE=${BASE_IMAGE}" \ + --tag "$CANDIDATE_IMAGE" "$BUILD_CONTEXT" >"$log_path" 2>&1 + sudo_host "$host" docker image inspect "$CANDIDATE_IMAGE" \ + --format '{{.Id}} {{.Size}}' >"$RESULT_ROOT/build/${host}.image.txt" +} + +smoke_image() { + local host="$1" + sudo_host "$host" docker run --rm --gpus device=0 \ + --entrypoint python3 "$CANDIDATE_IMAGE" -c \ + "import torch; from flashinfer.fused_moe import core; import sglang; assert torch.cuda.get_device_capability()[0] == 12; assert core.CUTLASS_FUSED_MOE_SUPPORTS_SITU is True; print(sglang.__file__, core.ActivationType.Situ)" \ + >"$RESULT_ROOT/build/${host}.smoke.txt" 2>&1 +} + +build_image_all_nodes() { + require_password + check_inputs + prepare_build_context + local host + for host in "${NODE_HOSTS[@]}"; do stage_context "$host"; done + local -a pids=() + for host in "${NODE_HOSTS[@]}"; do + build_one_node "$host" & + pids+=("$!") + done + local rc=0 pid + for pid in "${pids[@]}"; do wait "$pid" || rc=1; done + (( rc == 0 )) || { log "ERROR: candidate image build failed"; return 1; } + for host in "${NODE_HOSTS[@]}"; do smoke_image "$host"; done + log "candidate image built and smoke-tested on all four nodes" +} + +container_name() { + local rank="$1" + printf '%s_node%s' "$CONTAINER_PREFIX" "$rank" +} + +collect_gpu_snapshot() { + local label="$1" host rank=0 + for host in "${NODE_HOSTS[@]}"; do + sudo_host "$host" nvidia-smi \ + --query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \ + --format=csv,noheader,nounits \ + >"$RESULT_ROOT/gpu/${label}_node${rank}.csv" 2>&1 || true + rank=$((rank + 1)) + done +} + +collect_service_logs() { + local label="$1" host rank=0 name + for host in "${NODE_HOSTS[@]}"; do + name="$(container_name "$rank")" + sudo_host "$host" docker logs "$name" \ + >"$RESULT_ROOT/service/${label}_node${rank}.log" 2>&1 || true + rank=$((rank + 1)) + done +} + +stop_service() { + local host rank=0 name + for host in "${NODE_HOSTS[@]}"; do + name="$(container_name "$rank")" + sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true + rank=$((rank + 1)) + done +} + +start_node() { + local rank="$1" backend="$2" chunk="$3" + local host="${NODE_HOSTS[$rank]}" name bootstrap + name="$(container_name "$rank")" + bootstrap="export SGLANG_HOST_IP=174.1.60.$((rank + 1)); exec python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} --tp-size 32 --ep-size 32 --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend ${backend} --chunked-prefill-size ${chunk} --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}" + local -a cmd=( + docker run -d --name "$name" + --gpus all --network host --ipc=host --ulimit memlock=-1 + --device /dev/infiniband --shm-size 32g --entrypoint bash + -v "${MODEL_PATH}:${MODEL_PATH}:ro" + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + -e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0 + -e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3 + -e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7 + -e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + -e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1 + -e FLASHINFER_DISABLE_JIT=1 -e FLASHINFER_DISABLE_VERSION_CHECK=1 + "$CANDIDATE_IMAGE" -lc "$bootstrap" + ) + printf '%q ' "${cmd[@]}" >"$RESULT_ROOT/service/${backend}_chunk${chunk}_node${rank}.cmd.txt" + printf '\n' >>"$RESULT_ROOT/service/${backend}_chunk${chunk}_node${rank}.cmd.txt" + sudo_host "$host" "${cmd[@]}" >/dev/null +} + +wait_health() { + local backend="$1" chunk="$2" i + for ((i = 1; i <= HEALTH_WAIT_S; i++)); do + if curl --fail --silent --max-time 5 "http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then + log "service healthy backend=${backend} chunk=${chunk} wait_s=${i}" + return 0 + fi + if (( i % 30 == 0 )); then + log "waiting for service backend=${backend} chunk=${chunk} elapsed_s=${i}" + collect_service_logs "${backend}_chunk${chunk}_starting" + if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError' \ + "$RESULT_ROOT/service/${backend}_chunk${chunk}_starting_node"*.log; then + log "startup log contains an error; stopping early" + return 1 + fi + fi + sleep 1 + done + return 1 +} + +start_service() { + local backend="$1" chunk="$2" + stop_service + log "starting TP32/EP32 backend=${backend} chunk=${chunk} on 601-604" + start_node 1 "$backend" "$chunk" + start_node 2 "$backend" "$chunk" + start_node 3 "$backend" "$chunk" + sleep 5 + start_node 0 "$backend" "$chunk" + if ! wait_health "$backend" "$chunk"; then + collect_service_logs "${backend}_chunk${chunk}_startup_failed" + stop_service + return 1 + fi + collect_gpu_snapshot "${backend}_chunk${chunk}_healthy" + collect_service_logs "${backend}_chunk${chunk}_healthy" +} + +run_one_bench() { + local backend="$1" chunk="$2" concurrency="$3" repeat="$4" + local stem="${backend}_chunk${chunk}_c${concurrency}_r${repeat}" + local output_file="${RESULT_ROOT}/raw/${stem}.jsonl" + local bench_log="${RESULT_ROOT}/bench/${stem}.log" + rm -f "$output_file" + log "bench backend=${backend} chunk=${chunk} C=${concurrency} repeat=${repeat}/${REPEATS}" + sudo_host "$HEAD_HOST" docker run --rm --network host \ + -v "${MODEL_PATH}:${MODEL_PATH}:ro" \ + -v "${DATASET_PATH}:${DATASET_PATH}:ro" \ + -v "${RESULT_ROOT}:${RESULT_ROOT}" \ + -e PYTHONUNBUFFERED=1 --entrypoint python3 "$CANDIDATE_IMAGE" \ + -m "$BENCH_MODULE" \ + --backend sglang --host "$HEAD_HOST" --port "$PORT" \ + --dataset-name random --dataset-path "$DATASET_PATH" \ + --random-input-len "$INPUT_LEN" --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio 1.0 --num-prompts "$NUM_PROMPTS" \ + --max-concurrency "$concurrency" --request-rate 10000 \ + --warmup-requests "$WARMUP_REQUESTS" --output-file "$output_file" \ + --output-details --disable-tqdm >"$bench_log" 2>&1 + python3 - "$output_file" "$NUM_PROMPTS" <<'PY' +import json, sys +path, expected = sys.argv[1], int(sys.argv[2]) +with open(path, encoding="utf-8") as f: + rows = [json.loads(line) for line in f if line.strip()] +assert len(rows) == 1, (path, len(rows)) +assert rows[0].get("completed") == expected, rows[0].get("completed") +assert not any(rows[0].get("errors", [])), "benchmark contains request errors" +PY +} + +run_matrix() { + require_password + check_inputs + local backend chunk concurrency repeat label + collect_gpu_snapshot before_all + for chunk in "${CHUNKED_PREFILL_SIZES[@]}"; do + for backend in "${BACKENDS[@]}"; do + label="${backend}_chunk${chunk}" + if ! start_service "$backend" "$chunk"; then + log "ERROR: service failed for ${label}; preserving evidence and aborting" + return 1 + fi + for concurrency in "${CONCURRENCIES[@]}"; do + for ((repeat = 1; repeat <= REPEATS; repeat++)); do + run_one_bench "$backend" "$chunk" "$concurrency" "$repeat" + done + done + collect_service_logs "${label}_completed" + collect_gpu_snapshot "${label}_completed" + stop_service + sleep 5 + done + done + collect_gpu_snapshot after_all +} + +summarize() { + python3 - "$RESULT_ROOT" <<'PY' +import csv, json, re, statistics, sys +from pathlib import Path + +root = Path(sys.argv[1]) +pattern = re.compile(r"(.+)_chunk(\d+)_c(\d+)_r(\d+)\.jsonl$") +fields = [ + "backend", "chunked_prefill_size", "concurrency", "repeat", "completed", + "duration", "request_throughput", "input_throughput", "output_throughput", + "total_throughput", "mean_ttft_ms", "median_ttft_ms", "p95_ttft_ms", + "p99_ttft_ms", "mean_e2e_latency_ms", "median_e2e_latency_ms", + "p95_e2e_latency_ms", "p99_e2e_latency_ms", +] +rows = [] +for path in sorted((root / "raw").glob("*.jsonl")): + match = pattern.match(path.name) + if not match: + continue + with path.open(encoding="utf-8") as f: + data = next(json.loads(line) for line in f if line.strip()) + backend, chunk, concurrency, repeat = match.groups() + row = {key: data.get(key) for key in fields} + row.update( + backend=backend, + chunked_prefill_size=int(chunk), + concurrency=int(concurrency), + repeat=int(repeat), + ) + rows.append(row) + +with (root / "results.csv").open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + +metric_names = [ + "request_throughput", "input_throughput", "total_throughput", + "mean_ttft_ms", "median_ttft_ms", "p95_ttft_ms", "p99_ttft_ms", + "mean_e2e_latency_ms", "median_e2e_latency_ms", "p95_e2e_latency_ms", +] +groups = {} +for row in rows: + key = (row["backend"], row["chunked_prefill_size"], row["concurrency"]) + groups.setdefault(key, []).append(row) + +summary = [] +for (backend, chunk, concurrency), group in sorted(groups.items()): + item = { + "backend": backend, + "chunked_prefill_size": chunk, + "concurrency": concurrency, + "repeats": len(group), + "completed_each": [row["completed"] for row in group], + } + for metric in metric_names: + values = [float(row[metric]) for row in group if row.get(metric) is not None] + item[f"median_{metric}"] = statistics.median(values) if values else None + summary.append(item) + +comparisons = [] +index = {(x["backend"], x["chunked_prefill_size"], x["concurrency"]): x for x in summary} +for chunk in sorted({x["chunked_prefill_size"] for x in summary}): + for concurrency in sorted({x["concurrency"] for x in summary}): + base = index.get(("marlin", chunk, concurrency)) + cand = index.get(("flashinfer_mxfp4", chunk, concurrency)) + if not base or not cand: + continue + comp = {"chunked_prefill_size": chunk, "concurrency": concurrency} + for metric in ("input_throughput", "total_throughput"): + b = base[f"median_{metric}"] + c = cand[f"median_{metric}"] + comp[f"flashinfer_{metric}_change_pct"] = (c / b - 1.0) * 100.0 + for metric in ("median_ttft_ms", "p95_ttft_ms", "median_e2e_latency_ms"): + b = base[f"median_{metric}"] + c = cand[f"median_{metric}"] + comp[f"flashinfer_{metric}_change_pct"] = (c / b - 1.0) * 100.0 + comparisons.append(comp) + +payload = {"run_id": root.name, "summary": summary, "comparisons": comparisons} +(root / "summary.json").write_text(json.dumps(payload, indent=2), encoding="utf-8") +print(json.dumps(payload, indent=2)) +PY +} + +cleanup() { + collect_service_logs cleanup 2>/dev/null || true + stop_service 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +case "$ACTION" in + build-image) + build_image_all_nodes + ;; + run) + run_matrix + summarize + ;; + all) + build_image_all_nodes + run_matrix + summarize + ;; + summarize) + summarize + ;; + stop) + stop_service + ;; + -h|--help|help) + usage + ;; + *) + usage >&2 + exit 2 + ;; +esac + +trap - EXIT INT TERM +log "action=${ACTION} complete run_id=${RUN_ID} result_root=${RESULT_ROOT}"