Files
zl_bench/_seed_plugin.py
zhoulin 26b071c6e1 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.
2026-07-19 19:21:13 +00:00

43 lines
1.3 KiB
Python

"""pytest plugin: seed random/numpy/torch RNG for reproducible benchmark inputs.
Data-dependent kernels (sort, topk, nonzero, ...) have value-dependent latency;
a fixed seed makes every run generate byte-identical inputs, so the latency
delta between two runs reflects the change under test, not the input data.
"""
from __future__ import annotations
import sys
_SEED = 0
def pytest_configure(config):
seed = _SEED
# random.seed too: some kits (cutlass_scaled_mm) pick cases via random.shuffle.
import random
random.seed(seed)
seeded = ["random"]
try:
import numpy as _np
_np.random.seed(seed)
seeded.append("numpy")
except Exception:
pass
try:
import torch
except Exception as exc: # torch missing should never happen here, stay safe
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")
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)