[Test] Add Kimi-K3 Prefill PP baseline search

2026-08-21 14:19:45 +08:00 · 2026-08-21 14:19:45 +08:00 · d7381abe84
commit d7381abe84
parent 60f77cd4ef
2 changed files with 500 additions and 0 deletions
--- a/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search/README.md
+++ 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
+```
--- 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
@ -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