zl_bench/run_pytest.sh

#!/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
# 在 $BENCH_LEVEL 下的原生 shape 集。设了 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}"

# 上游 --level：core（默认）只用各算子的 core_shapes；comprehensive 额外并入
# 子类的 set_more_shapes()，并对同一 shape 再产出一次变体（BLAS 是转置侧）。
# 差距可以很大：mm 在 core 下是 5 个 shape，comprehensive 下并入 4 个大 K
# （N/K 到 152064）与 models_shapes/*.yaml 的模型 shape，共 29 个，且用例翻倍。
# 代价同步放大——sweep 是墙上时间的大头，shape 多时配合 PARALLEL_WARMUP_GPUS。
BENCH_LEVEL="${BENCH_LEVEL:-core}"
case "$BENCH_LEVEL" in
  core|comprehensive) ;;
  *) echo "!!! BENCH_LEVEL 只能是 core 或 comprehensive（当前：$BENCH_LEVEL）" >&2; exit 2 ;;
esac

# FlagGems 的 libtuner winner 存在 $FLAGGEMS_CACHE_DIR/config_cache 的 sqlite 里，
# 默认 ~/.flaggems，跨 run 长存。1=本次 run 用独立缓存目录，两侧各自冷调优
# （README「各自最优」口径）；0=沿用共享缓存（默认，快）。
#
# 共享缓存的风险是单向的：先跑的一侧把 winner 写进去，后跑的一侧直接命中，
# 于是它只编译那几个胜出变体。实测 mm 上后跑侧的 TTGIR 变体从 945 塌到 33、
# cvt 从 136 变 10，看起来像被测改动的效果，其实只是继承了对侧的 config。
# 走 REPLAY_FROM 的标准 A/B 不受此影响（config 本来就锁死），record 对 record
# 才会中招。无论取哪个值都会写进 run_identity.json，口径不再隐形。
ZL_ISOLATE_TUNER_CACHE="${ZL_ISOLATE_TUNER_CACHE:-0}"

# ============================== 执行逻辑 ==============================

# 前置检查。静默失败最贵：路径不存在时 pytest 收集不到 benchmark 目录，注册
# --level/--mode 的 conftest 就没加载，报错变成 "unrecognized arguments"，指向
# 完全错误的方向。每项检查自报名字，失败即停。
preflight() {
  local what="$1"; shift
  if ! "$@"; then
    echo "!!! PREFLIGHT_FAILED check=$what" >&2
    exit 3
  fi
}

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}}"

preflight "FlagGems directory exists: $FLAGGEMS_DIR" test -d "$FLAGGEMS_DIR"
# The node's file half; the ::test_<fn> half is left to pytest, which reports a
# missing function clearly ("not found: ...") while a missing FILE does not.
preflight "benchmark file exists: ${TEST_FILE%%::*}" test -f "${TEST_FILE%%::*}"
# The interpreter, not just the paths. triton and flag_gems normally come from a
# PYTHONPATH pointing at a compiler build and a source checkout, so a shell that
# merely reaches the right files can still be the wrong Python. Without this the
# failure surfaces per cell as `pytest exit 4` from a conftest ImportError, and a
# multi-round campaign burns every cell before anyone reads a log.
preflight "interpreter imports triton + flag_gems ($(command -v python))"   sh -c 'python -c "import triton, flag_gems" 2>/dev/null'
if [[ -n "$SHAPE_FILE" ]]; then
  preflight "shape file exists: $SHAPE_FILE" test -f "$SHAPE_FILE"
fi
if [[ -n "$REPLAY_FROM" ]]; then
  preflight "replay run directory exists: $REPLAY_FROM" test -d "$REPLAY_FROM"
fi

# 产物目录，可用 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: $BENCH_LEVEL" > "$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

# libtuner 缓存：隔离时用本 run 的私有目录，跑完随 Triton 缓存一起删（选中的
# config 已经在 autotune_records/ 里，DB 本身没有留存价值）。
TUNER_CACHE_ARGS=()
if [[ "$ZL_ISOLATE_TUNER_CACHE" == "1" ]]; then
  TUNER_CACHE_DIR="$OUT_DIR/.flaggems_cache"
  rm -rf "$TUNER_CACHE_DIR"; mkdir -p "$TUNER_CACHE_DIR"
  TUNER_CACHE_ARGS=(FLAGGEMS_CACHE_DIR="$TUNER_CACHE_DIR")
  TUNER_DESC=isolated
else
  TUNER_DESC="shared(${FLAGGEMS_CACHE_DIR:-~/.flaggems})"
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" \
  --bench-level "$BENCH_LEVEL" \
  --tuner-cache "$TUNER_DESC"

status=0
{
  echo "${C_BOLD}>>> op=$OP  backend=$ZL_BACKEND  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE  shape=$SHAPE_MODE  level=$BENCH_LEVEL  tuner=$TUNER_DESC${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" ${TUNER_CACHE_ARGS[@]+"${TUNER_CACHE_ARGS[@]}"} \
  USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
    "${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
    ${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
    --level "$BENCH_LEVEL" --mode kernel || status=$?

  rm -rf "$CACHE_DIR"
  [[ -n "${TUNER_CACHE_DIR:-}" ]] && rm -rf "$TUNER_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"