zhoulin
8ace0d790c
pytest plugin set + driver script for compiler A/B perf comparison: - reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay - per-shape ttgir dump of actually-used variants with readable naming - cudagraph-based timing with documented fallback semantics Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
109 lines
3.8 KiB
Bash
109 lines
3.8 KiB
Bash
#!/bin/bash
|
|
# Single-operator perf benchmark (manual debug).
|
|
# Edit OP/OP_FILE/INLINE_YAML below (or override via env: OP=... SHAPE_FILE=... ./run_pytest.sh).
|
|
# All artifacts go under runs/<op>_<timestamp>/: run.log, shapes.yaml,
|
|
# autotune_records/, ttgir/.
|
|
set -euo pipefail
|
|
|
|
OP="${OP:-fused_marlin_moe_mxfp4}"
|
|
OP_FILE="${OP_FILE:-fused_marlin_moe}" # benchmark file is test_<OP_FILE>.py
|
|
SHAPE_FILE="${SHAPE_FILE:-}"
|
|
|
|
# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name.
|
|
read -r -d '' INLINE_YAML <<'YAML' || true
|
|
fused_marlin_moe_mxfp4:
|
|
# 4 MoE models x 4 token counts (M=1,16,64,256) = 16 shapes
|
|
shapes:
|
|
# Mixtral (E=8)
|
|
- [1, 8, 4096, 14336, 2]
|
|
- [16, 8, 4096, 14336, 2]
|
|
- [64, 8, 4096, 14336, 2]
|
|
- [256, 8, 4096, 14336, 2]
|
|
# DeepSeek-V3 (E=256, H=7168)
|
|
- [1, 256, 7168, 2048, 8]
|
|
- [16, 256, 7168, 2048, 8]
|
|
- [64, 256, 7168, 2048, 8]
|
|
- [256, 256, 7168, 2048, 8]
|
|
# Qwen3 (E=512)
|
|
- [1, 512, 4096, 1024, 10]
|
|
- [16, 512, 4096, 1024, 10]
|
|
- [64, 512, 4096, 1024, 10]
|
|
- [256, 512, 4096, 1024, 10]
|
|
# DeepSeek-V4-Flash (E=256, H=4096)
|
|
- [1, 256, 4096, 2048, 6]
|
|
- [16, 256, 4096, 2048, 6]
|
|
- [64, 256, 4096, 2048, 6]
|
|
- [256, 256, 4096, 2048, 6]
|
|
shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
|
|
YAML
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
|
|
TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
|
|
|
|
OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
|
|
LOG_FILE="$OUT_DIR/run.log"
|
|
mkdir -p "$OUT_DIR"
|
|
|
|
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
|
PLUGINS=(
|
|
-p _seed_plugin
|
|
-p _shape_inject_plugin
|
|
-p _shape_iter_inject_plugin
|
|
-p _bespoke_shape_plugin
|
|
# -p _mm_cluster_fix_plugin
|
|
-p _autotune_record_plugin
|
|
-p _cudagraph_plugin
|
|
-p _ir_meta_plugin
|
|
)
|
|
# Autotune record/replay (see _autotune_record_plugin.py):
|
|
# - default: record chosen configs to $OUT_DIR/autotune_records/<op>.json
|
|
# - REPLAY_FROM=<previous runs/<op>_<ts> dir>: replay that run's recorded
|
|
# configs instead (for A/B runs where both sides must use the same config).
|
|
REPLAY_FROM="${REPLAY_FROM:-}"
|
|
if [[ -n "$REPLAY_FROM" ]]; then
|
|
AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records"
|
|
[[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \
|
|
echo ">>> warning: $REPLAY_FROM/autotune_records/$OP.json not found; replay will fall back to autotune" >&2
|
|
else
|
|
AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records"
|
|
mkdir -p "$OUT_DIR/autotune_records"
|
|
fi
|
|
export FLAGGEMS_PERF_CURRENT_OP="$OP"
|
|
|
|
# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file.
|
|
if [[ -z "$SHAPE_FILE" ]]; then
|
|
SHAPE_FILE="$(mktemp --suffix=.yaml)"
|
|
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
|
|
trap 'rm -f "$SHAPE_FILE"' EXIT
|
|
fi
|
|
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # archive the shape used
|
|
|
|
export PYTHONUNBUFFERED=1 # unbuffered live output
|
|
|
|
{
|
|
# _ir_meta_plugin dumps organized ttgir (only variants actually launched
|
|
# outside the autotune sweep) into TTGIR_DUMP_DIR at process exit, so the
|
|
# dump happens even if a CUDA crash kills the run. Layout/legend: see
|
|
# <dump>/naming.md and index.tsv (which also lists unused sweep losers).
|
|
CACHE_DIR="$OUT_DIR/.triton_cache"
|
|
rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR"
|
|
status=0
|
|
TRITON_CACHE_DIR="$CACHE_DIR" \
|
|
FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \
|
|
env "$AUTOTUNE_ENV" \
|
|
USE_FLAGTUNE=1 python -u -m pytest -s "$TEST_FILE" \
|
|
"${PLUGINS[@]}" \
|
|
--shape_file "$SHAPE_FILE" \
|
|
--level core --mode kernel || status=$?
|
|
|
|
rm -rf "$CACHE_DIR"
|
|
|
|
if (( status == 0 )); then
|
|
echo ">>> done. outputs in $OUT_DIR"
|
|
else
|
|
echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR"
|
|
fi
|
|
exit "$status"
|
|
} 2>&1 | tee "$LOG_FILE"
|