Adapt to upstream FlagGems repo split and record pre_hook configs for libtuner

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.

_autotune_record_plugin.py

@@ -18,6 +18,7 @@ FLAGGEMS_PERF_CURRENT_OP):
 from __future__ import annotations
 
 import atexit
+import importlib
 import json
 import os
 import sys
@@ -30,12 +31,17 @@ _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:
@@ -65,10 +71,16 @@ def _serialize_key(key: Tuple[Any, ...]) -> str:
     return json.dumps([repr(x) for x in key])
 
 
-def _serialize_config(cfg: Any) -> Optional[Dict[str, Any]]:
-    # Drop configs with a pre_hook (un-serializable callable) rather than replay
-    # without it and compute incorrectly.
-    if getattr(cfg, "pre_hook", None) is not None:
+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 {
@@ -141,7 +153,9 @@ def _record_run(original):
             cfg = self.configs[0]
         if cfg is None:
             return result
-        entry = _serialize_config(cfg)
+        # 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:
@@ -240,6 +254,7 @@ def _dump_record(record_dir: str) -> None:
 
 
 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()
@@ -252,26 +267,44 @@ def pytest_configure(config):
         return
 
     import triton.runtime.autotuner as _autotuner
-    wrap = _record_run if record_dir else _replay_run
-    _autotuner.Autotuner.run = wrap(_autotuner.Autotuner.run)
-    patched = ["Autotuner"]
-    try:
-        from flag_gems.utils.libentry import LibTuner
-        if "run" in LibTuner.__dict__:
-            LibTuner.run = wrap(LibTuner.__dict__["run"])
-            patched.append("LibTuner")
-    except Exception:
-        pass
+    _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)} ({'+'.join(patched)})",
+              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)} ({'+'.join(patched)})",
+              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)