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

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
 

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