diff --git a/aiter/ops/triton/_triton_kernels/attention/extend_attention.py b/aiter/ops/triton/_triton_kernels/attention/extend_attention.py index b3c0047459..d618431428 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" + ) # 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 52cd711495..36d77215aa 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,10 @@ def _hstu_attn_bwd( def _get_fwd_config( AUTOTUNE_Z: int, ): - if not hasattr(_get_fwd_config, "_config_dict"): - dev = arch_info.get_arch() - fpath = f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_fwd_config._config_dict = config + dev = arch_info.get_arch() + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json", + ) if AUTOTUNE_Z < 512: batch_key = "small_batch" @@ -885,23 +882,21 @@ def _get_fwd_config( else: batch_key = "large_batch" - return _get_fwd_config._config_dict[batch_key] + return config[batch_key] @functools.lru_cache(maxsize=1024) def _get_bwd_config( AUTOTUNE_Z: int, ): - if not hasattr(_get_bwd_config, "_config_dict"): - dev = arch_info.get_arch() - fpath = f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_bwd_config._config_dict = config + dev = arch_info.get_arch() + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json", + ) if AUTOTUNE_Z < 512: batch_key = "small_batch" else: batch_key = "large_batch" - return _get_bwd_config._config_dict[batch_key] + return config[batch_key] diff --git a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py index 16fed247e8..fc634572f4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/lean_atten.py +++ b/aiter/ops/triton/_triton_kernels/attention/lean_atten.py @@ -17,32 +17,24 @@ - """ -import functools -import json - import triton import triton.language as tl from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json # Support tensor in [B, Seqlen, H, d] format. Taking tensors in [B*Seqlen, H, d] as inputs -@functools.lru_cache(maxsize=1024) def _get_config(): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - config = _get_config._config_dict["any"] - return ( - config.copy() - ) # return a copy to avoid mutation of stored config in LRU cache + # No lru_cache here: load_config_json already caches the parse, and + # caching the .copy() would hand every caller the same mutable object. + dev = arch_info.get_arch() + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-LEANATTN-DEFAULT.json" + ) + return config["any"].copy() # fresh copy per call — safe for callers to mutate @triton.jit diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index 7dfb20283a..8bcc27b3e3 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,9 @@ def _get_config( has_pe: bool = False, head_dim_v: int | None = None, ): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - _get_config._config_dict = {} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict["default"] = config - fwd_cfg = _get_config._config_dict["default"]["fwd"] + dev = arch_info.get_arch() + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json") + fwd_cfg = config["fwd"] has_dropout_or_fp32 = enable_dropout or dtype == torch.float32 # TODO: pe + dropout is not tuned if has_pe and has_dropout_or_fp32 and "pe_dropout_or_fp32" in fwd_cfg: 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 8569d5afad..d137a81a59 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,6 @@ def _bwd_kernel_dkdvdq_noncausal( @functools.lru_cache(maxsize=1024) def _get_config(): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - _get_config._config_dict = {} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - return _get_config._config_dict["bkwd_fused"] + dev = arch_info.get_arch() + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json") + return config["bkwd_fused"] 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 0bb423d3ad..2b88a434e8 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,6 @@ def bwd_kernel_noncausal( @functools.lru_cache(maxsize=1024) def _get_config(): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - _get_config._config_dict = {} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - return _get_config._config_dict["bkwd_onekernel"] + dev = arch_info.get_arch() + config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json") + return config["bkwd_onekernel"] 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 4f0c38c702..cf1e33f4b3 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,7 @@ def _fwd_kernel_stage2( @functools.lru_cache(maxsize=1024) def _get_config(): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - _get_config._config_dict = {} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config - - return _get_config._config_dict + dev = arch_info.get_arch() + return load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json", + ) 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 14ab60d29e..2121b0f8c2 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,7 +618,6 @@ def _fused_gemm_a8w8_blockscale_split_cat_reduce( tl.store(c1_ptrs, y, mask=y_mask) -@functools.lru_cache(maxsize=1024) def _get_config( M: int, N: int, diff --git a/aiter/ops/triton/_triton_kernels/gmm.py b/aiter/ops/triton/_triton_kernels/gmm.py index 9e5abdfcc5..e8beabc884 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,16 @@ def get_config( "ptgmm", "nptgmm", }, f"'{gmm_type}' is an invalid GMM variant." - if not hasattr(get_config, "_config_dict"): - dev = arch_info.get_arch() - config_filename = f"{AITER_TRITON_CONFIGS_PATH}/{dev}-GMM.json" - assert os.path.exists(config_filename) and os.path.isfile( - config_filename - ), f"'{config_filename}' isn't an existent file." - with open(config_filename, "r") as config_file: - get_config._config_dict = json.load(config_file) - assert all( - gmm_type in get_config._config_dict - for gmm_type in ("gmm", "ptgmm", "nptgmm") - ), "Not all GMM variants are present in the configuration file." + dev = arch_info.get_arch() + config_dict = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-GMM.json") + assert all( + variant in config_dict for variant in ("gmm", "ptgmm", "nptgmm") + ), "Not all GMM variants are present in the configuration file." # TODO: Fine tune GMM kernels and use (M, K, N, G) shape to query the best # config in the dictionary. - assert ( - "default" in get_config._config_dict[gmm_type] - ), "Default configuration is absent." + assert "default" in config_dict[gmm_type], "Default configuration is absent." key = "accumulate" if accumulate else "default" - return get_config._config_dict[gmm_type][key] + return config_dict[gmm_type][key] # Common code shared by GMM and TGMM kernels. 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 408844e9ad..2d92737408 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,10 @@ def _routing_sigmoid_top1_kernel( @functools.lru_cache(maxsize=1024) def _get_config(M, N, K): - if not hasattr(_get_config, "_config_dict"): - dev = arch_info.get_arch() - _get_config._config_dict = {} - fpath = f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE_ROUTING_SIGMOID_TOPK1.json" - with open(fpath, "r") as file: - config = json.load(file) - _get_config._config_dict = config + dev = arch_info.get_arch() + config = load_config_json( + f"{AITER_TRITON_CONFIGS_PATH}/moe/{dev}-MOE_ROUTING_SIGMOID_TOPK1.json", + ) n_key = "N16" if N <= 16 else "N128" m_key = ( @@ -141,4 +137,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/lean_atten.py b/aiter/ops/triton/attention/lean_atten.py index 935bc842f7..d41a1d2fc4 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/attention/mla_decode_rope.py b/aiter/ops/triton/attention/mla_decode_rope.py index e9cd903f87..32fcead23d 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 72fb72000d..ccda98c87d 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -113,16 +113,19 @@ 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. - -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. +- Results are cached twice: `functools.lru_cache` on the full argument + tuple, plus a per-path cache of parsed JSON + (`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). + +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. --- @@ -139,9 +142,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`. @@ -192,7 +195,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 a0a550dcab..0000000000 --- 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 689eb50235..0000000000 --- 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 7a4a48cfb5..6d5ec6e3ca 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 629d62d9c8..222f27b6d1 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 372268c260..43f7876e18 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_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 1d1df58c1d..0000000000 --- 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 f9eba0dfe6..0000000000 --- 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 a0a550dcab..0000000000 --- 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 b5b4b7b66d..0000000000 --- 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 70042715e3..0000000000 --- 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 3faf5c1c93..b3c32442b0 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,13 +67,9 @@ def gemm_a16w16_atomic_( config, _ = _get_config(M, N, K) 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"] = "" + # 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/gemm/basic/gemm_afp4wfp4.py b/aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py index 7719207dd9..683327928f 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 591526e289..365e67fcf6 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 f7c8f6a5dc..94d9865f26 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.py b/aiter/ops/triton/gluon/gemm_a8w8.py index aa8dd8d849..da5d00f7e5 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,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.core import AITER_TRITON_CONFIGS_PATH 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() @@ -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 809c324889..251ed32b7a 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,39 +980,32 @@ 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)." - ) - _get_config._config_dict = {} - fpath = ( + 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-N={N}-K={K}.json", + required=False, + ) + # 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" ) - 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 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)) @@ -1023,18 +1014,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 = _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/gluon/gemm_afp4wfp4.py b/aiter/ops/triton/gluon/gemm_afp4wfp4.py index d860a923bd..da39234bb9 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,15 @@ 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.") + config, _ = get_gemm_config("GEMM-AFP4WFP4", M, N, K, backend="gluon") + return config def gemm_afp4wfp4( diff --git a/aiter/ops/triton/moe/moe_op_gemm_a8w4.py b/aiter/ops/triton/moe/moe_op_gemm_a8w4.py index 9080783db1..9f01acbd69 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,19 @@ 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 {} + 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/_triton/gemm_tune_check.py b/aiter/ops/triton/utils/_triton/gemm_tune_check.py index 47719395b9..b8033d88c3 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 5530189fc4..e4d0c99497 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" + ) # Tier 1: literal shape key. shapes = config_dict.get("shapes", {}) diff --git a/aiter/ops/triton/utils/core.py b/aiter/ops/triton/utils/core.py index 969355d944..13ea84aafa 100644 --- a/aiter/ops/triton/utils/core.py +++ b/aiter/ops/triton/utils/core.py @@ -1,5 +1,34 @@ +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 = 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()``). 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) + 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 73a83efb68..70d882e1cc 100644 --- a/aiter/ops/triton/utils/gemm_config_utils.py +++ b/aiter/ops/triton/utils/gemm_config_utils.py @@ -1,27 +1,19 @@ import copy import functools import itertools -import json import os 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 -""" +STANDARD_M_BOUNDS = (1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) def _dtype_dir(config_name: str) -> str: @@ -30,26 +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. - """ - 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, @@ -80,11 +52,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 +89,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, required=False) + 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", required=False + ) + 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 +122,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 797e2a217e..426bb5f68f 100644 --- a/aiter/ops/triton/utils/mhc_config_utils.py +++ b/aiter/ops/triton/utils/mhc_config_utils.py @@ -3,13 +3,44 @@ 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}", required=False + ) + 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,104 +80,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 - if cache_key not in get_mhc_config._config_cache: - get_mhc_config._config_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, - 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( - get_mhc_config._config_cache, - cache_key, - fpath_fallback, - "default", - fpath_should_exist=True, - ) - - 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: @@ -184,29 +138,8 @@ 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") 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 d237abf15b..c32f2bb92f 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", required=False + ) + 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 470e52528d..c7b5f65018 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,15 +11,16 @@ 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.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) @@ -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 b487ff4c7d..fd9a41e23c 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 df66d048d7..4ec6c83ce9 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 259de53205..09ef87ad8b 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)