diff --git a/README.md b/README.md index 7924998..5d9c49a 100644 --- a/README.md +++ b/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_.py::test_`)。设了且未指定 `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` | diff --git a/_run_identity.py b/_run_identity.py index 78af133..dd2406a 100644 --- a/_run_identity.py +++ b/_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 diff --git a/run_pytest.sh b/run_pytest.sh index 85f5ed8..d56d4ac 100644 --- a/run_pytest.sh +++ b/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"