Stabilize first-run cudagraph timing and colorize terminal output
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture (internal 5-iter warmup is too short), tag fallback markers with the failing phase; first-run latency no longer jitters run-to-run. - _device_guard_plugin (new): set GEMS_VENDOR via torch probe before importing flag_gems, avoiding its timeout-less nvidia-smi subprocess probe that can hang import in fork-broken environments. - _pretty_report_plugin (new) + _term_style (new): fold inputs identical across all result rows into a legend line, color status/plugin tags/markers on the live terminal; run.log is ANSI-stripped and keeps upstream SUCCESS/column wording for grep compatibility. - run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a config section with Chinese comments, add start/end banners. - README: document the warmup semantics, new plugins/env vars, and the A/B rule that both sides must use the same USE_FLAGTUNE.
This commit is contained in:
+82
-36
@@ -1,18 +1,27 @@
|
||||
#!/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/.
|
||||
# 单算子性能测试入口（手动调试用）。
|
||||
# 用法：改下方【配置区】，或用环境变量覆盖，例如：
|
||||
# OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
|
||||
# 产物统一落在 runs/<op>_<时间戳>/：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:-}"
|
||||
# ============================== 配置区 ==============================
|
||||
# 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。
|
||||
|
||||
# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name.
|
||||
# 被测算子：OP=测试函数名（去掉 test_ 前缀）；OP_FILE=benchmark 文件名
|
||||
# （对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>）
|
||||
OP="${OP:-fused_marlin_moe_mxfp4}"
|
||||
OP_FILE="${OP_FILE:-fused_marlin_moe}"
|
||||
|
||||
# FlagGems 仓库路径
|
||||
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
|
||||
|
||||
# shape 来源：SHAPE_FILE 非空则用该 yaml；为空则用下方内置 yaml。
|
||||
# yaml 顶层 key 必须是 op 名。
|
||||
SHAPE_FILE="${SHAPE_FILE:-}"
|
||||
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
|
||||
# 4 个 MoE 模型 x 4 档 token 数（M=1,16,64,256）= 16 组 shape
|
||||
shapes:
|
||||
# Mixtral (E=8)
|
||||
- [1, 8, 4096, 14336, 2]
|
||||
@@ -37,16 +46,20 @@ fused_marlin_moe_mxfp4:
|
||||
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}"
|
||||
# 调优空间：1=FlagTune 扩展空间（首跑全量搜索、慢）；0=普通 autotune（快速验证）
|
||||
USE_FLAGTUNE="${USE_FLAGTUNE:-1}"
|
||||
|
||||
OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
|
||||
LOG_FILE="$OUT_DIR/run.log"
|
||||
mkdir -p "$OUT_DIR"
|
||||
# autotune record/replay（见 _autotune_record_plugin.py）：
|
||||
# - 空：record 模式，本次选中的 config 记录到 $OUT_DIR/autotune_records/<op>.json；
|
||||
# - 指向某次历史 runs/<op>_<时间戳> 目录：replay 该次记录（A/B 两侧同 config）。
|
||||
REPLAY_FROM="${REPLAY_FROM:-}"
|
||||
|
||||
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
||||
# 终端颜色：always/never 强制开/关；为空则按 tty 自动判断（run.log 始终去色）
|
||||
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
|
||||
|
||||
# pytest 插件列表（可按需注释停用；各插件作用见 README「插件说明」）
|
||||
PLUGINS=(
|
||||
-p _device_guard_plugin
|
||||
-p _seed_plugin
|
||||
-p _shape_inject_plugin
|
||||
-p _shape_iter_inject_plugin
|
||||
@@ -54,55 +67,88 @@ PLUGINS=(
|
||||
# -p _mm_cluster_fix_plugin
|
||||
-p _autotune_record_plugin
|
||||
-p _cudagraph_plugin
|
||||
-p _pretty_report_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:-}"
|
||||
|
||||
# ============================== 执行逻辑 ==============================
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
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"
|
||||
|
||||
# 颜色决策：下方 tee 管道会让 python 侧看不到 tty，所以在这里判断，
|
||||
# 并通过 FLAGGEMS_PERF_COLOR 下传给各插件（_term_style.py 消费）。
|
||||
if [[ -z "$FLAGGEMS_PERF_COLOR" && -t 1 && -z "${NO_COLOR:-}" ]]; then
|
||||
FLAGGEMS_PERF_COLOR=always
|
||||
fi
|
||||
export FLAGGEMS_PERF_COLOR
|
||||
if [[ "$FLAGGEMS_PERF_COLOR" == "always" ]]; then
|
||||
C_RED=$'\033[31m'; C_GREEN=$'\033[32m'; C_YELLOW=$'\033[33m'
|
||||
C_DIM=$'\033[2m'; C_BOLD=$'\033[1m'; C_RESET=$'\033[0m'
|
||||
PYTEST_COLOR=(--color=yes)
|
||||
else
|
||||
C_RED='' C_GREEN='' C_YELLOW='' C_DIM='' C_BOLD='' C_RESET=''
|
||||
PYTEST_COLOR=()
|
||||
fi
|
||||
|
||||
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
||||
export FLAGGEMS_PERF_CURRENT_OP="$OP"
|
||||
export PYTHONUNBUFFERED=1 # 实时输出不缓冲
|
||||
|
||||
# record/replay 二选一，通过环境变量激活对应模式
|
||||
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
|
||||
echo "${C_YELLOW}>>> warning: $REPLAY_FROM/autotune_records/$OP.json 不存在；replay 将回退为现场 autotune${C_RESET}" >&2
|
||||
MODE_DESC="replay($REPLAY_FROM)"
|
||||
else
|
||||
AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records"
|
||||
mkdir -p "$OUT_DIR/autotune_records"
|
||||
MODE_DESC=record
|
||||
fi
|
||||
export FLAGGEMS_PERF_CURRENT_OP="$OP"
|
||||
|
||||
# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file.
|
||||
# 解析 shape 文件：SHAPE_FILE 为空时把内置 yaml 写到临时文件
|
||||
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
|
||||
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
|
||||
|
||||
status=0
|
||||
{
|
||||
# _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).
|
||||
echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
|
||||
echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
|
||||
|
||||
# _ir_meta_plugin 在进程退出时把"实际被使用"的变体的 ttgir 按 shape 整理落盘
|
||||
# （挂在 atexit 上，CUDA crash 后已编译部分仍可拿到）。
|
||||
# 目录结构与命名图例见 <dump>/naming.md 和 index.tsv（后者也记录落选的 sweep 变体）。
|
||||
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[@]}" \
|
||||
USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
|
||||
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
|
||||
--shape_file "$SHAPE_FILE" \
|
||||
--level core --mode kernel || status=$?
|
||||
|
||||
rm -rf "$CACHE_DIR"
|
||||
|
||||
if (( status == 0 )); then
|
||||
echo ">>> done. outputs in $OUT_DIR"
|
||||
echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}"
|
||||
else
|
||||
echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR"
|
||||
echo "${C_RED}>>> FAILED (pytest exit $status). partial outputs in $OUT_DIR${C_RESET}"
|
||||
fi
|
||||
exit "$status"
|
||||
} 2>&1 | tee "$LOG_FILE"
|
||||
} 2>&1 | tee "$LOG_FILE" || status=$?
|
||||
|
||||
# 终端保留颜色；落盘的 run.log 去掉 ANSI 转义，保证 grep/diff 面对纯文本
|
||||
# （Ctrl-C 中断时会跳过去色，仅影响观感）。
|
||||
sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG_FILE"
|
||||
exit "$status"
|
||||
|
||||
Reference in New Issue
Block a user