From b68463c937d6bcf106b1e0e28a8d6c41736cac13 Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Fri, 7 Aug 2026 16:06:32 +0000 Subject: [PATCH 1/6] unify the loading and add a UT for configs --- .../attention/extend_attention.py | 18 ++- .../attention/hstu_attention.py | 29 +++-- .../_triton_kernels/attention/lean_atten.py | 21 ++-- .../triton/_triton_kernels/attention/mha.py | 16 +-- .../attention/mha_fused_bwd.py | 17 +-- .../attention/mha_onekernel_bwd.py | 17 +-- .../attention/mla_decode_rope.py | 17 +-- aiter/ops/triton/_triton_kernels/gmm.py | 29 ++--- .../moe/moe_routing_sigmoid_top1_fused.py | 17 ++- aiter/ops/triton/attention/mla_decode_rope.py | 6 +- aiter/ops/triton/configs/CLAUDE.md | 8 +- aiter/ops/triton/gluon/gemm_a8w8.py | 26 ++--- .../ops/triton/gluon/gemm_a8w8_blockscale.py | 3 +- aiter/ops/triton/gluon/gemm_afp4wfp4.py | 25 ++--- aiter/ops/triton/utils/core.py | 25 +++++ aiter/ops/triton/utils/gemm_config_utils.py | 106 ++++++------------ aiter/ops/triton/utils/mhc_config_utils.py | 11 +- 17 files changed, 159 insertions(+), 232 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/extend_attention.py b/aiter/ops/triton/_triton_kernels/attention/extend_attention.py index b3c00474599..f3d327a3639 100644 --- a/aiter/ops/triton/_triton_kernels/attention/extend_attention.py +++ b/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", required=True + ) # 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"] diff --git a/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py b/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py index 52cd7114958..65ece9b8fc4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py +++ b/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,11 @@ 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", + required=True, + ) if AUTOTUNE_Z < 512: batch_key = "small_batch" @@ -885,23 +883,22 @@ 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", + required=True, + ) 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] diff --git a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py index 16fed247e85..d5dceca0df7 100644 --- a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py +++ b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py @@ -18,31 +18,26 @@ """ 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 + dev = arch_info.get_arch() + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json", required=True + ) + return config[ + "any" + ].copy() # return a copy to avoid mutation of the shared cached config @triton.jit diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index 7dfb20283a2..176eb295a1e 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha.py +++ b/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,11 @@ 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", required=True + ) + 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: diff --git a/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py b/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py index 8569d5afad0..e943e2e19a4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py +++ b/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,8 @@ 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", required=True + ) + return config["bkwd_fused"] diff --git a/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py b/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py index 0bb423d3ada..ff961cbd3c4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py @@ -2,7 +2,6 @@ # 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 @@ -10,7 +9,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.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,8 @@ 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", required=True + ) + return config["bkwd_onekernel"] diff --git a/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py b/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py index 4f0c38c7028..1420dc77387 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py +++ b/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,8 @@ 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", + required=True, + ) diff --git a/aiter/ops/triton/_triton_kernels/gmm.py b/aiter/ops/triton/_triton_kernels/gmm.py index 9e5abdfcc5a..e26cdcb1299 100644 --- a/aiter/ops/triton/_triton_kernels/gmm.py +++ b/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,18 @@ 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", required=True + ) + 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. diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py index 408844e9ada..fc3b8d011dc 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py +++ b/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,13 +126,11 @@ 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", + required=True, + ) n_key = "N16" if N <= 16 else "N128" m_key = ( @@ -141,4 +138,4 @@ def _get_config(M, N, K): 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] diff --git a/aiter/ops/triton/attention/mla_decode_rope.py b/aiter/ops/triton/attention/mla_decode_rope.py index e9cd903f879..32fcead23d3 100644 --- a/aiter/ops/triton/attention/mla_decode_rope.py +++ b/aiter/ops/triton/attention/mla_decode_rope.py @@ -23,6 +23,8 @@ # https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py # https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py +import copy + import torch import triton @@ -198,7 +200,9 @@ def decode_attention_fwd_grouped_rope( + f"k_pe_tokens={tuple(k_pe_tokens.shape) if k_pe_tokens is not None else None} cos_sin_cache={tuple(cos_sin_cache.shape) if cos_sin_cache is not None else None}" ) if config is None: - config = _get_config() + # _get_config() returns the shared cached dict and the launch helpers + # below write derived fields into its sub-configs — work on a copy. + config = copy.deepcopy(_get_config()) is_fp32: bool = any( float_tensor is not None and float_tensor.dtype == torch.float32 diff --git a/aiter/ops/triton/configs/CLAUDE.md b/aiter/ops/triton/configs/CLAUDE.md index 72fb72000d4..51356459ffb 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -113,9 +113,11 @@ Consequences to keep in mind: gluon default (currently gfx1250 `GEMM-AFP4WFP4`), lookup falls through to gluon. Adding `configs/gfx1250/triton/gemm/gemm_afp4wfp4/DEFAULT.json` later would change which file gfx1250 resolves to — verify that is intended. -- Results are cached per `(arch, config_name, backend)` via - `functools.lru_cache` plus `_config_cache`. Adding a file at runtime after a - lookup has happened has no effect; restart the process. +- Results are cached twice: `functools.lru_cache` on the full argument + tuple, plus a per-path cache of parsed JSON + (`_load_json`) that also caches negative results (missing files). Adding a + config file at runtime therefore has no effect; restart the process + (tooling may call `_load_json.cache_clear()` instead). Direct-path loaders bypass all of this. Grep for `f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything — `gluon/gemm_a8w8.py` diff --git a/aiter/ops/triton/gluon/gemm_a8w8.py b/aiter/ops/triton/gluon/gemm_a8w8.py index aa8dd8d8499..0c017f860b6 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8.py +++ b/aiter/ops/triton/gluon/gemm_a8w8.py @@ -1,6 +1,3 @@ -import functools -import json - import torch import triton from triton.experimental import gluon @@ -8,7 +5,7 @@ from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config from aiter.ops.triton.utils.device_info import get_num_xcds from aiter.ops.triton.utils.logger import AiterTritonLogger @@ -553,25 +550,18 @@ def _gemm_a8w8_preshuffled_kernel( gl.amd.cdna4.buffer_store(stored_value=c, ptr=c_ptr, offsets=c_offs, mask=c_mask) -@functools.lru_cache(maxsize=1024) def _get_config( M: int, N: int, K: int, ): - - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - if dev != "gfx950": - raise ValueError( - "Gluon implementation is not supported on this device (requires CDNA4)." - ) - fpath = f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - return _get_config._config_dict["any"] + if arch_info.get_arch() != "gfx950": + raise ValueError( + "Gluon implementation is not supported on this device (requires CDNA4)." + ) + # get_gemm_config caches internally and returns a fresh deep copy. + config, _ = get_gemm_config("GEMM-A8W8", M, N, K, backend="gluon") + return config def gemm_a8w8( diff --git a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py index 809c324889f..4de625992eb 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py @@ -993,13 +993,12 @@ def _get_config( raise ValueError( "Gluon implementation is not supported on this device (requires CDNA4)." ) - _get_config._config_dict = {} fpath = ( f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json" ) with open(fpath, "r") as file: config = json.load(file) - _get_config._config_dict["default"] = config + _get_config._config_dict = {"default": config} key = f"{N}_{K}" if key not in _get_config._config_dict: diff --git a/aiter/ops/triton/gluon/gemm_afp4wfp4.py b/aiter/ops/triton/gluon/gemm_afp4wfp4.py index d860a923bd4..05b5478901c 100644 --- a/aiter/ops/triton/gluon/gemm_afp4wfp4.py +++ b/aiter/ops/triton/gluon/gemm_afp4wfp4.py @@ -1,6 +1,3 @@ -import functools -import json - import torch import triton from triton.experimental import gluon @@ -8,7 +5,7 @@ from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -470,25 +467,17 @@ def _gemm_afp4wfp4_reduce_kernel( gl.store(c_out_ptrs, c) -@functools.lru_cache(maxsize=1024) def _get_config( M: int, N: int, K: int, ): - - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - if dev not in ["gfx950", "gfx1250"]: - raise ValueError("Gluon implementation is not supported on this device.") - fpath = ( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}/gluon/gemm/gemm_afp4wfp4/DEFAULT.json" - ) - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - return _get_config._config_dict["any"] + if arch_info.get_arch() not in ["gfx950", "gfx1250"]: + raise ValueError("Gluon implementation is not supported on this device.") + # get_gemm_config caches internally and returns a fresh deep copy, safe + # for the callers below that write derived fields into the config. + config, _ = get_gemm_config("GEMM-AFP4WFP4", M, N, K, backend="gluon") + return config def gemm_afp4wfp4( diff --git a/aiter/ops/triton/utils/core.py b/aiter/ops/triton/utils/core.py index 969355d944c..f59102b06f0 100644 --- a/aiter/ops/triton/utils/core.py +++ b/aiter/ops/triton/utils/core.py @@ -1,5 +1,30 @@ +import functools +import json import os this_dir = os.path.dirname(os.path.abspath(__file__)) AITER_TRITON_OPS_PATH = os.path.abspath(f"{this_dir}/../") AITER_TRITON_CONFIGS_PATH = os.path.abspath(f"{this_dir}/../configs") + +# This flag should be set to True, unless it is being used for debugging. +# When False, config JSON files are re-read on every call, so live edits to +# the JSON are picked up. +USE_LRU_CACHE = True + + +@functools.lru_cache(maxsize=None if USE_LRU_CACHE else 0) +def load_config_json(fpath: str, required: bool = False) -> dict | None: + """Load a config JSON file, cached per path (including negative results — + add config files before process start, or call + ``load_config_json.cache_clear()``). Returns None if the file doesn't + exist; with required=True raises FileNotFoundError instead, consistently + on every call (exceptions are never cached).""" + try: + with open(fpath, "r") as file: + return json.load(file) + except FileNotFoundError: + if required: + raise FileNotFoundError( + f"Required config file doesn't exist: {fpath}" + ) from None + return None diff --git a/aiter/ops/triton/utils/gemm_config_utils.py b/aiter/ops/triton/utils/gemm_config_utils.py index 73a83efb685..06c13513419 100644 --- a/aiter/ops/triton/utils/gemm_config_utils.py +++ b/aiter/ops/triton/utils/gemm_config_utils.py @@ -7,22 +7,15 @@ import triton from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +from aiter.ops.triton.utils.core import ( + AITER_TRITON_CONFIGS_PATH, + USE_LRU_CACHE, + load_config_json, +) # Standard bounds for M_LEQ_x keys (tuple for hashability with LRU cache) STANDARD_M_BOUNDS = (4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) -# This flag should be set to True, unless it is being used for debugging -USE_LRU_CACHE = True -""" -Cold start: 290.8928 ms -LRU Cache: ENABLED -Avg per call: 0.110 us -vs -LRU Cache: DISABLED -Avg per call: 2.503 us -""" - def _dtype_dir(config_name: str) -> str: """Nested-layout directory for a config family: @@ -39,6 +32,8 @@ def _load_config_file( ) -> bool: """ Helper function to load a config file and cache it. + + Retained for mhc_config_utils.py — the GEMM resolver uses _load_json(). """ if os.path.exists(fpath): with open(fpath, "r") as file: @@ -80,11 +75,7 @@ def _get_gemm_config_cached( and all(x < y for x, y in itertools.pairwise(bounds)) ), "When provided, bounds must be a non-empty tuple of strictly increasing positive numbers." - if not hasattr(_get_gemm_config_cached, "_config_cache"): - _get_gemm_config_cached._config_cache = {} - dev = arch_info.get_arch() - cache_key = f"{dev}_{config_name}" + (f"_{backend}" if backend else "") # Nested layout //gemm// (no arch prefix, default # named DEFAULT.json) first, then legacy flat gemm/ (arch-prefixed) for unmigrated configs. @@ -121,65 +112,31 @@ def _get_gemm_config_cached( cfg_dir, name_prefix, default_stem = _dir, _prefix, _stem break - if cache_key not in _get_gemm_config_cached._config_cache: - _get_gemm_config_cached._config_cache[cache_key] = {} - - # Load default config (must exist) - fpath = f"{cfg_dir}/{default_stem}.json" - _load_config_file( - _get_gemm_config_cached._config_cache, - cache_key, - fpath, - "default", - fpath_should_exist=True, - ) - - config_dict_key = "default" + # Load default config (must exist) + default_fpath = f"{cfg_dir}/{default_stem}.json" + config_dict = load_config_json(default_fpath) + if config_dict is None: + raise AssertionError(f"Required config file doesn't exist: {default_fpath}") - # Handle custom specialized filename (for fused kernels with multiple N dims) + # Specialized configs override the default; the first existing file wins. + # A custom specialized_filename (fused kernels with multiple N dims) + # bypasses the B/N/K candidates. + specialized_suffixes = [] if specialized_filename is not None: - spec_key = specialized_filename - if spec_key not in _get_gemm_config_cached._config_cache[cache_key]: - fpath = f"{cfg_dir}/{name_prefix}{config_name}-{specialized_filename}.json" - if _load_config_file( - _get_gemm_config_cached._config_cache, cache_key, fpath, spec_key - ): - config_dict_key = spec_key - else: - config_dict_key = spec_key - + specialized_suffixes = [specialized_filename] elif N is not None and K is not None: - # Try B-specialized config first: {config_name}-B={B}-N={N}-K={K}.json if B is not None: - bnk_key = f"{B}_{N}_{K}" - if bnk_key not in _get_gemm_config_cached._config_cache[cache_key]: - fpath = f"{cfg_dir}/{name_prefix}{config_name}-B={B}-N={N}-K={K}.json" - if _load_config_file( - _get_gemm_config_cached._config_cache, - cache_key, - fpath, - bnk_key, - ): - config_dict_key = bnk_key - else: - config_dict_key = bnk_key - - # Fall back to N/K-specialized config - if config_dict_key == "default": - nk_key = f"{N}_{K}" - if nk_key not in _get_gemm_config_cached._config_cache[cache_key]: - fpath = f"{cfg_dir}/{name_prefix}{config_name}-N={N}-K={K}.json" - if _load_config_file( - _get_gemm_config_cached._config_cache, - cache_key, - fpath, - nk_key, - ): - config_dict_key = nk_key - else: - config_dict_key = nk_key - - config_dict = _get_gemm_config_cached._config_cache[cache_key][config_dict_key] + specialized_suffixes.append(f"B={B}-N={N}-K={K}") + specialized_suffixes.append(f"N={N}-K={K}") + + is_tuned = False + for suffix in specialized_suffixes: + specialized_config = load_config_json( + f"{cfg_dir}/{name_prefix}{config_name}-{suffix}.json" + ) + if specialized_config is not None: + config_dict, is_tuned = specialized_config, True + break # use standard bounds unless custom bounds are passed search_bounds = bounds if bounds is not None else STANDARD_M_BOUNDS @@ -188,19 +145,20 @@ def _get_gemm_config_cached( for bound in search_bounds: key = f"M_LEQ_{bound}" if M <= bound and key in config_dict: - return dict(config_dict[key]), config_dict_key != "default" + return dict(config_dict[key]), is_tuned # Search for M_GEQ_x keys for bound in reversed(search_bounds): key = f"M_GEQ_{bound}" if M >= bound and key in config_dict: - return dict(config_dict[key]), config_dict_key != "default" + return dict(config_dict[key]), is_tuned if "any" in config_dict: return dict(config_dict["any"]), False raise KeyError( - f"No matching configuration found for M={M}, N={N}, K={K} in config '{config_name}'." + f"No matching configuration found for M={M}, N={N}, K={K}, B={B}, " + f"specialized_filename={specialized_filename!r} in config '{config_name}'." ) diff --git a/aiter/ops/triton/utils/mhc_config_utils.py b/aiter/ops/triton/utils/mhc_config_utils.py index 797e2a217ee..c7de4975142 100644 --- a/aiter/ops/triton/utils/mhc_config_utils.py +++ b/aiter/ops/triton/utils/mhc_config_utils.py @@ -61,14 +61,16 @@ def get_mhc_config( cache_key = f"{dev}_{actual_config_name}" - # Load default config with fallback for unsupported architectures + # Load default config with fallback for unsupported architectures. + # Load into a temp dict and commit only on success, so a failed load + # doesn't leave a stale empty entry that masks the error on later calls. if cache_key not in get_mhc_config._config_cache: - get_mhc_config._config_cache[cache_key] = {} + tmp_cache = {cache_key: {}} fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{actual_config_name}.json" # Try loading architecture-specific config first if not _load_config_file( - get_mhc_config._config_cache, + tmp_cache, cache_key, fpath, "default", @@ -79,12 +81,13 @@ def get_mhc_config( f"{AITER_TRITON_CONFIGS_PATH}/{fallback_dev}-{actual_config_name}.json" ) _load_config_file( - get_mhc_config._config_cache, + tmp_cache, cache_key, fpath_fallback, "default", fpath_should_exist=True, ) + get_mhc_config._config_cache[cache_key] = tmp_cache[cache_key] config_dict_key = "default" used_specialized = False From 7160dd7bf7e0d263162cb9bdcab8132f0a713e8c Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Fri, 7 Aug 2026 17:26:21 +0000 Subject: [PATCH 2/6] formatting --- aiter/ops/triton/gluon/gemm_a8w8.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/triton/gluon/gemm_a8w8.py b/aiter/ops/triton/gluon/gemm_a8w8.py index 0c017f860b6..da5d00f7e57 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8.py +++ b/aiter/ops/triton/gluon/gemm_a8w8.py @@ -5,8 +5,8 @@ from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd -from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config from aiter.ops.triton.utils.device_info import get_num_xcds +from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() From 8cc8711b2a752e0d591ef1c868270cf90099ffc0 Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Fri, 7 Aug 2026 18:44:43 +0000 Subject: [PATCH 3/6] config file and loader changes, remove unused configs --- .../_triton_kernels/attention/lean_atten.py | 6 +- .../gemm/basic/gemm_a16w16_gated.py | 2 +- .../fused_gemm_a8w8_blockscale_split_cat.py | 4 +- aiter/ops/triton/attention/lean_atten.py | 2 +- aiter/ops/triton/configs/CLAUDE.md | 32 ++-- ...fx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json | 80 --------- .../gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json | 74 -------- ..._PER_BATCHED_TENSOR_QUANT-N=128-K=512.json | 12 +- ..._PER_BATCHED_TENSOR_QUANT-N=512-K=128.json | 12 +- ...P_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json | 12 +- ..._PER_BATCHED_TENSOR_QUANT-N=128-K=512.json | 18 +- ..._PER_BATCHED_TENSOR_QUANT-N=512-K=128.json | 18 +- ...ER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json | 18 +- ...P_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json | 18 +- ...ED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json | 74 -------- ...ED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json | 74 -------- ...gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json | 80 --------- ...0-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json | 74 -------- .../gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json | 74 -------- .../triton/gemm/basic/gemm_a16w16_atomic.py | 7 - aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py | 6 +- ...oup_prequant_w_per_batched_tensor_quant.py | 1 - .../gemm/fused/fused_gemm_afp4wfp4_a16w16.py | 3 + .../ops/triton/gluon/gemm_a8w8_blockscale.py | 64 ++++--- aiter/ops/triton/moe/moe_op_gemm_a8w4.py | 12 +- .../triton/utils/_triton/gemm_tune_check.py | 8 +- aiter/ops/triton/utils/conv_config_utils.py | 43 +---- aiter/ops/triton/utils/gemm_config_utils.py | 25 +-- aiter/ops/triton/utils/mhc_config_utils.py | 158 +++++------------- aiter/ops/triton/utils/moe_config_utils.py | 32 ++-- .../triton/bench_moe_gemm_a8w8_blockscale.py | 10 +- op_tests/triton_tests/attention/test_la.py | 5 +- .../attention/test_mla_decode_rope.py | 5 +- .../triton_tests/gemm/basic/test_gemm_a8w8.py | 10 +- 34 files changed, 219 insertions(+), 854 deletions(-) delete mode 100644 aiter/ops/triton/configs/gemm/gfx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json delete mode 100644 aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json diff --git a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py index d5dceca0df7..431abd22155 100644 --- a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py +++ b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py @@ -17,7 +17,6 @@ - """ -import functools import triton import triton.language as tl @@ -29,15 +28,16 @@ # 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(): + # 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", required=True ) return config[ "any" - ].copy() # return a copy to avoid mutation of the shared cached config + ].copy() # fresh copy per call — safe for callers to mutate @triton.jit diff --git a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py index 04214385fb0..990d658c7e6 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py +++ b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py @@ -156,5 +156,5 @@ def _get_config( ): return get_gemm_config( - "GEMM-A16W16-gated", M, N, K, bounds=(64, 128, 256, 512, 2048) + "GEMM-A16W16-gated", M, N, K, bounds=(8, 16, 32, 64, 128, 256, 512, 2048) ) diff --git a/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py b/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py index 14ab60d29ec..c6749e8a08a 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py +++ b/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,13 +618,14 @@ 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, K: int, shuffle: bool = False, ) -> dict: + # No lru_cache: get_gemm_config caches internally and returns a fresh + # deep copy — memoizing here would hand callers a shared mutable dict. shuffle_suffix = "_PRESHUFFLED" if shuffle else "" config_name = f"GEMM-A8W8_BLOCKSCALE{shuffle_suffix}" diff --git a/aiter/ops/triton/attention/lean_atten.py b/aiter/ops/triton/attention/lean_atten.py index 935bc842f70..d41a1d2fc47 100644 --- a/aiter/ops/triton/attention/lean_atten.py +++ b/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 diff --git a/aiter/ops/triton/configs/CLAUDE.md b/aiter/ops/triton/configs/CLAUDE.md index 51356459ffb..839f0d32bfc 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -115,16 +115,20 @@ Consequences to keep in mind: would change which file gfx1250 resolves to — verify that is intended. - Results are cached twice: `functools.lru_cache` on the full argument tuple, plus a per-path cache of parsed JSON - (`_load_json`) that also caches negative results (missing files). Adding a - config file at runtime therefore has no effect; restart the process - (tooling may call `_load_json.cache_clear()` instead). - -Direct-path loaders bypass all of this. Grep for -`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything — `gluon/gemm_a8w8.py` -and `gluon/gemm_a8w8_blockscale.py` still build legacy `gemm/gluon/` paths by -hand and must be edited when their configs move. `gluon/gemm_afp4wfp4.py` -builds the nested `/gluon/gemm/gemm_afp4wfp4/DEFAULT.json` path by hand -and must be kept in sync with any future layout change. + (`utils/core.py::load_config_json`) that also caches negative results + (missing files). Adding a config file at runtime therefore has no effect; + restart the process (tooling may call `load_config_json.cache_clear()` + instead). +- `.github/scripts/select_triton_tests.py` only globs the legacy flat + `configs/gemm/` — configs in the nested layout are invisible to CI test + selection until it learns the `/` layout. + +Direct-path loaders bypass the resolver's directory probe. Grep for +`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything — +`gluon/gemm_a8w8_blockscale.py` still builds legacy `gemm/gluon/` paths by +hand (via `load_config_json`) and must be edited when its configs move. +`gluon/gemm_a8w8.py` and `gluon/gemm_afp4wfp4.py` go through +`get_gemm_config(backend="gluon")` and need no changes. --- @@ -141,9 +145,9 @@ Required top-level shape: ``` - `M_LEQ_x` is searched ascending over `STANDARD_M_BOUNDS = - (4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)`, then `M_GEQ_x` - descending, then `any`. A caller may override with `bounds=(...)`, which must - be strictly increasing positive ints. + (1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)`, then + `M_GEQ_x` descending, then `any`. A caller may override with + `bounds=(...)`, which must be strictly increasing positive ints. - `any` must exist unless every reachable `M` is covered by an explicit bound. - The deprecated `{"large": ..., "small": ...}` shape must not be introduced. - A `KeyError` at lookup time means no bound matched — usually a missing `any`. @@ -194,7 +198,7 @@ Dashes, underscores, and case all fold together, so new config names must stay distinct under that transform — `GEMM-FOO-BAR` and `GEMM-FOO_BAR` would collide. Config-name patterns: `GEMM-A{x}W{y}`, `BATCHED_GEMM-A{x}W{y}`, -`GEMM_PREQUANT-...`, `FUSED-GEMM-{op}`, `FF-A{x}W{y}-fused`; variant suffixes +`FUSED-GEMM-{op}`, `FF-A{x}W{y}-fused`; variant suffixes `_PRESHUFFLED`, `_BLOCKSCALE`. **`K` in AFP4WFP4 filenames is the logical K, i.e. `2 * K_bytes`.** The kernel diff --git a/aiter/ops/triton/configs/gemm/gfx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json b/aiter/ops/triton/configs/gemm/gfx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json deleted file mode 100644 index a0a550dcab8..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json +++ /dev/null @@ -1,80 +0,0 @@ -{ - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 6, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 2, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 2, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "any": { - "BLOCK_SIZE_M": 256, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 64, - "num_warps": 8, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": null, - "NUM_KSPLIT": 1 - } -} diff --git a/aiter/ops/triton/configs/gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json b/aiter/ops/triton/configs/gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json deleted file mode 100644 index 689eb502350..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json +++ /dev/null @@ -1,74 +0,0 @@ -{ - "M_LEQ_8": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "any": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": null, - "NUM_KSPLIT": 4 - } -} \ No newline at end of file diff --git a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json index 7a4a48cfb5a..6d5ec6e3ca8 100644 --- a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json +++ b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json @@ -7,7 +7,7 @@ "num_stages": 2, "waves_per_eu": 4, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_32": { @@ -18,7 +18,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_64": { @@ -29,7 +29,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_128": { @@ -40,7 +40,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_256": { @@ -51,7 +51,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "any": { @@ -62,7 +62,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json index 629d62d9c83..222f27b6d1c 100644 --- a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json +++ b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json @@ -7,7 +7,7 @@ "num_stages": 2, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_32": { @@ -18,7 +18,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_64": { @@ -29,7 +29,7 @@ "num_stages": 1, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_128": { @@ -40,7 +40,7 @@ "num_stages": 1, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_256": { @@ -51,7 +51,7 @@ "num_stages": 1, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "any": { @@ -62,7 +62,7 @@ "num_stages": 1, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json index 372268c2605..43f7876e181 100644 --- a/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json +++ b/aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json @@ -7,7 +7,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_32": { @@ -18,7 +18,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_64": { @@ -29,7 +29,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_128": { @@ -40,7 +40,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "M_LEQ_256": { @@ -51,7 +51,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" }, "any": { @@ -62,7 +62,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "kpack": 2, + "kpack": 1, "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json index 1124e1b843a..ed3d040ccdd 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json @@ -7,7 +7,8 @@ "num_stages": 2, "waves_per_eu": 4, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_32": { "BLOCK_SIZE_M": 16, @@ -17,7 +18,8 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_64": { "BLOCK_SIZE_M": 32, @@ -27,7 +29,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -37,7 +40,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": null + "cache_modifier": null, + "kpack": 1 }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -47,7 +51,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": null + "cache_modifier": null, + "kpack": 1 }, "any": { "BLOCK_SIZE_M": 32, @@ -57,6 +62,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json index cfc0b934e5c..d2bd791d13f 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json @@ -7,7 +7,8 @@ "num_stages": 2, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -17,7 +18,8 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_64": { "BLOCK_SIZE_M": 32, @@ -27,7 +29,8 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_128": { "BLOCK_SIZE_M": 32, @@ -37,7 +40,8 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_256": { "BLOCK_SIZE_M": 32, @@ -47,7 +51,8 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "any": { "BLOCK_SIZE_M": 32, @@ -57,6 +62,7 @@ "num_stages": 1, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json index 01ecea0d300..b2d817abd49 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json @@ -7,7 +7,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -17,7 +18,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_64": { "BLOCK_SIZE_M": 64, @@ -27,7 +29,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -37,7 +40,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -47,7 +51,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "any": { "BLOCK_SIZE_M": 32, @@ -57,6 +62,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json index e6f1bd4b192..302eb89b9d2 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json @@ -7,7 +7,8 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -17,7 +18,8 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_64": { "BLOCK_SIZE_M": 64, @@ -27,7 +29,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -37,7 +40,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -47,7 +51,8 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 }, "any": { "BLOCK_SIZE_M": 32, @@ -57,6 +62,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg" + "cache_modifier": ".cg", + "kpack": 1 } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json deleted file mode 100644 index 1d1df58c1d8..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json +++ /dev/null @@ -1,74 +0,0 @@ -{ - "M_LEQ_16": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 8, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "any": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "cache_modifier": null, - "NUM_KSPLIT": 1 - } -} diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json deleted file mode 100644 index f9eba0dfe69..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json +++ /dev/null @@ -1,74 +0,0 @@ -{ - "M_LEQ_16": { - "BLOCK_SIZE_M": 4, - "BLOCK_SIZE_N": 32, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 0, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 0, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 0, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 0, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "any": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 0, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - } -} diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json deleted file mode 100644 index a0a550dcab8..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json +++ /dev/null @@ -1,80 +0,0 @@ -{ - "M_LEQ_16": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 6, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 32, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 2, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 2, - "num_warps": 4, - "num_stages": 2, - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": ".cg", - "NUM_KSPLIT": 1 - }, - "any": { - "BLOCK_SIZE_M": 256, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 256, - "GROUP_SIZE_M": 64, - "num_warps": 8, - "num_stages": 2, - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "kpack": 1, - "cache_modifier": null, - "NUM_KSPLIT": 1 - } -} diff --git a/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json b/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json deleted file mode 100644 index b5b4b7b66d9..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json +++ /dev/null @@ -1,74 +0,0 @@ -{ - "M_LEQ_8": { - "BLOCK_SIZE_M": 4, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - }, - "any": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 14 - } -} diff --git a/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json b/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json deleted file mode 100644 index 70042715e3d..00000000000 --- a/aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json +++ /dev/null @@ -1,74 +0,0 @@ -{ - "M_LEQ_8": { - "BLOCK_SIZE_M": 4, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_32": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_64": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_128": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "M_LEQ_256": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "NUM_KSPLIT": 4 - }, - "any": { - "BLOCK_SIZE_M": 8, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 512, - "GROUP_SIZE_M": 1, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 2, - "matrix_instr_nonkdim": 16, - "cache_modifier": null, - "NUM_KSPLIT": 4 - } -} diff --git a/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py b/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py index 3faf5c1c931..bf0fd47ee78 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py +++ b/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py @@ -67,13 +67,6 @@ def gemm_a16w16_atomic_( else: config = deserialize_str(config) - # For compatability reasons, these keys may not exist in the config - # TODO: This needs to be embedded in the configs later - if "NUM_KSPLIT" not in config: - config["NUM_KSPLIT"] = 1 - if "cache_modifier" not in config: - config["cache_modifier"] = "" - if y is None: # atomic add requires 0 tensor if config["NUM_KSPLIT"] == 1: diff --git a/aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py b/aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py index 7719207dd99..683327928f4 100644 --- a/aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py +++ b/aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py @@ -456,11 +456,11 @@ def gemm_afp4wfp4_preshuffle( n16, _ = w_preshuf.shape N = n16 * 16 K_elems = 2 * K_bytes - # _get_config doubles K for config - 2 * K_bytes == K_elems - K_cfg = K_elems if config is None: - config, _ = _get_config(M, N, K_cfg, True) + # _get_config doubles K itself (logical K = 2 * K_bytes) — pass bytes, + # matching the non-preshuffled path. + config, _ = _get_config(M, N, K_bytes, True) config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32) if M < 32: diff --git a/aiter/ops/triton/gemm/batched/batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py b/aiter/ops/triton/gemm/batched/batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py index 591526e289e..365e67fcf67 100644 --- a/aiter/ops/triton/gemm/batched/batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py +++ b/aiter/ops/triton/gemm/batched/batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py @@ -88,7 +88,6 @@ def batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( if config is None: config, _ = _get_config(M, N, K) config["BLOCK_SIZE_K"] = group_size - config["kpack"] = 1 grid = lambda META: ( B, diff --git a/aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_a16w16.py b/aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_a16w16.py index f7c8f6a5dcc..94d9865f26a 100644 --- a/aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_a16w16.py +++ b/aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_a16w16.py @@ -13,6 +13,7 @@ _get_config, ) from aiter.ops.triton.gemm.basic.gemm_afp4wfp4 import get_splitk +from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH from aiter.ops.triton.utils.logger import AiterTritonLogger from aiter.utility.triton.triton_metadata_redirect import AOTMetadataContext @@ -64,6 +65,8 @@ def fused_gemm_afp4wfp4_a16w16( f"FUSED_GEMM_A8W8_BLOCKSCALE_A16W16: x_fp4={tuple(x_fp4.shape)} w_fp4={tuple(w_fp4.shape)} x_fp4_scale={tuple(x_fp4_scale.shape)} w_fp4_scale={tuple(w_fp4_scale.shape)} x_bf16={tuple(x_bf16.shape)} w_bf16={tuple(w_bf16.shape)}" ) + assert arch_info.is_fp4_avail(), "MXFP4 is not available on your device" + M, K = x_fp4.shape N_fp4, K = w_fp4.shape if is_fp4_preshuffled: diff --git a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py index 4de625992eb..9d2763d771a 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py @@ -2,8 +2,6 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import functools -import json -import os import torch import triton @@ -14,7 +12,7 @@ from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, 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 from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -982,38 +980,29 @@ def _gemm_a8w8_blockscale_reduce_kernel( @functools.lru_cache(maxsize=1024) -def _get_config( +def _get_config_cached( M: int, N: int, K: int, ): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - if int(dev.split("gfx")[1]) < 950: - raise ValueError( - "Gluon implementation is not supported on this device (requires CDNA4)." - ) - fpath = ( - f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json" + dev = arch_info.get_arch() + if not arch_info.is_gluon_avail(): + raise ValueError( + "Gluon implementation is not supported on this device (requires CDNA4)." ) - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = {"default": config} - - key = f"{N}_{K}" - if key not in _get_config._config_dict: - dev = arch_info.get_arch() - fpath = f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json" - if os.path.exists(fpath): - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict[key] = config - else: - key = "default" # fall back to default config + config_dict = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json", + required=True, + ) + specialized = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json" + ) + if specialized is not None: + config_dict = specialized # Config keys should be named M_LEQ_ or "any" bounds = [] - for setting in _get_config._config_dict[key]: + for setting in config_dict: potential_block_m = setting.replace("M_LEQ_", "") if potential_block_m.isnumeric(): bounds.append(int(potential_block_m)) @@ -1022,18 +1011,27 @@ def _get_config( # the kernel currently supports. Unsupported buckets are skipped (those # configs become live again once the kernel grows the corresponding # padded-LDS layouts), so we may fall through to "any". - config = _get_config._config_dict[key]["any"] + config = config_dict["any"] for bound in sorted(bounds): - if M > bound or f"M_LEQ_{bound}" not in _get_config._config_dict[key]: + if M > bound or f"M_LEQ_{bound}" not in config_dict: continue - candidate = _get_config._config_dict[key][f"M_LEQ_{bound}"] + candidate = config_dict[f"M_LEQ_{bound}"] if (candidate["BLOCK_SIZE_M"], candidate["BLOCK_SIZE_N"]) in _SUPPORTED_TILES: config = candidate break - config = ( - config.copy() - ) # avoid later inplace modification from interacting with cached config + return config + + +def _get_config( + M: int, + N: int, + K: int, +): + # Fresh copy per call, outside the lru boundary — the caller writes + # derived fields (SPLITK_BLOCK_SIZE here, GROUP_K/GROUP_N at the call + # site) into the returned dict. + config = dict(_get_config_cached(M, N, K)) block_size_k = config["BLOCK_SIZE_K"] num_k_blocks = triton.cdiv(K, block_size_k) diff --git a/aiter/ops/triton/moe/moe_op_gemm_a8w4.py b/aiter/ops/triton/moe/moe_op_gemm_a8w4.py index 9080783db18..b7554b74bb3 100644 --- a/aiter/ops/triton/moe/moe_op_gemm_a8w4.py +++ b/aiter/ops/triton/moe/moe_op_gemm_a8w4.py @@ -1,10 +1,7 @@ # adapted from triton_kernels package # original code https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/matmul_ogs.py -import functools import itertools -import json -import os import warnings import torch @@ -25,21 +22,16 @@ from aiter.ops.triton.moe.moe_routing.routing import RoutingData from aiter.ops.triton.moe.reduce import reduce_grouped from aiter.ops.triton.utils._triton.arch_info import get_arch -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 from aiter.ops.triton.utils.device_info import get_num_sms from aiter.ops.triton.utils.gemm_config_utils import pick_gemm_num_stages -@functools.lru_cache def _get_a8w4_dispatch(arch: str) -> dict: """Per-(block_m, N, K) dispatch table for moe_gemm_a8w4. Returns {} if no tuned file is shipped for this arch (caller uses the safe-default fallback). Mirrors get_moe_configs() in utils/moe_config_utils.py.""" - fpath = f"{AITER_TRITON_CONFIGS_PATH}/moe/{arch}-A8W4.json" - if os.path.exists(fpath): - with open(fpath, "r") as f: - return json.load(f) - return {} + return load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/moe/{arch}-A8W4.json") or {} def can_overflow_int32(tensor: torch.Tensor): diff --git a/aiter/ops/triton/utils/_triton/gemm_tune_check.py b/aiter/ops/triton/utils/_triton/gemm_tune_check.py index 47719395b94..b8033d88c3f 100644 --- a/aiter/ops/triton/utils/_triton/gemm_tune_check.py +++ b/aiter/ops/triton/utils/_triton/gemm_tune_check.py @@ -19,19 +19,21 @@ def gemm_tune_check( example 1: FP4 GEMM preshuffled weight scales for shape (16, 1280, 8192) from aiter.ops.triton.utils._triton.gemm_tune_check import gemm_tune_check - from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4_preshuffle + from aiter.ops.triton.gemm.basic.gemm_afp4wfp4 import gemm_afp4wfp4_preshuffle is_tunned = gemm_tune_check(gemm_afp4wfp4_preshuffle, N=1280, K=8192//2, M=16, shuffle=True) print(is_tunned) # return True or False example 2: FP8 GEMM blockscale for shape (16, 1280, 8192) from aiter.ops.triton.utils._triton.gemm_tune_check import gemm_tune_check - from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale + from aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale import gemm_a8w8_blockscale is_tunned = gemm_tune_check(gemm_a8w8_blockscale, N=1024, K=8192, M=16) print(is_tunned) # return True or False """ + # Public wrappers live at aiter.ops.triton..; their kernel + # modules mirror that path under aiter.ops.triton._triton_kernels. module_pth = func.__module__.split(".") - module_pth = module_pth[:-1] + ["_triton_kernels"] + module_pth[-1:] + module_pth = module_pth[:3] + ["_triton_kernels"] + module_pth[3:] module_pth = ".".join(module_pth) module = importlib.import_module(module_pth) _LOGGER.info(f"Function {func} found at {module}") diff --git a/aiter/ops/triton/utils/conv_config_utils.py b/aiter/ops/triton/utils/conv_config_utils.py index 5530189fc4f..f64d32ee1d8 100644 --- a/aiter/ops/triton/utils/conv_config_utils.py +++ b/aiter/ops/triton/utils/conv_config_utils.py @@ -3,13 +3,13 @@ import copy import functools -import json -import os from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH - -USE_LRU_CACHE = True +from aiter.ops.triton.utils.core import ( + AITER_TRITON_CONFIGS_PATH, + USE_LRU_CACHE, + load_config_json, +) STANDARD_M_BOUNDS: tuple[int, ...] = ( 4, @@ -56,21 +56,6 @@ def format_shape_key( ) -def _load_config_file( - cache_dict: dict, - cache_key: str, - fpath: str, - fpath_should_exist: bool = True, -) -> bool: - if os.path.exists(fpath): - with open(fpath, "r") as file: - cache_dict[cache_key] = json.load(file) - return True - elif fpath_should_exist: - raise AssertionError(f"Required config file doesn't exist: {fpath}") - return False - - @functools.lru_cache(maxsize=512 if USE_LRU_CACHE else 0) def _get_conv_config_cached( config_name: str, @@ -78,22 +63,10 @@ def _get_conv_config_cached( M: int | None, ) -> dict: """Three-tier walk: literal shape entry -> M_LEQ bucket -> 'any'.""" - if not hasattr(_get_conv_config_cached, "_file_cache"): - _get_conv_config_cached._file_cache = {} - dev = arch_info.get_arch() - file_cache_key = f"{dev}_{config_name}" - - if file_cache_key not in _get_conv_config_cached._file_cache: - fpath = f"{AITER_TRITON_CONFIGS_PATH}/conv/{dev}-{config_name}.json" - _load_config_file( - _get_conv_config_cached._file_cache, - file_cache_key, - fpath, - fpath_should_exist=True, - ) - - config_dict = _get_conv_config_cached._file_cache[file_cache_key] + config_dict = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/conv/{dev}-{config_name}.json", required=True + ) # Tier 1: literal shape key. shapes = config_dict.get("shapes", {}) diff --git a/aiter/ops/triton/utils/gemm_config_utils.py b/aiter/ops/triton/utils/gemm_config_utils.py index 06c13513419..adf9d8357e4 100644 --- a/aiter/ops/triton/utils/gemm_config_utils.py +++ b/aiter/ops/triton/utils/gemm_config_utils.py @@ -1,7 +1,6 @@ import copy import functools import itertools -import json import os import triton @@ -14,7 +13,7 @@ ) # Standard bounds for M_LEQ_x keys (tuple for hashability with LRU cache) -STANDARD_M_BOUNDS = (4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) +STANDARD_M_BOUNDS = (1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) def _dtype_dir(config_name: str) -> str: @@ -23,28 +22,6 @@ def _dtype_dir(config_name: str) -> str: return config_name.lower().replace("-", "_") -def _load_config_file( - cache_dict: dict, - cache_key: str, - fpath: str, - config_key: str, - fpath_should_exist: bool = False, -) -> bool: - """ - Helper function to load a config file and cache it. - - Retained for mhc_config_utils.py — the GEMM resolver uses _load_json(). - """ - if os.path.exists(fpath): - with open(fpath, "r") as file: - config = json.load(file) - cache_dict[cache_key][config_key] = config - return True - elif fpath_should_exist: - raise AssertionError(f"Required config file doesn't exist: {fpath}") - return False - - @functools.lru_cache(maxsize=1024 if USE_LRU_CACHE else 0) def _get_gemm_config_cached( config_name: str, diff --git a/aiter/ops/triton/utils/mhc_config_utils.py b/aiter/ops/triton/utils/mhc_config_utils.py index c7de4975142..9671dc06545 100644 --- a/aiter/ops/triton/utils/mhc_config_utils.py +++ b/aiter/ops/triton/utils/mhc_config_utils.py @@ -3,13 +3,42 @@ import functools import glob -import json import os import re from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH -from aiter.ops.triton.utils.gemm_config_utils import USE_LRU_CACHE, _load_config_file +from aiter.ops.triton.utils.core import ( + AITER_TRITON_CONFIGS_PATH, + USE_LRU_CACHE, + load_config_json, +) + +_FALLBACK_DEV = "gfx942" + + +def _load_with_fallback(dev: str, fname: str, required: bool = False) -> dict | None: + """Load ``{dev}-{fname}``, falling back to the gfx942 copy for arches + without tuned MHC configs (may be suboptimal).""" + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{fname}") + if config is None: + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{_FALLBACK_DEV}-{fname}", required=required + ) + return config + + +@functools.lru_cache(maxsize=None if USE_LRU_CACHE else 0) +def _c_thresholds(dev: str, actual_config_name: str) -> tuple[int, ...]: + """C values that have a specialized config file (arch-specific plus the + gfx942 fallback), sorted ascending.""" + thresholds = set() + for d in {dev, _FALLBACK_DEV}: + pattern = f"{AITER_TRITON_CONFIGS_PATH}/{d}-{actual_config_name}-C=*.json" + for fpath in glob.glob(pattern): + match = re.search(r"-C=(\d+)\.json$", os.path.basename(fpath)) + if match: + thresholds.add(int(match.group(1))) + return tuple(sorted(thresholds)) @functools.lru_cache(maxsize=1024 if USE_LRU_CACHE else 0) @@ -49,107 +78,27 @@ def get_mhc_config( ValueError: If mode is invalid or missing when required KeyError: If no matching config found """ - if not hasattr(get_mhc_config, "_config_cache"): - get_mhc_config._config_cache = {} - dev = arch_info.get_arch() - fallback_dev = "gfx942" if mode is None or mode != "sinkhorn": raise ValueError(f"mode must be 'sinkhorn', got '{mode}'") actual_config_name = f"{config_name}_{mode.upper()}" - cache_key = f"{dev}_{actual_config_name}" - - # Load default config with fallback for unsupported architectures. - # Load into a temp dict and commit only on success, so a failed load - # doesn't leave a stale empty entry that masks the error on later calls. - if cache_key not in get_mhc_config._config_cache: - tmp_cache = {cache_key: {}} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{actual_config_name}.json" - - # Try loading architecture-specific config first - if not _load_config_file( - tmp_cache, - cache_key, - fpath, - "default", - fpath_should_exist=False, - ): - # Fallback to gfx942 configs if architecture-specific config doesn't exist - fpath_fallback = ( - f"{AITER_TRITON_CONFIGS_PATH}/{fallback_dev}-{actual_config_name}.json" - ) - _load_config_file( - tmp_cache, - cache_key, - fpath_fallback, - "default", - fpath_should_exist=True, - ) - get_mhc_config._config_cache[cache_key] = tmp_cache[cache_key] - - config_dict_key = "default" + # Default config (must exist for the arch or the gfx942 fallback) + config_dict = _load_with_fallback(dev, f"{actual_config_name}.json", required=True) used_specialized = False - # Try C-specific config (threshold matching: largest C <= input C) - c_thresholds_key = f"{cache_key}_c_thresholds" - - # Discover available C-specific config files once per cache_key - if c_thresholds_key not in get_mhc_config._config_cache: - c_thresholds = [] - - # Check architecture-specific C configs - pattern = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{actual_config_name}-C=*.json" - for fpath in glob.glob(pattern): - basename = os.path.basename(fpath) - match = re.search(r"-C=(\d+)\.json$", basename) - if match: - c_thresholds.append(int(match.group(1))) - - # Also check fallback architecture C configs - if dev != fallback_dev: - pattern_fallback = f"{AITER_TRITON_CONFIGS_PATH}/{fallback_dev}-{actual_config_name}-C=*.json" - for fpath in glob.glob(pattern_fallback): - basename = os.path.basename(fpath) - match = re.search(r"-C=(\d+)\.json$", basename) - if match: - c_val = int(match.group(1)) - if c_val not in c_thresholds: - c_thresholds.append(c_val) - - c_thresholds.sort() - get_mhc_config._config_cache[c_thresholds_key] = c_thresholds - - # Find largest C threshold <= input C - for c_threshold in reversed(get_mhc_config._config_cache[c_thresholds_key]): + # C-specific config: largest discovered threshold <= input C wins + for c_threshold in reversed(_c_thresholds(dev, actual_config_name)): if C >= c_threshold: - c_key = f"C_{c_threshold}" - if c_key not in get_mhc_config._config_cache[cache_key]: - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{actual_config_name}-C={c_threshold}.json" - # Try architecture-specific C config first, fallback to gfx942 if needed - if not _load_config_file( - get_mhc_config._config_cache, - cache_key, - fpath, - c_key, - fpath_should_exist=False, - ): - fpath_fallback = f"{AITER_TRITON_CONFIGS_PATH}/{fallback_dev}-{actual_config_name}-C={c_threshold}.json" - _load_config_file( - get_mhc_config._config_cache, - cache_key, - fpath_fallback, - c_key, - fpath_should_exist=False, - ) - if c_key in get_mhc_config._config_cache[cache_key]: - config_dict_key = c_key + specialized = _load_with_fallback( + dev, f"{actual_config_name}-C={c_threshold}.json" + ) + if specialized is not None: + config_dict = specialized used_specialized = True break - config_dict = get_mhc_config._config_cache[cache_key][config_dict_key] - # Extract M_LEQ_x keys and their thresholds, sorted ascending m_leq_keys = [] for key in config_dict: @@ -187,29 +136,10 @@ def get_mhc_post_config(M: int, C: int) -> dict: Picks the largest ``C_ <= C``, else ``"default"``. """ - if M <= 1024 and C == 4096: - return { - "BLOCK_M": 8, - "BLOCK_C": 512, - "num_warps": 8, - "num_stages": 1, - "waves_per_eu": 3, - } - - if not hasattr(get_mhc_post_config, "_file_cache"): - get_mhc_post_config._file_cache = {} - dev = arch_info.get_arch() - if dev not in get_mhc_post_config._file_cache: - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHC_POST.json" - if not os.path.exists(fpath): - raise FileNotFoundError( - f"Required MHC_POST config file doesn't exist: {fpath}" - ) - with open(fpath, "r") as f: - get_mhc_post_config._file_cache[dev] = json.load(f) - - cfg = get_mhc_post_config._file_cache[dev] + cfg = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHC_POST.json", required=True + ) c_thresholds = sorted( int(k[2:]) for k in cfg if k.startswith("C_") and k[2:].isdigit() diff --git a/aiter/ops/triton/utils/moe_config_utils.py b/aiter/ops/triton/utils/moe_config_utils.py index d237abf15bd..f771a8c767f 100644 --- a/aiter/ops/triton/utils/moe_config_utils.py +++ b/aiter/ops/triton/utils/moe_config_utils.py @@ -2,15 +2,13 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import functools -import json -import os import warnings from typing import Any import torch from ._triton import arch_info -from .core import AITER_TRITON_CONFIGS_PATH +from .core import AITER_TRITON_CONFIGS_PATH, USE_LRU_CACHE, load_config_json M_THRESHOLD_SMALL = 256 M_THRESHOLD_MEDIUM = 1024 @@ -41,7 +39,7 @@ def get_config_dtype_str( return None -@functools.lru_cache +@functools.lru_cache(maxsize=1024 if USE_LRU_CACHE else 0) def get_moe_configs(dtype: str | None) -> dict[int, Any] | None: """ Return optimized configurations for the fused MoE kernel. @@ -55,12 +53,11 @@ def get_moe_configs(dtype: str | None) -> dict[int, Any] | None: # directory dtype_str = "DEFAULT" if dtype is None else dtype dev = arch_info.get_arch() - config_file_path = f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE-{dtype_str}.json" - - if os.path.exists(config_file_path): - with open(config_file_path) as f: - # If a configuration has been found, return it - return {key: val for key, val in json.load(f).items()} + configs = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE-{dtype_str}.json" + ) + if configs is not None: + return configs # If no optimized configuration is available, we will use the default # configuration @@ -84,14 +81,13 @@ def get_optimal_moe_config( ) # print(f"dtype_str={dtype_str}") configs = get_moe_configs(dtype_str) - if configs is not None: - if configs: - if M < M_THRESHOLD_SMALL: - config = configs["small_M"] - elif M < M_THRESHOLD_MEDIUM: - config = configs["medium_M"] - else: - config = configs["large_M"] + if configs: + if M < M_THRESHOLD_SMALL: + config = configs["small_M"] + elif M < M_THRESHOLD_MEDIUM: + config = configs["medium_M"] + else: + config = configs["large_M"] else: # default config config = { diff --git a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py index 470e52528dc..b8515eb0676 100644 --- a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py +++ b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py @@ -11,10 +11,11 @@ import torch import triton.profiler as proton -from aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16w16 import ( - _get_config, -) from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 +from aiter.ops.triton.utils.gemm_config_utils import ( + compute_splitk_params, + get_gemm_config, +) from aiter.ops.triton.moe.moe_op_gemm_a8w8_blockscale import ( moe_gemm_a8w8_blockscale, ) @@ -215,7 +216,8 @@ def num_blocks(length, block): fpath = Path(tempfile.mktemp()) M, K = xg.shape K, N = wg.shape - config, _ = _get_config(M, N, K) + config, _ = get_gemm_config("GEMM-A16W16", M, N, K) + compute_splitk_params(config, K) config["BLOCK_SIZE_M"] = min(config["BLOCK_SIZE_M"], 128) config["BLOCK_SIZE_N"] = min(config["BLOCK_SIZE_N"], 128) config["BLOCK_SIZE_K"] = min(config["BLOCK_SIZE_K"], 128) diff --git a/op_tests/triton_tests/attention/test_la.py b/op_tests/triton_tests/attention/test_la.py index b487ff4c7d2..fd9a41e23ce 100644 --- a/op_tests/triton_tests/attention/test_la.py +++ b/op_tests/triton_tests/attention/test_la.py @@ -456,10 +456,7 @@ def test_persistent_lean_attention_outer( ): torch.manual_seed(20) - config = _get_config( - batch_size=batch, - causal=causal, - ) + config = _get_config() sm_count = arch_info.get_num_sms() # Long seqlen (>512K) can hit memory access fault. Suspect compiler issue diff --git a/op_tests/triton_tests/attention/test_mla_decode_rope.py b/op_tests/triton_tests/attention/test_mla_decode_rope.py index df66d048d72..4ec6c83ce91 100644 --- a/op_tests/triton_tests/attention/test_mla_decode_rope.py +++ b/op_tests/triton_tests/attention/test_mla_decode_rope.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import copy from typing import Any import pytest @@ -245,7 +246,9 @@ def ref_compute_full_fwd( def get_config(dtype: torch.dtype): - config: dict[str, Any] = _get_config() + # _get_config() returns the shared cached dict and the launch helpers + # write derived fields into its sub-configs — work on a copy. + config: dict[str, Any] = copy.deepcopy(_get_config()) base_config_key: str = "fwd_grouped_kernel_stage1_rope" fp32_config_key: str = f"{base_config_key}_fp32" config_key: str = ( diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py index 259de53205d..09ef87ad8b3 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py @@ -6,7 +6,6 @@ import torch.nn.functional as F from aiter.ops.shuffle import shuffle_weight -from aiter.ops.triton._triton_kernels.gemm.basic.gemm_a8w8 import _get_config from aiter.ops.triton.gemm.basic.gemm_a8w8 import gemm_a8w8 as triton_gemm_a8w8 from aiter.ops.triton.gluon.gemm_a8w8 import ( gemm_a8w8 as gluon_gemm_a8w8, @@ -15,7 +14,10 @@ gemm_a8w8_preshuffle as gluon_gemm_a8w8_preshuffle, ) from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.gemm_config_utils import compute_splitk_params +from aiter.ops.triton.utils.gemm_config_utils import ( + compute_splitk_params, + get_gemm_config, +) from aiter.ops.triton.utils.types import get_fp8_dtypes, str_to_torch_dtype DEVICE_ARCH = arch_info.get_arch() @@ -300,7 +302,7 @@ def test_gemm_splitk(in_dtype, out_dtype, m, n, k, num_ksplit, has_bias): if not has_bias: bias = None - config, _ = _get_config(m, n, k) + config, _ = get_gemm_config("GEMM-A8W8", m, n, k) config["NUM_KSPLIT"] = num_ksplit compute_splitk_params(config, k) @@ -351,7 +353,7 @@ def test_gemm_splitk_skip_reduce(in_dtype, out_dtype, m, n, k, num_ksplit): output=False, ) - config, _ = _get_config(m, n, k) + config, _ = get_gemm_config("GEMM-A8W8", m, n, k) config["NUM_KSPLIT"] = num_ksplit compute_splitk_params(config, k) From 6cdf0daa6b7b1accc5d76e07545082c4787887fe Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Fri, 7 Aug 2026 19:03:56 +0000 Subject: [PATCH 4/6] format --- aiter/ops/triton/_triton_kernels/attention/lean_atten.py | 5 +---- .../_triton_kernels/gemm/basic/gemm_a16w16_gated.py | 2 +- .../gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py | 2 -- aiter/ops/triton/gluon/gemm_afp4wfp4.py | 2 -- .../triton/bench_moe_gemm_a8w8_blockscale.py | 8 ++++---- 5 files changed, 6 insertions(+), 13 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py index 431abd22155..650ae5a42df 100644 --- a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py +++ b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py @@ -17,7 +17,6 @@ - """ - import triton import triton.language as tl @@ -35,9 +34,7 @@ def _get_config(): config = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json", required=True ) - return config[ - "any" - ].copy() # fresh copy per call — safe for callers to mutate + return config["any"].copy() # fresh copy per call — safe for callers to mutate @triton.jit diff --git a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py index 990d658c7e6..04214385fb0 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py +++ b/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16_gated.py @@ -156,5 +156,5 @@ def _get_config( ): return get_gemm_config( - "GEMM-A16W16-gated", M, N, K, bounds=(8, 16, 32, 64, 128, 256, 512, 2048) + "GEMM-A16W16-gated", M, N, K, bounds=(64, 128, 256, 512, 2048) ) diff --git a/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py b/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py index c6749e8a08a..2121b0f8c2b 100644 --- a/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py +++ b/aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py @@ -624,8 +624,6 @@ def _get_config( K: int, shuffle: bool = False, ) -> dict: - # No lru_cache: get_gemm_config caches internally and returns a fresh - # deep copy — memoizing here would hand callers a shared mutable dict. shuffle_suffix = "_PRESHUFFLED" if shuffle else "" config_name = f"GEMM-A8W8_BLOCKSCALE{shuffle_suffix}" diff --git a/aiter/ops/triton/gluon/gemm_afp4wfp4.py b/aiter/ops/triton/gluon/gemm_afp4wfp4.py index 05b5478901c..da39234bb9b 100644 --- a/aiter/ops/triton/gluon/gemm_afp4wfp4.py +++ b/aiter/ops/triton/gluon/gemm_afp4wfp4.py @@ -474,8 +474,6 @@ def _get_config( ): if arch_info.get_arch() not in ["gfx950", "gfx1250"]: raise ValueError("Gluon implementation is not supported on this device.") - # get_gemm_config caches internally and returns a fresh deep copy, safe - # for the callers below that write derived fields into the config. config, _ = get_gemm_config("GEMM-AFP4WFP4", M, N, K, backend="gluon") return config diff --git a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py index b8515eb0676..c7b5f650185 100644 --- a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py +++ b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py @@ -12,15 +12,15 @@ import triton.profiler as proton from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 -from aiter.ops.triton.utils.gemm_config_utils import ( - compute_splitk_params, - get_gemm_config, -) from aiter.ops.triton.moe.moe_op_gemm_a8w8_blockscale import ( moe_gemm_a8w8_blockscale, ) from aiter.ops.triton.moe.moe_routing.routing import routing from aiter.ops.triton.utils._triton.arch_info import get_arch +from aiter.ops.triton.utils.gemm_config_utils import ( + compute_splitk_params, + get_gemm_config, +) # Default group_m, group_n, group_k group_shape = (128, 128, 128) From 60097333e7421c55ccfc422cadc7e8b2709b2ddf Mon Sep 17 00:00:00 2001 From: Satya Nikhil Kodukula Date: Mon, 10 Aug 2026 14:52:56 +0000 Subject: [PATCH 5/6] address comments --- .../attention/extend_attention.py | 2 +- .../attention/hstu_attention.py | 2 -- .../_triton_kernels/attention/lean_atten.py | 2 +- .../triton/_triton_kernels/attention/mha.py | 4 +--- .../_triton_kernels/attention/mha_fused_bwd.py | 4 +--- .../attention/mha_onekernel_bwd.py | 4 +--- .../attention/mla_decode_rope.py | 1 - aiter/ops/triton/_triton_kernels/gmm.py | 4 +--- .../moe/moe_routing_sigmoid_top1_fused.py | 1 - aiter/ops/triton/configs/CLAUDE.md | 3 --- ...W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json | 18 ++++++------------ ...W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json | 18 ++++++------------ ...PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json | 18 ++++++------------ ...UP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json | 18 ++++++------------ .../triton/gemm/basic/gemm_a16w16_atomic.py | 4 ++++ aiter/ops/triton/gluon/gemm_a8w8_blockscale.py | 6 +++--- aiter/ops/triton/moe/moe_op_gemm_a8w4.py | 5 ++++- aiter/ops/triton/utils/conv_config_utils.py | 2 +- aiter/ops/triton/utils/core.py | 12 ++++++++---- aiter/ops/triton/utils/gemm_config_utils.py | 4 ++-- aiter/ops/triton/utils/mhc_config_utils.py | 8 ++++---- aiter/ops/triton/utils/moe_config_utils.py | 2 +- 22 files changed, 57 insertions(+), 85 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/extend_attention.py b/aiter/ops/triton/_triton_kernels/attention/extend_attention.py index f3d327a3639..d618431428a 100644 --- a/aiter/ops/triton/_triton_kernels/attention/extend_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/extend_attention.py @@ -322,7 +322,7 @@ def _fwd_kernel( def _get_config(HEAD_SIZE, dtype): dev = arch_info.get_arch() config = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION.json", required=True + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION.json" ) # HEAD_SIZE 192 = 128 head and 64 pe head dim diff --git a/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py b/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py index 65ece9b8fc4..36d77215aa7 100644 --- a/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py +++ b/aiter/ops/triton/_triton_kernels/attention/hstu_attention.py @@ -873,7 +873,6 @@ def _get_fwd_config( dev = arch_info.get_arch() config = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json", - required=True, ) if AUTOTUNE_Z < 512: @@ -893,7 +892,6 @@ def _get_bwd_config( dev = arch_info.get_arch() config = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json", - required=True, ) if AUTOTUNE_Z < 512: diff --git a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py index 650ae5a42df..fc634572f45 100644 --- a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py +++ b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py @@ -32,7 +32,7 @@ def _get_config(): # 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", required=True + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json" ) return config["any"].copy() # fresh copy per call — safe for callers to mutate diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index 176eb295a1e..8bcc27b3e3a 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha.py @@ -953,9 +953,7 @@ def _get_config( head_dim_v: int | None = None, ): dev = arch_info.get_arch() - config = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json", required=True - ) + 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 diff --git a/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py b/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py index e943e2e19a4..d137a81a597 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py @@ -1062,7 +1062,5 @@ def _bwd_kernel_dkdvdq_noncausal( @functools.lru_cache(maxsize=1024) def _get_config(): dev = arch_info.get_arch() - config = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json", required=True - ) + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json") return config["bkwd_fused"] diff --git a/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py b/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py index ff961cbd3c4..2b88a434e85 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py @@ -1769,7 +1769,5 @@ def bwd_kernel_noncausal( @functools.lru_cache(maxsize=1024) def _get_config(): dev = arch_info.get_arch() - config = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json", required=True - ) + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json") return config["bkwd_onekernel"] diff --git a/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py b/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py index 1420dc77387..cf1e33f4b39 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py +++ b/aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py @@ -406,5 +406,4 @@ def _get_config(): dev = arch_info.get_arch() return load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json", - required=True, ) diff --git a/aiter/ops/triton/_triton_kernels/gmm.py b/aiter/ops/triton/_triton_kernels/gmm.py index e26cdcb1299..e8beabc8844 100644 --- a/aiter/ops/triton/_triton_kernels/gmm.py +++ b/aiter/ops/triton/_triton_kernels/gmm.py @@ -32,9 +32,7 @@ def get_config( "nptgmm", }, f"'{gmm_type}' is an invalid GMM variant." dev = arch_info.get_arch() - config_dict = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-GMM.json", required=True - ) + 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." diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py index fc3b8d011dc..2d92737408f 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py @@ -129,7 +129,6 @@ def _get_config(M, N, K): dev = arch_info.get_arch() config = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE_ROUTING_SIGMOID_TOPK1.json", - required=True, ) n_key = "N16" if N <= 16 else "N128" diff --git a/aiter/ops/triton/configs/CLAUDE.md b/aiter/ops/triton/configs/CLAUDE.md index 839f0d32bfc..ccda98c87d8 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -119,9 +119,6 @@ Consequences to keep in mind: (missing files). Adding a config file at runtime therefore has no effect; restart the process (tooling may call `load_config_json.cache_clear()` instead). -- `.github/scripts/select_triton_tests.py` only globs the legacy flat - `configs/gemm/` — configs in the nested layout are invisible to CI test - selection until it learns the `/` layout. Direct-path loaders bypass the resolver's directory probe. Grep for `f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything — diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json index ed3d040ccdd..1124e1b843a 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json @@ -7,8 +7,7 @@ "num_stages": 2, "waves_per_eu": 4, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_32": { "BLOCK_SIZE_M": 16, @@ -18,8 +17,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_64": { "BLOCK_SIZE_M": 32, @@ -29,8 +27,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -40,8 +37,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": null, - "kpack": 1 + "cache_modifier": null }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -51,8 +47,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": null, - "kpack": 1 + "cache_modifier": null }, "any": { "BLOCK_SIZE_M": 32, @@ -62,7 +57,6 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json index d2bd791d13f..cfc0b934e5c 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json @@ -7,8 +7,7 @@ "num_stages": 2, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -18,8 +17,7 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_64": { "BLOCK_SIZE_M": 32, @@ -29,8 +27,7 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_128": { "BLOCK_SIZE_M": 32, @@ -40,8 +37,7 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_256": { "BLOCK_SIZE_M": 32, @@ -51,8 +47,7 @@ "num_stages": 1, "waves_per_eu": 1, "matrix_instr_nonkdim": 32, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "any": { "BLOCK_SIZE_M": 32, @@ -62,7 +57,6 @@ "num_stages": 1, "waves_per_eu": 6, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json index b2d817abd49..01ecea0d300 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json @@ -7,8 +7,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -18,8 +17,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_64": { "BLOCK_SIZE_M": 64, @@ -29,8 +27,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -40,8 +37,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -51,8 +47,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "any": { "BLOCK_SIZE_M": 32, @@ -62,7 +57,6 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json index 302eb89b9d2..e6f1bd4b192 100644 --- a/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json +++ b/aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json @@ -7,8 +7,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_32": { "BLOCK_SIZE_M": 32, @@ -18,8 +17,7 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_64": { "BLOCK_SIZE_M": 64, @@ -29,8 +27,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_128": { "BLOCK_SIZE_M": 64, @@ -40,8 +37,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "M_LEQ_256": { "BLOCK_SIZE_M": 64, @@ -51,8 +47,7 @@ "num_stages": 2, "waves_per_eu": 1, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" }, "any": { "BLOCK_SIZE_M": 32, @@ -62,7 +57,6 @@ "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, - "cache_modifier": ".cg", - "kpack": 1 + "cache_modifier": ".cg" } } diff --git a/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py b/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py index bf0fd47ee78..b3c32442b06 100644 --- a/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py +++ b/aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py @@ -10,6 +10,7 @@ _get_config, ) from aiter.ops.triton.utils.common_utils import deserialize_str, serialize_dict +from aiter.ops.triton.utils.gemm_config_utils import add_default_gemm_config_params from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -66,6 +67,9 @@ def gemm_a16w16_atomic_( config, _ = _get_config(M, N, K) else: config = deserialize_str(config) + # Caller-supplied configs may omit NUM_KSPLIT/cache_modifier; backfill + # with the canonical defaults (the shipped JSONs always carry both). + add_default_gemm_config_params(config) if y is None: # atomic add requires 0 tensor diff --git a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py index 9d2763d771a..14574e031a1 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py @@ -992,10 +992,10 @@ def _get_config_cached( ) config_dict = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json", - required=True, ) specialized = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json" + f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json", + required=False, ) if specialized is not None: config_dict = specialized @@ -1031,7 +1031,7 @@ def _get_config( # Fresh copy per call, outside the lru boundary — the caller writes # derived fields (SPLITK_BLOCK_SIZE here, GROUP_K/GROUP_N at the call # site) into the returned dict. - config = dict(_get_config_cached(M, N, K)) + config = _get_config_cached(M, N, K).copy() block_size_k = config["BLOCK_SIZE_K"] num_k_blocks = triton.cdiv(K, block_size_k) diff --git a/aiter/ops/triton/moe/moe_op_gemm_a8w4.py b/aiter/ops/triton/moe/moe_op_gemm_a8w4.py index b7554b74bb3..9f01acbd697 100644 --- a/aiter/ops/triton/moe/moe_op_gemm_a8w4.py +++ b/aiter/ops/triton/moe/moe_op_gemm_a8w4.py @@ -31,7 +31,10 @@ def _get_a8w4_dispatch(arch: str) -> dict: """Per-(block_m, N, K) dispatch table for moe_gemm_a8w4. Returns {} if no tuned file is shipped for this arch (caller uses the safe-default fallback). Mirrors get_moe_configs() in utils/moe_config_utils.py.""" - return load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/moe/{arch}-A8W4.json") or {} + dispatch = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/moe/{arch}-A8W4.json", required=False + ) + return dispatch if dispatch is not None else {} def can_overflow_int32(tensor: torch.Tensor): diff --git a/aiter/ops/triton/utils/conv_config_utils.py b/aiter/ops/triton/utils/conv_config_utils.py index f64d32ee1d8..e4d0c994979 100644 --- a/aiter/ops/triton/utils/conv_config_utils.py +++ b/aiter/ops/triton/utils/conv_config_utils.py @@ -65,7 +65,7 @@ def _get_conv_config_cached( """Three-tier walk: literal shape entry -> M_LEQ bucket -> 'any'.""" dev = arch_info.get_arch() config_dict = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/conv/{dev}-{config_name}.json", required=True + f"{AITER_TRITON_CONFIGS_PATH}/conv/{dev}-{config_name}.json" ) # Tier 1: literal shape key. diff --git a/aiter/ops/triton/utils/core.py b/aiter/ops/triton/utils/core.py index f59102b06f0..13ea84aafae 100644 --- a/aiter/ops/triton/utils/core.py +++ b/aiter/ops/triton/utils/core.py @@ -13,12 +13,16 @@ @functools.lru_cache(maxsize=None if USE_LRU_CACHE else 0) -def load_config_json(fpath: str, required: bool = False) -> dict | None: +def load_config_json(fpath: str, required: bool = True) -> dict | None: """Load a config JSON file, cached per path (including negative results — add config files before process start, or call - ``load_config_json.cache_clear()``). Returns None if the file doesn't - exist; with required=True raises FileNotFoundError instead, consistently - on every call (exceptions are never cached).""" + ``load_config_json.cache_clear()``). Raises FileNotFoundError if the file + doesn't exist, consistently on every call (exceptions are never cached); + pass required=False for probe/fallback lookups to get None instead. + + The returned dict is the shared cached object — copy before mutating: + a shallow ``.copy()`` suffices for flat bucket dicts (scalar values), + ``copy.deepcopy`` when nested sub-dicts will be mutated.""" try: with open(fpath, "r") as file: return json.load(file) diff --git a/aiter/ops/triton/utils/gemm_config_utils.py b/aiter/ops/triton/utils/gemm_config_utils.py index adf9d8357e4..70d882e1ccf 100644 --- a/aiter/ops/triton/utils/gemm_config_utils.py +++ b/aiter/ops/triton/utils/gemm_config_utils.py @@ -91,7 +91,7 @@ def _get_gemm_config_cached( # Load default config (must exist) default_fpath = f"{cfg_dir}/{default_stem}.json" - config_dict = load_config_json(default_fpath) + config_dict = load_config_json(default_fpath, required=False) if config_dict is None: raise AssertionError(f"Required config file doesn't exist: {default_fpath}") @@ -109,7 +109,7 @@ def _get_gemm_config_cached( is_tuned = False for suffix in specialized_suffixes: specialized_config = load_config_json( - f"{cfg_dir}/{name_prefix}{config_name}-{suffix}.json" + f"{cfg_dir}/{name_prefix}{config_name}-{suffix}.json", required=False ) if specialized_config is not None: config_dict, is_tuned = specialized_config, True diff --git a/aiter/ops/triton/utils/mhc_config_utils.py b/aiter/ops/triton/utils/mhc_config_utils.py index 9671dc06545..426bb5f68f5 100644 --- a/aiter/ops/triton/utils/mhc_config_utils.py +++ b/aiter/ops/triton/utils/mhc_config_utils.py @@ -19,7 +19,9 @@ def _load_with_fallback(dev: str, fname: str, required: bool = False) -> dict | None: """Load ``{dev}-{fname}``, falling back to the gfx942 copy for arches without tuned MHC configs (may be suboptimal).""" - config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{fname}") + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-{fname}", required=False + ) if config is None: config = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/{_FALLBACK_DEV}-{fname}", required=required @@ -137,9 +139,7 @@ def get_mhc_post_config(M: int, C: int) -> dict: Picks the largest ``C_ <= C``, else ``"default"``. """ dev = arch_info.get_arch() - cfg = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHC_POST.json", required=True - ) + cfg = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHC_POST.json") c_thresholds = sorted( int(k[2:]) for k in cfg if k.startswith("C_") and k[2:].isdigit() diff --git a/aiter/ops/triton/utils/moe_config_utils.py b/aiter/ops/triton/utils/moe_config_utils.py index f771a8c767f..c32f2bb92fb 100644 --- a/aiter/ops/triton/utils/moe_config_utils.py +++ b/aiter/ops/triton/utils/moe_config_utils.py @@ -54,7 +54,7 @@ def get_moe_configs(dtype: str | None) -> dict[int, Any] | None: dtype_str = "DEFAULT" if dtype is None else dtype dev = arch_info.get_arch() configs = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE-{dtype_str}.json" + f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE-{dtype_str}.json", required=False ) if configs is not None: return configs From 0d828a4590d3aa75b0fdc3adaabc659cee32181d Mon Sep 17 00:00:00 2001 From: Satya Nikhil Kodukula Date: Mon, 10 Aug 2026 15:00:18 +0000 Subject: [PATCH 6/6] address comment --- aiter/ops/triton/gluon/gemm_a8w8_blockscale.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py index 14574e031a1..251ed32b7a4 100644 --- a/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py +++ b/aiter/ops/triton/gluon/gemm_a8w8_blockscale.py @@ -985,20 +985,23 @@ def _get_config_cached( N: int, K: int, ): - dev = arch_info.get_arch() if not arch_info.is_gluon_avail(): raise ValueError( "Gluon implementation is not supported on this device (requires CDNA4)." ) + + dev = arch_info.get_arch() + + # Try specialized config first. config_dict = load_config_json( - f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json", - ) - specialized = load_config_json( f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json", required=False, ) - if specialized is not None: - config_dict = specialized + # Fall back to the general config (must exist). + if config_dict is None: + config_dict = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json" + ) # Config keys should be named M_LEQ_ or "any" bounds = []