zhoulin
95895cea0e
Upstream moved from single /workspace/FlagGems-dev to /workspace/dev/FlagGems (flag_gems) plus /workspace/dev/FlagGems-vllm (flaggems_vllm). Patch LibTuner per actually-imported package at collection finish instead of hardcoding flag_gems, and update the default FLAGGEMS_DIR. Also stop dropping pre_hook configs for libtuner kernels: hopper mm's TMA configs carry a pre_hook even with USE_FLAGTUNE=0, so record silently skipped them and replay fell back. Upstream LibTuner.run now re-attaches the hook by kwargs match on cache read (same path as its own ConfigCache round-trip), so pre_hook-less injection is safe there; plain triton Autotuner keeps the drop.
311 lines
12 KiB
Python
311 lines
12 KiB
Python
"""pytest plugin: record or replay autotune configs across benchmark runs.
|
|
|
|
When comparing two runs (e.g. two compiler builds), replay the first run's
|
|
recorded configs in the second: the autotuner's choice is pinned, so the
|
|
latency delta reflects the change under test rather than the autotuner picking
|
|
a different config under different IR.
|
|
|
|
Modes (mutually exclusive, selected by env var; op identity from
|
|
FLAGGEMS_PERF_CURRENT_OP):
|
|
- FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR: run full autotune, dump the chosen Config
|
|
per (kernel, key) to <dir>/<op>.json at session end.
|
|
- FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR: pre-populate Autotuner.cache from
|
|
<dir>/<op>.json, skipping the sweep. Missing key / non-compiling config /
|
|
missing json fall back to full autotune and emit an AUTOTUNE_REPLAY_FALLBACK
|
|
marker into the log.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import importlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
|
|
_RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR"
|
|
_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR"
|
|
_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP"
|
|
|
|
# Upstream split into two repos with different package names; either may host
|
|
# LibTuner-decorated kernels depending on which repo's benchmark is running.
|
|
_LIBTUNER_PKGS = ("flag_gems", "flaggems_vllm")
|
|
|
|
# Module-global aggregates, accumulated as run() is called and dumped at exit.
|
|
# Lock guards merges in case a kernel uses threads internally.
|
|
_record_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
|
_record_lock = threading.Lock()
|
|
_replay_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
|
_dump_done = False
|
|
_wrap = None # set in pytest_configure when record/replay is active
|
|
|
|
|
|
def _op_name() -> str:
|
|
return os.environ.get(_OP_ENV, "").strip() or "default"
|
|
|
|
|
|
def _record_path(record_dir: str) -> Path:
|
|
return Path(record_dir) / f"{_op_name()}.json"
|
|
|
|
|
|
def _kernel_id(tuner: Any) -> str:
|
|
# Stable across runs: source module + qualified name. Shapes/keys split
|
|
# by the inner key dict.
|
|
name = (
|
|
getattr(tuner, "__name__", None)
|
|
or getattr(getattr(tuner, "base_fn", None), "__name__", None)
|
|
or getattr(getattr(tuner, "fn", None), "__name__", None)
|
|
or "unknown_kernel"
|
|
)
|
|
module = getattr(getattr(tuner, "fn", None), "__module__", "") or ""
|
|
return f"{module}::{name}"
|
|
|
|
|
|
def _serialize_key(key: Tuple[Any, ...]) -> str:
|
|
# Per-element repr, JSON-encoded into a flat string usable as a JSON key;
|
|
# stable across runs given seeded RNG + yaml-fixed shapes.
|
|
return json.dumps([repr(x) for x in key])
|
|
|
|
|
|
def _serialize_config(cfg: Any, allow_pre_hook: bool = False) -> Optional[Dict[str, Any]]:
|
|
# pre_hook is an un-serializable callable. For plain triton Autotuner drop
|
|
# such configs rather than replay without the hook and compute incorrectly.
|
|
# For LibTuner callers pass allow_pre_hook=True: upstream's own ConfigCache
|
|
# round-trips configs pre_hook-less and LibTuner.run re-attaches the hook by
|
|
# matching all_kwargs against self.configs, and replay-injected configs go
|
|
# through that same path (hopper mm's TMA configs carry a pre_hook even
|
|
# with USE_FLAGTUNE=0 — dropping them would record nothing for those
|
|
# kernels and silently defeat replay).
|
|
if getattr(cfg, "pre_hook", None) is not None and not allow_pre_hook:
|
|
return None
|
|
try:
|
|
return {
|
|
"kwargs": dict(cfg.kwargs),
|
|
"num_warps": cfg.num_warps,
|
|
"num_stages": cfg.num_stages,
|
|
"num_ctas": cfg.num_ctas,
|
|
"maxnreg": cfg.maxnreg,
|
|
}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _deserialize_config(entry: Dict[str, Any]) -> Any:
|
|
# Late import (triton loads during collection, after plugin registration).
|
|
from triton import Config
|
|
return Config(
|
|
kwargs=dict(entry["kwargs"]),
|
|
num_warps=entry.get("num_warps", 4),
|
|
num_stages=entry.get("num_stages", 3),
|
|
num_ctas=entry.get("num_ctas", 1),
|
|
maxnreg=entry.get("maxnreg"),
|
|
)
|
|
|
|
|
|
def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[Tuple[Any, ...]]:
|
|
# Reproduce Autotuner/LibTuner key computation; return None on any internals
|
|
# shift (caller then just loses replay for that call).
|
|
try:
|
|
arg_names = tuner.arg_names
|
|
nargs = dict(zip(arg_names, args))
|
|
all_args = {**nargs, **kwargs}
|
|
_args = {k: v for k, v in all_args.items() if k in arg_names}
|
|
if hasattr(tuner, "get_key"):
|
|
return tuner.get_key(_args)
|
|
key = [_args[name] for name in tuner.keys if name in _args]
|
|
for _, arg in _args.items():
|
|
if hasattr(arg, "dtype"):
|
|
key.append(str(arg.dtype))
|
|
return tuple(key)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
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. Yellow on the live terminal only —
|
|
# 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):
|
|
# Read back this call's own key after run(), rather than diffing cache keys.
|
|
# LibTuner.cache is backed by a persistent sqlite DB (flag_gems libcache), so
|
|
# on a warm DB the key is already present and a before/after diff comes up
|
|
# empty -- which silently recorded nothing for every libtuner kernel.
|
|
# Single-config kernels bypass the cache write, so fall back to configs[0].
|
|
def runner(self, *args, **kwargs):
|
|
result = original(self, *args, **kwargs)
|
|
key = _compute_key(self, args, kwargs)
|
|
if key is None:
|
|
return result
|
|
cfg = None
|
|
try:
|
|
cfg = self.cache[key]
|
|
except Exception:
|
|
cfg = None
|
|
if cfg is None and len(getattr(self, "configs", []) or []) == 1:
|
|
cfg = self.configs[0]
|
|
if cfg is None:
|
|
return result
|
|
# get_key marks LibTuner (both flag_gems and flaggems_vllm); plain
|
|
# triton Autotuner has no pre_hook re-attach on cache read, LibTuner does.
|
|
entry = _serialize_config(cfg, allow_pre_hook=hasattr(self, "get_key"))
|
|
if entry is not None:
|
|
kid = _kernel_id(self)
|
|
with _record_lock:
|
|
_record_map.setdefault(kid, {})[_serialize_key(key)] = entry
|
|
return result
|
|
return runner
|
|
|
|
|
|
def _replay_run(original):
|
|
# Inject recorded configs before benchmarking; on key-missing or a recorded
|
|
# config failing to compile/launch, fall back to full autotune (+ marker).
|
|
def runner(self, *args, **kwargs):
|
|
injected_key = None
|
|
if len(getattr(self, "configs", []) or []) > 1:
|
|
kid = _kernel_id(self)
|
|
bucket = _replay_map.get(kid)
|
|
if bucket:
|
|
key = _compute_key(self, args, kwargs)
|
|
# Overwrite unconditionally: LibTuner.cache is backed by a
|
|
# persistent sqlite DB, so gating on `key not in self.cache`
|
|
# would skip injection whenever that DB is warm -- leaving the
|
|
# run silently self-tuned instead of replaying.
|
|
if key is not None:
|
|
rec = bucket.get(_serialize_key(key))
|
|
if rec is None:
|
|
_emit_marker("key_missing")
|
|
else:
|
|
try:
|
|
cfg = _deserialize_config(rec)
|
|
self.cache[key] = cfg
|
|
injected_key = key
|
|
except Exception as exc:
|
|
_emit_marker(f"deserialize_{type(exc).__name__}")
|
|
try:
|
|
return original(self, *args, **kwargs)
|
|
except Exception as exc:
|
|
if injected_key is not None:
|
|
# Recorded config failed under the current build: drop it and
|
|
# let run() autotune. Only Autotuner.cache is a plain dict;
|
|
# LibTuner.cache is a sqlite-backed ConfigCache with no
|
|
# __delitem__, so eviction is impossible there -- say so in the
|
|
# marker instead of retrying with the same bad config and
|
|
# reporting a clean fallback.
|
|
evicted = True
|
|
try:
|
|
del self.cache[injected_key]
|
|
except Exception:
|
|
evicted = False
|
|
# _no_evict means the retry below re-reads the same recorded
|
|
# config, so it is not a clean "fell back to live autotune":
|
|
# treat those measurement points as unverified.
|
|
suffix = "" if evicted else "_no_evict"
|
|
_emit_marker(f"compile_{type(exc).__name__}{suffix}")
|
|
return original(self, *args, **kwargs)
|
|
raise
|
|
return runner
|
|
|
|
|
|
def _load_replay_dir(replay_dir: str) -> int:
|
|
path = _record_path(replay_dir)
|
|
if not path.is_file():
|
|
_emit_marker("record_missing")
|
|
return 0
|
|
try:
|
|
data = json.loads(path.read_text())
|
|
except Exception as exc:
|
|
_emit_marker(f"load_{type(exc).__name__}")
|
|
return 0
|
|
if not isinstance(data, dict):
|
|
return 0
|
|
count = 0
|
|
for kid, bucket in data.items():
|
|
if not isinstance(bucket, dict):
|
|
continue
|
|
_replay_map[kid] = dict(bucket)
|
|
count += len(bucket)
|
|
return count
|
|
|
|
|
|
def _dump_record(record_dir: str) -> None:
|
|
global _dump_done
|
|
if _dump_done:
|
|
return
|
|
_dump_done = True
|
|
path = _record_path(record_dir)
|
|
try:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with _record_lock:
|
|
payload = {kid: dict(bucket) for kid, bucket in _record_map.items()}
|
|
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.
|
|
from _term_style import tag
|
|
print(f"{tag('[autotune-record-plugin]')} dump failed: {exc}",
|
|
file=sys.stderr, flush=True)
|
|
|
|
|
|
def pytest_configure(config):
|
|
global _wrap
|
|
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(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
|
|
"REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
|
|
return
|
|
if not record_dir and not replay_dir:
|
|
return
|
|
|
|
import triton.runtime.autotuner as _autotuner
|
|
_wrap = _record_run if record_dir else _replay_run
|
|
_autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run)
|
|
|
|
if record_dir:
|
|
# 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"{tag('[autotune-record-plugin]')} recording autotune configs to "
|
|
f"{_record_path(record_dir)} (Autotuner)",
|
|
file=sys.stderr, flush=True)
|
|
else:
|
|
loaded = _load_replay_dir(replay_dir)
|
|
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
|
|
f"{_record_path(replay_dir)} (Autotuner)",
|
|
file=sys.stderr, flush=True)
|
|
|
|
|
|
def pytest_collection_finish(session):
|
|
# LibTuner is patched here rather than in pytest_configure: only after
|
|
# collection (which imports the benchmark module) do we know which FlagGems
|
|
# package is actually in use, and importing the unused one just to patch it
|
|
# would initialize a second runtime in this process for nothing.
|
|
if _wrap is None:
|
|
return
|
|
from _term_style import tag
|
|
patched = []
|
|
for pkg in _LIBTUNER_PKGS:
|
|
if pkg not in sys.modules:
|
|
continue
|
|
try:
|
|
libentry = importlib.import_module(f"{pkg}.utils.libentry")
|
|
tuner_cls = libentry.LibTuner
|
|
if "run" in tuner_cls.__dict__:
|
|
tuner_cls.run = _wrap(tuner_cls.__dict__["run"])
|
|
patched.append(f"LibTuner[{pkg}]")
|
|
except Exception as exc:
|
|
print(f"{tag('[autotune-record-plugin]')} warning: LibTuner patch "
|
|
f"failed for {pkg}: {exc}", file=sys.stderr, flush=True)
|
|
print(f"{tag('[autotune-record-plugin]')} libtuner coverage: "
|
|
f"{', '.join(patched) or 'none (no FlagGems package imported?)'}",
|
|
file=sys.stderr, flush=True)
|