[Bugfix][Triton] Set the _get_config memo only after the config JSON loads

[Ragua1]

Motivation

When a Triton op has no config JSON for the running architecture, the first call reports the real error and every call after it reports a misleading one.

call 1: FileNotFoundError: .../aiter/ops/triton/configs/gfx1101-MHA-DEFAULT.json
call 2: KeyError: 'default'

The second message is the one users see and report, because retries, later layers and test suites rarely stop at the first exception. It names an internal dict key instead of the missing file, so it sends people looking in the wrong place.

This is not specific to one architecture or to a missing file: it happens on any arch, for any op whose config file cannot be read or parsed.

Technical Details

Seven _get_config helpers memoize their parsed config on the function object and guard the load with hasattr:

@functools.lru_cache(maxsize=1024)
def _get_config(enable_dropout, dtype, has_pe=False):
    if not hasattr(_get_config, "_config_dict"):
        dev = arch_info.get_arch()
        _get_config._config_dict = {}          # published before the load
        fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json"
        with open(fpath, "r") as file:         # may raise
            config = json.load(file)
        _get_config._config_dict["default"] = config
    fwd_cfg = _get_config._config_dict["default"]["fwd"]

functools.lru_cache does not cache exceptions, so the next call re-enters the body — but
hasattr is now True, the load block is skipped, and the caller reads an empty dict.

Affected sites and what the second call reports today:

File	Second call raises
aiter/ops/triton/_triton_kernels/attention/mha.py	KeyError: 'default'
aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py	KeyError: 'bkwd_fused'
aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py	KeyError: 'bkwd_onekernel'
aiter/ops/triton/_triton_kernels/attention/extend_attention.py	KeyError: 'default'
aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py	nothing — returns {}
aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py	KeyError: 'N16'
aiter/ops/triton/gluon/gemm_a8w8_blockscale.py	KeyError: 'default'

mla_decode_rope is the one worth a second look: it hands an empty config back to the caller without raising, so the failure surfaces later as a missing kernel parameter.

The fix publishes the memo only once there is something to publish. Two shapes, +2/-9 in total:

• Five of the sites already end with a whole-dict assignment (_get_config._config_dict = config), which means the earlier = {} is dead code on the success path — it is overwritten three lines later. Those lines are simply removed.
• mha.py and gluon/gemm_a8w8_blockscale.py build into the dict, so the two statements fold into one: _get_config._config_dict = {"default": config}.

No behaviour changes on any path where the config loads.

This is the same family as #2169 (merged), which stopped get_gemm_config from leaking cached state into callers. That one was about a caller mutating the cached dict; this one is about an exception leaving a half-built memo behind.

Test Plan

New test: op_tests/triton_tests/test_config_load_failure.py, parametrised over the six sites whose failure path is arch-independent. For each one it points AITER_TRITON_CONFIGS_PATH at pytest's tmp_path and asserts that two consecutive calls raise the same FileNotFoundError.

It needs no tuned config and no particular GPU — the load fails for whichever architecture is running, so it exercises the fix on CDNA as well. It also clears the memo in a finally block so it cannot affect later tests in the same session.

gluon/gemm_a8w8_blockscale.py is fixed but not covered by the test: its load sits behind a gfx < 950 check, so whether the failure path is reachable depends on the running architecture, and a test whose outcome depends on that does not belong in the suite.

Also run: a standalone reproducer that calls each of the seven _get_config helpers twice in a fresh subprocess, with an option to override the reported architecture so the missing-config condition can be produced on any GPU, including ones that ship every config.

Test Result

Environment: Radeon RX 7800 XT (gfx1101), Windows, torch 2.10.0+rocm7.14.0a20260611, triton 3.7.1, AITER_TRITON_ONLY (implicit on win32). CDNA hardware was not available to me.

New test, with the fix:

$ python -m pytest op_tests/triton_tests/test_config_load_failure.py -q
......                                                                   [100%]
6 passed in 2.79s

Same test with the fix reverted — 5 sites raise the wrong error, 1 raises none:

$ git checkout -- <the 7 files> && python -m pytest op_tests/triton_tests/test_config_load_failure.py -q
E       KeyError: 'default'
aiter\ops\triton\_triton_kernels\attention\mha.py:946: KeyError
E       KeyError: 'bkwd_fused'
aiter\ops\triton\_triton_kernels\attention\mha_fused_bwd.py:1073: KeyError
E       KeyError: 'bkwd_onekernel'
aiter\ops\triton\_triton_kernels\attention\mha_onekernel_bwd.py:1780: KeyError
E       KeyError: 'default'
aiter\ops\triton\_triton_kernels\attention\extend_attention.py:336: KeyError
E           Failed: DID NOT RAISE FileNotFoundError
E       KeyError: 'N16'
6 failed in 2.83s

Reproducer, all seven sites, before and after. --arch overrides the reported architecture; the gfx942 run is the control, since those config files do exist:

before, real arch (gfx1101, no configs):        7/7 sites diverge on the second call
before, --arch gfx9999 (no configs):            7/7 sites diverge on the second call
before, --arch gfx942  (configs present):       0/7   <- control
after,  real arch (gfx1101, no configs):        0/7
after,  --arch gfx942  (configs present):       0/7

End to end through the public API, flash_attn_func called twice in one process:

before:  call 1: FileNotFoundError: ...configs/gfx1101-MHA-DEFAULT.json
         call 2: KeyError: 'default'
after:   call 1: FileNotFoundError: ...configs/gfx1101-MHA-DEFAULT.json
         call 2: FileNotFoundError: ...configs/gfx1101-MHA-DEFAULT.json

Style checks on the changed files: black --check clean (and clean repo-wide, 972 files, with black 26.5.1); ruff check with the CI-pinned 0.16.0 reports All checks passed!.

I could not run the existing MHA op_tests on this platform: op_tests/triton_tests/attention/test_mha.py imports aiter.test_mha_common, which does from aiter import dtypes, and that attribute is not bound on the AITER_TRITON_ONLY path. Supplying it from outside then fails on the missing aiter.jit.module_aiter_core prebuilt, which sends architecture detection to rocminfo. That is why the new test is written not to depend on those helpers.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4447 --add-label <label>

[Boss2002n]

These will be addressed with #4613 and #4630