[Test] Add Kimi-K3 Prefill PP baseline search
This commit is contained in:
Zhiyi Hong
parent
60f77cd4ef
commit
d7381abe84
@ -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
|
||||||
|
```
|
||||||
@ -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
|
||||||
Loading…
x
Reference in New Issue
Block a user