Skip to content

[Cache] Memoize kernel config lookups to avoid per-launch file access - #1123

Open
EastZeus wants to merge 1 commit into
fla-org:mainfrom
EastZeus:cache-memoize-config-lookup
Open

[Cache] Memoize kernel config lookups to avoid per-launch file access#1123
EastZeus wants to merge 1 commit into
fla-org:mainfrom
EastZeus:cache-memoize-config-lookup

Conversation

@EastZeus

@EastZeus EastZeus commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

Summary

Addresses #1072.

With the FLA autotune cache enabled, load_cached_config performs a Path.exists() stat and re-validates the config JSON on every kernel launch whose autotune key misses the in-memory autotuner cache. In large multi-rank trainings this shows up as thousands of small open/read/stat operations per second against the shared node-local disk, and host-side gaps before kernel launches (see the profiling in #1072, where chunk_gated_delta_rule_fwd_h roughly doubled its wall time).

This PR routes all lookups through a new memoized load_kernel_config_file helper that caches, per path: the existence check, the raw JSON, and the validated KernelConfigFile. In the steady state a kernel launch performs no filesystem access at all.

Design notes:

  • Missing files are memoized as negative entries. This matches the existing behavior of load_config_file, which already caches a missing file as None (as documented in scripts/utils/autotune_export.py).
  • ALWAYS mode bypasses the memoization through __wrapped__ and keeps its re-read-on-every-call debugging contract (edit a config JSON, next kernel call picks it up).
  • A side effect worth noting: a malformed config file now logs its validation warning once per process instead of once per launch.

Blast radius: host-side lookup path only, no kernel or numerics changes. The one behavior change is that config files created mid-process are no longer picked up outside ALWAYS mode; previously that was already true for file contents (memoized by load_config_file) but not for the existence check, so only the narrow case "file did not exist at first lookup, created later in the same process" changes, and ALWAYS mode retains the old behavior for that workflow.

Test plan

  • Added three host-side tests to tests/ops/test_cache.py (no CUDA required): DEFAULT mode returns default_config, a missing file is memoized as a negative entry, and ALWAYS mode re-reads edits.
  • Local machine is CPU-only, so I verified the module logic by loading fla/ops/utils/cache.py standalone and running the equivalent scenarios, plus STRICT exact match, STRICT miss, and the FULL mode legacy raw-config fallback, all passing. Patching Path.exists with a counter shows 100 repeated lookups for a missing config perform 0 stats after this change (previously 100).
  • find_dependent_tests flags most of the ops suite (everything imports the cache module), so CI coverage is broad by construction; the directly relevant file is tests/ops/test_cache.py.
  • ruff check clean with the repo config.

Benchmark / NCU (kernel changes only)

Neutral for kernels themselves; this removes host-side filesystem work from the launch path. I could not reproduce the multi-rank IO saturation from #1072 locally (no GPU cluster), but the stat-count measurement above shows the per-launch filesystem access dropping to zero, which is the mechanism behind the reported IOPS.

Breaking changes

None intended. See blast radius note above for the one narrow mid-process file-creation case, which ALWAYS mode still supports.

Checklist

  • I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
  • I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
  • This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
  • Dependent tests pass locally or in CI; new behavior is covered by tests where applicable. (New tests added; they are host-side and should pass on any CI runner, but I could not run the triton-importing suite locally.)
  • Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable). (Not a kernel change.)

load_cached_config stats the config file and revalidates its JSON on
every call whose autotune key misses the in-memory autotuner cache. In
large multi-rank trainings with the FLA cache enabled this shows up as
thousands of small open/read/stat operations per second against a shared
node-local disk, degrading GPU utilization (fla-org#1072).

Route all lookups through a new load_kernel_config_file helper that
memoizes the existence check, the raw JSON and the validated
KernelConfigFile by path, so steady state kernel launches perform no
filesystem access at all. Missing files are memoized as negative
entries, matching the existing behavior of load_config_file. ALWAYS
mode bypasses the memoization through __wrapped__ and keeps its
re-read-on-every-call debugging contract.

Add host-side tests covering the DEFAULT mode lookup, the memoized
negative lookup, and the ALWAYS mode re-read contract.
@zhiyuan1i zhiyuan1i added the needs-verification Lacks real execution evidence (CI skipped / no before-after data) label Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

needs-verification Lacks real execution evidence (CI skipped / no before-after data)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants