feat: expose the upstream benchmark level as BENCH_LEVEL

run_pytest.sh hardcoded --level core, so the only shapes ever reachable were
each operator's core_shapes entry. FlagGems merges a subclass's
set_more_shapes() on top of that at --level comprehensive, and for BLAS that
addition is where the large GEMMs live: mm goes from 5 shapes to 29, taking in
four large-K tuples plus models_shapes/*.yaml, with N/K up to 152064. The
comprehensive path also yields each shape a second time (transposed, for BLAS),
so the case count doubles. None of that was reachable without editing the driver.

BENCH_LEVEL defaults to core, so existing runs are unchanged, and rejects
anything other than core/comprehensive rather than passing it through to pytest.
The chosen level is echoed in the run header, written into the archived
shapes.yaml, and recorded in run_identity.json, because an A/B whose two sides
used different levels is not comparable and is otherwise invisible afterwards.

Sweep cost scales with the shape count, so comprehensive on a BLAS operator
wants PARALLEL_WARMUP_GPUS.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

README.md

@@ -53,7 +53,7 @@ softmax:
 | `OP` | `fused_marlin_moe_mxfp4` | 测试函数名（`test_` 之后的部分）；设了 `TEST_NODE` 时由节点覆盖 |
 | `OP_FILE` | `fused_marlin_moe` | benchmark 文件名（`test_` 与 `.py` 之间的部分）；设了 `TEST_NODE` 时由节点覆盖 |
 | `TEST_NODE` | 空 | FlagGems pytest 节点（`…/benchmark/test_<op_file>.py::test_<op>`）。设了且未指定 `SHAPE_FILE` 时默认不传 `--shape_file` |
-| `SHAPE_FILE` | 空 | shape yaml 路径；空且无 `TEST_NODE` 时用脚本内置 yaml；`-`/`upstream`/`none` 表示沿用上游 `--level core` |
+| `SHAPE_FILE` | 空 | shape yaml 路径；空且无 `TEST_NODE` 时用脚本内置 yaml；`-`/`upstream`/`none` 表示沿用上游在 `BENCH_LEVEL` 下的原生 shape 集 |
 | `SHAPE_MODE` | 空 | `inline` / `file` / `upstream`；通常不用手设，`TEST_NODE` 会把空 `SHAPE_FILE` 推成 `upstream` |
 | `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径。测试文件缺失时回退 `/workspace/FlagGems`；`TEST_NODE` 也可直接带出仓库根 |
 | `ZL_BACKEND` | `cuda` | 设备后端：`cuda` / `musa` / `rocm` / `ascend`（别名 `nvidia`/`amd`/`npu` 等；未知名直接失败） |
@@ -65,6 +65,7 @@ softmax:
 |-----|------|------|
 | `REPLAY_FROM` | 空（record 模式） | 指向某次历史 run 目录，replay 其 autotune 选择。做 A/B 时 B 侧必须设（见「A/B 对比测试标准流程」） |
 | `USE_FLAGTUNE` | `0` | `0` 走普通 autotune（跳过搜索、快速验证）；`1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）。A/B 两侧须取同值 |
+| `BENCH_LEVEL` | `core` | 透传上游 `--level`。`comprehensive` 相对 `core` 做两件事：① 在 `set_shapes()` 里并入子类 `set_more_shapes()`（`mm` 因此从 5 个 shape 变 29 个，含 N/K 到 152064 的大 K 与 `models_shapes/*.yaml` 的模型 shape）；② 在 `get_input_iter()` 里对每个 shape 再产出一次变体（BLAS 是转置侧），用例数翻倍。注意 ① 只在 `SHAPE_MODE=upstream` 下可见——传了 `SHAPE_FILE` 时 `_shape_inject_plugin` 会在上游合并之后再覆盖一次、把 shape 集拉回 yaml 内容（BLAS 那类 arity 不一致的情形除外，插件会退让）；② 不经插件，任何 shape 来源下都生效。sweep 成本同步放大，shape 多时配合 `PARALLEL_WARMUP_GPUS`。A/B 两侧须取同值 |
 | `DTYPES` | 空（上游默认 dtype 扫描） | 空格分隔的 dtype 白名单（如 `"bfloat16 float16"`），逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` |
 | `PARALLEL_WARMUP_GPUS` | 空（串行） | 设为 `N`（≥2）时把 autotune sweep 分片到 N 卡并行，测量仍单卡串行；不设或 `<2` 则完全串行（见「加速：多卡并行 sweep」） |
 | `ZL_ENABLE_CUDAGRAPH` | `auto` | `auto`/`1` 仅在 CUDA 上走 Graph 计时；非 CUDA 即使设 `1` 也拒绝并保持 `timer=direct` |

_run_identity.py

@@ -70,6 +70,7 @@ def main(argv=None) -> int:
     p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
     p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
     p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
+    p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
     args = p.parse_args(argv)
     init_identity(
         args.out_dir,
@@ -81,6 +82,7 @@ def main(argv=None) -> int:
         device=args.device,
         test_file=args.test_file,
         shape_mode=args.shape_mode,
+        bench_level=args.bench_level,
     )
     return 0
 

run_pytest.sh

@@ -20,7 +20,7 @@ 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。
+# 在 $BENCH_LEVEL 下的原生 shape 集。设了 TEST_NODE 且未指定 SHAPE_FILE 时默认 upstream。
 SHAPE_FILE="${SHAPE_FILE:-}"
 SHAPE_MODE="${SHAPE_MODE:-}"
 read -r -d '' INLINE_YAML <<'YAML' || true
@@ -64,6 +64,17 @@ FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
 # 非 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
+
 # ============================== 执行逻辑 ==============================
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -139,7 +150,7 @@ 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"
+    printf '%s\n' "source: upstream" "level: $BENCH_LEVEL" > "$OUT_DIR/shapes.yaml"
   fi
 else
   if [[ -z "$SHAPE_FILE" ]]; then
@@ -159,11 +170,12 @@ python "$SCRIPT_DIR/_run_identity.py" "$OUT_DIR" \
   --autotune-mode "$MODE_DESC" \
   --device "${ZL_DEVICE:-}" \
   --test-file "$TEST_FILE" \
-  --shape-mode "$SHAPE_MODE"
+  --shape-mode "$SHAPE_MODE" \
+  --bench-level "$BENCH_LEVEL"
 
 status=0
 {
-  echo "${C_BOLD}>>> op=$OP  backend=$ZL_BACKEND  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE  shape=$SHAPE_MODE${C_RESET}"
+  echo "${C_BOLD}>>> op=$OP  backend=$ZL_BACKEND  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE  shape=$SHAPE_MODE  level=$BENCH_LEVEL${C_RESET}"
   echo "${C_DIM}>>> node=$TEST_FILE${C_RESET}"
   echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
 
@@ -178,7 +190,7 @@ status=0
   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=$?
+    --level "$BENCH_LEVEL" --mode kernel || status=$?
 
   rm -rf "$CACHE_DIR"