[Cache] Memoize kernel config lookups to avoid per-launch file access - #1123
Open
EastZeus wants to merge 1 commit into
Open
[Cache] Memoize kernel config lookups to avoid per-launch file access#1123EastZeus wants to merge 1 commit into
EastZeus wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Addresses #1072.
With the FLA autotune cache enabled,
load_cached_configperforms aPath.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 smallopen/read/statoperations per second against the shared node-local disk, and host-side gaps before kernel launches (see the profiling in #1072, wherechunk_gated_delta_rule_fwd_hroughly doubled its wall time).This PR routes all lookups through a new memoized
load_kernel_config_filehelper that caches, per path: the existence check, the raw JSON, and the validatedKernelConfigFile. In the steady state a kernel launch performs no filesystem access at all.Design notes:
load_config_file, which already caches a missing file asNone(as documented inscripts/utils/autotune_export.py).ALWAYSmode 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).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
ALWAYSmode; previously that was already true for file contents (memoized byload_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, andALWAYSmode retains the old behavior for that workflow.Test plan
tests/ops/test_cache.py(no CUDA required):DEFAULTmode returnsdefault_config, a missing file is memoized as a negative entry, andALWAYSmode re-reads edits.fla/ops/utils/cache.pystandalone and running the equivalent scenarios, plusSTRICTexact match,STRICTmiss, and theFULLmode legacy raw-config fallback, all passing. PatchingPath.existswith a counter shows 100 repeated lookups for a missing config perform 0 stats after this change (previously 100).find_dependent_testsflags most of the ops suite (everything imports the cache module), so CI coverage is broad by construction; the directly relevant file istests/ops/test_cache.py.ruff checkclean 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
ALWAYSmode still supports.Checklist