Stabilize first-run cudagraph timing and colorize terminal output

- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture
  (internal 5-iter warmup is too short), tag fallback markers with the
  failing phase; first-run latency no longer jitters run-to-run.
- _device_guard_plugin (new): set GEMS_VENDOR via torch probe before
  importing flag_gems, avoiding its timeout-less nvidia-smi subprocess
  probe that can hang import in fork-broken environments.
- _pretty_report_plugin (new) + _term_style (new): fold inputs identical
  across all result rows into a legend line, color status/plugin
  tags/markers on the live terminal; run.log is ANSI-stripped and keeps
  upstream SUCCESS/column wording for grep compatibility.
- run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a
  config section with Chinese comments, add start/end banners.
- README: document the warmup semantics, new plugins/env vars, and the
  A/B rule that both sides must use the same USE_FLAGTUNE.

README.md

@@ -4,7 +4,7 @@
 
 1. **可复现**：固定随机种子、指定 shape、固定 autotune config，把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量；
 2. **可解释**：自动按 shape 收集每次运行实际使用的 ttgir（带可读的变体命名），供 IR 级 diff；
-3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没。
+3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没；capture 前按 warmup 预算显式预热吸收 autotune/JIT，首轮即稳态、多轮一致。
 
 ## 快速开始
 
@@ -28,7 +28,7 @@ softmax:
 
 跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。
 
-脚本内固定了 `--level core --mode kernel` 和 `USE_FLAGTUNE=1`（FlagTune 扩展调优空间，见"为什么需要 replay"一节）；需要改动这些口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
+脚本内固定了 `--level core --mode kernel`；`USE_FLAGTUNE` 默认 `1`（FlagTune 扩展调优空间，见"为什么需要 replay"一节），可用环境变量覆盖（`USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune，跳过全量搜索，快速验证时用）；需要改动其它口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
 
 ### 环境变量一览
 
@@ -39,6 +39,8 @@ softmax:
 | `SHAPE_FILE`   | 空（用脚本内置 yaml）       | shape yaml 路径                                   |
 | `FLAGGEMS_DIR` | `/workspace/FlagGems-dev` | FlagGems 仓库路径                                 |
 | `REPLAY_FROM`  | 空（record 模式）          | 指向某次历史 run 目录，replay 其 autotune 选择（见下） |
+| `USE_FLAGTUNE` | `1`                       | `1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）；`0` 走普通 autotune，跳过搜索、快速验证 |
+| `FLAGGEMS_PERF_COLOR` | 空（按 tty 自动判断） | `always`/`never` 强制开/关终端颜色；`run.log` 始终为去色纯文本 |
 
 ## 输出目录结构
 
@@ -46,7 +48,7 @@ softmax:
 
 ```
 runs/softmax_20260716-031752/
-├── run.log                  # 完整日志（含 SUCCESS 行的 latency/speedup 表）
+├── run.log                  # 完整日志（含 SUCCESS 行的 latency/speedup 表；已去 ANSI 色的纯文本）
 ├── shapes.yaml              # 本次实际使用的 shape（存档）
 ├── autotune_records/        # Triton autotune 选中的 config（供 replay）
 │   └── softmax.json
@@ -77,13 +79,17 @@ bash run_pytest.sh                          # -> runs/<op>_<ts_A>/
 REPLAY_FROM=$PWD/runs/<op>_<ts_A> bash run_pytest.sh
 ```
 
+A/B 两侧的 `USE_FLAGTUNE` 必须取同值：该开关决定调优空间，两侧不一致时 replay 会大量 fallback，libtuner 持久缓存也各自独立命中，"同 config"前提不再成立。
+
 对比 `run.log` 的 latency 表看性能差异；diff 两侧 `ttgir/<shape>/<kernel>/` 下的同名文件看 IR 差异。
 
 replay 的兜底行为：B 侧遇到记录中没有的 key、或记录的 config 在新编译器下编译失败时，自动回退到现场 autotune 并在 `run.log` 打 `AUTOTUNE_REPLAY_FALLBACK reason=...` 标记——出现该标记的测量点不再满足"同 config"前提，解读时注意。另外 replay 模式的 run 目录不产生 `autotune_records/`，后续 run 的 `REPLAY_FROM` 应始终指向最初 record 的那次 A 侧目录，不要链式指向 replay 产物。
 
-### cudagraph 回退与测量精度
+### cudagraph 计时的 warmup、回退与精度
 
-部分算子本身不支持 CUDA graph capture（测量函数内含 host 同步、动态显存分配、不合法的流操作等），这类算子会自动回退到普通 do_bench 计时，`run.log` 中打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记（含具体原因）。
+**capture 前的显式 warmup**：`do_bench_cudagraph` 自带的内部预热只有 5 次迭代，对首跑要触发 autotune 编译（尤其含 FlagTune 扩展空间）、libtuner 选择、lazy JIT 的算子远远不够——这些一次性开销若漏进被捕获的 graph 或第一个计时迭代，测出的 latency 会 run-to-run 抖动（M=1 多 kernel 路径最明显，实测首轮可低到稳态的 ~1/3）。`_cudagraph_plugin` 因此在 capture 前显式预热：先跑一次并丢弃（吸收 autotune/JIT 编译），再按调用方传入的 warmup 时间预算（`Config.warm_up`）循环稳态预热，然后才 capture。预热次数按稳态单次耗时换算，并夹在 5–200 次之间——亚毫秒 kernel 的实际预热时长因此低于名义预算（1000ms），实测足够；若换新算子仍见首轮抖动，优先调大 `_warmup_before_capture` 里的次数上限。这样首轮即稳态、多轮一致（实测同一 M=1 shape 两轮 speedup 差 <0.1%）。
+
+部分算子本身不支持 CUDA graph capture（测量函数内含 host 同步、动态显存分配、不合法的流操作等），这类算子会自动回退到普通 do_bench 计时，`run.log` 中打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记（含失败阶段与具体原因；`phase=warmup` 表示预热阶段就失败了，并非 graph capture 被拒）。
 
 回退本身不影响 A/B 公平性（两侧同一算子回退行为一致），但**回退口径的测量误差更大**：do_bench 每次迭代都走完整的 Python → launch 路径，kernel 越小，launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上，且行间波动更明显。解读这类算子的结果时：
 
@@ -109,12 +115,14 @@ libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_confi
 
 | 插件 | 作用 | 何时关闭 |
 |-----|------|---------|
+| `_device_guard_plugin` | 在 import flag_gems 之前直接用 torch 探测到 NVIDIA 卡就设 `GEMS_VENDOR=nvidia`，跳过 flag_gems 启动时 `nvidia-smi` 子进程探测（该探测在部分 fork 环境下会挂在 `wait4` 上导致 import 卡死） | 一般无需关：已设 `GEMS_VENDOR`/`FLAGGEMS_VENDOR` 等 env 时自动跳过，非 NVIDIA 卡上自动 no-op |
 | `_seed_plugin` | 固定 random/numpy/torch 种子，数据相关算子（sort/topk 等）输入逐字节一致 | 不关 |
 | `_shape_inject_plugin` | 让 shape yaml 覆盖子类硬编码的 `set_shapes()` | 不关 |
 | `_shape_iter_inject_plugin` | 覆盖在 `get_input_iter` 里硬编码 shape 的类（conv/pool 等） | 不关 |
 | `_bespoke_shape_plugin` | 覆盖特殊输入构造的算子（upsample/flash_mla/cutlass 等，按类名注册） | 测这些算子之外可关 |
 | `_autotune_record_plugin` | record/replay Triton autotune 选择（由 RECORD/REPLAY 环境变量二选一激活） | 不关 |
-| `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`（kernel 纯耗时；内部先做跨流同步，修过一个间歇性 illegal instruction）。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
+| `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`（kernel 纯耗时；内部先做跨流同步，修过一个间歇性 illegal instruction）。**capture 前按 warmup 时间预算显式预热**（先一次丢弃跑吸收 autotune/JIT 编译，再稳态预热），消除首轮抖动。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
+| `_pretty_report_plugin` | 结果表整理 + 着色：全表相同的输入折叠成表头下一行图例，每行 Size Detail 只留随行变化的部分（shape 保持 `torch.Size([...])` 原样，MoE 类算子单行从 ~400 字符缩到一屏内）；SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文，`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时注释掉 |
 | `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化；launch 钩子统计每个 (kernel, shape) 的实际使用；退出时按上述结构 dump ttgir | 不关 |
 | `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认注释；测 fp16 mm 崩溃时打开 |
 
@@ -126,11 +134,11 @@ libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_confi
 
 **Q: 第一次跑某算子特别慢？**
 
-libtuner 算子首跑要做全量 sweep（mm 约 50 分钟），winner 持久化到 `~/.flaggems/config_cache/` 后，之后同 shape 秒级命中。
+两个来源：（1）默认 `USE_FLAGTUNE=1`，首跑要在 FlagTune 扩展空间做全量搜索（fused_marlin_moe_mxfp4 单进程可达数分钟甚至十几分钟，期间 GPU 满载、`run.log` 停在测试名不动属正常，不是卡死）；（2）libtuner 算子首跑的全量 sweep（mm 约 50 分钟）。winner 持久化到 `~/.flaggems/config_cache/` 后同 shape 秒级命中（但该缓存会因源码/开关/GPU 变化失效，失效后又需重搜）。快速验证（不追求"各自最优 config"、只看功能/稳定性）时直接 `USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune，单 shape 通常 10 秒级出结果。
 
 **Q: `--warmup/--iter` 要设吗？**
 
-不用。cudagraph 计时路径下 warmup 参数本就不生效（内部自带预热），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。
+不用。cudagraph 计时路径下，传入的 warmup 时间预算会被用来在 capture 前显式预热（先吸收 autotune/JIT 编译再稳态预热，稳定首轮，见"cudagraph 计时的 warmup、回退与精度"），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。
 
 **Q: ttgir 目录里某个 shape 少了文件？**
 

_autotune_record_plugin.py

@@ -115,8 +115,10 @@ def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> O
 
 def _emit_marker(reason: str) -> None:
     # Marker lands in run.log (grep it to audit replay coverage); flush so it
-    # appears before the kernel runs.
+    # appears before the kernel runs. Yellow on the live terminal only —
-    print(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", flush=True)
+    # run_pytest.sh strips ANSI from run.log.
+    from _term_style import YELLOW, paint
+    print(paint(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", YELLOW), flush=True)
 
 
 def _record_run(original):
@@ -235,16 +237,19 @@ def _dump_record(record_dir: str) -> None:
         path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
     except Exception as exc:
         # Don't break the run if dump fails; the perf row is still produced.
-        print(f"[autotune-record-plugin] dump failed: {exc}", file=sys.stderr, flush=True)
+        from _term_style import tag
+        print(f"{tag('[autotune-record-plugin]')} dump failed: {exc}",
+              file=sys.stderr, flush=True)
 
 
 def pytest_configure(config):
+    from _term_style import tag
     record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
     replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
     if record_dir and replay_dir:
         # Mutually exclusive: recording while replaying is pointless. Surface loudly.
-        print("[autotune-record-plugin] error: both RECORD_DIR and REPLAY_DIR set; "
+        print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
-              "ignoring both (no-op)", file=sys.stderr, flush=True)
+              "REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
         return
     if not record_dir and not replay_dir:
         return
@@ -265,11 +270,11 @@ def pytest_configure(config):
         # atexit (not sessionfinish): persist whatever was recorded even if an op
         # crash kills the session; a later replay run falls back for missing keys.
         atexit.register(_dump_record, record_dir)
-        print(f"[autotune-record-plugin] recording autotune configs to "
+        print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
               f"{_record_path(record_dir)} ({'+'.join(patched)})",
               file=sys.stderr, flush=True)
     else:
         loaded = _load_replay_dir(replay_dir)
-        print(f"[autotune-record-plugin] replaying {loaded} recorded entries from "
+        print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
               f"{_record_path(replay_dir)} ({'+'.join(patched)})",
               file=sys.stderr, flush=True)

_bespoke_shape_plugin.py

@@ -27,7 +27,9 @@ def _load_yaml(path: str) -> dict:
         with open(path, "r") as f:
             return _yaml.safe_load(f) or {}
     except Exception as exc:
-        print(f"[bespoke-shape-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        from _term_style import tag
+        print(f"{tag('[bespoke-shape-plugin]')} failed to load {path}: {exc}",
+              file=sys.stderr)
         return {}
 
 
@@ -146,5 +148,7 @@ def pytest_collection_finish(session):
         if cls.__name__ == "CutlassScaledMMBenchmark" and _patch_cutlass(cls):
             covered.append("CutlassScaledMMBenchmark(mnk)")
 
-    print(f"[bespoke-shape-plugin] yaml-driven inputs for: {', '.join(covered) or 'none'}",
+    from _term_style import tag
+    print(f"{tag('[bespoke-shape-plugin]')} yaml-driven inputs for: "
+          f"{', '.join(covered) or 'none'}",
           file=sys.stderr, flush=True)

_cudagraph_plugin.py

@@ -5,8 +5,10 @@ line out to time with plain do_bench). Monkey-patches do_bench →
 do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
 loaded but before any benchmark calls it. Kernels that can't be graph-captured
 fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
-into the log. The warmup kwarg is dropped (do_bench_cudagraph warms up
+into the log. Before capturing, fn is explicitly warmed up for the caller's
-internally).
+warmup time budget (do_bench_cudagraph's own 5-iter warmup is too short to
+settle autotune/JIT for the MoE kernels, which makes the first captured
+measurement unstable run-to-run).
 """
 
 from __future__ import annotations
@@ -37,10 +39,55 @@ def _is_fatal_cuda_error(exc: Exception) -> bool:
     return any(m in msg for m in markers)
 
 
+def _warmup_before_capture(fn, warmup_ms, grad_to_none):
+    """Warm fn up before cudagraph capture.
+
+    do_bench_cudagraph only warms up 5 iterations internally — too few for the
+    MoE kernels, where the first calls trigger Triton autotune config
+    compilation, libtuner selection and lazy JIT/init. If that work leaks into
+    the captured graph or the first timed iteration the latency is unstable
+    run-to-run (most visibly on the small-M multi-kernel path). Absorb the
+    one-off compile with a throwaway run, then loop for the caller's warmup time
+    budget so the graph is captured at steady state.
+
+    The loop is clamped to 5..200 iterations: the cap keeps sub-ms kernels from
+    spinning through a large budget (FlagGems' Config.warm_up defaults to
+    1000 ms) — 200 iterations settles them fine in practice — and the floor
+    covers kernels whose single run exceeds the budget. If first-run jitter
+    reappears on a new op, raising the cap is the first knob to try.
+    """
+    import time
+
+    import torch
+
+    def _run():
+        if grad_to_none is not None:
+            for x in grad_to_none:
+                x.grad = None
+        fn()
+
+    # First call absorbs autotune/JIT compilation (can take seconds); discard it.
+    _run()
+    torch.cuda.synchronize()
+    # Size the warmup loop from one timed steady-state run. dt_ms includes
+    # launch overhead, so n errs low for tiny kernels; warmup needn't be exact.
+    t0 = time.perf_counter()
+    _run()
+    torch.cuda.synchronize()
+    dt_ms = (time.perf_counter() - t0) * 1000.0
+    budget = warmup_ms if (warmup_ms and warmup_ms > 0) else 25.0
+    n = int(budget / dt_ms) if dt_ms > 0 else 25
+    n = max(5, min(n, 200))
+    for _ in range(n):
+        _run()
+    torch.cuda.synchronize()
+
+
 def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
                       quantiles=None, return_mode="mean"):
     """Drop-in replacement for triton.testing.do_bench using cudagraph capture."""
     op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
+    phase = "sync"  # which step rejected cudagraph, for the fallback marker
     try:
         # do_bench_cudagraph runs fn on a fresh side stream without syncing
         # with the default stream first; input tensors produced just before
@@ -49,6 +96,10 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
         # device trap on garbage indices). Sync before switching streams.
         import torch
         torch.cuda.synchronize()
+        phase = "warmup"
+        # Settle autotune/JIT before capture (rationale: _warmup_before_capture).
+        _warmup_before_capture(fn, warmup, grad_to_none)
+        phase = "capture"
         return _tt.do_bench_cudagraph(
             fn,
             rep=rep,
@@ -57,13 +108,16 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
             return_mode=return_mode,
         )
     except Exception as exc:
-        # Re-raise genuine device failures; only graph-capture rejections fall back.
+        # Re-raise genuine device failures; anything else falls back to plain
+        # do_bench, with the failing phase (sync/warmup/capture) in the marker.
         if _is_fatal_cuda_error(exc):
             raise
         reason = type(exc).__name__
         detail = " ".join(str(exc).split())[:160]
+        from _term_style import YELLOW, paint
         print(
-            f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}",
+            paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} phase={phase} "
+                  f"reason={reason} detail={detail!r}", YELLOW),
             flush=True,
         )
         return _ORIGINAL_DO_BENCH(
@@ -77,6 +131,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
 
 
 def pytest_configure(config):
+    from _term_style import tag
     _tt.do_bench = _patched_do_bench
-    print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph",
+    print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph",
           file=sys.stderr, flush=True)

_device_guard_plugin.py

@@ -0,0 +1,45 @@
+"""Guard against FlagGems' device-detection hang (benchmark-side fix).
+
+FlagGems' `runtime.backend.device_finder` falls back to
+`subprocess.run(nvidia-smi)` **without a timeout**. Under this conda env's
+forked/inconsistent subprocess (`_posixsubprocess` symbol mismatch), that
+child can hang indefinitely, leaving `import flag_gems` stuck in `wait4`
+(seen as run_pytest "卡住 with no result", GPU 0%).
+
+This plugin runs at import time — before any test module imports flag_gems —
+detects the vendor via torch (no subprocess), and sets GEMS_VENDOR so
+device_finder takes its env fast-path (`_get_vendor_from_env`) and never
+reaches the hang-prone subprocess probe. No change to the FlagGems repo.
+
+Side effect: get_device_properties initializes the CUDA context very early in
+the pytest process. Fine for the current single-process runs; revisit if
+fork-based parallelism (e.g. pytest-xdist) is ever introduced.
+"""
+import os
+import sys
+
+_VENDOR_ENV_KEYS = ("GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND")
+
+
+def _guard_device_vendor():
+    # Respect an explicit choice if the user already set one.
+    if any(k in os.environ for k in _VENDOR_ENV_KEYS):
+        return
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            name = torch.cuda.get_device_properties(0).name.upper()
+            if "NVIDIA" in name:
+                os.environ["GEMS_VENDOR"] = "nvidia"
+                from _term_style import tag
+                print(f"{tag('[device-guard-plugin]')} set GEMS_VENDOR=nvidia "
+                      "(skip flag_gems nvidia-smi subprocess probe, avoids import hang)",
+                      file=sys.stderr, flush=True)
+    except Exception as e:  # pragma: no cover
+        print(f"[device-guard-plugin] torch vendor probe failed ({e}); "
+              "leaving detection to flag_gems",
+              file=sys.stderr, flush=True)
+
+
+_guard_device_vendor()

_ir_meta_plugin.py

@@ -114,7 +114,8 @@ def _compile_listener(*, src, metadata, metadata_group, times, cache_hit):
                       f, indent=1, sort_keys=True)
         os.replace(tmp, path)
     except Exception as exc:  # never break compilation over a metadata dump
-        print(f"[ir-meta-plugin] sidecar dump failed: {exc}", file=sys.stderr)
+        from _term_style import tag
+        print(f"{tag('[ir-meta-plugin]')} sidecar dump failed: {exc}", file=sys.stderr)
 
 
 # --- shape + launch tracking ----------------------------------------------
@@ -202,7 +203,8 @@ def pytest_collection_finish(session):
         if own is not None:
             cls.get_input_iter = _wrap_input_iter(own)
             wrapped += 1
-    print(f"[ir-meta-plugin] shape tracking wrapped on {wrapped} Benchmark classes",
+    from _term_style import tag
+    print(f"{tag('[ir-meta-plugin]')} shape tracking wrapped on {wrapped} Benchmark classes",
           file=sys.stderr, flush=True)
 
 
@@ -312,7 +314,9 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
 
     cache, dest = pathlib.Path(cache_dir), pathlib.Path(dump_dir)
     if not cache.is_dir():
-        print(f"[ir-meta-plugin] no cache dir {cache}; nothing to dump", file=sys.stderr)
+        from _term_style import tag
+        print(f"{tag('[ir-meta-plugin]')} no cache dir {cache}; nothing to dump",
+              file=sys.stderr)
         return
     dest.mkdir(parents=True, exist_ok=True)
 
@@ -386,11 +390,12 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
 
     shapes = sorted({r[0] for r in rows if r[0] != "-"})
     losers = sum(1 for r in rows if r[0] == "-")
+    from _term_style import DIM, paint
     print(f">>> [Dump] ttgir -> {dest} ({copied} files across {len(shapes)} shapes; "
           f"{losers} unused variants index-only; legend: naming.md)", flush=True)
     for s in shapes:
         n = sum(1 for r in rows if r[0] == s)
-        print(f">>> [Dump]   {s}: {n}", flush=True)
+        print(paint(f">>> [Dump]   {s}: {n}", DIM), flush=True)
 
 
 # --- registration -----------------------------------------------------------
@@ -419,10 +424,12 @@ def pytest_configure(config):
             _autotuner.Autotuner._bench = _wrap_bench(_autotuner.Autotuner._bench)
             atexit.register(_dump_ttgir, cache_dir, dump_dir)
         else:
-            print("[ir-meta-plugin] warning: dump dir set but TRITON_CACHE_DIR "
+            from _term_style import tag
+            print(f"{tag('[ir-meta-plugin]')} warning: dump dir set but TRITON_CACHE_DIR "
                   "is not; ttgir dump disabled", file=sys.stderr, flush=True)
 
-    print("[ir-meta-plugin] compilation listener registered"
+    from _term_style import tag
+    print(f"{tag('[ir-meta-plugin]')} compilation listener registered"
           + (f"; shape/launch tracking on, ttgir dump -> {dump_dir}"
              if _dump_enabled else ""),
           file=sys.stderr, flush=True)

_mm_cluster_fix_plugin.py

@@ -92,22 +92,26 @@ def _warmup_load_backend():
             torch.mm(a, b)
             torch.cuda.synchronize()
     except (ImportError, AttributeError) as e:
-        print(f"[mm-cluster-fix-plugin] warmup skipped (benign): {e!r}",
+        from _term_style import tag
+        print(f"{tag('[mm-cluster-fix-plugin]')} warmup skipped (benign): {e!r}",
               file=sys.stderr, flush=True)
     except Exception as e:  # noqa: BLE001
         # e.g. a CUDA error — do NOT hide it; the harness needs to see it.
-        print(f"[mm-cluster-fix-plugin] WARNING: warmup mm failed unexpectedly: {e!r}",
+        from _term_style import tag
+        print(f"{tag('[mm-cluster-fix-plugin]')} WARNING: warmup mm failed "
+              f"unexpectedly: {e!r}",
               file=sys.stderr, flush=True)
 
 
 def pytest_configure(config):
+    from _term_style import tag
     _patch_loaded()
     _warmup_load_backend()
     _patch_loaded()
     guarded = _guarded_modules()
     if guarded:
         print(
-            "[mm-cluster-fix-plugin] guarded cluster_remote_mm_scenario "
+            f"{tag('[mm-cluster-fix-plugin]')} guarded cluster_remote_mm_scenario "
             f"(odd N-tile unmasked OOB) in: {guarded}",
             file=sys.stderr,
             flush=True,
@@ -125,7 +129,7 @@ def pytest_configure(config):
     else:
         # Non-Hopper: the cluster path never runs, so an unpatched state is fine.
         print(
-            "[mm-cluster-fix-plugin] no cluster_remote_mm_scenario found; "
+            f"{tag('[mm-cluster-fix-plugin]')} no cluster_remote_mm_scenario found; "
             "non-Hopper device, guard not needed (no-op).",
             file=sys.stderr,
             flush=True,

_pretty_report_plugin.py

@@ -0,0 +1,148 @@
+"""pytest plugin: readable, colorized benchmark result table.
+
+Replaces FlagGems' BenchmarkResult.__str__ (benchmark/consts.py), whose rows
+print every input's torch.Size(...) on every line — for the MoE ops that is 11
+tensors and ~400 chars per row, ~90% of it identical across rows. This plugin:
+
+- folds inputs identical across all rows into one legend line above the table,
+  so each row's Size Detail keeps only what varies between rows;
+- shapes keep their original torch.Size([...]) spelling (no compaction);
+- colors the table title and Status (SUCCESS green / FAILED red);
+- keeps the literal "SUCCESS"/"FAILED" words and upstream column titles, so
+  run.log greps and __str__-wrapping hooks (vllm column renames) still work.
+
+Color only reaches the live terminal: run_pytest.sh strips ANSI from run.log,
+and _term_style disables color for non-tty stdout unless forced.
+"""
+
+from __future__ import annotations
+
+import sys
+
+from _term_style import BOLD, CYAN, DIM, GREEN, RED, paint, tag
+
+
+def _shape_rows(metrics_list) -> list:
+    """One list of per-input cell strings per metrics row (upstream spelling)."""
+    rows = []
+    for m in metrics_list:
+        sd = m.shape_detail
+        if isinstance(sd, (list, tuple)):
+            rows.append([str(e) for e in sd])
+        else:
+            rows.append([str(sd) if sd is not None else "N/A"])
+    return rows
+
+
+def _factor_static(rows):
+    """Fold columns identical across all rows into a legend.
+
+    Only kicks in when it actually helps: >=2 rows of equal arity >=3, with
+    >=2 static positions and at least one varying position left. Returns
+    (legend_cells_or_None, per_row_varying_strings_or_None).
+    """
+    if len(rows) >= 2:
+        arities = {len(r) for r in rows}
+        if len(arities) == 1:
+            n = next(iter(arities))
+            if n >= 3:
+                static = {i for i in range(n) if len({r[i] for r in rows}) == 1}
+                if 2 <= len(static) < n:
+                    legend = [rows[0][i] for i in sorted(static)]
+                    kept = [
+                        ", ".join(r[i] for i in range(n) if i not in static)
+                        for r in rows
+                    ]
+                    return legend, kept
+    return None, None
+
+
+def _num_cell(value, fmt: str, width: int) -> str:
+    s = f"{value:{fmt}}" if value is not None else "N/A"
+    return f"{s:>{width}}"
+
+
+def _pretty_str(self) -> str:
+    metrics = self.result or []
+    title = (
+        "\n"
+        + paint(f"Operator: {self.op_name}", BOLD, CYAN)
+        + paint(f"  (dtype={self.dtype}, mode={self.mode}, level={self.level})", DIM)
+        + "\n"
+    )
+    if not metrics:
+        return title + "(no results)\n"
+
+    # Same optional-column conditions as upstream.
+    with_tflops = bool(metrics[0].tflops)
+    with_gbps = metrics[0].gbps is not None
+
+    legend, varying = _factor_static(_shape_rows(metrics))
+    legend_line = ""
+    if legend:
+        legend_line = (
+            paint(
+                f"{len(legend)} inputs identical across all rows: "
+                + ", ".join(legend),
+                DIM,
+            )
+            + "\n"
+        )
+        shape_cells = varying
+    else:
+        # No folding possible: print the full shape_detail, upstream style.
+        shape_cells = [
+            str(m.shape_detail) if m.shape_detail is not None else "N/A"
+            for m in metrics
+        ]
+
+    cols = [
+        ("Status", 8, "<"),
+        ("Torch Latency (ms)", 19, ">"),
+        ("Gems Latency (ms)", 18, ">"),
+        ("Gems Speedup", 13, ">"),
+    ]
+    if with_tflops:
+        cols.append(("TFLOPS", 13, ">"))
+    if with_gbps:
+        cols.append(("Torch GBPS", 12, ">"))
+        cols.append(("Gems GBPS", 12, ">"))
+    head_plain = (
+        "  ".join(f"{name:{align}{width}}" for name, width, align in cols)
+        + "  Size Detail"
+    )
+    header = paint(head_plain, BOLD) + "\n" + paint("-" * len(head_plain), DIM) + "\n"
+
+    lines = []
+    for m, shape_cell in zip(metrics, shape_cells):
+        ok = m.error_msg is None
+        cells = [paint(f"{'SUCCESS' if ok else 'FAILED':<8}", GREEN if ok else RED)]
+        cells.append(_num_cell(m.latency_base, ".6f", 19))
+        cells.append(_num_cell(m.latency, ".6f", 18))
+        cells.append(_num_cell(m.speedup, ".3f", 13))
+        if with_tflops:
+            cells.append(_num_cell(m.tflops, ".3f", 13))
+        if with_gbps:
+            cells.append(_num_cell(m.gbps_base, ".3f", 12))
+            cells.append(_num_cell(m.gbps, ".3f", 12))
+        lines.append("  ".join(cells) + "  " + shape_cell)
+
+    return title + legend_line + header + "\n".join(lines) + "\n"
+
+
+def pytest_configure(config):
+    try:
+        from benchmark import consts as fg_consts
+    except Exception as exc:
+        print(
+            f"{tag('[pretty-report-plugin]')} disabled "
+            f"(cannot import benchmark.consts: {exc})",
+            file=sys.stderr, flush=True,
+        )
+        return
+    fg_consts.BenchmarkResult.__str__ = _pretty_str
+    print(
+        f"{tag('[pretty-report-plugin]')} BenchmarkResult table -> folded + color "
+        "(SUCCESS/FAILED words and column titles unchanged)",
+        file=sys.stderr, flush=True,
+    )

_seed_plugin.py

@@ -27,13 +27,16 @@ def pytest_configure(config):
     try:
         import torch
     except Exception as exc:  # torch missing should never happen here, stay safe
-        print(f"[seed-plugin] torch unavailable, seeded {'+'.join(seeded)} only: {exc}",
+        from _term_style import tag
+        print(f"{tag('[seed-plugin]')} torch unavailable, "
+              f"seeded {'+'.join(seeded)} only: {exc}",
               file=sys.stderr)
         return
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
     seeded.append("torch")
-    print(f"[seed-plugin] manual_seed({seed}) for {'+'.join(seeded)} "
+    from _term_style import tag
+    print(f"{tag('[seed-plugin]')} manual_seed({seed}) for {'+'.join(seeded)} "
           "— reproducible benchmark inputs/cases",
           file=sys.stderr, flush=True)

_shape_inject_plugin.py

@@ -26,7 +26,9 @@ def _load_yaml(path: str) -> dict:
         with open(path, "r") as f:
             return _yaml.safe_load(f) or {}
     except Exception as exc:
-        print(f"[shape-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        from _term_style import tag
+        print(f"{tag('[shape-inject-plugin]')} failed to load {path}: {exc}",
+              file=sys.stderr)
         return {}
 
 
@@ -76,8 +78,9 @@ def _patch_init_user_config():
             cur_homogeneous = len(cur_arities) == 1
             yaml_homogeneous = len(yaml_arities) == 1
             if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities:
+                from _term_style import tag
                 print(
-                    f"[shape-inject-plugin] skip override for op_name={self.op_name!r}: "
+                    f"{tag('[shape-inject-plugin]')} skip override for op_name={self.op_name!r}: "
                     f"subclass produced homogeneous arity {next(iter(cur_arities))}, "
                     f"yaml has homogeneous arity {next(iter(yaml_arities))} "
                     f"— subclass normalization preserved",
@@ -93,7 +96,8 @@ def _patch_init_user_config():
 
 
 def pytest_configure(config):
+    from _term_style import tag
     _patch_init_user_config()
-    print("[shape-inject-plugin] Benchmark.init_user_config patched: "
+    print(f"{tag('[shape-inject-plugin]')} Benchmark.init_user_config patched: "
           "yaml shapes now win over hardcoded subclass set_shapes()",
           file=sys.stderr, flush=True)

_shape_iter_inject_plugin.py

@@ -31,7 +31,9 @@ def _load_yaml(path: str) -> dict:
         with open(path, "r") as f:
             return _yaml.safe_load(f) or {}
     except Exception as exc:
-        print(f"[shape-iter-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        from _term_style import tag
+        print(f"{tag('[shape-iter-inject-plugin]')} failed to load {path}: {exc}",
+              file=sys.stderr)
         return {}
 
 
@@ -95,6 +97,7 @@ def pytest_collection_finish(session):
             cls.get_input_iter = _make_patched(own)
             patched += 1
 
-    print(f"[shape-iter-inject-plugin] get_input_iter redirected on {patched} "
+    from _term_style import tag
+    print(f"{tag('[shape-iter-inject-plugin]')} get_input_iter redirected on {patched} "
           f"hardcoded-shape Benchmark classes; yaml shapes now win",
           file=sys.stderr, flush=True)

_term_style.py

@@ -0,0 +1,61 @@
+"""Shared ANSI color helper for zl_bench plugins and scripts.
+
+Color policy, decided once at import (env is fixed before python starts):
+1. FLAGGEMS_PERF_COLOR=always|never wins — run_pytest.sh sets `always` when
+   *its* stdout is a tty, because the tee pipeline hides the tty from python;
+2. FORCE_COLOR forces on (or off when set to 0/false, node semantics), taking
+   precedence over NO_COLOR;
+3. NO_COLOR (non-empty) forces off;
+4. otherwise: on iff stdout is a tty.
+
+run.log stays plain either way: run_pytest.sh strips ANSI from the log file
+after the run.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+_FALSY = ("", "0", "false", "no", "off", "none")
+
+
+def _color_enabled() -> bool:
+    mode = os.environ.get("FLAGGEMS_PERF_COLOR", "").strip().lower()
+    if mode in ("always", "1", "yes", "on"):
+        return True
+    if mode in ("never", "0", "no", "off"):
+        return False
+    force = os.environ.get("FORCE_COLOR")
+    if force is not None:
+        return force.strip().lower() not in _FALSY
+    if os.environ.get("NO_COLOR"):
+        return False
+    try:
+        return sys.stdout.isatty()
+    except Exception:
+        return False
+
+
+COLOR = _color_enabled()
+
+RESET = "\033[0m"
+BOLD = "\033[1m"
+DIM = "\033[2m"
+RED = "\033[31m"
+GREEN = "\033[32m"
+YELLOW = "\033[33m"
+CYAN = "\033[36m"
+BRIGHT_BLUE = "\033[94m"
+
+
+def paint(text: str, *codes: str) -> str:
+    """Wrap text in ANSI codes when color is enabled; identity otherwise."""
+    if not COLOR or not codes:
+        return text
+    return "".join(codes) + text + RESET
+
+
+def tag(label: str) -> str:
+    """Bright-blue '[xxx-plugin]' prefix for plugin log lines."""
+    return paint(label, BRIGHT_BLUE)

run_pytest.sh

@@ -1,18 +1,27 @@
 #!/bin/bash
-# Single-operator perf benchmark (manual debug).
+# 单算子性能测试入口（手动调试用）。
-# Edit OP/OP_FILE/INLINE_YAML below (or override via env: OP=... SHAPE_FILE=... ./run_pytest.sh).
+# 用法：改下方【配置区】，或用环境变量覆盖，例如：
-# All artifacts go under runs/<op>_<timestamp>/: run.log, shapes.yaml,
+#   OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
-# autotune_records/, ttgir/.
+# 产物统一落在 runs/<op>_<时间戳>/：run.log、shapes.yaml、autotune_records/、ttgir/。
 set -euo pipefail
 
-OP="${OP:-fused_marlin_moe_mxfp4}"
+# ============================== 配置区 ==============================
-OP_FILE="${OP_FILE:-fused_marlin_moe}"   # benchmark file is test_<OP_FILE>.py
+# 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。
-SHAPE_FILE="${SHAPE_FILE:-}"
 
-# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name.
+# 被测算子：OP=测试函数名（去掉 test_ 前缀）；OP_FILE=benchmark 文件名
+# （对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>）
+OP="${OP:-fused_marlin_moe_mxfp4}"
+OP_FILE="${OP_FILE:-fused_marlin_moe}"
+
+# FlagGems 仓库路径
+FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
+
+# shape 来源：SHAPE_FILE 非空则用该 yaml；为空则用下方内置 yaml。
+# yaml 顶层 key 必须是 op 名。
+SHAPE_FILE="${SHAPE_FILE:-}"
 read -r -d '' INLINE_YAML <<'YAML' || true
 fused_marlin_moe_mxfp4:
-  # 4 MoE models x 4 token counts (M=1,16,64,256) = 16 shapes
+  # 4 个 MoE 模型 x 4 档 token 数（M=1,16,64,256）= 16 组 shape
   shapes:
     # Mixtral (E=8)
     - [1,   8, 4096, 14336,  2]
@@ -37,16 +46,20 @@ fused_marlin_moe_mxfp4:
   shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
 YAML
 
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# 调优空间：1=FlagTune 扩展空间（首跑全量搜索、慢）；0=普通 autotune（快速验证）
-FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
+USE_FLAGTUNE="${USE_FLAGTUNE:-1}"
-TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
 
-OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
+# autotune record/replay（见 _autotune_record_plugin.py）：
-LOG_FILE="$OUT_DIR/run.log"
+# - 空：record 模式，本次选中的 config 记录到 $OUT_DIR/autotune_records/<op>.json；
-mkdir -p "$OUT_DIR"
+# - 指向某次历史 runs/<op>_<时间戳> 目录：replay 该次记录（A/B 两侧同 config）。
+REPLAY_FROM="${REPLAY_FROM:-}"
 
-export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
+# 终端颜色：always/never 强制开/关；为空则按 tty 自动判断（run.log 始终去色）
+FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
+
+# pytest 插件列表（可按需注释停用；各插件作用见 README「插件说明」）
 PLUGINS=(
+  -p _device_guard_plugin
   -p _seed_plugin
   -p _shape_inject_plugin
   -p _shape_iter_inject_plugin
@@ -54,55 +67,88 @@ PLUGINS=(
   # -p _mm_cluster_fix_plugin
   -p _autotune_record_plugin
   -p _cudagraph_plugin
+  -p _pretty_report_plugin
   -p _ir_meta_plugin
 )
-# Autotune record/replay (see _autotune_record_plugin.py):
+
-# - default: record chosen configs to $OUT_DIR/autotune_records/<op>.json
+# ============================== 执行逻辑 ==============================
-# - REPLAY_FROM=<previous runs/<op>_<ts> dir>: replay that run's recorded
+
-#   configs instead (for A/B runs where both sides must use the same config).
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPLAY_FROM="${REPLAY_FROM:-}"
+TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
+
+OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
+LOG_FILE="$OUT_DIR/run.log"
+mkdir -p "$OUT_DIR"
+
+# 颜色决策：下方 tee 管道会让 python 侧看不到 tty，所以在这里判断，
+# 并通过 FLAGGEMS_PERF_COLOR 下传给各插件（_term_style.py 消费）。
+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 PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
+export FLAGGEMS_PERF_CURRENT_OP="$OP"
+export PYTHONUNBUFFERED=1  # 实时输出不缓冲
+
+# 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 ">>> warning: $REPLAY_FROM/autotune_records/$OP.json not found; replay will fall back to autotune" >&2
+    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
-export FLAGGEMS_PERF_CURRENT_OP="$OP"
 
-# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file.
+# 解析 shape 文件：SHAPE_FILE 为空时把内置 yaml 写到临时文件
 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"  # archive the shape used
+cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml"  # 存档本次实际使用的 shape
-
-export PYTHONUNBUFFERED=1  # unbuffered live output
 
+status=0
 {
-  # _ir_meta_plugin dumps organized ttgir (only variants actually launched
+  echo "${C_BOLD}>>> op=$OP  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
-  # outside the autotune sweep) into TTGIR_DUMP_DIR at process exit, so the
+  echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
-  # dump happens even if a CUDA crash kills the run. Layout/legend: see
+
-  # <dump>/naming.md and index.tsv (which also lists unused sweep losers).
+  # _ir_meta_plugin 在进程退出时把"实际被使用"的变体的 ttgir 按 shape 整理落盘
+  # （挂在 atexit 上，CUDA crash 后已编译部分仍可拿到）。
+  # 目录结构与命名图例见 <dump>/naming.md 和 index.tsv（后者也记录落选的 sweep 变体）。
   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" \
-  USE_FLAGTUNE=1 python -u -m pytest -s "$TEST_FILE" \
+  USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
-    "${PLUGINS[@]}" \
+    "${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
     --shape_file "$SHAPE_FILE" \
     --level core --mode kernel || status=$?
 
   rm -rf "$CACHE_DIR"
 
   if (( status == 0 )); then
-    echo ">>> done. outputs in $OUT_DIR"
+    echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}"
   else
-    echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR"
+    echo "${C_RED}>>> FAILED (pytest exit $status). partial outputs in $OUT_DIR${C_RESET}"
   fi
   exit "$status"
-} 2>&1 | tee "$LOG_FILE"
+} 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"