From d7381abe8462e80c256a51261e248c883acff71e Mon Sep 17 00:00:00 2001 From: Zhiyi Hong <2497491955@qq.com> Date: Fri, 21 Aug 2026 14:19:45 +0800 Subject: [PATCH] [Test] Add Kimi-K3 Prefill PP baseline search --- .../README.md | 149 ++++++++ .../run_pp_baseline_search.sh | 351 ++++++++++++++++++ 2 files changed, 500 insertions(+) create mode 100644 experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/README.md create mode 100755 experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/run_pp_baseline_search.sh diff --git a/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/README.md b/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/README.md new file mode 100644 index 0000000..f784c57 --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/README.md @@ -0,0 +1,149 @@ +# Kimi-K3 Prefill PP baseline search + +## Goal + +Select the pipeline-parallel configuration for the next MoE A2A experiment on +four RTX 6000D nodes. The comparison fixes every workload and backend variable +except PP/TP: + +- SGLang image: `local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5` +- Model: Kimi-K3 +- GPUs: 32 GPUs on 601-604 +- MoE: FlashInfer MXFP4, EP4, A2A disabled +- Workload: 16K input, one output token, C=8/16, 40 requests +- Chunked prefill: 8K +- Repeats: three per point + +`run_pp_baseline_search.sh` is the only entry point. It uses the upstream +SGLang PP implementation and does not apply a runtime PP patch. + +## Result + +Values are the medians of three runs. PP1 is the directly comparable existing +run; PP2/4/8 come from the new search. + +| PP / TP / EP | C | Input TPS | TTFT p50 | TTFT p95 | Input TPS vs PP1 | TTFT p50 vs PP1 | +|---|---:|---:|---:|---:|---:|---:| +| 1 / 32 / 4 | 8 | 3,257.96 | 39.19 s | 41.62 s | baseline | baseline | +| 2 / 16 / 4 | 8 | 5,015.95 | 25.22 s | 26.69 s | +53.96% | -35.65% | +| 4 / 8 / 4 | 8 | 7,155.20 | 17.28 s | 18.62 s | +119.62% | -55.90% | +| 8 / 4 / 4 | 8 | **8,809.50** | **13.47 s** | **16.45 s** | **+170.40%** | **-65.64%** | +| 1 / 32 / 4 | 16 | 3,260.14 | 78.39 s | 80.88 s | baseline | baseline | +| 2 / 16 / 4 | 16 | 5,018.38 | 50.70 s | 52.14 s | +53.93% | -35.32% | +| 4 / 8 / 4 | 16 | 7,152.58 | 34.84 s | 35.95 s | +119.39% | -55.56% | +| 8 / 4 / 4 | 16 | **8,824.12** | **27.31 s** | **29.77 s** | **+170.67%** | **-65.16%** | + +All 18 new samples completed 40/40 requests. No OOM, traceback, NCCL failure, +request error, or engine failure was found. The three repeats are stable: PP8 +Input TPS varies by less than 0.7% at both concurrency points. + +## Memory + +The table reports the highest per-GPU memory observed in the post-benchmark +snapshot across all four nodes. + +| PP / TP / EP | Peak GPU memory | Remaining from 85,651 MiB | +|---|---:|---:| +| 1 / 32 / 4 | 81,355 MiB | 4,296 MiB | +| 2 / 16 / 4 | 80,335 MiB | 5,316 MiB | +| 4 / 8 / 4 | 82,417 MiB | 3,234 MiB | +| 8 / 4 / 4 | 79,407 MiB | 6,244 MiB | + +PP stages do not have identical model/state allocations, so memory is less +balanced than PP1. PP8 nevertheless has the best worst-rank headroom in this +run, which matters because A2A backends allocate additional communication +buffers. + +## Why PP helps here + +The existing PP1 Nsight profile measured a median 57.27% of the Prefill window +inside exposed NCCL work, with about 557 AllReduce calls and effectively no +compute/communication overlap. Its TP32 collectives span all four nodes. + +Changing PP also changes the TP communication domain: + +- PP2/TP16: each TP group spans two nodes. +- PP4/TP8: each TP group fits on one node; only pipeline activation transfers + cross node boundaries. +- PP8/TP4: two pipeline stages fit per node and each TP/EP4 group is node-local. + +This replaces frequent cross-node TP AllReduce with a smaller number of +pipeline activation transfers. At C=8/16 there is enough request-level work to +fill the deeper pipeline, so the communication saving is larger than the +pipeline bubble cost. + +The current run verifies the rank topology, successful NCCL P2P communicators, +and absence of communication failures. It does not contain per-PP Nsight or HCA +bandwidth samples; the 57.27% NCCL attribution belongs specifically to the +directly comparable PP1 profile and is used to explain, not fabricate, the PP8 +timeline. + +## Why older PP tests could be slower + +The old configurations were not an apples-to-apples PP-only comparison. In +particular, an old PP4 profile set: + +```bash +--pp-max-micro-batch-size 1 +``` + +Current SGLang leaves this value unset and computes: + +```python +pp_max_micro_batch_size = max(max_running_requests // pp_size, 1) +``` + +The scheduler then limits admission with +`pp_max_micro_batch_size - running_bs`. Setting it to one therefore permits +only one running request per PP scheduler, serializing much of a C=8/16 load +and exposing pipeline bubbles. Older tests also mixed Marlin, 16K chunks, +EP1/EP8, PP16, older images, and a runtime PP-group patch. They cannot be used +to conclude that PP itself is slower. + +## Decision + +Use **PP8 / TP4 / EP4** as the primary configuration for the next A2A +compatibility and memory smoke test. Keep **PP4 / TP8 / EP4** as the fallback: +it is slower and has less memory headroom here, but provides a useful check if +an A2A backend imposes TP/PP constraints not exercised by the baseline. + +Do not carry forward `--pp-max-micro-batch-size 1`. Leave it unset unless a +separate controlled experiment establishes a reason to cap it. + +## Reproduce and evidence + +Run from 601: + +```bash +cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search +RUN_ID=kimi3-pp-$(date +%Y%m%d-%H%M%S) \ + SUDO_PASSWORD_FILE=/path/to/password-file \ + bash run_pp_baseline_search.sh run +``` + +New PP2/4/8 evidence: + +```text +results/kimi3-pp-upstream-20260821-120208/ + summary.csv + results.csv + metadata/ + bench/ + raw/ + gpu/ + service/ +``` + +Existing PP1 evidence: + +```text +../kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/ + results/kimi3-ep4-moe-full-20260818-151349/ +``` + +Existing PP1 communication profile: + +```text +../kimi3_pro6000_sglang_prefill_communication_profile/ + results/kimi3-prefill-comm-20260820-143749/nsys_analysis.json +``` diff --git a/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/run_pp_baseline_search.sh b/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/run_pp_baseline_search.sh new file mode 100755 index 0000000..0ed19ed --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/run_pp_baseline_search.sh @@ -0,0 +1,351 @@ +#!/usr/bin/env bash +# Search Kimi-K3 prefill pipeline-parallel baselines on 601-604. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +# shellcheck source=/dev/null +source "${ROOT_DIR}/scripts/common/lib.sh" + +ACTION="${1:-run}" +EXPERIMENT="kimi3_pro6000_sglang_pp_baseline_search" +RUN_ID="${RUN_ID:-kimi3-pp-baseline-$(date '+%Y%m%d-%H%M%S')}" +RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}" + +MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}" +SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-kimi-k3}" +DOCKER_IMAGE="${DOCKER_IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}" +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) +SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10) +DIST_PORT="${DIST_PORT:-20000}" + +EP_SIZE=4 +PP_SIZES=(2 4 8) +CONCURRENCIES=(8 16) +INPUT_LEN=16384 +OUTPUT_LEN=1 +CHUNKED_PREFILL_SIZE=8192 +NUM_PROMPTS=40 +REPEATS="${REPEATS:-3}" +WARMUP_REQUESTS=2 +HEALTH_WAIT_S="${HEALTH_WAIT_S:-2400}" +CONTAINER_PREFIX="${EXPERIMENT}" + +mkdir -p "${RESULT_ROOT}"/{service,raw,bench,gpu,metadata} +log_init "${RESULT_ROOT}/orchestrator.log" + +usage() { + printf 'Usage: SUDO_PASSWORD=... %s {dry-run|run|summarize|stop}\n' "$0" +} + +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 + [[ -n "${SUDO_PASSWORD:-}" ]] || { + echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE" >&2 + exit 2 + } +} + +is_head() { + [[ "$1" == "$HEAD_HOST" ]] +} + +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}" +} + +container_name() { + printf '%s_node%s' "$CONTAINER_PREFIX" "$1" +} + +tp_for_pp() { + local pp="$1" + (( 32 % pp == 0 )) || { echo "ERROR: PP=$pp does not divide 32" >&2; return 1; } + echo $((32 / pp)) +} + +preflight() { + require_password + [[ -d "$MODEL_PATH" ]] || { echo "ERROR: missing model $MODEL_PATH" >&2; exit 2; } + local host + for host in "${NODE_HOSTS[@]}"; do + sudo_host "$host" docker image inspect "$DOCKER_IMAGE" >/dev/null + sudo_host "$host" test -e /dev/infiniband/uverbs0 + done + { + echo "run_id=$RUN_ID" + echo "image=$DOCKER_IMAGE" + echo "model=$MODEL_PATH" + echo "pp_sizes=${PP_SIZES[*]}" + echo "tp_sizes=16 8 4" + echo "ep_size=$EP_SIZE" + echo "shape=${INPUT_LEN}->${OUTPUT_LEN}" + echo "chunked_prefill_size=$CHUNKED_PREFILL_SIZE" + echo "concurrencies=${CONCURRENCIES[*]}" + echo "num_prompts=$NUM_PROMPTS" + echo "repeats=$REPEATS" + } >"${RESULT_ROOT}/metadata/manifest.env" +} + +stop_service() { + local rank host name + for rank in 0 1 2 3; do + host="${NODE_HOSTS[$rank]}" + name="$(container_name "$rank")" + sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true + done +} + +collect_service_logs() { + local label="$1" rank host name + for rank in 0 1 2 3; do + host="${NODE_HOSTS[$rank]}" + name="$(container_name "$rank")" + sudo_host "$host" docker logs "$name" \ + >"${RESULT_ROOT}/service/${label}_node${rank}.log" 2>&1 || true + done +} + +collect_gpu_snapshot() { + local label="$1" rank host + for rank in 0 1 2 3; do + host="${NODE_HOSTS[$rank]}" + 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 + done +} + +start_node() { + local rank="$1" pp="$2" tp="$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 ${tp} --pp-size ${pp} --ep-size ${EP_SIZE} --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend flashinfer_mxfp4 --chunked-prefill-size ${CHUNKED_PREFILL_SIZE} --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 + "$DOCKER_IMAGE" -lc "$bootstrap" + ) + printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/service/pp${pp}_node${rank}.cmd.txt" + printf '\n' >>"${RESULT_ROOT}/service/pp${pp}_node${rank}.cmd.txt" + sudo_host "$host" "${cmd[@]}" >/dev/null +} + +wait_health() { + local pp="$1" 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 PP=${pp} wait_s=${i}" + return 0 + fi + if (( i % 30 == 0 )); then + log "waiting for PP=${pp} service elapsed_s=${i}" + collect_service_logs "pp${pp}_starting" + if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError|AssertionError' \ + "${RESULT_ROOT}/service/pp${pp}_starting_node"*.log; then + log "ERROR: PP=${pp} startup log contains a fatal error" + return 1 + fi + fi + sleep 1 + done + return 1 +} + +verify_server() { + local pp="$1" tp="$2" info="${RESULT_ROOT}/metadata/pp${pp}_server_info.json" + curl --fail --silent "http://${HEAD_HOST}:${PORT}/get_server_info" >"$info" + python3 - "$info" "$pp" "$tp" "$EP_SIZE" <<'PY' +import json, sys + +path, pp, tp, ep = sys.argv[1], *map(int, sys.argv[2:]) +info = json.load(open(path, encoding="utf-8")) +expected = {"pp_size": pp, "tp_size": tp, "ep_size": ep, + "moe_runner_backend": "flashinfer_mxfp4"} +bad = [f"{k}={info.get(k)!r}, expected={v!r}" for k, v in expected.items() + if info.get(k) != v] +if bad: + raise SystemExit("server configuration mismatch: " + "; ".join(bad)) +print(expected) +PY +} + +start_service() { + local pp="$1" tp + tp="$(tp_for_pp "$pp")" + stop_service + log "starting PP=${pp} TP=${tp} EP=${EP_SIZE} on 601-604" + start_node 1 "$pp" "$tp" + start_node 2 "$pp" "$tp" + start_node 3 "$pp" "$tp" + sleep 5 + start_node 0 "$pp" "$tp" + if ! wait_health "$pp"; then + collect_service_logs "pp${pp}_startup_failed" + stop_service + return 1 + fi + verify_server "$pp" "$tp" + collect_service_logs "pp${pp}_healthy" + collect_gpu_snapshot "pp${pp}_healthy" +} + +run_one_bench() { + local pp="$1" concurrency="$2" repeat="$3" + local stem="pp${pp}_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 PP=${pp} C=${concurrency} repeat=${repeat}/${REPEATS}" + sudo_host "$HEAD_HOST" docker run --rm --network host \ + -v "${MODEL_PATH}:${MODEL_PATH}:ro" \ + -v "${RESULT_ROOT}:${RESULT_ROOT}" \ + -e PYTHONUNBUFFERED=1 --entrypoint python3 "$DOCKER_IMAGE" \ + -m sglang.benchmark.serving \ + --backend sglang --host "$HEAD_HOST" --port "$PORT" \ + --tokenizer "$MODEL_PATH" --dataset-name random-ids \ + --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]) +rows = [json.loads(line) for line in open(path, encoding="utf-8") 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 +} + +summarize() { + python3 - "$RESULT_ROOT" <<'PY' +import csv, json, re, statistics, sys +from pathlib import Path + +root = Path(sys.argv[1]) +pattern = re.compile(r"pp(\d+)_c(\d+)_r(\d+)\.jsonl$") +metrics = [ + "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 + data = next(json.loads(line) for line in path.open(encoding="utf-8") if line.strip()) + pp, concurrency, repeat = map(int, match.groups()) + row = {"pp_size": pp, "tp_size": 32 // pp, "ep_size": 4, + "concurrency": concurrency, "repeat": repeat, + "completed": data.get("completed")} + row.update({name: data.get(name) for name in metrics}) + rows.append(row) + +fields = list(rows[0]) if rows else [] +with (root / "results.csv").open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fields) + if fields: + writer.writeheader() + writer.writerows(rows) + +groups = {} +for row in rows: + groups.setdefault((row["pp_size"], row["concurrency"]), []).append(row) +summary = [] +for (pp, concurrency), group in sorted(groups.items()): + item = {"pp_size": pp, "tp_size": 32 // pp, "ep_size": 4, + "concurrency": concurrency, "repeats": len(group)} + for name in metrics: + values = [float(row[name]) for row in group if row.get(name) is not None] + item[f"median_{name}"] = statistics.median(values) if values else None + summary.append(item) + +summary_fields = list(summary[0]) if summary else [] +with (root / "summary.csv").open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=summary_fields) + if summary_fields: + writer.writeheader() + writer.writerows(summary) +print(f"wrote {len(rows)} result rows and {len(summary)} summary rows") +PY +} + +run_matrix() { + preflight + collect_gpu_snapshot before_all + local pp concurrency repeat + for pp in "${PP_SIZES[@]}"; do + if ! start_service "$pp"; then + log "ERROR: PP=${pp} failed to start; preserving evidence and continuing" + continue + fi + for concurrency in "${CONCURRENCIES[@]}"; do + for ((repeat = 1; repeat <= REPEATS; repeat++)); do + run_one_bench "$pp" "$concurrency" "$repeat" + done + done + collect_service_logs "pp${pp}_completed" + collect_gpu_snapshot "pp${pp}_completed" + stop_service + sleep 5 + done + collect_gpu_snapshot after_all + summarize +} + +dry_run() { + local pp tp + for pp in "${PP_SIZES[@]}"; do + tp="$(tp_for_pp "$pp")" + echo "PP=${pp} TP=${tp} EP=${EP_SIZE} ${INPUT_LEN}->${OUTPUT_LEN} C=${CONCURRENCIES[*]} chunk=${CHUNKED_PREFILL_SIZE}" + done + bash -n "$0" +} + +case "$ACTION" in + dry-run) dry_run ;; + run) + trap stop_service EXIT INT TERM + run_matrix + ;; + summarize) summarize ;; + stop) + require_password + stop_service + ;; + *) usage; exit 2 ;; +esac