Add FlagGems single-op perf benchmark harness

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>
2026-07-16 03:52:07 +00:00
parent 242d74af0e
commit 8ace0d790c
11 changed files with 1562 additions and 1 deletions
@@ -0,0 +1,108 @@
+#!/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"