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:
@@ -118,7 +118,9 @@ bash ab_fold_test.sh
|
||||
| Triton `@triton.autotune` | 仅进程内存 | 每次进程重新 sweep | 计时噪声可能让 A/B 选中**不同 config** → 用 REPLAY_FROM 固定 |
|
||||
| FlagGems `@libtuner`(含 FlagTune 扩展空间) | `~/.flaggems/config_cache/*.db`(sqlite,跨进程持久) | **不失效**(表名只含 kernel 源码与 config 空间的 hash),A/B 自动命中同一 winner | 反向风险:B 侧沿用 A 侧选的旧 winner,测的是"旧 config 下的编译器差异"而非"各自最优" |
|
||||
|
||||
libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_configs.yaml / expand yaml / `USE_FLAGTUNE` 开关变化、Triton 大版本或 GPU 型号变化。如果需要"各自最优"口径(让 B 侧重新 sweep),删除 sqlite 中对应表,或设 `FLAGGEMS_DB_URL` 指向一次性文件。**两种口径都合理,报告结论时注明用的哪种。**
|
||||
libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_configs.yaml / expand yaml / `USE_FLAGTUNE` 开关变化、Triton 大版本或 GPU 型号变化。如果需要"各自最优"口径(让两侧各自重新 sweep),**删掉 `~/.flaggems/config_cache/TunedConfig_*.db` 或设 `FLAGGEMS_DB_URL` 指向一次性文件**——注意"各自 fresh tune"不等于"不设 `REPLAY_FROM`":缓存已热时两侧都会直接命中同一 winner,看似独立调优实则同 config。**两种口径都合理,报告结论时注明用的哪种。**
|
||||
|
||||
同一层缓存也决定了 `_autotune_record_plugin` 的实现方式:record 不能靠"cache 新增了哪个 key"来判断本次选中的 config(缓存一热就走 cached 分支、不写新 key,键集差集恒为空),改为 `run()` 之后直接读回本次调用的 `self.cache[key]`;replay 同理不能加"key 不在 cache 里才注入"的前置条件,否则注入被跳过、该 run 表面在 replay 实际在自调优。想确认 record 真的覆盖到目标 kernel,查 `runs/<run>/autotune_records/<op>.json` 里有无对应 kernel 条目——漏记时该文件照样生成,只是少了 libtuner 那几个。
|
||||
|
||||
## 插件说明
|
||||
|
||||
|
||||
+21
-33
@@ -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:
|
||||
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
|
||||
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
|
||||
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
|
||||
_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")
|
||||
|
||||
Reference in New Issue
Block a user