unify the loading and add a UT for configs

aiter/ops/triton/_triton_kernels/attention/extend_attention.py

@@ -18,7 +18,6 @@
 """
 
 import functools
-import json
 
 import torch
 import triton
@@ -28,7 +27,7 @@
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 _fwd_kernel_extend_repr = make_kernel_repr(
     "_fwd_kernel",
@@ -321,16 +320,13 @@ def _fwd_kernel(
 
 @functools.lru_cache(maxsize=1024)
 def _get_config(HEAD_SIZE, dtype):
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
+    dev = arch_info.get_arch()
+    config = load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION.json"
+    )
 
     # HEAD_SIZE 192 = 128 head and 64 pe head dim
     if (HEAD_SIZE > 192) or dtype == torch.float32:
-        return _get_config._config_dict["large_head_or_fp32"]
+        return config["large_head_or_fp32"]
 
-    return _get_config._config_dict["default"]
+    return config["default"]

aiter/ops/triton/_triton_kernels/attention/hstu_attention.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import functools
-import json
 
 # @manual=//triton:triton
 import triton
@@ -24,7 +23,7 @@
 
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 try:
     from triton.language.extra.libdevice import (
@@ -871,12 +870,10 @@ def _hstu_attn_bwd(
 def _get_fwd_config(
     AUTOTUNE_Z: int,
 ):
-    if not hasattr(_get_fwd_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_fwd_config._config_dict = config
+    dev = arch_info.get_arch()
+    config = load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json",
+    )
 
     if AUTOTUNE_Z < 512:
         batch_key = "small_batch"
@@ -885,23 +882,21 @@ def _get_fwd_config(
     else:
         batch_key = "large_batch"
 
-    return _get_fwd_config._config_dict[batch_key]
+    return config[batch_key]
 
 
 @functools.lru_cache(maxsize=1024)
 def _get_bwd_config(
     AUTOTUNE_Z: int,
 ):
-    if not hasattr(_get_bwd_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_bwd_config._config_dict = config
+    dev = arch_info.get_arch()
+    config = load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json",
+    )
 
     if AUTOTUNE_Z < 512:
         batch_key = "small_batch"
     else:
         batch_key = "large_batch"
 
-    return _get_bwd_config._config_dict[batch_key]
+    return config[batch_key]

aiter/ops/triton/_triton_kernels/attention/lean_atten.py

@@ -17,32 +17,24 @@
     -
 """
 
-import functools
-import json
-
 import triton
 import triton.language as tl
 
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 # Support tensor in [B, Seqlen, H, d] format. Taking tensors in [B*Seqlen, H, d] as inputs
 
 
-@functools.lru_cache(maxsize=1024)
 def _get_config():
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
-
-    config = _get_config._config_dict["any"]
-    return (
-        config.copy()
-    )  # return a copy to avoid mutation of stored config in LRU cache
+    # No lru_cache here: load_config_json already caches the parse, and
+    # caching the .copy() would hand every caller the same mutable object.
+    dev = arch_info.get_arch()
+    config = load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json"
+    )
+    return config["any"].copy()  # fresh copy per call — safe for callers to mutate
 
 
 @triton.jit

aiter/ops/triton/_triton_kernels/attention/mha.py

@@ -2,7 +2,6 @@
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import functools
-import json
 
 import torch
 import triton
@@ -15,7 +14,7 @@
     remap_workgroup_spatial,
     remap_xcd,
 )
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 
 @triton.jit
@@ -953,14 +952,9 @@ def _get_config(
     has_pe: bool = False,
     head_dim_v: int | None = None,
 ):
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict["default"] = config
-    fwd_cfg = _get_config._config_dict["default"]["fwd"]
+    dev = arch_info.get_arch()
+    config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
+    fwd_cfg = config["fwd"]
     has_dropout_or_fp32 = enable_dropout or dtype == torch.float32
     # TODO: pe + dropout is not tuned
     if has_pe and has_dropout_or_fp32 and "pe_dropout_or_fp32" in fwd_cfg:

aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py

@@ -2,7 +2,6 @@
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import functools
-import json
 
 import triton
 import triton.language as tl
@@ -11,7 +10,7 @@
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
 from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 # This function computes delta given output Out and gradient DO
 # Here is the I/O shape:
@@ -1062,12 +1061,6 @@ def _bwd_kernel_dkdvdq_noncausal(
 
 @functools.lru_cache(maxsize=1024)
 def _get_config():
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
-
-    return _get_config._config_dict["bkwd_fused"]
+    dev = arch_info.get_arch()
+    config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
+    return config["bkwd_fused"]

aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py

@@ -2,15 +2,14 @@
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import functools
-import json
 
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 # NOTE: triton fails to import tl.constexprs so create them here for the file
 DROPOUT_USE_PYTORCH = False
@@ -1769,12 +1768,6 @@ def bwd_kernel_noncausal(
 
 @functools.lru_cache(maxsize=1024)
 def _get_config():
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
-
-    return _get_config._config_dict["bkwd_onekernel"]
+    dev = arch_info.get_arch()
+    config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
+    return config["bkwd_onekernel"]

aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py

@@ -24,7 +24,6 @@
 # https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
 
 import functools
-import json
 
 import triton
 import triton.language as tl
@@ -33,7 +32,7 @@
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
 from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 _fwd_grouped_kernel_stage1_rope_repr = make_kernel_repr(
     "_fwd_grouped_kernel_stage1_rope",
@@ -404,12 +403,7 @@ def _fwd_kernel_stage2(
 
 @functools.lru_cache(maxsize=1024)
 def _get_config():
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
-
-    return _get_config._config_dict
+    dev = arch_info.get_arch()
+    return load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json",
+    )

aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
-import functools
 
 import triton
 import triton.language as tl
@@ -619,7 +618,6 @@ def _fused_gemm_a8w8_blockscale_split_cat_reduce(
     tl.store(c1_ptrs, y, mask=y_mask)
 
 
-@functools.lru_cache(maxsize=1024)
 def _get_config(
     M: int,
     N: int,

aiter/ops/triton/_triton_kernels/gmm.py

@@ -7,8 +7,6 @@
 
 # Python standard library
 import functools
-import json
-import os.path
 
 # Triton
 import triton
@@ -18,7 +16,7 @@
 from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd
 
 # AITER
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 # Kernel config.
 # ------------------------------------------------------------------------------
@@ -33,25 +31,16 @@ def get_config(
         "ptgmm",
         "nptgmm",
     }, f"'{gmm_type}' is an invalid GMM variant."
-    if not hasattr(get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        config_filename = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-GMM.json"
-        assert os.path.exists(config_filename) and os.path.isfile(
-            config_filename
-        ), f"'{config_filename}' isn't an existent file."
-        with open(config_filename, "r") as config_file:
-            get_config._config_dict = json.load(config_file)
-            assert all(
-                gmm_type in get_config._config_dict
-                for gmm_type in ("gmm", "ptgmm", "nptgmm")
-            ), "Not all GMM variants are present in the configuration file."
+    dev = arch_info.get_arch()
+    config_dict = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-GMM.json")
+    assert all(
+        variant in config_dict for variant in ("gmm", "ptgmm", "nptgmm")
+    ), "Not all GMM variants are present in the configuration file."
     # TODO: Fine tune GMM kernels and use (M, K, N, G) shape to query the best
     #       config in the dictionary.
-    assert (
-        "default" in get_config._config_dict[gmm_type]
-    ), "Default configuration is absent."
+    assert "default" in config_dict[gmm_type], "Default configuration is absent."
     key = "accumulate" if accumulate else "default"
-    return get_config._config_dict[gmm_type][key]
+    return config_dict[gmm_type][key]
 
 
 # Common code shared by GMM and TGMM kernels.

aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py

@@ -2,14 +2,13 @@
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import functools
-import json
 
 import triton
 import triton.language as tl
 
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
-from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
+from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
 
 _routing_sigmoid_top1_repr = make_kernel_repr(
     "_routing_sigmoid_top1_kernel",
@@ -127,18 +126,15 @@ def _routing_sigmoid_top1_kernel(
 
 @functools.lru_cache(maxsize=1024)
 def _get_config(M, N, K):
-    if not hasattr(_get_config, "_config_dict"):
-        dev = arch_info.get_arch()
-        _get_config._config_dict = {}
-        fpath = f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE_ROUTING_SIGMOID_TOPK1.json"
-        with open(fpath, "r") as file:
-            config = json.load(file)
-        _get_config._config_dict = config
+    dev = arch_info.get_arch()
+    config = load_config_json(
+        f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE_ROUTING_SIGMOID_TOPK1.json",
+    )
 
     n_key = "N16" if N <= 16 else "N128"
     m_key = (
         "xlarge"
         if M >= 8192
         else "large" if M >= 4096 else "medium" if M >= 2048 else "small"
     )
-    return _get_config._config_dict[n_key][m_key]
+    return config[n_key][m_key]

aiter/ops/triton/attention/lean_atten.py

@@ -83,7 +83,7 @@ def persistent_lean_attention(
         f"LEAN_ATTEN: q={tuple(q.shape)}  k={tuple(k.shape)}  v={tuple(v.shape)} Mp={tuple(Mp.shape)} Lp={tuple(Lp.shape)}  Op={tuple(Op.shape)}"
     )
     if config is None:
-        config = _get_config(causal=causal, batch_size=batch_size)
+        config = _get_config()
     sm_count = arch_info.get_num_sms()
     total_programs = (
         program_count