forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
c8ea586d0f
Everything about which accelerator a run targets was previously spelled CUDA: device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES, the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in run_pytest.sh with no way to drop the CUDA-only members. Running the same operator corpus on another vendor's device therefore meant editing the drivers. The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with the usual aliases): - backends/ owns per-backend device enumeration, visibility env keys and capability flags, so nothing above it names a vendor; - plugin_plan.py derives the -p list from those capabilities instead of a hardcoded array, which is what lets a non-CUDA run omit the graph-capture and cluster-fix plugins rather than watch them fail; - _timing_plugin selects the timer from the same capabilities: CUDA keeps graph capture with its per-shape fallback, other backends stay on upstream do_bench and record timer=direct, because capture-then-fallback on a device without graph support only produces noise; - _device_setup_plugin replaces the NVIDIA-only vendor guard; - providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to FlagGems' own core_shapes; - _run_identity.py writes run_identity.json next to each run, recording backend, device, timer, autotune mode and the resolved test node. Two runs whose calibers differ are otherwise indistinguishable after the fact. run_batch gains --backend and --devices (--gpus kept as an alias) and carries backend/device/timer into ops_status.csv. test_backends.py covers the registry, alias resolution, visibility env and plugin planning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
196 lines
7.7 KiB
Bash
196 lines
7.7 KiB
Bash
#!/bin/bash
|
||
# 单算子性能测试入口。改配置区或用同名环境变量覆盖：
|
||
# 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
|
||
|
||
# ============================== 配置区 ==============================
|
||
# 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。
|
||
|
||
# 被测算子，对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
|
||
# 也可直接传 FlagGems pytest 节点，例如：
|
||
# TEST_NODE=/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
|
||
OP="${OP:-fused_marlin_moe_mxfp4}"
|
||
OP_FILE="${OP_FILE:-fused_marlin_moe}"
|
||
|
||
# 上游已拆分为两个仓库：主仓 /workspace/dev/FlagGems（包名 flag_gems）与
|
||
# /workspace/dev/FlagGems-vllm（包名 flaggems_vllm）。测 vllm 仓库算子时覆盖本变量即可。
|
||
# NV 产品路径也常见 /workspace/FlagGems；TEST_NODE 或缺失测试文件时会回退到该路径。
|
||
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}"
|
||
TEST_NODE="${TEST_NODE:-${ZL_TEST_NODE:-}}"
|
||
|
||
# shape 来源：非空则用该 yaml；SHAPE_MODE=upstream 则不传 --shape_file，沿用 FlagGems
|
||
# `--level core` 的原生 core_shapes。设了 TEST_NODE 且未指定 SHAPE_FILE 时默认 upstream。
|
||
SHAPE_FILE="${SHAPE_FILE:-}"
|
||
SHAPE_MODE="${SHAPE_MODE:-}"
|
||
read -r -d '' INLINE_YAML <<'YAML' || true
|
||
fused_marlin_moe_mxfp4:
|
||
shapes:
|
||
- [1, 256, 4096, 256, 6]
|
||
- [2, 256, 4096, 256, 6]
|
||
- [4, 256, 4096, 256, 6]
|
||
- [8, 256, 4096, 256, 6]
|
||
- [16, 256, 4096, 256, 6]
|
||
- [32, 256, 4096, 256, 6]
|
||
- [64, 256, 4096, 256, 6]
|
||
- [128, 256, 4096, 256, 6]
|
||
- [256, 256, 4096, 256, 6]
|
||
- [512, 256, 4096, 256, 6]
|
||
- [1024, 256, 4096, 256, 6]
|
||
- [2048, 256, 4096, 256, 6]
|
||
- [4096, 256, 4096, 256, 6]
|
||
- [8192, 256, 4096, 256, 6]
|
||
- [16384, 256, 4096, 256, 6]
|
||
- [32768, 256, 4096, 256, 6]
|
||
shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
|
||
YAML
|
||
|
||
# 调优空间：0=普通 autotune（默认，快速验证）；1=FlagTune 扩展空间（首跑全量搜索、慢）
|
||
USE_FLAGTUNE="${USE_FLAGTUNE:-0}"
|
||
|
||
# 可选：限制 dtype 集（空格分隔，如 "bfloat16 float16"）。空=上游默认 dtype 扫描。
|
||
# 算子不支持指定 dtype 时上游会报 "can't be supported by this op"（批量驱动据此降级重试）。
|
||
DTYPES="${DTYPES:-}"
|
||
|
||
# 空=record 模式，把本次选中的 config 记入 autotune_records/<op>.json；
|
||
# 指向某次历史 run 目录则 replay 其记录，用于 A/B 两侧锁同一套 config。
|
||
REPLAY_FROM="${REPLAY_FROM:-}"
|
||
|
||
# always/never 强制开关终端颜色，空则按 tty 判断（run.log 始终去色）
|
||
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
|
||
|
||
# 设备后端：cuda / musa / rocm / ascend（别名见 backends/__init__.py）。
|
||
# 插件列表由 plugin_plan.py 按后端能力生成，不再在此硬编码 -p。
|
||
# 非 CUDA 默认走上游 do_bench，不会加载 _cudagraph_plugin。
|
||
ZL_BACKEND="${ZL_BACKEND:-cuda}"
|
||
|
||
# ============================== 执行逻辑 ==============================
|
||
|
||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
||
eval "$(python "$SCRIPT_DIR/providers/flaggems.py" --emit-shell \
|
||
--parse-node "$TEST_NODE" \
|
||
--op "$OP" \
|
||
--op-file "$OP_FILE" \
|
||
--flaggems-dir "$FLAGGEMS_DIR" \
|
||
--shape-file "$SHAPE_FILE" \
|
||
--shape-mode "$SHAPE_MODE")"
|
||
: "${TEST_FILE:=$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}}"
|
||
|
||
# 产物目录，可用 OUT_DIR 指定。时间戳只到秒，并发启动会撞名，故用不带 -p 的
|
||
# mkdir 抢占（已存在即失败），撞上就退让到 -2、-3……
|
||
if [[ -n "${OUT_DIR:-}" ]]; then
|
||
mkdir -p "$OUT_DIR"
|
||
else
|
||
mkdir -p "$SCRIPT_DIR/runs"
|
||
BASE="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
|
||
OUT_DIR="$BASE"
|
||
n=1
|
||
until mkdir "$OUT_DIR" 2>/dev/null; do
|
||
n=$((n + 1))
|
||
OUT_DIR="${BASE}-${n}"
|
||
(( n > 99 )) && { echo "!!! 无法创建产物目录（$BASE 及后缀均被占用）" >&2; exit 1; }
|
||
done
|
||
fi
|
||
LOG_FILE="$OUT_DIR/run.log"
|
||
|
||
# tee 管道会让 python 侧检测不到 tty，故在 shell 层判断后经环境变量下传
|
||
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 FLAGGEMS_PERF_CURRENT_OP="$OP"
|
||
export PYTHONUNBUFFERED=1 # 实时输出不缓冲
|
||
export ZL_BACKEND
|
||
|
||
# 按 ZL_BACKEND / ZL_ENABLE_* 生成 -p 列表（见 plugin_plan.py）
|
||
read -r -a PLUGINS <<< "$(python "$SCRIPT_DIR/plugin_plan.py" --backend "$ZL_BACKEND" --pytest-args)"
|
||
|
||
# DTYPES 非空时逐个转为上游 --dtypes 选项（action=append，每个 dtype 一次）
|
||
DTYPE_ARGS=()
|
||
for _dt in $DTYPES; do
|
||
DTYPE_ARGS+=(--dtypes "$_dt")
|
||
done
|
||
|
||
# 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 "${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
|
||
|
||
# 未指定 shape 文件时：inline 用 INLINE_YAML；upstream 不传 --shape_file。
|
||
SHAPE_ARGS=()
|
||
if [[ "${PASS_SHAPE_FILE:-1}" == "0" ]]; then
|
||
if [[ -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" ]]; then
|
||
cp -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" "$OUT_DIR/shapes.yaml"
|
||
else
|
||
printf '%s\n' "source: upstream" "level: core" > "$OUT_DIR/shapes.yaml"
|
||
fi
|
||
else
|
||
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" # 存档本次实际使用的 shape
|
||
SHAPE_ARGS=(--shape_file "$SHAPE_FILE")
|
||
fi
|
||
|
||
python "$SCRIPT_DIR/_run_identity.py" "$OUT_DIR" \
|
||
--backend "$ZL_BACKEND" \
|
||
--op "$OP" \
|
||
--op-file "$OP_FILE" \
|
||
--flaggems-dir "$FLAGGEMS_DIR" \
|
||
--autotune-mode "$MODE_DESC" \
|
||
--device "${ZL_DEVICE:-}" \
|
||
--test-file "$TEST_FILE" \
|
||
--shape-mode "$SHAPE_MODE"
|
||
|
||
status=0
|
||
{
|
||
echo "${C_BOLD}>>> op=$OP backend=$ZL_BACKEND mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE shape=$SHAPE_MODE${C_RESET}"
|
||
echo "${C_DIM}>>> node=$TEST_FILE${C_RESET}"
|
||
echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
|
||
|
||
# 每次用独立的 Triton 缓存目录，跑完即删：保证编译过程可复现，且 ttgir 落盘
|
||
# 只包含本次的变体（_ir_meta_plugin 在 atexit 里按 shape 整理）。
|
||
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=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
|
||
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
|
||
${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
|
||
--level core --mode kernel || status=$?
|
||
|
||
rm -rf "$CACHE_DIR"
|
||
|
||
if (( status == 0 )); then
|
||
echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}"
|
||
else
|
||
echo "${C_RED}>>> FAILED (pytest exit $status). partial outputs in $OUT_DIR${C_RESET}"
|
||
fi
|
||
exit "$status"
|
||
} 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"
|