diff --git a/aiter/ops/triton/attention/kda.py b/aiter/ops/triton/attention/kda.py index e9616104e12..99577be061d 100644 --- a/aiter/ops/triton/attention/kda.py +++ b/aiter/ops/triton/attention/kda.py @@ -8,10 +8,7 @@ import triton from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.config_utils import ( - AITER_TRITON_CONFIGS_PATH, - load_config_json, -) +from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -39,9 +36,8 @@ def get_kda_config( overrides: dict | None = None, fused: bool = False, ) -> dict: - tuned = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{_ARCH}-KDA_DECODE-DEFAULT.json" - ) + cfg_dir = resolve_config_dir("attention", "KDA_DECODE", backend="gluon") + tuned = load_config_json(f"{cfg_dir}/DEFAULT.json") num_seq_heads = num_seqs * HV aligned = K % 32 == 0 and V % 32 == 0 if fused: @@ -143,8 +139,8 @@ def fused_recurrent_kda( pad_slot_guard: paged only: state slots <= 0 mark padded sequences config: tuning dict (BV, SK, num_warps, num_buffers, use_tdm_store, use_tdm_load, use_tdm_fused_load); None resolves a tuned bucket - from ``configs/-KDA_DECODE-DEFAULT.json``, and omitted - keys inherit the resolved bucket's published values. + from ``configs//gluon/attention/kda_decode/DEFAULT.json``, + and omitted keys inherit the resolved bucket's published values. Returns: (o, final_state): o is [B, T, HV, V] in v's dtype. final_state is diff --git a/aiter/ops/triton/configs/gfx1250-KDA_DECODE-DEFAULT.json b/aiter/ops/triton/configs/gfx1250/gluon/attention/kda_decode/DEFAULT.json similarity index 100% rename from aiter/ops/triton/configs/gfx1250-KDA_DECODE-DEFAULT.json rename to aiter/ops/triton/configs/gfx1250/gluon/attention/kda_decode/DEFAULT.json diff --git a/aiter/ops/triton/configs/gfx950-KDA_DECODE-DEFAULT.json b/aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json similarity index 100% rename from aiter/ops/triton/configs/gfx950-KDA_DECODE-DEFAULT.json rename to aiter/ops/triton/configs/gfx950/gluon/attention/kda_decode/DEFAULT.json diff --git a/op_tests/triton_tests/attention/test_kda.py b/op_tests/triton_tests/attention/test_kda.py index 0bcddaa2ab7..ea1316dc2e0 100644 --- a/op_tests/triton_tests/attention/test_kda.py +++ b/op_tests/triton_tests/attention/test_kda.py @@ -2,13 +2,13 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import json +import os import sys import pytest import torch import torch.nn.functional as F -from aiter.ops.triton.attention import kda as kda_module from aiter.ops.triton.attention.kda import ( fused_recurrent_kda, fused_recurrent_kda_packed_decode, @@ -167,15 +167,20 @@ def _replay_into_snapshots(pool, indices, lens, BV=32): @pytest.fixture def get_config(monkeypatch, tmp_path): - name = f"{arch}-KDA_DECODE-DEFAULT.json" - with open(f"{triton_core.AITER_TRITON_CONFIGS_PATH}/{name}") as f: + # The shipped family lives at //gluon/attention/kda_decode/; + # a variant is written to the same relative path under tmp_path and the + # config root is pointed there, so the wrapper resolves the variant. + cfg_dir = triton_core.resolve_config_dir("attention", "KDA_DECODE", backend="gluon") + with open(f"{cfg_dir}/DEFAULT.json") as f: shipped = json.load(f) + relative = os.path.relpath(cfg_dir, triton_core.AITER_TRITON_CONFIGS_PATH) def _get(**variant): merged = {b: {**values, **variant} for b, values in shipped.items()} - (tmp_path / name).write_text(json.dumps(merged)) + variant_dir = tmp_path / relative + variant_dir.mkdir(parents=True, exist_ok=True) + (variant_dir / "DEFAULT.json").write_text(json.dumps(merged)) monkeypatch.setattr(triton_core, "AITER_TRITON_CONFIGS_PATH", str(tmp_path)) - monkeypatch.setattr(kda_module, "AITER_TRITON_CONFIGS_PATH", str(tmp_path)) monkeypatch.setattr(sys.modules[__name__], "_TUNED", dict(variant)) return _get