From 3b1ee5d0f9cc3167dd38b54faff503db7946c8c2 Mon Sep 17 00:00:00 2001 From: Satya Nikhil Kodukula Date: Fri, 28 Aug 2026 15:12:16 +0000 Subject: [PATCH] [Triton] Update config docs for the single nested layout The config tree is now one nested layout, `configs//// /`, resolved by a deterministic path builder; the flat arch-prefixed directories and every fallback that reached them are gone. Six docs still described the mid-migration world. - configs/CLAUDE.md: rewritten as the rulebook for the layout that exists. One layout section, the `resolve_config_dir()` / `load_config_json()` contract (validation, fail-closed assertions, caching incl. negatives, shared-object return), the layered loader-module table, and a section per family: GEMM order-of-operations and file contents, the MOE dispatch key schemes per backend, the conv four-tier walk, MHC's C-threshold discovery and documented gfx942 fallback, and the pinned-tile loader. Adding a config and seeding a new arch replace the migration playbook. The planned `get_moe_config()` design section is dropped -- `get_moe_dispatch()` shipped. - README.md: config sections rewritten to match; `utils/core.py` references point at `utils/config_utils.py`; MOE naming row and key schemes corrected. - .github/instructions: placement rules restated for one layout, plus new review rules for MOE dispatch tables, kernel-level `backend` defaults, the layered `utils/` split, and arch seeding. - conv/DESIGN.md, conv/README.md: seven flat `configs/conv/{arch}-CONV-*.json` references updated to the nested paths, and section 9 now names the shared resolver. - tunning/README.md: one copy destination instead of a migrated-vs-legacy split. Corrections found while checking every claim against the tree: the `kpack` rule is scoped to gfx950 (the RDNA trees do still carry it), the MHC C-threshold glob unions the running arch with gfx942, the six-bucket MOE rule is guidance for newly tuned shapes rather than a description of current coverage, and the AOT-guard note now matches what each of the two paths actually does. --- .../aiter-ops-triton.instructions.md | 125 ++-- aiter/ops/triton/README.md | 133 ++-- aiter/ops/triton/configs/CLAUDE.md | 607 ++++++++++-------- aiter/ops/triton/conv/DESIGN.md | 28 +- aiter/ops/triton/conv/README.md | 20 +- .../triton/utils/_triton/tunning/README.md | 14 +- 6 files changed, 529 insertions(+), 398 deletions(-) diff --git a/.github/instructions/aiter-ops-triton.instructions.md b/.github/instructions/aiter-ops-triton.instructions.md index e903b3f224..1bf14f7a79 100644 --- a/.github/instructions/aiter-ops-triton.instructions.md +++ b/.github/instructions/aiter-ops-triton.instructions.md @@ -11,11 +11,18 @@ relevant rule — reviewers may not know these conventions yet. Always prefer reusing existing code over adding new code. Before a PR adds a helper, kernel, or utility, the existing ones should have been checked: -`utils/` (config loading, shuffling, arch info, logging, `kernel_repr`), -`_triton_kernels/common/` (shared split-K reduce), and existing kernels and -test helpers. Flag new code that duplicates functionality already in the -tree, even partially — the fix is to extend or import the existing -implementation, not to add a parallel copy. +`utils/` (`config_utils` and the per-family `*_config_utils` loaders, +shuffling, arch info, logging, `kernel_repr`), `_triton_kernels/common/` +(shared split-K reduce), and existing kernels and test helpers. Flag new code +that duplicates functionality already in the tree, even partially — the fix is +to extend or import the existing implementation, not to add a parallel copy. + +`utils/` is layered on purpose: `config_utils.py` holds the shared core +(`resolve_config_dir`, `load_config_json`, the path constants) and each family +keeps its own loader module (`gemm_config_utils`, `conv_config_utils`, +`mhc_config_utils`, `moe_config_utils`, `tuned_config_utils`) on top of it. +Flag a function given a second home — a re-export, a wrapper that only +forwards to another module, or a copy of a core helper inside a family module. ## Folder structure and imports @@ -46,35 +53,39 @@ and tuned JSON in `configs/`. Flag: ## Tuned configs: JSON placement and naming -The config tree is mid-migration from a legacy flat layout -(`configs/gemm/-.json`) to a nested layout -(`configs/////`, e.g. -`configs/gfx950/triton/gemm/gemm_afp4wfp4/DEFAULT.json`). The legacy layout is -deprecated. Flag: - -- A new GEMM config JSON added under legacy `configs/gemm/` when the family - already has a nested `//gemm//` directory, or a brand - new family added to the legacy layout instead of the nested one. -- A config family split across the two layouts (e.g. `DEFAULT.json` nested but - new specialized files in `configs/gemm/`, or vice versa). The directory - probe in `get_gemm_config()` picks ONE directory — files in the losing - directory are silently ignored. Families move wholesale or not at all. +Every tuned config lives in one nested layout: +`configs/////`, e.g. +`configs/gfx950/triton/gemm/gemm_afp4wfp4/DEFAULT.json`. `` is `gemm`, +`moe`, `conv`, `mhc`, `attention`, `gmm` or `fusions`; `` is +`config_name.lower().replace("-", "_")`. The flat arch-prefixed directories +and every fallback that reached them are gone. Flag: + +- A config JSON added outside `configs/////` — a + re-created `configs/gemm/`, `configs/moe/` or `configs/conv/` directory, or + a loose file at the top of `configs/`. Nothing resolves there any more. - An arch prefix on a filename inside `configs//...` (wrong: - `configs/gfx950/triton/gemm/x/gfx950-GEMM-X.json`), or a nested default file - named anything other than exactly `DEFAULT.json`. -- A specialized file added to a nested `/` directory that contains no - `DEFAULT.json` — the resolver probes only for the default, so the whole - directory is invisible. -- Any MOE config created or moved under `//moe/`. No MOE - resolver for the nested layout exists yet; MOE configs stay in - `configs/moe/` with the arch prefix (see `configs/CLAUDE.md` §5). -- Deleting or renaming `.gitkeep` placeholder directories under `configs/`. + `configs/gfx950/triton/gemm/x/gfx950-GEMM-X.json`), or a default file named + anything other than exactly `DEFAULT.json`. +- A specialized file added to a `/` directory that contains no + `DEFAULT.json`, for a family whose loader requires the default — the load + raises for every shape, not just the unspecialized ones. +- A family's files split across two `/` directories that differ only + by the `_dtype_dir()` fold (`GEMM-FOO-BAR` and `GEMM-FOO_BAR` collide; two + spellings of one family must not both exist). - A config file that is both moved and content-edited in the same commit — - migrations must be pure `git mv` renames, content changes in a follow-up. -- `kpack` in a config file for gfx950 or a newer arch. Triton's AMD backend - deprecates `kpack` starting from gfx950 — it warns and force-overrides - `kpack = 1` there, and the parameter is slated for removal. Only gfx942 - configs may still carry `kpack`. + moves must be pure `git mv` renames, content changes in a follow-up. +- A new `.gitkeep` under `configs/`. A `/` directory is created + populated; the few `.gitkeep` files left from the migration are inert + leftovers, not placeholders to maintain. +- A new arch directory seeded from another arch without the copy being + byte-identical and called out in the commit message. The one seeding rule in + force is gfx950 → gfx1250, triton only — never into a gluon directory, never + backwards into gfx950. +- `kpack` newly added to a gfx950 config. Triton's AMD backend deprecates + `kpack` on CDNA4 — it warns and force-overrides `kpack = 1` there, and the + parameter is slated for removal. The gfx950 tree is clean of it; gfx942 may + still carry it, and existing RDNA (gfx1151/gfx1201/gfx1250) entries predate + the rule, so flag additions rather than the entries already there. - Checked-in files under `configs/gemm/aot/` or `configs/paged_mqa_logits/aot/` — these are runtime AOT caches, never committed. @@ -90,12 +101,27 @@ Flag, inside GEMM-family config JSON: `BLOCK_SIZE_K`, `GROUP_SIZE_M`, `num_warps`, `num_stages`, `waves_per_eu`, `matrix_instr_nonkdim`, `cache_modifier`, `NUM_KSPLIT`. (Loader backfill of `NUM_KSPLIT`/`cache_modifier` is a last resort, not a license to omit.) -- MOE-scheme keys (`small_M`/`medium_M`/`large_M`) in a GEMM config or GEMM - `M_LEQ_x`/`M_GEQ_y` keys in a MOE config — the schemes must not mix. +- MOE dispatch keys (`bm_n_k`) in a GEMM config or GEMM + `M_LEQ_x`/`M_GEQ_y` keys in a MOE dispatch table — the schemes must not mix. - For `*AFP4WFP4*` specialized filenames: `K` must be the logical K (`2 * K_bytes`) — the wrapper doubles K before lookup, so a file named by the packed byte width will never resolve. +And inside MOE dispatch tables: + +- A newly tuned gluon dispatch shape that does not carry all six `m2bucket` + suffixes (`tiny`, `small`, `medium`, `medium2`, `large`, `xlarge`) — a + missing bucket falls through to `bm_any` and silently loses that + shape's tuning for that M range. (Existing entries are unevenly covered; + flag new gaps, not the ones already shipped.) +- A gluon dispatch file with no `bm_any` tier for a `block_m` it + otherwise covers: that tier is the last resort for an unmeasured shape. +- A `BLOCK_SIZE_M` / `block_m` key inside a dispatch entry — `block_m` is the + dispatch key (routing decides it), not a tunable. +- Triton-shaped entry keys (`BLOCK_SIZE_N`, `num_stages`, ...) in a gluon + dispatch file or gluon-shaped keys (`block_n`, `num_buffers`, + `persistent_iters`) in a triton one — the two paths read disjoint params. + ## Python-side config hygiene These rules apply equally to Triton and Gluon wrappers and kernels. Tuning @@ -105,6 +131,12 @@ values for either backend live in JSON, never in Python. Flag: inline dict literals with `BLOCK_SIZE_*`/`num_warps`/`waves_per_eu` keys, arch-conditional tuning constants, or hardcoded fallback configs. The fix is always in the JSON file, not the Python. +- A kernel-level `_get_config()` whose `backend` parameter defaults to `None` + (or any value outside `("triton", "gluon")`) — `None` is not a backend and + `resolve_config_dir()` asserts on it, so the kernel raises the moment a + caller omits the argument. Kernel-level helpers default to `"triton"`; + public wrappers that expose `backend: str | None = None` must normalize it + before calling down. - A new or modified Triton or Gluon GEMM `_get_config()` that does anything beyond calling `get_gemm_config(...)` with the appropriate backend selection (plus `compute_splitk_params()` for split-K kernels), and that does not @@ -125,13 +157,24 @@ values for either backend live in JSON, never in Python. Flag: re-tune. - Raw config-file reads — `json.load(open(...))` or function-attribute caches like `_get_config._config_dict` — instead of - `aiter.ops.triton.utils.core.load_config_json` (which caches per path, - including negative results) or the resolvers `get_gemm_config` / - `get_tuned_kernel_config`. All hand-rolled loaders were deliberately - removed; do not add them back. -- New hand-built config paths (`f"{AITER_TRITON_CONFIGS_PATH}/..."`) where - `get_gemm_config()` / `get_tuned_kernel_config()` would work — hand-built - paths break silently when the family migrates to the nested layout. + `aiter.ops.triton.utils.config_utils.load_config_json` (which caches per + path, including negative results) or a family loader. All hand-rolled + loaders were deliberately removed; do not add them back. +- Mutating the dict returned by `load_config_json()` — it is the shared cached + object. Copy first (`dict(...)` for flat entries, `copy.deepcopy` for nested + ones); the family loaders already copy on the caller's behalf. +- New hand-built config paths (`f"{AITER_TRITON_CONFIGS_PATH}/..."`) where a + family loader or `resolve_config_dir()` would work — a hand-built path is a + second place the layout is encoded, and it skips the argument validation + that makes a wrong value fail closed. +- A second MOE config reader. `utils/moe_config_utils.py::get_moe_dispatch` is + the only MOE fetcher; flag any new MOE path built by hand, any direct + `load_config_json` on a `moe/` file, and any reintroduced per-wrapper MOE + loader. +- A new arch- or backend-fallback chain inside a loader (try this arch, then + that one; try triton, then gluon). Resolution is deterministic. MHC's gfx942 + fallback is the one documented exception and it goes through the `arch=` + override, not through a probe. ## Weight & scale shuffling — must come from `utils/shuffle.py` diff --git a/aiter/ops/triton/README.md b/aiter/ops/triton/README.md index b35762fafb..80dc8c799e 100644 --- a/aiter/ops/triton/README.md +++ b/aiter/ops/triton/README.md @@ -52,53 +52,80 @@ import from the categorized path** (`aiter.ops.triton.gemm.basic.gemm_a16w16`). ## Tuned configs -### Two layouts are live; only one accepts new files +### One layout, one path builder -The config tree is mid-migration from a flat, arch-prefixed layout to a nested -one. The legacy flat layout is **deprecated — treat it as read-only history**: +Every tuned JSON file lives in a single nested tree, keyed by architecture, +backend, op and config family: ```text -# Target layout (all new GEMM configs go here) configs/////DEFAULT.json configs/////-.json # gfx950 triton gemm gemm_afp4wfp4 -# gluon moe - -# Legacy layout (deprecated, pending removal) -configs/gemm/-[-].json -configs/moe/-MOE-.json +# gfx1250 gluon moe a8w4 ``` +`` is one of `gemm`, `moe`, `conv`, `mhc`, `attention`, `gmm`, `fusions`. +The flat, arch-prefixed directories (`configs/gemm/`, `configs/moe/`, +`configs/conv/`, the loose files at the top of `configs/`) and the fallback +code that reached them are gone. + Rules that follow from the layout: - `` is `config_name.lower().replace("-", "_")` (`GEMM-AFP4WFP4` → - `gemm_afp4wfp4`, see `gemm_config_utils._dtype_dir()`). New config names must + `gemm_afp4wfp4`, see `config_utils._dtype_dir()`). New config names must stay distinct under that fold — `GEMM-FOO-BAR` and `GEMM-FOO_BAR` collide. - Files inside `configs//...` carry **no arch prefix**; the default file is named exactly `DEFAULT.json`. -- A `/` directory without a `DEFAULT.json` is **invisible** — the - resolver probes only for the default file, so its specialized files are - silently ignored. Never split a config family across the two layouts; a - family migrates wholesale or not at all (`git mv`, 100% rename similarity, - content changes in a separate commit). -- `//moe/` directories exist but are `.gitkeep` placeholders. - **MOE has no nested-layout resolver yet** — MOE configs stay in - `configs/moe/` with the arch prefix until `get_moe_config()` lands - (design in `configs/CLAUDE.md` §5). -- `kpack` is deprecated starting from gfx950: the Triton AMD backend warns - and force-overrides `kpack = 1` there, and the parameter is slated for - removal. Configs for gfx950 and newer must not carry it; only gfx942 - configs may still set it. +- A required `DEFAULT.json` that is missing raises naming that exact path. + There is no probe to fall through, so a misplaced file surfaces as an error + instead of silently resolving somewhere else. +- Moves and renames go in a pure `git mv` commit (100% rename similarity), + with content changes in a follow-up. +- `kpack` is deprecated on CDNA4: the Triton AMD backend warns and + force-overrides `kpack = 1` on gfx950, and the parameter is slated for + removal. No gfx950 config carries it and none should; gfx942 configs still + may. Existing `kpack` entries in the RDNA trees predate the rule — do not + add new ones. + +All of it is built by one function in `utils/config_utils.py`: + +```python +resolve_config_dir(op, config_name, backend="triton", arch=None) -> str +``` + +It **builds** the path — no probing, no candidate list, no cross-backend or +cross-arch search — and validates every argument against a whitelist, so a bad +value fails closed with an `AssertionError` rather than resolving to an +escaped or wrong directory. `backend` is declared by the caller (gluon kernels +and gluon dispatch paths pass `"gluon"`; everything else takes the `"triton"` +default), because the two backends take disjoint config params and borrowing +across them would be a bug. `arch=` overrides the running architecture only +where a loader deliberately retries elsewhere — today just MHC's documented +gfx942 fallback. + +`config_utils.py` is the shared core; each family keeps its own small loader +module on top of it, and every function has exactly one home: + +| Module | Entry points | +| ------ | ------------ | +| `utils/config_utils.py` | `resolve_config_dir`, `load_config_json`, path constants | +| `utils/gemm_config_utils.py` | `get_gemm_config`, `compute_splitk_params`, `add_default_gemm_config_params`, `pick_gemm_num_stages` | +| `utils/conv_config_utils.py` | `get_conv_config` + the shape-key formatters and table probes | +| `utils/mhc_config_utils.py` | `get_mhc_config`, `get_mhc_post_config` | +| `utils/moe_config_utils.py` | `get_moe_dispatch` — the only MOE config fetcher | +| `utils/tuned_config_utils.py` | `get_tuned_kernel_config` | + +Attention and GMM kernels read their single `DEFAULT.json` straight off the +core (`resolve_config_dir()` + `load_config_json()`); a family module earns +its place once a family grows real selection logic. ### How GEMM configs resolve — `get_gemm_config()` All GEMM-family kernels load configs through one function, `utils/gemm_config_utils.py::get_gemm_config(config_name, M, N=None, K=None, -bounds=None, specialized_filename=None, backend=None, B=None)`. It probes -candidate directories for the default file and takes the first hit -(`backend=None` → `/triton/gemm/` → `/gluon/gemm/` → legacy -`configs/gemm/`), then reads specialized files from that same directory. -It returns `(config, is_tuned)`: +bounds=None, specialized_filename=None, backend="triton", B=None)`. It reads +`//gemm//DEFAULT.json` (required), overlays a +specialized file when one matches, and returns `(config, is_tuned)`: - the config is a fresh deep copy, safe to mutate; - `is_tuned` is `True` only when a specialized (`N=…-K=…`, `B=…-N=…-K=…`, or @@ -119,6 +146,11 @@ def _get_config(M: int, N: int, K: int): return compute_splitk_params(config, K), is_tuned ``` +A kernel-level `_get_config()` that takes a `backend` argument defaults it to +`"triton"`, never to `None` — `None` is not a backend and the resolver rejects +it. Public wrappers that expose `backend: str | None = None` normalize it +(typically `"gluon"` on gfx1250, `"triton"` elsewhere) before calling down. + Split-K kernels also share one common second-stage reduce — `_gemm_splitk_reduce_kernel` (and `_batched_gemm_splitk_reduce_kernel`) in `_triton_kernels/common/splitk_reduce.py` — rather than carrying a per-kernel @@ -152,20 +184,27 @@ Gluon. New split-K kernels import it from there. No `setdefault(...)` blocks, no inline config dict literals, no arch-conditional constants, no hardcoded fallback configs in `.py` files. If a -value is missing at runtime, the fix is in the JSON. MOE still has legacy -Python fallbacks (`get_optimal_moe_config()`, the `moe_op_gemm_a8w4` dispatch -tiers) — fix them as you touch them, and do not add more. +value is missing at runtime, the fix is in the JSON. The MOE a8w4 triton +dispatch still ends in a Python default after its tuned lookup and its +same-`(N, K)` proxy — that last tier shrinks as coverage grows; do not add +more like it. ### Loading is unified — `load_config_json()` -Every config-file read goes through `utils/core.py::load_config_json(fpath, -required=...)`. It is cached per path **including negative results**: -a config file added at runtime is not picked up without -`load_config_json.cache_clear()` or a process restart. Do not hand-roll -`json.load(open(...))` or function-attribute caches, and prefer the resolvers -(`get_gemm_config` / `get_tuned_kernel_config`) over hand-built -`f"{AITER_TRITON_CONFIGS_PATH}/..."` paths — hand-built paths break silently -when a family migrates and must be grepped for during every migration. +Every config-file read goes through +`utils/config_utils.py::load_config_json(fpath, required=...)`. +`required=True` (the default) raises `FileNotFoundError` naming the exact +nested path; `required=False` returns `None` for genuinely optional tables and +the caller handles it. It is cached per path **including negative results**: a +config file added at runtime is not picked up without +`load_config_json.cache_clear()` or a process restart. The dict it returns is +the shared cached object — copy before mutating (the family loaders already do +this for their callers). + +Do not hand-roll `json.load(open(...))` or function-attribute caches, and +prefer the family loaders over hand-built +`f"{AITER_TRITON_CONFIGS_PATH}/..."` paths — a hand-built path is a second +place the layout is encoded, and it goes stale silently. Kernels that carry a Python autotune search space (opt-in tuning) pin their single default tile per arch via @@ -181,19 +220,21 @@ kernel_name, fallback, backend)`, which reads the nested-layout | Batched GEMM | `BATCHED_GEMM-A{x}W{y}`, specialized `-B={B}-N={N}-K={K}` | | Fused ops | `FUSED-GEMM-{operation}` | | Feed-forward | `FF-A{x}W{y}-fused` | -| MOE | `MOE-` (`DEFAULT`, `FP8_W8A8`, `MX_FP4`, ...) | +| MOE | `A8W4`, `A4W4` — dispatch tables, one `DEFAULT.json` per backend | - **`K` in AFP4WFP4 filenames is the logical K, i.e. `2 * K_bytes`** — the wrapper doubles K before calling `get_gemm_config`. Tuning output named by the packed byte width will never be found. -- MOE files use `small_M` / `medium_M` / `large_M` (thresholds 256 / 1024 in - `moe_config_utils.py`). Never mix that scheme with the GEMM +- MOE dispatch tables are keyed `bm_n_k` on the triton path and + `bm_n_k_` (plus a `bm_any` tier) on the + gluon path — a different scheme per backend, which is why `backend` is a + caller-declared argument. Never mix a MOE dispatch key with the GEMM `M_LEQ_x`/`M_GEQ_y` scheme in either direction. - `configs/gemm/aot/` and `configs/paged_mqa_logits/aot/` are runtime AOT caches, not tuning configs — never check them in or migrate them. -For migrating a family into the nested layout, follow the playbook in -`configs/CLAUDE.md` §6 step by step. For the manual tuning flow, see +For adding a config, seeding a new arch, and the per-family key schemes, follow +`configs/CLAUDE.md` (§5 and §6). For the manual tuning flow, see `utils/_triton/tunning/README.md`. --- @@ -253,8 +294,8 @@ decide, and use `return_layout=True` instead of hardcoding swizzle labels. ## Architecture naming -Key behavior off GPU architecture identifiers, never product names — in config -filenames (`gfx950-...`), directory names (`configs/gfx950/...`), and code: +Key behavior off GPU architecture identifiers, never product names — in +directory names (`configs/gfx950/...`) and in code: ```python DEVICE_ARCH = arch_info.get_arch() diff --git a/aiter/ops/triton/configs/CLAUDE.md b/aiter/ops/triton/configs/CLAUDE.md index af4f49cdf8..f69911382b 100644 --- a/aiter/ops/triton/configs/CLAUDE.md +++ b/aiter/ops/triton/configs/CLAUDE.md @@ -1,137 +1,184 @@ # Triton kernel configs — rules for automated edits -Scope: **GEMM and MOE configs only.** Read this before adding, moving, -renaming, or tuning a GEMM or MOE JSON file under -`aiter/ops/triton/configs/`, or before touching -`utils/gemm_config_utils.py` or `utils/moe_config_utils.py`. +Scope: **every tuned JSON file under `aiter/ops/triton/configs/`**, and the +loader modules that read them (`aiter/ops/triton/utils/config_utils.py` and +the per-family `*_config_utils.py` modules). Read this before adding, moving, +renaming, or tuning a config file. -Out of scope, do not touch without an explicit request: `configs/conv/`, -`configs/hstu_attn/`, and the flat attention / GMM / MHC / MLA files at the top -of `configs/`. They have their own loaders and are unaffected by anything here. - -The tree is **mid-migration** from a flat, arch-prefixed layout to a nested -`////` layout. Both layouts are live, but **the legacy flat -layout is deprecated and will be removed** — treat it as read-only history, not -as a place to add things. +There is now exactly **one layout**. The flat, arch-prefixed directories +(`configs/gemm/`, `configs/moe/`, `configs/conv/`, the loose attention / GMM / +MHC files at the top of `configs/`) are gone, and so is the fallback code that +used to reach them. Nothing resolves outside the nested tree. Two non-negotiables: 1. **Tuning values live in JSON, never in Python.** No `setdefault`, no inline dict literals, no arch-conditional constants, no hardcoded fallback configs. If a value is missing, fix the JSON. -2. **New configs go in the target layout** unless their family is still in the - legacy directory (see §6). - -`GEMM-AFP4WFP4` (gfx950 triton, gfx950/gfx1250 gluon) and -`GEMM-AFP4WFP4_PRESHUFFLED` (gfx950/gfx1250 triton) are the migrated -families; `GEMM-AFP4WFP4` is the worked reference — copy its shape when in -doubt. +2. **Every config read goes through a loader** in `utils/`, which goes through + `resolve_config_dir()` + `load_config_json()`. No `json.load(open(...))`, no + function-attribute caches, no hand-built `f"{AITER_TRITON_CONFIGS_PATH}/…"` + paths. --- -## 1. Layouts - -### Target layout (use for all new configs) +## 1. The layout ``` configs/////DEFAULT.json configs/////-.json ``` -| Segment | Values | -| ----------- | --------------------------------------------------------- | -| `` | `gfx942`, `gfx950`, `gfx1250`, `gfx1151`, `gfx1200`, `gfx1201` | -| `` | `triton` or `gluon` | -| `` | `gemm` or `moe` | -| `` | `config_name.lower().replace("-", "_")` — `GEMM-AFP4WFP4` → `gemm_afp4wfp4`. The transform lives in `gemm_config_utils._dtype_dir()` | +| Segment | Values | +| ----------- | ----------------------------------------------------------- | +| `` | `gfx942`, `gfx950`, `gfx1100`, `gfx1151`, `gfx1200`, `gfx1201`, `gfx1250` | +| `` | `triton` or `gluon` | +| `` | `gemm`, `moe`, `conv`, `mhc`, `attention`, `gmm`, `fusions` | +| `` | `config_name.lower().replace("-", "_")` — `GEMM-AFP4WFP4` → `gemm_afp4wfp4`. The transform is `config_utils._dtype_dir()` | | filename | **no arch prefix** — the arch is the directory. The default is literally `DEFAULT.json`; specialized files keep the `-` stem | ``` configs/gfx950/triton/gemm/gemm_afp4wfp4/DEFAULT.json configs/gfx950/triton/gemm/gemm_afp4wfp4/GEMM-AFP4WFP4-N=8192-K=8192.json -configs/gfx950/gluon/gemm/gemm_afp4wfp4/DEFAULT.json configs/gfx1250/gluon/gemm/gemm_afp4wfp4/DEFAULT.json +configs/gfx1250/gluon/moe/a8w4/DEFAULT.json +configs/gfx942/triton/mhc/mhc_fused_sinkhorn/MHC_FUSED_SINKHORN-C=128.json +configs/gfx1201/triton/conv/conv_3x3_nhwc/DEFAULT.json ``` -The `//moe/` directories exist but are empty, held open with -`.gitkeep`. Keep them. **No MOE config has been migrated and no MOE resolver -understands the nested layout yet** — see §5. - -### Legacy layout — deprecated, pending removal +Regenerate rather than trusting any listing in this file: ``` -configs/gemm/-[-].json -configs/gemm/gluon/-[-].json -configs/moe/-MOE-.json -configs/moe/-A8W4.json -configs/moe/-MOE_ROUTING_SIGMOID_TOPK1.json +git ls-tree -r --name-only HEAD aiter/ops/triton/configs/ ``` -Regenerate rather than trusting this listing: -`git ls-tree -r --name-only HEAD aiter/ops/triton/configs/` - -Still authoritative for every family not yet migrated. For GEMM it is reached -through the fallback chain in §2; for MOE it is the *only* path that works. -The GEMM fallback is temporary — anything left in `configs/gemm/` when the -legacy candidates are dropped from `gemm_config_utils.py` will stop resolving. +A few `.gitkeep` files survive from the migration in directories that later +got real content or never got any. They are inert: no loader looks for them, +and nothing breaks if one is deleted along with an otherwise-empty directory. +Do not add new ones — a `/` directory is created populated. --- -## 2. GEMM resolution order — `get_gemm_config()` +## 2. Resolution — `resolve_config_dir()` -`utils/gemm_config_utils.py` picks a directory by probing for the *default* -config file (`DEFAULT.json` in the nested layout, `-.json` -in legacy) in order and taking the first hit. Specialized files are then read -from that same directory. +```python +resolve_config_dir(op, config_name, backend="triton", arch=None) -> str +``` -**`backend=None`** (what every caller uses today): +in `utils/config_utils.py` is the single path builder. It **builds** the +directory; it never probes, never searches, and has no fallback chain: -1. `configs//triton/gemm//DEFAULT.json` -2. `configs//gluon/gemm//DEFAULT.json` -3. `configs/gemm/-.json` *(legacy)* +``` +{AITER_TRITON_CONFIGS_PATH}/{arch}/{backend}/{op}/{_dtype_dir(config_name)} +``` -**`backend="triton"|"gluon"`**: +- `arch` defaults to `arch_info.get_arch()`. The `arch=` argument is an + explicit override for loaders that deliberately retry under another + architecture (MHC's gfx942 fallback, §5.3) — not a search order. +- `backend` is **declared by the caller** and defaults to `"triton"`. Gluon + kernels and gluon dispatch paths pass `"gluon"`. There is no cross-backend + search: the two backends take disjoint config params, so a config tuned for + the other backend is not usable, and silently borrowing one would be a bug, + not a convenience. +- Every argument becomes a path component, so each is validated against a + whitelist and the function **fails closed** with an `AssertionError`: + + | Argument | Pattern | Why | + | -------- | ------- | --- | + | `op` | `[a-z][a-z0-9_]*` | directory name in the tree | + | `config_name` | `[A-Za-z0-9][A-Za-z0-9_-]*` | folded into `` | + | `backend` | one of `("triton", "gluon")` | | + | `arch=` override | `[a-z][a-z0-9_]*` | programmer-written literal | + | running arch | `[A-Za-z0-9][A-Za-z0-9_.+:-]*` | driver-derived; tolerates vendor formats but must stay path-safe | + +- Existence is **not** checked. Whether a given file inside the directory has + to exist is the loader's decision, expressed through `load_config_json()`. + +### `load_config_json()` -1. `configs///gemm//DEFAULT.json` -2. `configs/gemm//-.json` *(legacy)* -3. `configs/gemm/-.json` *(legacy)* +```python +load_config_json(fpath, required=True) -> dict | None +``` -If nothing matches, the last legacy candidate is used and the missing-default -assertion fires there — so error messages still point at `configs/gemm/`. +- `required=True` (the default) raises `FileNotFoundError` naming the exact + nested path when the file is missing. That message is the error every + missing required table should produce — do not wrap it in a vaguer one. +- `required=False` returns `None`, for genuinely optional tables (probes, + arch-specific extras, per-path fallbacks). Callers must handle `None` + explicitly. +- Cached per path with `functools.lru_cache`, **including negative results**. + Adding a config file at runtime therefore has no effect: restart the + process, or call `load_config_json.cache_clear()` from tooling. Exceptions + are never cached, so a missing required file raises consistently on every + call. +- The returned dict is the **shared cached object**. Copy before mutating — a + shallow `.copy()` for flat bucket dicts, `copy.deepcopy` when nested + sub-dicts get mutated. The family loaders already do this for their callers. -The legacy candidates are marked `# TODO(satya): legacy, remove` and are -scheduled for deletion. Do not write new code that depends on them resolving. +--- -Consequences to keep in mind: +## 3. Loader modules + +`utils/config_utils.py` is the shared core: the config-tree paths, +`load_config_json()`, `_dtype_dir()` and `resolve_config_dir()`. Each op +family keeps its own small module built on that core. Every function has +exactly one home; there is no facade or re-export layer. + +| Module | Entry points | Reads | +| ------ | ------------ | ----- | +| `utils/config_utils.py` | `resolve_config_dir`, `load_config_json`, `AITER_TRITON_CONFIGS_PATH`, `AITER_TRITON_OPS_PATH`, `USE_LRU_CACHE` | — (core) | +| `utils/gemm_config_utils.py` | `get_gemm_config`, `add_default_gemm_config_params`, `compute_splitk_params`, `pick_gemm_num_stages`, `STANDARD_M_BOUNDS` | `//gemm//` | +| `utils/conv_config_utils.py` | `get_conv_config`, `has_conv_config`, `has_exact_conv_config`, `conv_config_uses_exact_routes`, `format_shape_key`, `format_prepack_shape_key`, `CONV_STANDARD_M_BOUNDS` | `/triton/conv//` | +| `utils/mhc_config_utils.py` | `get_mhc_config`, `get_mhc_post_config`, `hip_post_dispatch_block` | `/triton/mhc//` (gfx942 fallback) | +| `utils/moe_config_utils.py` | `get_moe_dispatch` | `//moe//` | +| `utils/tuned_config_utils.py` | `get_tuned_kernel_config` | `////DEFAULT.json` | + +Attention and GMM kernels have no family module: they call +`resolve_config_dir()` + `load_config_json()` directly from their kernel file, +which is fine for a single `DEFAULT.json` read with no selection logic. + +Adding a family module is the right move only when a family grows real +selection logic (bucket walks, specialized-file discovery, fallbacks). Until +then, two lines against the core beat a module. + +Two places still interpolate `AITER_TRITON_CONFIGS_PATH` by hand instead of +calling the resolver: `tuned_config_utils._get_tuned_kernel_entry()` and +`fusions/fused_clamp_act_mul.py::_get_config()` (whose gfx950 fallback is +exactly the resolver's `arch=` override). Both land on the same directory the +resolver would build, but they re-encode the layout and skip the argument +validation. Move them onto `resolve_config_dir()` when you touch them; do not +add a third. -- **A directory is chosen as a unit.** The unit is the family's `/` - directory. Splitting a config family across `/triton/gemm//` - and legacy `configs/gemm/` silently drops the specialized files in whichever - directory loses the probe. Move a family wholesale or not at all. Worse: a - `/` directory with specialized files but **no `DEFAULT.json` is - invisible** — the probe keys only on `DEFAULT.json` and falls through to - legacy, ignoring everything in the directory. -- **`backend=None` prefers `triton` over `gluon`.** On an arch with only a - 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 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_afp4wfp4.py` goes through `get_gemm_config(backend="gluon")` and -needs no changes. +## 4. GEMM — `get_gemm_config()` ---- +```python +get_gemm_config(config_name, M, N=None, K=None, bounds=None, + specialized_filename=None, backend="triton", B=None) + -> (config: dict, is_tuned: bool) +``` + +Order of operations: + +1. `//gemm//DEFAULT.json` — **must exist**, else + `AssertionError` naming the path. +2. Specialized file, first hit wins: + `{config_name}-B={B}-N={N}-K={K}.json` (when `B` is given), then + `{config_name}-N={N}-K={K}.json`; or `{config_name}-{specialized_filename}.json` + when the caller passes one (fused kernels with several N dims), which + bypasses the B/N/K candidates. +3. Inside the chosen file: `M_LEQ_x` ascending, then `M_GEQ_x` descending, + then `"any"`. `KeyError` if nothing matches. + +`is_tuned` is `True` only when a specialized file was hit — `False` for the +default file and for `"any"`. **Do not discard it**; it is how callers and +tuning tooling detect a shape running on untuned numbers. Call sites that +legitimately ignore it use `config, _ = _get_config(...)`. -## 3. GEMM config file contents +The returned config is a fresh deep copy, safe to mutate. -Required top-level shape: +### File contents ```json { @@ -141,259 +188,249 @@ Required top-level shape: } ``` -- `M_LEQ_x` is searched ascending over `STANDARD_M_BOUNDS = - (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. +- `M_LEQ_x` is searched over `STANDARD_M_BOUNDS = + (1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)`. 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`. + A `KeyError` at lookup time usually means it is missing. +- The deprecated `{"large": …, "small": …}` shape must not be introduced. +- Each `M_*` entry carries at minimum: -Each `M_*` entry carries at minimum: + ``` + BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, + num_warps, num_stages, waves_per_eu, matrix_instr_nonkdim, + cache_modifier, NUM_KSPLIT + ``` -``` -BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, -num_warps, num_stages, waves_per_eu, matrix_instr_nonkdim, -cache_modifier, NUM_KSPLIT -``` - -`add_default_gemm_config_params()` backfills `NUM_KSPLIT=1` and -`cache_modifier=None` as a last resort, and `compute_splitk_params()` derives -`SPLITK_BLOCK_SIZE` and may clamp `BLOCK_SIZE_K` / `NUM_KSPLIT`. Neither is a -license to omit keys. - -`get_gemm_config()` returns `(config, is_tuned)`. `is_tuned` is `True` only when -a specialized (N/K-, B-, or `specialized_filename`-keyed) file was hit, `False` -for the default file or `any`. Do not discard it. - -### The JSON is the only place tuning values live + `add_default_gemm_config_params()` backfills `NUM_KSPLIT=1` and + `cache_modifier=None` as a last resort, and `compute_splitk_params()` + derives `SPLITK_BLOCK_SIZE` and may clamp `BLOCK_SIZE_K` / `NUM_KSPLIT`. + Neither is a license to omit keys. -A `_get_config()` should do nothing but call `get_gemm_config()` and return: +### `_get_config()` stays a thin wrapper ```python -def _get_config(M: int, N: int, K: int): - return get_gemm_config("GEMM-AFP4WFP4", M, N, K) +def _get_config(M: int, N: int, K: int, backend: str = "triton"): + return get_gemm_config("GEMM-A16W16", M, N, K, backend=backend) ``` -`_triton_kernels/gemm/basic/gemm_afp4wfp4.py` carried a block of `setdefault` -calls and it was deleted — it masked incomplete config files with values nobody -had tuned, and made the effective config un-inspectable from the JSON. +A kernel-level `_get_config()` that takes a `backend` argument **defaults it +to `"triton"`**, never to `None`: `None` is not a backend, and +`resolve_config_dir()` rejects it. Public wrappers that expose +`backend: str | None = None` normalize it (typically to `"gluon"` on gfx1250, +`"triton"` elsewhere) before calling down. ---- - -## 4. GEMM naming +`_triton_kernels/gemm/basic/gemm_afp4wfp4.py` once carried a block of +`setdefault` calls; it was deleted, because it masked incomplete config files +with values nobody had tuned and made the effective config un-inspectable from +the JSON. -| Kind | Target layout | Legacy layout | -| ----------------- | ----------------------------------------------- | ---------------------------------------------------- | -| Default | `gemm_a16w16/DEFAULT.json` | `gfx950-GEMM-A16W16.json` | -| N/K specialized | `gemm_a16w16/GEMM-A16W16-N=256-K=7168.json` | `gfx950-GEMM-A16W16-N=256-K=7168.json` | -| Batched (B, N, K) | `batched_gemm_a16w16/BATCHED_GEMM-A16W16-B=4-N=1024-K=4096.json` | `gfx1250-BATCHED_GEMM-A16W16-B=4-N=1024-K=4096.json` | -| Custom suffix | `fused_gemm_afp4wfp4_a16w16/FUSED-GEMM-AFP4WFP4-A16W16-N4=512-N16=256-K=7168.json` | same, arch-prefixed | +### Naming -The `` directory name is `config_name.lower().replace("-", "_")`. -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. +| Kind | Path | +| ----------------- | ---------------------------------------------------------------- | +| Default | `gemm_a16w16/DEFAULT.json` | +| N/K specialized | `gemm_a16w16/GEMM-A16W16-N=256-K=7168.json` | +| Batched (B, N, K) | `batched_gemm_a16w16/BATCHED_GEMM-A16W16-B=4-N=1024-K=4096.json` | +| Custom suffix | `fused_gemm_afp4wfp4_a16w16/FUSED-GEMM-AFP4WFP4-A16W16-N4=512-N16=256-K=7168.json` | Config-name patterns: `GEMM-A{x}W{y}`, `BATCHED_GEMM-A{x}W{y}`, -`FUSED-GEMM-{op}`, `FF-A{x}W{y}-fused`; variant suffixes -`_PRESHUFFLED`, `_BLOCKSCALE`. +`FUSED-GEMM-{op}`, `FF-A{x}W{y}-fused`; variant suffixes `_PRESHUFFLED`, +`_BLOCKSCALE`. + +Dashes, underscores and case all fold together in ``, so new config +names must stay distinct under that transform — `GEMM-FOO-BAR` and +`GEMM-FOO_BAR` would collide on one directory. **`K` in AFP4WFP4 filenames is the logical K, i.e. `2 * K_bytes`.** The kernel -does `K = 2 * K` before calling `get_gemm_config`. Tuning output that names +does `K = 2 * K` before calling `get_gemm_config()`. Tuning output that names files by the packed byte width will never be found. --- -## 5. MOE configs - -MOE does **not** go through `get_gemm_config()`. There is no probe order, no -nested-layout support, and no `is_tuned` signal. Four independent loaders read -`configs/moe/` directly, each with its own schema: - -| Loader | File | Schema | -| ------ | ---- | ------ | -| `utils/moe_config_utils.py::get_moe_configs` | `moe/-MOE-.json` | `small_M` / `medium_M` / `large_M` | -| `moe/moe_op_gemm_a8w4.py::_get_a8w4_dispatch` | `moe/-A8W4.json` | `bm_n_k` | -| `moe/moe_op_gemm_a4w4.py::_get_a4w4_dispatch` | `moe/-A4W4.json` | `bm_n_k_` | -| `_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py` | `moe/-MOE_ROUTING_SIGMOID_TOPK1.json` | `N16` → `small` / `medium` / … | +## 5. The other families -`` comes from `get_config_dtype_str()`: `DEFAULT`, `FP8_W8A8`, -`INT8_W8A16`, `INT8_W8A8`, `INT4_W4A16`, `MX_FP4`. +### 5.1 MOE — `get_moe_dispatch()` -`small_M` / `medium_M` / `large_M` split on `M_THRESHOLD_SMALL = 256` and -`M_THRESHOLD_MEDIUM = 1024`, both module constants in `moe_config_utils.py`. -This is **not** the GEMM `M_LEQ_x` / `M_GEQ_y` scheme — do not mix them. - -`A4W4` feeds the **gluon path only** (`get_kernel_config_gluon`); a4w4's triton -path still computes its config in Python. Its `` is a *third* M scheme — -`m2bucket()` in `moe_op_gemm_a4w4.py`, splitting on 8 / 32 / 128 / 256 / 512 into -`tiny` / `small` / `medium` / `medium2` / `large` / `xlarge`. Lookup is two tiers: -`bm_n_k_`, then `bm_any`. Since a missing bucket -falls all the way through to `_any` and loses the shape's tuning, a tuned shape -must supply **all six** buckets, even where the values repeat. +```python +get_moe_dispatch(config_name, arch, backend) -> dict +``` -### MOE is the main offender for tuning values in Python +is the **only** MOE config fetcher. It reads +`//moe//DEFAULT.json` and returns `{}` when this arch +and backend ship no tuned file, so callers fall through to their own defaults +instead of crashing. `arch` is passed in by callers that already resolved it +(it keys the cache); `resolve_config_dir()` reads the same value. + +The returned dict is the shared cached object — **read-only**. + +The two dispatch paths key the same family differently, which is exactly why +`backend` is a caller-declared argument rather than something the loader +guesses: + +| Family | Backend | Key | Entry keys | +| ------ | ------- | --- | ---------- | +| `A8W4` | `triton` | `bm_n_k` | `BLOCK_SIZE_N`, `BLOCK_SIZE_K`, `num_warps`, `num_stages`, `waves_per_eu`, `matrix_instr_nonkdim` | +| `A8W4` | `gluon` | `bm_n_k_`, then `bm_any` | `block_n`, `block_k`, `num_buffers`, `num_warps`, `persistent_iters` | +| `A4W4` | `gluon` | `bm_n_k_`, then `bm_any` | `block_n`, `block_k`, `num_buffers`, `num_warps` | + +- Entries omit `BLOCK_SIZE_M` / `block_m` on purpose: `block_m` is the + dispatch **key**, decided by routing, not a tunable. +- `` is `m2bucket()`, splitting M on 8 / 32 / 128 / 256 / 512 into + `tiny` / `small` / `medium` / `medium2` / `large` / `xlarge`. + `moe_op_gemm_a8w4.py` and `moe_op_gemm_a4w4.py` currently carry identical + copies of it — if you touch one, they belong in a shared home, not diverged. + A missing bucket falls all the way through to `bm_any`, which + loses that shape's tuning entirely — so a newly tuned shape should carry all + six, even where values repeat. Coverage in the shipped tables is uneven + (`a4w4` is complete; most `a8w4` gluon shapes cover only the M range they + were measured over), which is a gap to close, not a pattern to copy. +- The `bm_any` tier must exist in every gluon dispatch file — it is + the last resort for an unmeasured shape. +- `moe_op_gemm_a8w4.py`'s triton path additionally derives a proxy from a + tuned entry with the same `(N, K)` under a different `block_m` before + reaching its Python default. That proxy reads tuned numbers out of JSON; the + final default does not, and is the tier to delete as coverage grows. + +### 5.2 Conv — `get_conv_config()` -Fix these as you touch them; do not add more: +```python +get_conv_config(config_name, shape_key=None, M=None, variants=()) -> dict +``` -- `get_optimal_moe_config()` returns a hardcoded dict (`BLOCK_SIZE_M: 256`, - `BLOCK_SIZE_N: 256`, …) when no config file exists, behind a - `warnings.warn`. A missing config silently runs untuned values. -- `moe_op_gemm_a8w4.py` has a three-tier Python fallback: exact - `bm_n_k` hit → any-`block_m` proxy with matching `(N, K)` → a gfx942-gated - shape heuristic → a conservative default. Only the first tier reads tuned - numbers from JSON. +reads `/triton/conv//DEFAULT.json` (conv is triton-only) and +walks four tiers, first hit wins: -### Planned: `get_moe_config()` — design, not yet implemented +1. `shapes_[shape_key]` — optional variant-specific pin + (`shapes_nchw`, `shapes_nhwc`). +2. `shapes[shape_key]` — generic exact-shape pin. +3. `M_LEQ_` — bucket walk over `CONV_STANDARD_M_BOUNDS` + (`4 … 262144`), on `M_total` for the GEMM-like kernels, `T` for Winograd. +4. `"any"` — global fallback. -**Status: design only. Nothing below exists in the tree yet.** Do not assume -`get_moe_config` is importable; do not move MOE JSON files in anticipation of -it. If you are asked to implement it, follow this shape. +Shape keys are built by `format_shape_key()` +(`N=…,C=…,H=…,W=…,K=…,R=…,S=…,sh=…,sw=…,ph=…,pw=…,dh=…,dw=…`) and +`format_prepack_shape_key()` (`N=…,C=…,H=…,W=…,CB=…`) — never hand-formatted at +a call site. `route_exact_only: true` in a file restricts routing to exact +entries, read through `conv_config_uses_exact_routes()`. -The unification is **path resolution only.** The three MOE schemas stay as they -are — the loader finds and parses the file, each caller keeps interpreting its -own structure. Converging the schemas is a separate, later decision (it would -touch every MOE config file and require re-validating dispatch on every arch). +`has_conv_config()` probes whether the running arch ships an optional table at +all (it passes `required=False`); the other entry points require the file. +Families: `CONV-1X1`, `CONV-3X3-NHWC`, `CONV-3X3-CBLOCKED`, `CONV-3X3-NCHW`, +`CONV-GENERAL`, `CONV-PREPACK`, `CONV-WINO-F4X3-{INPUT,GEMM,OUTPUT}`. -**Step 1 — extract the shared probe.** The candidate-directory logic currently -inlined in `_get_gemm_config_cached()` becomes a helper in `utils/`, parameterised -on ``: +### 5.3 MHC — `get_mhc_config()` ```python -def resolve_config_dir(op: str, config_name: str, backend: str | None = None, - legacy_dir: str | None = None) -> tuple[str, str]: - """Return (cfg_dir, name_prefix) for the first candidate whose default - file exists: DEFAULT.json when name_prefix is empty (the nested layout, - dir from _dtype_dir()), else .json. Falls back - to the last candidate so the missing-file assertion names a legacy path.""" +get_mhc_config(config_name, M, C, mode=None) -> (config, used_specialized) +get_mhc_post_config(M, C) -> dict ``` -Candidates for `op="moe"`, mirroring §2: +`mode` is required in practice: anything other than `"sinkhorn"` — including +the `None` default — raises `ValueError`. -- `backend=None` → `/triton/moe//DEFAULT.json`, - `/gluon/moe//DEFAULT.json`, then legacy `configs/moe/` with - the `-` prefix -- `backend=...` → `//moe//DEFAULT.json`, then legacy - `configs/moe/` prefixed +The family directory is `/triton/mhc//`, with `` derived +from `f"{config_name}_{mode.upper()}"` — `MHC_FUSED` + `sinkhorn` → +`mhc_fused_sinkhorn`. Selection is C first, then M: -`` comes from the unprefixed stem via the same transform: -`MOE-FP8_W8A8` → `moe_fp8_w8a8`, `A8W4` → `a8w4`, -`MOE_ROUTING_SIGMOID_TOPK1` → `moe_routing_sigmoid_topk1`. +- C: the largest `-C=` file threshold `<= C` wins. Available thresholds + are **discovered by globbing** `MHC_FUSED_SINKHORN-C=*.json` — in the running + arch's directory *and* gfx942's, unioned — so a new specialized file becomes + reachable just by being added (subject to the load cache), and a threshold + that exists only under gfx942 is still a candidate on other arches. +- M: within the selected file, the largest `M_LEQ_ <= M`, else `"any"`. -`gemm_config_utils.py` is then refactored onto the same helper with `op="gemm"` -and `legacy_dir="gemm"` — behaviour-identical, and the legacy candidates stay -tagged `# TODO(satya): legacy, remove` so both ops retire together. +**Arch fallback:** an arch with no MHC directory falls back to the `gfx942` +files, via the `arch=` override on `resolve_config_dir()`. This is a +deliberate, documented exception — it keeps MHC running (possibly +suboptimally) on untuned hardware. Do not copy the pattern into other +families without the same explicit justification. -**Step 2 — the loader.** +`get_mhc_post_config()` reads `mhc_post/DEFAULT.json` and picks the largest +`C_ <= C`, else `"default"`. `hip_post_dispatch_block()` mirrors +`MHC_POST_KERNEL_DISPATCH` in `csrc/kernels/mhc_kernels.cu` and belongs next +to it in review. + +### 5.4 Pinned autotune tiles — `get_tuned_kernel_config()` ```python -def get_moe_config(config_name: str, backend: str | None = None) -> dict | None: - """Load a MOE config by name. Returns the parsed JSON (a deep copy, safe to - mutate), or None if no file exists for this arch. Schema interpretation is - the caller's job — MOE files are not uniform.""" +get_tuned_kernel_config(op, config_name, kernel_name, fallback, backend="triton") + -> triton.Config ``` -- `config_name` is the unprefixed stem: `MOE-FP8_W8A8`, `MOE-DEFAULT`, `A8W4`, - `MOE_ROUTING_SIGMOID_TOPK1`. -- Cache with `functools.lru_cache` and deep-copy on return, exactly as - `get_gemm_config()` does — callers mutate configs. -- Return `None` rather than raising: unlike GEMM, a missing MOE config is - currently normal (only gfx942 and gfx950 ship any). -- No `is_tuned` flag. MOE has no default-vs-specialized distinction to report. - -**Step 3 — port the three loaders** onto it, one PR each, without changing -schemas or file locations: - -| Loader | Call becomes | -| ------ | ------------ | -| `moe_config_utils.py::get_moe_configs` | `get_moe_config(f"MOE-{dtype_str}")` | -| `moe_op_gemm_a8w4.py::_get_a8w4_dispatch` | `get_moe_config("A8W4") or {}` | -| `moe_routing_sigmoid_top1_fused.py` | `get_moe_config("MOE_ROUTING_SIGMOID_TOPK1")` | - -**Step 4 — delete the hardcoded Python fallbacks.** Blocked on shipping a -`-MOE-DEFAULT.json` for the arches that have none — today only gfx942 and -gfx950 have MOE configs at all, so gfx1250/gfx1151/gfx1200/gfx1201 hit the -hardcoded dict in `get_optimal_moe_config()`. Once every supported arch has a -file, that dict and the `warnings.warn` come out. +For kernels whose autotune search space lives in Python and only need **one +pinned tile per device**. It reads `//DEFAULT.json` for the running +arch and looks up `kernel_name` inside it, returning `fallback` (with a +warning) when the arch publishes no entry. -**Step 5 — only now `git mv`** MOE files into `//moe/` and drop -the arch prefix, following §6. +The `fallback` must be **launchable anywhere**, not fastest somewhere: the same +tile can fit in 16 KB of LDS on one arch and overflow another's 64 KB. An +unmeasured device stays on the fallback until a measured entry is published. -Moving MOE JSON before step 2 lands would silently break resolution, and the -Python fallbacks would swallow the breakage instead of surfacing it. +File shape: `{"": {tile keys…, "num_warps": n, "num_stages": n}}`. --- -## 6. Migration playbook (GEMM) - -One family = one `` × `` × ``, including every -specialized file. `GEMM-AFP4WFP4` is the worked example — diff it against its -legacy form if a step is ambiguous. - -1. **Scope it.** `ls configs/gemm/-*.json` — the default plus - every specialized file. All of them move together. -2. **Find every reader.** Grep for the config name and for - `AITER_TRITON_CONFIGS_PATH` in `aiter/ops/triton/`. Hand-built paths must be - rewritten; `get_gemm_config()` callers need no change. -3. **Move with `git mv`** so the change reviews as a rename. Strip the - `-` prefix; the default file becomes `DEFAULT.json`, specialized - files keep the `-` stem: - ``` - git mv configs/gemm/gfx950-GEMM-FOO.json \ - configs/gfx950/triton/gemm/gemm_foo/DEFAULT.json - git mv configs/gemm/gfx950-GEMM-FOO-N=1-K=2.json \ - configs/gfx950/triton/gemm/gemm_foo/GEMM-FOO-N=1-K=2.json - ``` - `mkdir` the `/` directory first; it needs no `.gitkeep` (it is - created populated). Create `//{gemm,moe}/` with a `.gitkeep` - if absent. -4. **Do not edit contents in the same commit.** Keep renames at 100% similarity; - content changes go in a follow-up commit. -5. **Update docs.** `aiter/ops/triton/README.md` ("How GEMM configs resolve", - "Config naming") and `aiter/ops/triton/utils/_triton/tunning/README.md` - (the copy step under "Verify performance") both describe the two layouts — - keep them current if a migration changes what they say. -6. **Pull any tuning values still hardcoded in Python into the JSON.** A - migrated family must be fully described by its config files. -7. **Verify** on the target arch: config resolves, `is_tuned` is `True` for a - shape that has a specialized file, and numerics are unchanged. -8. Leave the `# TODO(satya): legacy, remove` markers in `gemm_config_utils.py` - until `configs/gemm/` is empty. Deleting the legacy fallback is the final - step of the migration, not an intermediate one. - -### Adding a *new* tuned config (no migration) - -- **GEMM**, new arch/backend combination or a family already migrated → - target layout, inside the family's `/` directory (`DEFAULT.json` - for the default; specialized files keep the config-name stem, no arch - prefix). -- **GEMM**, family still in `configs/gemm/` → add to `configs/gemm/` with the - arch prefix, and consider migrating the whole family in the same PR. Never - create a lone nested file for a family whose default lives in legacy; the - directory probe picks one directory and ignores the other. -- **MOE** → `configs/moe/` with the arch prefix, matching the schema of the - loader that will read it. The nested layout is not wired up for MOE. +## 6. Adding a config + +1. **Name the family.** Pick the ``, check `` does not + collide with an existing directory under that arch/backend/op. +2. **Put the default in place**: `////DEFAULT.json`. + No arch prefix in the filename. A `/` directory whose required + default is missing raises at first lookup naming that exact path — that is + the intended failure, not something to paper over with a fallback. +3. **Specialized files** keep the `-` stem next to the default: + `GEMM-A16W16-N=256-K=7168.json`, `MHC_FUSED_SINKHORN-C=128.json`. +4. **Pull any tuning value still hardcoded in Python into the JSON.** A family + must be fully described by its config files. +5. **Verify on the target arch**: the config resolves, `is_tuned` is `True` + for a shape that has a specialized file, and numerics are unchanged. +6. If the change moves or renames files, keep the commit a pure `git mv` + (100% rename similarity) and put content edits in a follow-up commit. + +### Seeding a new architecture + +An arch that ships no file for a family gets whatever that family's loader +does on a miss — a hard error for GEMM/conv/MHC required tables, `{}` for the +MOE dispatch, the `fallback` tile for `get_tuned_kernel_config()`. Where that +is not acceptable, seed the directory with a **byte-identical copy** from the +closest measured arch and say so in the commit message. + +The one seeding rule currently in force: **gfx950 → gfx1250, triton only.** +Never seed a gluon directory from another arch (gluon configs carry +arch-specific tile and buffer counts), and never seed backwards into gfx950. +A seed is a placeholder that unblocks the caller-declared backend policy — it +is not tuning, and it should be replaced by measured numbers. ### Do not -- Rename or delete `.gitkeep` placeholder directories. -- Put an arch prefix on a file inside `/...`, or name a nested default - anything other than `DEFAULT.json`. +- Put an arch prefix on a file inside `/…`, or name a default anything + other than `DEFAULT.json`. - Put tuning values in `.py` files — no `setdefault`, no inline dicts, no arch-conditional constants, no hardcoded fallback configs. -- Mix the GEMM `M_LEQ_x`/`M_GEQ_y` scheme with the MOE - `small_M`/`medium_M`/`large_M` scheme. -- Move MOE configs into `//moe/` before a resolver exists. +- Mix key schemes: GEMM `M_LEQ_x`/`M_GEQ_y`, MOE `bm…` dispatch keys, conv + shape keys and MHC `C_`/`M_LEQ_` all belong to their own families. +- Reintroduce a probe, a candidate list, or a cross-backend/cross-arch search + in `resolve_config_dir()`. A miss is an error with a path in it. +- Add `kpack` to a gfx950 config. Triton's AMD backend deprecates `kpack` on + CDNA4 — it warns and force-overrides `kpack = 1` there, and the parameter is + slated for removal. No gfx950 config carries it today; keep it that way. + gfx942 configs still may. The RDNA trees (gfx1151, gfx1201, gfx1250) do + carry `kpack` in places — those entries predate the rule, so do not add new + ones and drop them when you retune the family. ### Not tuning configs Two AOT code paths build directories under this tree at runtime that are **not checked in and out of scope**: -- `configs/gemm/aot/_M=…-N=…-K=…` — `gemm/fused/fused_gemm_afp4wfp4_a16w16.py`, +- `configs/gemm/aot/_M=…-N=…-K=…` — + `gemm/fused/fused_gemm_afp4wfp4_a16w16.py`, `gemm/fused/fused_gemm_afp4wfp4_mul_add.py` - `configs/paged_mqa_logits/aot/` — `attention/pa_mqa_logits.py` -Both are guarded by `use_aot and os.path.exists(...)` and hold compiled-kernel -metadata, not tuning parameters. Do not create, migrate, or document them as -config directories. +Both hold compiled-kernel metadata, not tuning parameters: the GEMM ones are +written only under `use_aot and os.path.exists(...)`, and the +`paged_mqa_logits` one only on the AOT gluon branch. Do not create, migrate, +or document them as config directories. diff --git a/aiter/ops/triton/conv/DESIGN.md b/aiter/ops/triton/conv/DESIGN.md index 158ee22181..cd07bcd18e 100644 --- a/aiter/ops/triton/conv/DESIGN.md +++ b/aiter/ops/triton/conv/DESIGN.md @@ -152,7 +152,7 @@ def conv2d(x, w_oihw, bias=None, stride=(1,1), padding=(0,0), dilation=(1,1), The semi-public method-specific functions (`conv2d_nchw_cblocked`, `conv2d_winograd_f4x3_cblocked`, …) take an internal `block_k=64` channel-block and padding granularity. It is intentionally not surfaced on the public -`conv2d` — every shipped config in `configs/conv/` was tuned with 64, so the +`conv2d` — every shipped conv config was tuned with 64, so the parameter has no good user story. `x.dtype` is validated at entry: anything other than `torch.float16` / @@ -247,7 +247,7 @@ into the user's chosen output layout. ### 5.0 A platform note on `num_stages` -Every shipped config in `configs/conv/` pins **`num_stages=1`**. That is +Every shipped conv config pins **`num_stages=1`**. That is deliberate. The `num_stages > 1` Triton knob is meant to lower to a software-pipelined @@ -281,7 +281,7 @@ The kernel fuses this with the index unwrap `m → (n, p, q)` and a Highlights: - **Tile shape:** `BLOCK_M × BLOCK_N × BLOCK_K`, set per-arch in - `aiter/ops/triton/configs/conv/{arch}-CONV-1X1.json` and loaded at launch + `aiter/ops/triton/configs/{arch}/triton/conv/conv_1x1/DEFAULT.json` and loaded at launch via `get_conv_config("CONV-1X1")`. `num_stages=1` always — see 5.0. - **L2 cache swizzle.** Tiles are reordered into super-groups of `GROUP_SIZE_M` along the `M` axis so each weight (`N`-axis) tile is reused @@ -312,7 +312,7 @@ NHWC-native 3×3 with **K-major weight layout** `W3[K_out, 9, C_pad]`: emits one vectorized load per row. - **Same L2 swizzle as 5.1**: workgroups are reordered into super-groups of `GROUP_SIZE_M` along the `M` axis (set to 4 or 8 in - `configs/conv/{arch}-CONV-3X3-NHWC.json`) so each weight (`N`-axis) tile + `configs/{arch}/triton/conv/conv_3x3_nhwc/DEFAULT.json`) so each weight (`N`-axis) tile stays hot in L2 across `GROUP_SIZE_M` consecutive workgroups. ### 5.3 `_conv2d_3x3_cblocked_kernel` @@ -559,7 +559,7 @@ The main input and packed-weight layouts are: | `[36, K_out, C_pad]` (Winograd weight) | `winograd_f4x3_*` | `prepack_winograd_filter_f4x3` | LRU 256, `_PACK_CACHE_WINOGRAD_F4X3` | The fused NCHWc launcher in `_launch.py` loads its parameters from -`configs/conv/{arch}-CONV-PREPACK.json`. It uses the canonical key +`configs/{arch}/triton/conv/conv_prepack/DEFAULT.json`. It uses the canonical key `N=...,C=...,H=...,W=...,CB=...`; exact pins cover tuned shapes, then architecture-specific `M_LEQ_*` buckets and `any` cover unseen shapes by spatial size. @@ -647,12 +647,23 @@ If a new kernel needs the Winograd bump, mark it `is_winograd=True` in Each kernel fetches its launch parameters (tile sizes, `num_warps`, etc.) at launch time via `get_conv_config(...)` in `aiter/ops/triton/utils/conv_config_utils.py`, which reads the per-arch JSON in -`configs/conv/{arch}-CONV-{KERNEL}.json`. There is **no runtime autotune on the +`configs/{arch}/triton/conv//DEFAULT.json`, `` being the config +name lowercased with dashes folded to underscores (`CONV-3X3-NHWC` → +`conv_3x3_nhwc`). There is **no runtime autotune on the hot path** — the search already happened offline and its winners are frozen in JSON, so first-call latency and CI compile time stay predictable. Candidate search spaces are maintained by offline tuning tooling, outside the production package. +The directory itself is built by `resolve_config_dir("conv", config_name, +backend="triton")` in `utils/config_utils.py` — the same deterministic path +builder every op family uses. Conv is triton-only, so it always declares +`backend="triton"`; there is no probing and no cross-arch search, and an arch +that ships no file for a required table raises with that exact path in the +message. Optional tables (the direct-NCHW route) are probed instead with +`has_conv_config()`, which loads with `required=False` and reports whether the +running arch ships one at all. + ### The JSON is a keyed lookup, not a list of layers A common confusion: "ResNet-50 has 53 conv layers, but the JSON only has ~23 @@ -730,8 +741,9 @@ Concretely, to add (say) a `winograd_f6x3` variant: 1. **Implement the kernel.** New file under `aiter/ops/triton/_triton_kernels/conv/`. Add a `_get_config()` helper that calls `get_conv_config("CONV-")`, - and ship `{arch}-CONV-.json` files under `aiter/ops/triton/configs/conv/` - for each supported arch. + and ship a `DEFAULT.json` under + `aiter/ops/triton/configs/{arch}/triton/conv/conv_/` for each supported + arch. 2. **Add a launch wrapper** in `_launch.py` (`_launch_winograd_f6x3`) that sets up the grid and turns Python ints into Triton constexprs. 3. **Add a public function** in `conv2d.py` diff --git a/aiter/ops/triton/conv/README.md b/aiter/ops/triton/conv/README.md index 73705f8aff..30bd42750d 100644 --- a/aiter/ops/triton/conv/README.md +++ b/aiter/ops/triton/conv/README.md @@ -208,15 +208,17 @@ Tested on ROCm 7.2 / PyTorch `2.9.1+gitff65f5b` / Triton 3.7 (commit `23f4e522d` ### Tuning -Per-kernel configs ship as JSON under `aiter/ops/triton/configs/conv/`, one -file per `(arch, kernel)` — e.g. `gfx1201-CONV-3X3-NHWC.json` and -`gfx1201-CONV-PREPACK.json`. The loader walks four tiers: layout-specific shape -pin → generic shape pin → `M_LEQ_x` bucket → `"any"` fallback. No -runtime autotune in the hot path, so CI compile time stays predictable and -the first call hits no tuning tax. - -RDNA configs are available for gfx1100, gfx1151, gfx1200, and gfx1201. The -optional direct-NCHW table exists for gfx1100, gfx1151, and gfx1201. gfx1100 +Per-kernel configs ship as JSON in the nested config tree, one `DEFAULT.json` +per `(arch, kernel)` under `aiter/ops/triton/configs//triton/conv/` — +e.g. `gfx1201/triton/conv/conv_3x3_nhwc/DEFAULT.json` and +`gfx1201/triton/conv/conv_prepack/DEFAULT.json`. The loader walks four tiers: +layout-specific shape pin → generic shape pin → `M_LEQ_x` bucket → `"any"` +fallback. No runtime autotune in the hot path, so CI compile time stays +predictable and the first call hits no tuning tax. + +Conv config trees ship for gfx1100, gfx1151, gfx1200, gfx1201 and gfx1250 on +the RDNA side, and for gfx942 and gfx950 on CDNA. The optional direct-NCHW +table exists for gfx1100, gfx1151, and gfx1201. gfx1100 and gfx1151 use exact-shape routing, so shapes not measured faster than the complete NCHWc path remain on the cblocked fallback; gfx1201 uses the configurable spatial crossover. diff --git a/aiter/ops/triton/utils/_triton/tunning/README.md b/aiter/ops/triton/utils/_triton/tunning/README.md index 14f3ed3b3d..524800835c 100644 --- a/aiter/ops/triton/utils/_triton/tunning/README.md +++ b/aiter/ops/triton/utils/_triton/tunning/README.md @@ -60,18 +60,14 @@ Example 3: **Verify performance** -To verify that your tunned JSON config files actually is performant and can be correctly picked up by AITER, first you have to copy the generated JSON config files into the config tree. Where they go depends on whether the config family has been migrated to the nested layout yet (`/aiter/ops/triton/configs/CLAUDE.md` is the authoritative rulebook): +To verify that your tunned JSON config files actually is performant and can be correctly picked up by AITER, first you have to copy the generated JSON config files into the config tree. Every family lives in one nested layout, `configs/////` (`/aiter/ops/triton/configs/CLAUDE.md` is the authoritative rulebook). Files there carry **no arch prefix** — the arch is the directory — and the default file is named exactly `DEFAULT.json`, so drop the arch prefix when copying: -- Migrated families (e.g. `GEMM-AFP4WFP4`, `GEMM-AFP4WFP4_PRESHUFFLED`) live in `configs///gemm//`. Files there carry **no arch prefix** and the default file is named exactly `DEFAULT.json`, so drop the arch prefix when copying: + cp GEMM-AFP4WFP4_PRESHUFFLED-N=7168-K=2048.json \ + /aiter/ops/triton/configs/gfx950/triton/gemm/gemm_afp4wfp4_preshuffled/ - cp GEMM-AFP4WFP4_PRESHUFFLED-N=7168-K=2048.json \ - /aiter/ops/triton/configs/gfx950/triton/gemm/gemm_afp4wfp4_preshuffled/ +`` is the config name lowercased with dashes folded to underscores (`GEMM-AFP4WFP4_PRESHUFFLED` → `gemm_afp4wfp4_preshuffled`), and `` is `triton` unless you tuned the gluon kernel — the two backends read separate directories and never fall back to each other. -- Families still in the legacy flat layout keep the arch prefix and go to `configs/gemm/`: - - cp *.json /aiter/ops/triton/configs/gemm/ - -Two gotchas: a `/` directory is invisible to the resolver unless it contains a `DEFAULT.json`, and config reads are cached per path (including missing files), so restart the Python process after copying for the new files to be picked up. +Two gotchas: a family's `DEFAULT.json` must be in place before any specialized file resolves, and config reads are cached per path (including missing files), so restart the Python process after copying for the new files to be picked up. then, you can run, for example,