80 lines
2.5 KiB
Bash
Executable File
80 lines
2.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -Eeuo pipefail
|
|
|
|
ROOT=${ROOT:-/data/wxy}
|
|
RUN_ID=${RUN_ID:-h3-sdpa-kernel-matrix-$(date +%Y%m%d-%H%M%S)}
|
|
OUT=${OUT:-$ROOT/profile_results/$RUN_ID}
|
|
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
|
|
BENCH=${BENCH:-$ROOT/h3_profile/h3_sdpa_kernel_bench.py}
|
|
|
|
mkdir -p "$OUT"/{jobs,system,metadata}
|
|
cp "$0" "$BENCH" "$OUT/metadata/"
|
|
date -Ins > "$OUT/metadata/date.txt"
|
|
nvidia-smi -q > "$OUT/metadata/nvidia-smi-q.txt"
|
|
nvidia-smi topo -m > "$OUT/metadata/topology.txt"
|
|
"$PYTHON" -m pip freeze > "$OUT/metadata/pip-freeze.txt"
|
|
|
|
declare -a PIDS=()
|
|
cleanup(){
|
|
local rc=$? p
|
|
trap - EXIT INT TERM
|
|
for p in "${PIDS[@]:-}"; do kill -TERM "$p" 2>/dev/null || true; done
|
|
exit "$rc"
|
|
}
|
|
trap cleanup EXIT INT TERM
|
|
|
|
nvidia-smi dmon -s pucvmet -d 1 -o DT > "$OUT/system/nvidia-dmon.log" 2>&1 &
|
|
DMON_PID=$!
|
|
|
|
backends=(flash cudnn efficient auto)
|
|
for gpu in 0 1 2 3; do
|
|
backend=${backends[$gpu]}
|
|
CUDA_VISIBLE_DEVICES=$gpu PYTHONUNBUFFERED=1 "$PYTHON" "$BENCH" \
|
|
--scenario F3 --backend "$backend" --sequence 41778 \
|
|
--output "$OUT/jobs/gpu${gpu}-F3-${backend}.json" \
|
|
> "$OUT/jobs/gpu${gpu}-F3-${backend}.log" 2>&1 &
|
|
PIDS+=("$!")
|
|
done
|
|
for gpu in 4 5 6 7; do
|
|
backend=${backends[$((gpu-4))]}
|
|
CUDA_VISIBLE_DEVICES=$gpu PYTHONUNBUFFERED=1 "$PYTHON" "$BENCH" \
|
|
--scenario RVA --backend "$backend" --sequence 81547 \
|
|
--output "$OUT/jobs/gpu${gpu}-RVA-${backend}.json" \
|
|
> "$OUT/jobs/gpu${gpu}-RVA-${backend}.log" 2>&1 &
|
|
PIDS+=("$!")
|
|
done
|
|
|
|
failed=0
|
|
for p in "${PIDS[@]}"; do wait "$p" || failed=1; done
|
|
PIDS=()
|
|
kill -TERM "$DMON_PID" 2>/dev/null || true
|
|
wait "$DMON_PID" 2>/dev/null || true
|
|
|
|
"$PYTHON" - "$OUT" <<'PY'
|
|
import json, pathlib, sys
|
|
root = pathlib.Path(sys.argv[1])
|
|
rows = []
|
|
for path in sorted((root / "jobs").glob("*.json")):
|
|
payload = json.loads(path.read_text())
|
|
for result in payload.get("results", []):
|
|
rows.append({
|
|
"scenario": payload["scenario"],
|
|
"backend": payload["backend"],
|
|
"layout": result["layout"],
|
|
"median_ms": result["warm"]["median_ms"],
|
|
"cold_median_ms": result["cold"]["median_ms"],
|
|
"cold_over_warm": result["cold_over_warm"],
|
|
"effective_tflops": result["effective_tflops"],
|
|
"cuda_graph_ms": result["cuda_graph"].get("median_ms"),
|
|
"q_stride": result["q_stride"],
|
|
})
|
|
(root / "summary.json").write_text(json.dumps(rows, indent=2) + "\n")
|
|
print(json.dumps(rows, indent=2))
|
|
PY
|
|
|
|
if ((failed)); then
|
|
touch "$OUT/FAILED"
|
|
exit 1
|
|
fi
|
|
touch "$OUT/DONE"
|