Fix autotune record/replay silently skipping libtuner kernels

LibTuner.cache is a view onto flag_gems' persistent sqlite config DB
(~/.flaggems/config_cache/TunedConfig_*.db), which survives across runs.
Both record and replay assumed a cold, in-process cache:

- record captured the chosen config by diffing cache keys before/after
  run(). On a warm DB the key is already present, LibTuner.run takes the
  cached branch without writing it again, so the diff was always empty
  and nothing was recorded for any libtuner kernel. Only @triton.autotune
  kernels (in-process cache) made it into the json -- e.g. a
  fused_marlin_moe_mxfp4 run recorded moe_sum_kernel alone, missing both
  MXFP4 GEMMs.
- replay only injected when the key was absent from the cache, so a warm
  DB skipped injection entirely: the run reported "replaying N entries"
  while actually self-tuning.

Record now reads back this call's own self.cache[key] after run();
replay overwrites unconditionally. Verified on fused_marlin_moe_mxfp4:
recorded entries 1 -> 6 (both GEMMs present), replay injects all 6 with
zero AUTOTUNE_REPLAY_FALLBACK and reproduces latency.

README: note that "fresh tune per side" requires dropping the sqlite DB
(not merely omitting REPLAY_FROM), and how to verify record coverage.
This commit is contained in:
2026-07-29 00:38:14 +00:00
parent f98ec10fd0
commit 7e0e8648f1
2 changed files with 29 additions and 39 deletions
+26 -38
View File
@@ -122,46 +122,30 @@ def _emit_marker(reason: str) -> None:
def _record_run(original):
# Snapshot self.cache before/after run() to capture the chosen config (works
# for both Autotuner and LibTuner). Single-config kernels skip the cache
# write, so record configs[0] explicitly for uniform replay.
# 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):
try:
keys_before = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
except Exception:
keys_before = set()
result = original(self, *args, **kwargs)
try:
keys_after = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
except Exception:
keys_after = set()
new_keys = keys_after - keys_before
if not new_keys:
# No new entry: record configs[0] for single-config kernels (cache
# write bypassed); otherwise nothing to record (disk-cache hit).
if len(getattr(self, "configs", []) or []) == 1:
key = _compute_key(self, args, kwargs)
if key is not None and key not in self.cache:
cfg = self.configs[0]
entry = _serialize_config(cfg)
if entry is not None:
kid = _kernel_id(self)
with _record_lock:
bucket = _record_map.setdefault(kid, {})
bucket[_serialize_key(key)] = entry
key = _compute_key(self, args, kwargs)
if key is None:
return result
kid = _kernel_id(self)
with _record_lock:
bucket = _record_map.setdefault(kid, {})
for k in new_keys:
try:
cfg = self.cache[k]
except Exception:
continue
entry = _serialize_config(cfg)
if entry is None:
continue
bucket[_serialize_key(k)] = entry
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
entry = _serialize_config(cfg)
if entry is not None:
kid = _kernel_id(self)
with _record_lock:
_record_map.setdefault(kid, {})[_serialize_key(key)] = entry
return result
return runner
@@ -176,7 +160,11 @@ def _replay_run(original):
bucket = _replay_map.get(kid)
if bucket:
key = _compute_key(self, args, kwargs)
if key is not None and key not in self.cache:
# 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")