forked from zhoulin/zl_bench
feat: make the harness backend-agnostic
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>
This commit is contained in:
AharrypotterandClaude Opus 5 committed
1 parent
35a928a516
commit
c8ea586d0f
24 files changed
+1448
-159
No files matched your search
+53
-25
@@ -8,15 +8,21 @@ 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，否则用 INLINE_YAML。顶层 key 必须是 op 名。
|
||||
# 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:
|
||||
@@ -53,25 +59,23 @@ REPLAY_FROM="${REPLAY_FROM:-}"
|
||||
# always/never 强制开关终端颜色，空则按 tty 判断（run.log 始终去色）
|
||||
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
|
||||
|
||||
# 各插件作用见 README「插件说明」；注释掉某行即停用该插件
|
||||
PLUGINS=(
|
||||
-p _device_guard_plugin
|
||||
-p _parallel_warmup_plugin
|
||||
-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 _pretty_report_plugin
|
||||
-p _ir_meta_plugin
|
||||
)
|
||||
# 设备后端：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)"
|
||||
TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
|
||||
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……
|
||||
@@ -104,9 +108,12 @@ else
|
||||
PYTEST_COLOR=()
|
||||
fi
|
||||
|
||||
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
|
||||
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=()
|
||||
@@ -126,17 +133,38 @@ else
|
||||
MODE_DESC=record
|
||||
fi
|
||||
|
||||
# 未指定 shape 文件时，把 INLINE_YAML 落到临时文件供 pytest 读取
|
||||
if [[ -z "$SHAPE_FILE" ]]; then
|
||||
SHAPE_FILE="$(mktemp --suffix=.yaml)"
|
||||
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
|
||||
trap 'rm -f "$SHAPE_FILE"' EXIT
|
||||
# 未指定 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
|
||||
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
|
||||
|
||||
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 mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
|
||||
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 落盘
|
||||
@@ -149,7 +177,7 @@ status=0
|
||||
env "$AUTOTUNE_ENV" \
|
||||
USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
|
||||
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
|
||||
--shape_file "$SHAPE_FILE" ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
|
||||
${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
|
||||
--level core --mode kernel || status=$?
|
||||
|
||||
rm -rf "$CACHE_DIR"
|
||||
|
||||
Reference in new issue
Block a user