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>

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"