diff --git a/fla/ops/utils/cache.py b/fla/ops/utils/cache.py index 40ab2dbc18..801de34936 100644 --- a/fla/ops/utils/cache.py +++ b/fla/ops/utils/cache.py @@ -268,6 +268,31 @@ def read_config_file(config_file: Path) -> dict[str, Any] | None: return load_config_file(config_file) +@cache +def load_kernel_config_file(config_file: Path) -> "tuple[dict[str, Any] | None, KernelConfigFile | None]": + """Read, parse and validate a kernel config file, memoized by path. + + Returns a ``(raw config dict, validated KernelConfigFile)`` pair, with both set + to ``None`` when the file is missing or malformed. + + The existence check and the validation are memoized along with the file contents, + so in the steady state a kernel launch performs no filesystem access at all. + Without this, every launch whose autotune key misses the in-memory cache would + stat the config file (and re-validate its JSON) again, which showed up as + thousands of small file operations per second in large multi-rank trainings + (see issue #1072). + + ALWAYS mode should bypass the memoization via ``load_kernel_config_file.__wrapped__`` + so that config file edits are picked up on the next kernel call. + """ + if not config_file.exists(): + return None, None + config_data = read_config_file(config_file) + if config_data is None: + return None, None + return config_data, KernelConfigFile.from_dict(config_file, config_data) + + def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None: """ Load cached best config for a kernel from FLA configs directory. @@ -296,13 +321,10 @@ def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None config_dir = get_fla_config_dir() config_file = config_dir / f"{kernel_name}.json" - if not config_file.exists(): - return None - - config_data = read_config_file(config_file) - if config_data is None: - return None - config = KernelConfigFile.from_dict(config_file, config_data) + if FLA_CACHE_MODE is FlaCacheMode.ALWAYS: + config_data, config = load_kernel_config_file.__wrapped__(config_file) + else: + config_data, config = load_kernel_config_file(config_file) if config is None: return None diff --git a/tests/ops/test_cache.py b/tests/ops/test_cache.py index ed44274631..ee78e06cb0 100644 --- a/tests/ops/test_cache.py +++ b/tests/ops/test_cache.py @@ -5,11 +5,14 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +import json + import pytest import torch import triton import triton.language as tl +from fla.ops.utils import cache as cache_mod from fla.ops.utils.cache import fla_cache_autotune from fla.utils import device @@ -58,3 +61,49 @@ def test_fla_cache_autotune_handles_none_restore_value(): y2 = torch.full((M,), 7, dtype=torch.int32, device=device) _optional_restore_kernel[(triton.cdiv(M, 128),)](None, y2, M) assert torch.equal(y2, torch.zeros(M, dtype=torch.int32, device=device)) + + +def test_load_cached_config_default_mode(tmp_path, monkeypatch): + """DEFAULT mode returns the top-level default_config of the kernel's json file.""" + monkeypatch.setattr(cache_mod, "FLA_CACHE_MODE", cache_mod.FlaCacheMode.DEFAULT) + monkeypatch.setenv("FLA_CONFIG_DIR", str(tmp_path)) + + kernel = "cache_test_kernel_default" + default_config = {"kwargs": {"BLOCK": 64}, "num_warps": 4, "num_stages": 2, "num_ctas": 1} + (tmp_path / f"{kernel}.json").write_text(json.dumps({"default_config": default_config})) + + assert cache_mod.load_cached_config(kernel) == default_config + + +def test_load_kernel_config_file_memoizes_missing_file(tmp_path, monkeypatch): + """A missing config file is looked up on disk once, then served from memory. + + Repeated lookups for kernels without a config file previously performed one + filesystem stat per kernel launch, which showed up as excessive small-file + IO in multi-rank trainings. + """ + monkeypatch.setattr(cache_mod, "FLA_CACHE_MODE", cache_mod.FlaCacheMode.FULL) + monkeypatch.setenv("FLA_CONFIG_DIR", str(tmp_path)) + + kernel = "cache_test_kernel_missing" + assert cache_mod.load_cached_config(kernel) is None + + # outside ALWAYS mode the negative result is memoized, matching the + # existing behavior of load_config_file for files created mid-run + (tmp_path / f"{kernel}.json").write_text(json.dumps({"default_config": {"kwargs": {}}})) + assert cache_mod.load_cached_config(kernel) is None + + +def test_load_cached_config_always_mode_rereads(tmp_path, monkeypatch): + """ALWAYS mode picks up config file edits on the next lookup.""" + monkeypatch.setattr(cache_mod, "FLA_CACHE_MODE", cache_mod.FlaCacheMode.ALWAYS) + monkeypatch.setenv("FLA_CONFIG_DIR", str(tmp_path)) + + kernel = "cache_test_kernel_always" + config_file = tmp_path / f"{kernel}.json" + + config_file.write_text(json.dumps({"default_config": {"kwargs": {"BLOCK": 64}}})) + assert cache_mod.load_cached_config(kernel)["kwargs"]["BLOCK"] == 64 + + config_file.write_text(json.dumps({"default_config": {"kwargs": {"BLOCK": 128}}})) + assert cache_mod.load_cached_config(kernel)["kwargs"]["BLOCK"] == 128