Repository navigation
feat(ck-tile): enable GEMM bridges on gfx1250 (MI400) — universal + grouped + multi-D - #10921
Merged
ozturkosu merged 5 commits intoAug 31, 2026
Merged
Conversation
This was referenced Aug 18, 2026
ozturkosu
marked this pull request as ready for review
August 18, 2026 05:36
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
Contributor
Author
gfx1250 (MI400) CI-config validation — this PR's four bridge opsAll four ops in this PR were run against their CI config on gfx1250 / MI400 (ROCm 10.1) and pass:
Each merged bridge pinned the CDNA MFMA |
ozturkosu
requested review from
ThruptiRajLakshmanaGowda,
illsilin and
yraparti
August 18, 2026 06:21
ozturkosu
force-pushed
the
users/muozturk/ck/gemm-bridges-gfx1250-enable
branch
from
August 18, 2026 06:44
83e69e6 to
ff22c88
Compare
ozturkosu
added a commit
that referenced
this pull request
Aug 18, 2026
…fx1250 (MI400) The batched-GEMM (#9306) and batched-contraction (#9328) TE->Dispatcher bridges combined in merged #10612 were CDNA-only (gfx90a/gfx942/gfx950, MFMA). Both keep their OWN _SUPPORTED_ARCHES/_SUPPORTED_ARCHS tuple (they do not inherit the shared gemm_utils gate), and both ship MFMA CI configs, so neither runs on gfx1250 as merged. This folds both variants onto the shared gfx1250 foundation (base branch users/muozturk/ck/gemm-bridges-gfx1250-enable, PR #10921): - batched_gemm_utils / batched_contraction_utils: add "gfx1250" to the per-variant supported-arch tuple (arch gate; codegen is arch-agnostic). - batched_gemm: gfx1250 CI config (WMMA warp_tile 16x16x32). - batched_contraction: gfx1250 CI configs (default + bridge, WMMA 16x16x32); the merged configs pin MFMA 32x32x64 / 32x32x16 which do not run on gfx1250. - both --arch help + config help mention gfx1250. - CPU-only tests: gfx1250 in supported arches, arch resolve/detect accepts gfx1250, gfx1250 CI configs present with WMMA 16x16x32, and the contraction gfx1250 sweep expands to valid WMMA kernels. Validation: 66 CPU tests pass (incl. the new gfx1250 tests). On-GPU validation on MI400/gfx1250 is pending node availability (both Helios nodes were unreachable / GPU-less at authoring time).
Contributor
Author
Combined-tree validation (single checkout of this branch) on MI400 / gfx1250Fresh single-checkout build + on-device validation of this PR's four GEMM-bridge ops, all from one clean checkout of
Per-op results (fp16)
* The 16 Overall: all four ops build and run correctly on gfx1250 from a single combined-branch checkout. |
Contributor
Author
|
This is PR 1 of 3 in the recommended review order. Review order & dependencies — gfx1250 (MI400) GEMM-bridge enablement seriesThese PRs enable the merged CK GEMM bridges on gfx1250 (MI400, RDNA4/WMMA). They all target 1️⃣ #10921 — foundation: gemm_universal + grouped + multi-D + multi-ABD. 2️⃣ #10928 — batched GEMM + batched contraction. 3️⃣ #10927 — grouped AQuant/ABQuant/BQuant. Merge mechanicsAfter #10921 merges to Separate, not in this order (need C++/kernel follow-ups, kept as drafts): #10916 (mx — GPU blocked by a gfx950-only |
Contributor
Author
📋 Merge order for the gfx1250 (MI400) bridge-enablement seriesThis is the coordination note for landing the gfx1250 enablement PRs. There are two independent tracks plus follow-ups. Track A — core GEMM bridges (must be ordered; they share files)1️⃣ #10921 — FIRST. Foundation: gemm_universal + grouped + multi-D + multi-ABD plus the shared gfx1250 enablement ( 2️⃣ #10928 — batched GEMM + batched contraction. After #10921. GPU-validated.
Track B — block-scale quant bridges (independent)#10439 — block-scale AQuant/ABQuant/BQuant/RowColQuant/TensorQuant. ✅ Independent of Track A (verified: 0 file overlap with #10921). It introduces its own bridges and its gfx1250 enablement is self-contained ( Follow-ups — not ready (need C++/kernel work)
TL;DR |
This was referenced Aug 18, 2026
ozturkosu
changed the title
Contributor
Author
|
Removed the multi-ABD gfx1250 enablement from this PR (commit multi-ABD gfx1250 support now lives in its own PR: #11048 (stacked on this PR; do not merge before this one). Title updated to drop multi-ABD. |
Contributor
Author
|
Validated on gfx1250 (MI400) — Helios
Clean on MI400 — LGTM. |
Contributor
Author
|
Interim fix for #11161 pushed here (aac4891). The compv3+intrawave 8-warp defect (#11161) is confirmed on this PR's universal GEMM on gfx1250 (it's a core-pipeline bug — the compv3 intrawave pipeline is MFMA/wave64-scheduled and miscompiles on WMMA/wave32 at 8 warps, max_rel 0.14–0.87 vs CPU). Until the pipeline is ported to wave32, this commit gates that kernel class off the gfx1250 sweep so it can't emit wrong results. Verification of the gate: the previously-failing config now expands to 0 kernels on gfx1250 ( A full post-fix |
Contributor
Author
Post-gate
|
ozturkosu
force-pushed
the
users/muozturk/ck/gemm-bridges-gfx1250-enable
branch
from
August 24, 2026 09:14
aac4891 to
837dca8
Compare
Contributor
Author
Full
|
| op | kernels | measurements | pass | fail | max_rel (median / max) |
|---|---|---|---|---|---|
gemm_universal |
1,640 | 4,920 | 4,920 | 0 | 5.08e-4 / 5.11e-4 |
grouped_gemm |
1,640 | 4,920 | 4,920 | 0 | 5.08e-4 / 5.11e-4 |
gemm_multi_d |
6,560 | 19,680 | 19,680 | 0 | 5.68e-4 / 6.63e-4 |
29,520 measurements, 0 failures, all bit-accurate for fp16.
I also confirmed the interim gate is active and covers all three ops — there are zero compv3 + intrawave + 8-warp kernels in any of the three result sets (all three flow through the same gemm_utils.expand_sweep; the gate drops that class: universal/grouped expand 6144→5632 configs). That class is the #11161 correctness defect; the real pipeline fix is #11175, and once it merges I'll revert the gate so those kernels re-enter the gfx1250 sweep.
Result CSVs: /ck/mi400val/PR10921/sweep_aug24/{universal,grouped,multi_d}.csv.
Contributor
Author
Full dtype × layout correctness matrix on MI400 (gfx1250) — interim updateFollowing up on the earlier fp16 validation: I'm now running the complete dtype × layout matrix for this PR's ops on gfx1250 (full
Totals so far
fp8/bf8 verified at their quantization-appropriate tolerances (6e-2 / 1e-1); fp16/bf16 at 2e-2. Note: this runs against the branch as-is (interim gate on), so the |
Contributor
Author
|
| dtype \ layout | rcr | rrr | crr | ccr |
|---|---|---|---|---|
| fp16 | 1640 (4920) | 1296 (3888) | 1080 (3240) | 1296 (3888) |
| bf16 | 1640 (4920) | 1296 (3888) | 1080 (3240) | 1296 (3888) |
| fp8 | 2984 (8952) | 2288 (6864) | 2146 (6438) | 2284 (6852) |
| bf8 | 2984 (8952) | 2288 (6864) | 2146 (6438) | 2284 (6852) |
| kernels | result rows | fail | |
|---|---|---|---|
| total | 30,028 | 90,084 | 0 |
Notes:
- Run against this branch as-is (interim gate on), so the
compv3+intrawave+8-warpclass is excluded — that's the compv3 intrawave GEMM pipeline returns wrong results on gfx1250 (MI400) for 8-warp blocks (2x4x1/4x2x1) #11161 defect, fixed for real by fix(ck-tile): correct compv3 intrawave 8-warp schedule on gfx1250 (wave32) #11175 (validated separately there). - fp8/bf8 are bit-accurate at their quantization-appropriate tolerances (rel ≈ 3.9e-2 / 7.6e-2).
grouped_gemm is also complete (16/16, 0 fail) and gemm_multi_d is running now — I'll post the combined final table for all three of this PR's ops once multi_d finishes.
Contributor
Author
✅ #10921 — full dtype × layout correctness matrix on MI400 (gfx1250) — COMPLETEAll three ops this PR enables are now validated across every supported dtype × layout on gfx1250. Method per combo: expand the full Grand total: 81,304 kernels · 243,912 GPU-vs-CPU checks · 0 failures
Notes:
Every op × dtype × layout on gfx1250 is correct. Per-combo CSVs archived. |
illsilin
pushed a commit
that referenced
this pull request
Aug 27, 2026
…ve32) (#11175) ## Summary Fixes the `compv3` intrawave GEMM pipeline producing **numerically wrong results** on **gfx1250 (MI400, RDNA4 / WMMA / wave32)** for **8-warp block arrangements** (`2x4x1` / `4x2x1`). max relative error 0.14–0.87 vs an fp32 CPU reference (not tolerance). Fixes #11161. ## Root cause The `NumWarps == 8` special-cased hot-loop/tail schedule in `BaseGemmPipelineAgBgCrCompV3` (`BlockHasHotloop` / `GetBlockLoopTailNum` / `TailHandler`) was written for **wave64 512-thread** blocks (gfx9xx / MFMA). On gfx1250 an 8-warp block is **wave32 = 256 threads** — the same thread count as a 4-warp wave64 block — so it must follow the standard (≤4-warp) schedule. The wave64 8-warp path mis-selects the loop tail (`TailNumber::Even` for even loop counts), which makes the intrawave RUN path execute **one extra `block_gemm` on a non-existent K-tile** → garbage accumulation. (For `num_loop == 3` the 8-warp `BlockHasHotloop` uses `> 3` instead of `> 2`, skipping the hot loop entirely.) ## Fix Add an arch-guarded compile-time predicate and switch the three checks on it: ```cpp #if defined(__gfx1250__) || defined(__GFX12__) static constexpr bool Use8WarpSchedule = false; #else static constexpr bool Use8WarpSchedule = (Problem::BlockGemmShape::NumWarps == 8); #endif ``` - On **gfx942 / gfx950** (wave64): `Use8WarpSchedule == (NumWarps == 8)` — **byte-for-byte identical, zero regression**. - On **gfx1250 / GFX12** (wave32): the 8-warp special case is disabled, so those blocks use the standard schedule. All users of these functions are `CK_TILE_DEVICE` and `TailHandler::scenarios[]` compiles in the same device pass, so the `__gfx1250__`/`__GFX12__` guard is host/device consistent (verified `hipcc --offload-arch=gfx1250` predefines both `__gfx1250__` and `__GFX12__`). ## Validation (on-device, gfx1250 MI400, GPU output vs fp32 CPU reference) | config | before | after | |--------|--------|-------| | compv3+intrawave **2x4x1** (8-warp) | MISMATCH 0.326 | **VERIFY rel=5.08e-04** | | compv3+intrawave **4x2x1** (8-warp) | wrong | **VERIFY rel=5.08e-04** | | compv3+intrawave **1x8x1 / 8x1x1** (8-warp) | wrong | **VERIFY rel=5.08e-04** | | compv3+intrawave **2x2x1 / 1x4x1 / 4x1x1** (4-warp control) | correct | **VERIFY** (no regression) | Validated across `num_loop` = 3, 4, 5, 10, 11, 12 (odd + even, hot-loop and no-hot-loop) and both fp16 (~5e-4) and bf16 (~3.8e-3, normal tolerance). ## Notes This is the **real** fix for #11161. An interim gate that simply drops these kernels from the gfx1250 sweep was applied to the gfx1250 enablement PRs (#10921 / #11048 / #10928); once this lands, that gate can be removed so the (now-correct) 8-warp `compv3 + intrawave` kernels re-enter the sweep.
The regular-GEMM TE->Dispatcher bridge (#8997 fp16/bf16, #8998 fp8/bf8/int8) was CDNA-only (gfx90a/gfx942/gfx950, MFMA). This enables it on gfx1250 (MI400, RDNA4/WMMA): - gemm_utils: add gfx1250 to _SUPPORTED_ARCHES + wave-combo table; add an OCP fp8/bf8 host codec (float8_e4m3fn / float8_e5m2) selected by arch (gfx950/gfx12xx -> OCP, gfx942 -> FNUZ), since gfx1250's device fp8_t is OCP. - gemm_validation_utils: gfx1250 warp-configuration combos. - arch_specs_generated: gfx1250 fp8/bf8 WMMA warp tile (16x16x64). - gemm_full_benchmark: allow --dtype fp8/bf8. - configs: gfx1250 CI configs (fp16/bf16 warp_tile 16x16x32; fp8 16x16x64). - tests: CPU-only gfx1250/OCP-fp8 codec + arch tests. Validated on MI400/gfx1250 (ROCm 10.1): fp16/bf16 and fp8 kernels build and pass --verify (fp8 max_rel ~3.6e-2 vs fp32 ref, i.e. E4M3 quantization only).
Fold the grouped (#9000), multi-D (#9308), and multi-ABD (#9305) variant enablement onto the shared gfx1250 foundation so gemm_universal + these three "full" variants ship in one PR: - grouped: gfx1250 CI config (WMMA 16x16x32) + CPU tests - multi-D: gfx1250 CI config (WMMA 16x16x32) + CPU tests + --arch help - multi-ABD: gfx1250 CI config (WMMA 16x16x32, fp16-only) + CPU tests Each merged variant pinned the CDNA MFMA 32x32x16 tile in its CI config, which does not run on gfx1250 WMMA; the new configs use 16x16x32. No variant codegen change was needed beyond the shared arch enablement. Validated on MI400/gfx1250 (ROCm 10.1): grouped 64/64 fp16 OK; multi-D 80 measurements VERIFY (max_rel ~5e-4); multi-ABD 32/32 fp16 VERIFY. 63 CPU tests pass.
multi-ABD failed some correctness cases in MI400/gfx1250 validation, so it is being split into its own PR while gemm_universal + grouped + multi-D (all of which passed MI400 validation) stay here. Removes the multi-ABD gfx1250 CI config and its gfx1250 CPU test class; grouped and multi-D gfx1250 configs/tests are untouched.
#11161) The compv3 intrawave GEMM pipeline is hand-scheduled for the MFMA/CDNA (wave64) instruction cadence and miscompiles for 8-warp blocks (2x4x1 / 4x2x1) on gfx1250 (WMMA / wave32), producing wrong results -- validated on-device: max_rel 0.14-0.87 vs an fp32 CPU reference, while <=4-warp compv3 and all compv4/mem/interwave kernels are bit-accurate. Drop that kernel class from the gfx1250 sweep in expand_sweep until the pipeline is ported to wave32. Covers universal, multi-ABD and batched GEMM (they share gemm_utils.expand_sweep). See #11161.
ozturkosu
force-pushed
the
users/muozturk/ck/gemm-bridges-gfx1250-enable
branch
from
August 27, 2026 22:49
e322ced to
525b70c
Compare
yraparti
approved these changes
Aug 31, 2026
yraparti
approved these changes
Aug 31, 2026
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 31, 2026
=?UTF-8?q?feat(ck-tile):=20enable=20GEMM=20bridges=20on?= =?UTF-8?q?=20gfx1250=20(MI400)=20=E2=80=94=20universal=20+=20grouped=20+?= =?UTF-8?q?=20multi-D=20(#10921)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : AICK-2098 JIRA ID : AICK-2099 JIRA ID : AICK-2100 JIRA ID : AICK-2137 ## Summary Enables the **regular (gemm_universal)**, **grouped**, **multi-D**, and **multi-ABD** GEMM Tile-Engine → Dispatcher bridges on **gfx1250 (MI400, RDNA4 / WMMA)**, for fp16/bf16 (+ fp8 on the regular path). This is the consolidation of the previously-stacked PRs #10904 / #10914 / #10915 / #10917 into one. Python + JSON only (no C/C++) — no clang-format gate. The two *partial* variants (**stream-K #10919**, **microscaling mx #10916**) are **not** included — they need C++/kernel follow-ups (see below) and remain separate. ## Do the merged bridges (Universal/grouped/multi-D) support gfx1250? — No. As merged, none of these bridges run on gfx1250. They share a CDNA-only arch gate (`_SUPPORTED_ARCHES = gfx90a/gfx942/gfx950`) and most pinned the MFMA `32x32x16` warp tile in their CI config, which is invalid on gfx1250's RDNA4 WMMA engine (fp16/bf16 WMMA is `16x16x32`, fp8 is `16x16x64`). ## Changes **Shared enablement (was #10904)** - `dispatcher/python/gemm_utils.py`: `gfx1250` in `_SUPPORTED_ARCHES` + gfx1250 wave combos; OCP fp8/bf8 codec (`float8_e4m3fn`/`float8_e5m2`) selected by arch (gfx950/gfx12xx → OCP, gfx942 → FNUZ). - `tile_engine/ops/gemm/gemm_validation_utils.py`: gfx1250 warp-configuration combos. - `dispatcher/codegen/arch_specs_generated.py`: gfx1250 fp8/bf8 WMMA `[[16,16,64]]` (fp16/bf16 `[[16,16,32]]` already present). - `tile_engine/ops/gemm/gemm_full_benchmark.py`: `SUPPORTED_DTYPES` += fp8, bf8. - `tile_engine/ops/gemm/configs/`: gfx1250 regular CI configs (fp16/bf16 16x16x32; fp8 16x16x64). **Variant configs + tests** - grouped: `grouped_gemm/configs/default_ci_config_gfx1250.json`. - multi-D: `gemm_multi_d/configs/default_ci_config_gfx1250.json` + `--arch` help. - multi-ABD: `gemm_multi_abd/configs/default_ci_config_gfx1250.json` (fp16-only, upstream design). - CPU-only tests across `dispatcher/tests/` (gfx1250 arch, OCP fp8 codec, grouped/multi-D/multi-ABD configs). 63 tests pass. ## Validation (MI400 / gfx1250, ROCm 10.1) | Bridge | Result | |---|---| | gemm_universal | 640/640 fp16+bf16 VERIFY, 0 fail; fp8 up to ~625 TFLOP/s (max_rel ~3.8e-2 = E4M3 quant) | | grouped | 64/64 fp16 OK (~178–217 TFLOP/s) | | multi-D | 80 measurements VERIFY (max_rel ~5e-4) | | multi-ABD | 32/32 fp16 VERIFY (fp16-only by upstream design) | ## Notes - `arch_specs_generated.py` is auto-generated from `arch_specs.json`, which has **no gfx1250 entry** on develop (the gfx1250 fp16/bf16 rows were hand-added to the generated file). This PR follows that precedent; a follow-up should add a full gfx1250 stanza to `arch_specs.json`. - fp8 `--verify` needs a looser tol (~0.05–0.15) than the fp16 default (2e-2), since the reference is fp32 and fp8 carries E4M3 quantization. - **Not included (separate, partial):** stream-K (#10919 — kernel hangs on RDNA4: cross-workgroup K-reduction coherency relies on CDNA atomics) and microscaling mx (#10916 — GPU `.so` build `static_assert`s gfx950-only via `preShuffleScaleBuffer_gfx950`; ck_tile already ships `preShuffleScaleBuffer_gfx1250`). Both need C++ follow-ups. Supersedes #10904, #10914, #10915, #10917.
ozturkosu
added a commit
that referenced
this pull request
Aug 31, 2026
…fx1250 (MI400) (#10928) JIRA ID : AICK-2104 JIRA ID : AICK-2136 ## Summary Enables the **batched GEMM (#9306)** and **batched contraction (#9328)** bridges (merged together in #10612) on **gfx1250 (MI400, RDNA4/WMMA)**. Targets `develop` directly. Python + JSON only (no C/C++). > Recreated against `develop` (replaces the earlier stacked PR #10922). The diff therefore also includes the shared gfx1250 enablement that also appears in #10921 (gemm_utils `_SUPPORTED_ARCHES`/wave tables, arch_specs gfx1250, driver dtypes) — if #10921 lands first, this will rebase down to just the batched-specific files. ## Does merged #10612 support gfx1250? — No. The batched variants keep their **own** arch tuples (`batched_gemm_utils._SUPPORTED_ARCHES`, `batched_contraction_utils._SUPPORTED_ARCHS` = gfx90a/gfx942/gfx950) and their CI configs pin CDNA MFMA tiles (contraction `default` 32x32x64, `bridge` 32x32x16) that don't run on gfx1250 WMMA. ## Changes - Add `gfx1250` to both batched utils' supported-arch tuples. - gfx1250 WMMA `16x16x32` CI configs: `batched_gemm/configs/default_ci_config_gfx1250.json`; `batched_contraction/configs/default_ci_config_gfx1250.json` + `bridge_default_ci_config_gfx1250.json`. - `--arch`/config help in both `*_full_benchmark.py`. - CPU tests (gfx1250 arch + configs) in `test_batched_bridge.py` / `test_batched_contraction_bridge.py`. ## Validation — MI400 / gfx1250 (a07-1, ROCm 10.1, gpu0) - **batched_gemm** (fp16/rcr, --verify): **32/32 Successful, 0 failed** (verify ~5e-4). - **batched_contraction**: `bridge` config **1/1 verified** (max_rel 2.6e-4); `default` config **8/12 built, 8/8 verified** (the 4 non-builds are CompV3/CompV4+interwave pipeline incompatibilities, not gfx1250 runtime bugs). - 66 batched CPU tests pass. Refs: merged #10612 (#9306 + #9328); shared gfx1250 enablement also in #10921. JIRA ID : AICK-2104 JIRA ID : AICK-2136 --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
Contributor
TheRock Submodule Bump ActivityNewest first
|
ozturkosu
added a commit
that referenced
this pull request
Sep 1, 2026
…fx1250 (MI400) (#10927) JIRA ID : AICK-2107 JIRA ID : AICK-2140 ## Summary Enables three merged grouped **quant** GEMM TE->dispatcher bridges on gfx1250 (MI400): - `grouped_gemm_aquant` + `grouped_gemm_abquant` (originally from #9253) - `grouped_gemm_bquant` (originally from #9166) Stacked on the shared gfx1250 foundation in **#10921** (base branch `users/muozturk/ck/gemm-bridges-gfx1250-enable`). ## Requires C++? YES (small) The **3 ctypes-lib runtime arch allow-lists** rejected gfx1250 and had to be extended — that is the only C++ change. Each of `grouped_gemm_{aquant,abquant,bquant}_ctypes_lib.cpp` had a runtime gate of the form `if(arch != gfx950 && arch != gfx942 && arch != gfx90a) return -1;`; added `&& ... != gfx1250` and updated the `(supported: ...)` message. Each file is clang-format-18 clean. ## Correct gfx1250 config — FlatMM `warp_tile_k=128` (on-GPU determined) The original revision of this PR pinned the gfx12 WMMA tile (`warp_tile_k=16`) on the theory that gfx1250 needs WMMA. **On-GPU validation on MI400/gfx1250 (fp32-dequant reference) proved that wrong**: `warp_tile_k=16` returns **all-zeros** for every grouped-quant variant. The empirically-correct gfx1250 config for the fp8/bf8 quant path is the **FlatMM tile, `warp_tile_k=128`** — the same tile the gfx9 fp8/bf8 FlatMM decode uses. | variant | dtype | gfx1250 config warp_tile | on-GPU status | max_rel_err | |---|---|---|---|---| | aquant | fp8 | 16x16x**128** | **PASS** | ~3e-4 | | aquant | bf8 | 16x16x**128** | **PASS** | ~1e-4 | | abquant (compv3) | fp8 | 16x16x**128** | **PASS** | ~3e-4 | | abquant (compv3) | bf8 | 16x16x**128** | **PASS** | ~3e-4 | | bquant | fp8 | 16x16x**128** (= stock) | **PASS** | ~5e-4 | | bquant | bf8 | 16x16x**128** (= stock) | **PASS** | ~5e-4 | | a/bquant | fp8i4 / bf8i4 | — | **deferred** (does not compile on gfx1250) | — | - **aquant / abquant**: gfx9 stock uses `warp_tile_k=32` (plain MFMA), which is WRONG on gfx1250 (aquant max_rel=1.0; abquant all-zero). The gfx1250 helpers now pin `warp_tile_k=128`. - **bquant**: the stock fp8/bf8 config is already `warp_tile_k=128` and is correct on gfx1250 as-is, so the gfx1250 helper simply delegates to the stock config. - **fp8i4 / bf8i4** (packed-int4 weights): do **not** compile on gfx1250 at any `warp_tile_k` (no gfx12 instruction for the packed-int4 quant path). The C++ arch-gate stays open, but on-GPU correctness is **deferred pending kernel work**. No gfx1250 config helper is shipped for i4 (it would emit a non-buildable kernel). ## Changes - **C++ (arch gates only):** add `gfx1250` to the runtime allow-list + `(supported: ...)` string in `grouped_gemm_{aquant,abquant,bquant}_ctypes_lib.cpp`; clang-format-18 applied. - **Python utils:** gfx1250 helpers pin the FlatMM tile `warp_tile_k=128`: - aquant: `default_{fp8,bf8}_config_gfx1250` - bquant: `default_{fp8,bf8}_config_gfx1250` (delegate to stock) - abquant: `default_{fp8,bf8}_compv3_config_gfx1250` (compv3, `transpose_c=false`) - i4 gfx1250 helpers removed (unbuildable on gfx1250). - **CPU tests:** each `TestGfx1250Configs` now asserts `warp_tile_k=128` / `16x16x128` names and the correct pipeline (aquant `mem`; bquant/abquant `compv3`; abquant not eightwaves/transpose_c). ## Validation - **CPU + codegen:** `pytest dispatcher/tests/test_grouped_gemm_{aquant,bquant,abquant}_utils.py` → **176/176 pass**. - **On-GPU (MI400/gfx1250, a07-1, `HIP_VISIBLE_DEVICES=0`):** all 6 fp8/bf8 configs PASS against an fp32-dequant reference (see table + validation comment). Uses the real corrected helper functions. ## Dependencies / follow-ups - **Base:** #10921 (shared gfx1250 GEMM-bridge foundation). - **Originating PRs:** #9253 (aquant/abquant), #9166 (bquant). - **Follow-up:** gfx1250 packed-int4 (fp8i4/bf8i4) grouped-quant kernel support. _Note: no `Co-Authored-By` trailer per team policy._ --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
ozturkosu
added a commit
that referenced
this pull request
Sep 2, 2026
JIRA ID : AICK-2101 JIRA ID : AICK-2133 ## Summary This PR enables the **multi-ABD** (`a2_b2_d2`) GEMM TE→dispatcher bridge (#9305) on **gfx1250 / MI400**. It was split out of #10921 so the other three ops (`gemm_universal`, `grouped_gemm`, `gemm_multi_d`) weren't held back. Comparison of this PR's original file set vs `develop` (branch head `56e7bc3`): **So the substantive multi-ABD gfx1250 enablement is already on `develop`; the remaining unique delta in this PR is a single config file (`example_problems.json`).** ## Motivation (why the enablement was needed) gfx1250 runs **WMMA**, not MFMA: the valid fp16/bf16 warp tile is **16x16x32**, whereas the CDNA MFMA **32x32x16** tile in `default_ci_config.json` instantiates zero runnable kernels there. ## Supported surface (gfx1250) The `a2_b2_d2` variant is intentionally narrow — the codegen, ctypes lib, and runner are all **fp16-only** and **rcrr-only** (`expand_sweep` rejects other dtypes/layouts). So **fp16 × rcrr is the complete supported matrix**. ## Validation (on-device, MI400 / gfx1250) Old-TE has no gfx1250 back end, so this is **correctness-only** (GPU output vs an fp32 CPU reference), not a performance baseline. Full `default_config` sweep of the fp16/rcrr surface, 3 tile-divisible shapes (768³ / 1536³ / 3072³): | op | supported matrix | kernels | result rows | fail | |----|------------------|--------:|------------:|-----:| | `gemm_multi_abd` | fp16 × rcrr (only) | 1,640 | 4,920 | **0** | Every kernel in the supported surface is correct on gfx1250 (rel err ≈ 4e-4). ## Note on #11161 (compv3 + intrawave + 8-warp) The full sweep initially surfaced 552 failures, all in the `compv3 + intrawave + 8-warp (2x4x1 / 4x2x1)` class — the **#11161** defect (a wave64 assumption in the `comp_v3` pipeline, wrong on gfx1250 wave32). An interim gate excludes that class so the supported surface reports clean; the real fix, **#11175** (`Use8WarpSchedule`), is already on `develop`. ## Notes / dependencies - Built on #10921's shared gfx1250 foundation (now on `develop` via the sibling series). - Supersedes the earlier closed attempt **#10917**. - Sibling gfx1250 bridge PRs: #10921 (universal / grouped / multi-D), #10927 (grouped quant), #10928 (batched GEMM + contraction), #11043 (rowcol / tensor quant), #10934 (contraction multi-ABD). --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
This was referenced Sep 16, 2026
ozturkosu
added a commit
that referenced
this pull request
Sep 17, 2026
JIRA ID : AICK-2249 ## Summary gfx1250 is labelled **RDNA4** in several comments and one CI config string. The product generation is **CDNA5**. This corrects the label. Comments and one JSON `_comment` only — no code, no behaviour change. Raised in review of #12173, where the same value was corrected in `arch_specs.json`. Split out because renaming terminology across the tree does not belong in an LDS-capacity change. ## Why the RDNA4 label was there It was not arbitrary. gfx1250 has no MFMA units and multiplies through **WMMA**, which is an RDNA-family instruction, and the `gfx12xx` numbering points the same way. The observation about the *instruction path* was right; using it as a *generation* label was not. `arch.hpp` already draws exactly this distinction: > GFX1250 is its own standalone family. Although it **shares the RDNA architecture** with the GFX12 family, its MMA builtins and data-type ABI differ, so it must not be treated as a GFX12-family device. Sharing an ISA lineage is not the same as belonging to that generation. So there are two separate claims in play, and only the second one holds. ## Changes Where a generation label is meant, it now says CDNA5. Where only the instruction path matters, the generation is dropped rather than replaced — a WMMA path is a WMMA path regardless of which generation the part belongs to. | file | before | after | |---|---|---| | `dispatcher/tests/test_gemm_utils.py` | `gfx1250 (MI400 / RDNA4-WMMA) enablement` | `gfx1250 (MI400 / CDNA5, WMMA) enablement` | | `dispatcher/tests/test_batched_bridge.py` | same | same | | `dispatcher/tests/test_batched_contraction_bridge.py` | same | same | | `dispatcher/tests/test_multi_d_bridge.py` | `it runs the RDNA4 WMMA path` | `it runs the WMMA path` | | `dispatcher/codegen/grouped_conv/grouped_config_rules_full.py` | `such as rdna4/gfx1250` | `such as gfx1250` | | `tile_engine/.../bridge_default_ci_config_gfx1250.json` | `(MI400, RDNA4/WMMA)` | `(MI400, CDNA5/WMMA)` | ## Deliberately not changed - **`include/ck_tile/core/arch/arch.hpp`** and **`.../mma/scale/wmma/selector.hpp`** describe the ISA lineage and the selector strategy as RDNA. Both are accurate — the first is the quote above, the second documents an RDNA-style selector. Renaming these would replace a correct statement with a wrong one. - **`dispatcher/codegen/arch_specs_generated.py`** still reads `rdna4`. It is generated, and its source is the gfx1250 entry in `arch_specs.json` that #12173 adds; the value is corrected there. Hand-editing a generated file is the exact trap this codebase keeps hitting, so it is not done here. That last point means the tree is briefly inconsistent until #12173 lands. The affected field is metadata: nothing branches on its value, and its only consumer is a `None` check that rejects unknown architectures in strict mode. ## Test plan - [x] `test_gemm_utils`, `test_batched_bridge`, `test_batched_contraction_bridge`, `test_multi_d_bridge`, `test_grouped_conv_codegen` all pass - [x] The edited CI config still parses as JSON - [x] No `gfx1250`/`MI400` reference to RDNA4 remains outside the two files listed above as intentionally kept ## Note for the reviewer The label originated in #10921, the gfx1250 bridge enablement PR, and spread from there; AICK-2082 uses the same wording. So this is a retroactive terminology correction rather than a one-off typo. If **CDNA5 is not the right generation either**, say so before this merges — I have taken it from review feedback on #12173 and have not independently confirmed it. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
ozturkosu
added a commit
that referenced
this pull request
Sep 26, 2026
…bridge (#12544) JIRA ID: AICK-2427 JIRA ID: AICK-2137 JIRA ID: AICK-2136 JIRA ID: AICK-2142 JIRA ID: AICK-2429 Related PRs: #12508 #10921 #10927 #10928 #11043 #11048 #10916 ## Summary This PR enables the gfx1250-only GEMM pipelines in the tile_engine (TE) / dispatcher bridge: - `comp_tdm` (TDM V1) and `comp_tdm_v2` (TDM V2); - `comp_async`; - the TDM epilogue. They are enabled for GEMM universal and batched GEMM, whose gfx1250 bridge enablement is already merged. GEMM multi-ABD and the grouped quant ops only gain reject-guards: their kernels have no async/TDM pipeline yet. Most of the change is in `tile_engine/` and `dispatcher/`. It also fixes two CK-Tile kernel defects in pipelines this PR enables (see Kernel fixes), with gtest coverage. On gfx942 and gfx950 the generated bridge code is byte-identical. ## Motivation The gfx1250 bridge enablement PRs (#10921, #10928, #11048) exposed only the pipelines shared with gfx942/gfx950. The TDM and async pipelines exist in CK-Tile for gfx1250, and #10916 already uses them for mx GEMM, but the bridge could not generate or benchmark them for the other GEMM ops. ## Design note - **TDM pipelines** (`comp_tdm`, `comp_tdm_v2`): - They are intrawave only, always use the TDM epilogue, and are non-persistent. - They require `pad_m = pad_n = pad_k = false`. TDM clips A/B/E in hardware, and right-padding N/K inflates the bound-clip extents, which gives wrong results on unaligned shapes. This matches the CK-Tile TDM gtests, which pin padding off. - The TE builders and dispatcher codegen reject any other combination with a reason string. - **`comp_async` on gfx1250:** - It is emitted only for row-major A / col-major B. Other layouts hit the transpose-load `static_assert` with the gfx1250 WMMA tiles. - For non-MX ops it is also emitted only fully padded (`pad_m = pad_n = pad_k = true`). On develop, unpadded `comp_async` on gfx1250 has a kernel defect: the async K-prefetch reads past the A/B extent, and the `TailNumber::Two` path has no LDS fence. The kernel fix goes into #12508; this restriction can be relaxed once it merges. - fp8/bf8 `comp_async` on gfx1250 requires `warp_tile_k >= 128`. - mx GEMM is unaffected. - **TDM LDS budget:** the gfx1250 TDM pipelines share the `comp_async` LDS budget (aliases in `arch_specs.json`, emitted by `generate_arch_specs.py`). The LDS check adds the TDM 16-byte-per-row-group descriptor padding and checks that the TDM output tile fits. The tile_engine and the dispatcher ArchFilter use the same rule. Non-MX `comp_async` GEMM on gfx1250 uses the same base LDS descriptor, so it gets the same row padding in both gates (without it, fp16/bf16 64x256x256, 192x128x256 and 256x64x256 fail to link at 337856 B > 327680 B); MX `comp_async` keeps its own rule from #10916. - **GEMM multi-ABD:** reject-guards only for `comp_async`, `comp_tdm`, `comp_tdm_v2` and the TDM epilogue. - **Grouped quant ops** (`aquant`, `bquant`, `abquant`, `rowcolquant`, `tensorquant`): - These files only gain reject-guards for `comp_async`, `comp_tdm`, `comp_tdm_v2` and the TDM epilogue. Nothing is enabled for them. - The quant GEMM pipelines have no TDM/async implementation yet, so an unguarded config would fail to compile instead of being filtered. - **Runtime reject = SKIP.** The dispatcher returns status `-1` when a kernel's `IsSupportedArgument` rejects a problem (for example TDM alignment at `N % 8 != 0`). It is exposed as `GemmResult.unsupported` (also on grouped and multi-D results). The gfx1250 tests and `test_gemm_search_space` count it as a skip, not a failure, and a lane still fails if nothing was verified. - **Trait parsing** now goes through one shared helper (`trait_parse.py` / `gemm_trait_parse.cmake`), so the CMake and Python sides split kernel trait strings the same way. ### Deferred ops grouped GEMM, GEMM multi-D, batched contraction and contraction multi-ABD are not enabled here. Their kernels need the core changes in #12508 before they can host a TDM pipeline. A follow-up PR will enable them after #12508 merges. ## Kernel fixes Found by a gfx1250 tuning sweep with `--verify` (bf16 `rcr`, 5608x4096x2048): - **`comp_async` with `persistent=true` wrote only part of the output (9-36% of tiles non-zero).** - Cause: `GemmPipelineAgBgCrCompAsync` did not expose `UsePersistentKernel`. The host launched a persistent-sized grid, but the device took the one-block-per-tile entry, so most tiles were never written and out-of-range tile indices could fault. - Fix: the pipeline now exposes `UsePersistentKernel` for plain GEMM problems. MX problems stay non-persistent, because their host path launches a split-K grid that the persistent loop does not read. - The same pipeline is used on gfx950. The gfx950 persistent comp_async tests pass with the fix. - **`comp_tdm` / `comp_tdm_v2` returned all-NaN when the TDM policy picks `sub_tile_num > 1`.** - The TDM policy picks `sub_tile_num > 1` from its VGPR estimate. In the sweep this happened at 4 waves with 128x128x256 and 128x256x128 tiles. - Cause: the hot loop in `block_gemm_areg_breg_creg_v1.hpp` ran over the full `KIterPerWarp` instead of one K sub-tile (`KIterPerWarp / KSubTileNum`), so it read past the sub-tile register window. - Fix: the loop now runs over one K sub-tile. With `KSubTileNum == 1` it is unchanged; the compv4 code object is byte-identical. - **`comp_tdm_v2`:** the 4-wave `static_assert` is now written as `NumWarps == 4`. It is equivalent to the old block-size form, and the pipeline body hardcodes four wave roles. - **Tests:** - A Persistent variant is added to the comp_async WMMA gtest list. - New `test_ck_tile_gemm_pipeline_tdm_subtile_wmma`: fp16/bf16, V1/V2, with sub-tile 2 cases plus sub-tile 1 controls. - Both tests fail on the old kernel and pass with the fix. - **Scope:** `comp_async_v2` (sub-tiling policy) shares the block GEMM fix but has no gtest here and is not emitted by the bridge. ## Changes | Area | Change | |---|---| | `dispatcher/codegen/arch_filter.py`, `codegen_common.py`, `unified_gemm_codegen.py` | gfx1250 pipeline/epilogue emission and validity rules (TDM pad/scheduler/persistent, `comp_async` layout + pad) | | `dispatcher/codegen/arch_specs_generated.py`, `include/.../arch_specs_generated.hpp`, `kernel_key.hpp` | New pipeline/epilogue enum values and names | | `dispatcher/codegen/unified_grouped_gemm_*quant_codegen.py` | Reject-guards only | | `dispatcher/codegen/arch_specs.json`, `generate_arch_specs.py` | TDM LDS budget aliases (generated, not hand-edited) | | `dispatcher/python/gemm_utils.py` | Pipeline names, `GemmResult.unsupported` for status `-1` | | `dispatcher/tests/CMakeLists.txt` | Registers the gfx1250 pipeline, gate-equivalence, arch-specs regeneration and result-status tests | | `tile_engine/ops/gemm/gemm_instance_builder.py`, `gemm_validation_utils.py` | Shared gfx1250 validity checks | | `tile_engine/ops/gemm/gemm_universal/*`, `batched_gemm/*`, `gemm_multi_abd/*` | Builders, CMake, gfx1250 default and CI configs | | `tile_engine/ops/gemm/grouped_gemm_quant/*` | Reject-guards only | | `dispatcher/codegen/codegen_common.py`, `tile_engine/ops/gemm/gemm_validation_utils.py` | One shared async/TDM reject guard per layer; the grouped quant codegens and the multi-ABD / grouped quant builders bind it to their op name with `functools.partial` instead of keeping a local copy | | `tile_engine/ops/gemm/trait_parse.py`, `gemm_trait_parse.cmake` | Shared trait-string parsing | | `include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp`, `pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp`, `pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v2.hpp` | Kernel fixes (see above) | | `test/ck_tile/gemm/*` | Persistent comp_async gtest variant, new TDM sub-tile gtest | | `test_*.py` (dispatcher and tile_engine) | Unit tests for the gates, configs and guards | ## Testing | Check | Result | |---|---| | gfx942 / gfx950 generated-code byte identity (before vs after, all ops) | Identical | | TDM V1/V2 native on gfx1250, pad false | Correct (18/18) | | `comp_async` native on gfx1250, fully padded | Correct | | clang-format-18 on the C++ headers | Clean | | Dispatcher unit tests (targeted gfx1250 + new tests) | 109 passed | | tile_engine gfx1250 tests | 92/92 passed | | ArchFilter vs tile_engine gate equivalence (gfx1250 enumeration) | Identical | | gfx1250 build + `--verify` of every comp_async / comp_tdm / comp_tdm_v2 default_config instance (7323 kernels x 9 shapes) | 0 wrong results; the LDS gate rejects exactly the 90 instances that fail to link, 0 false rejects | | Generated arch specs reproducible from `arch_specs.json` | Identical | | ruff check (F, E9) on the changed Python | Clean | | gfx1250 end-to-end bridge run (TE + dispatcher, universal / batched / multi-ABD) | 0 failures. TE 1188 pass / 1102 skip; dispatcher universal 620/840 and batched 94/96 verified, the rest are runtime-unsupported skips; multi-ABD 18/22 with the 4 negative cases correctly rejected; codegen gate vs emitted kernels and ArchFilter vs codegen: 0 mismatches | | gfx1250 kernel-fix gtests | TDM sub-tile 24/24, comp_async WMMA 14/14 (incl. Persistent), persistent WMMA 14/14 | | gfx1250 tile_engine re-run of the reported configs (bf16 `rcr`, 5 shapes incl. 5608x4096x2048) | comp_async 40/40 pass (20 with `persistent=true`); comp_tdm / comp_tdm_v2 72/72 pass, including 128x128x256 and 128x256x128 at 2x2x1 | | Mutation check (gfx1250, old kernel code restored) | Persistent cases fail without the comp_async fix; all sub-tile 2 cases fail without the block GEMM fix; controls pass | | gfx950 regression (comp_async, eight-waves, persistent async, MX async rcr/rrr/crr/ccr) | All pass | --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
ozturkosu
added a commit
that referenced
this pull request
Oct 16, 2026
…ched contraction kernels (#12508) JIRA ID: AICK-2082 JIRA ID: AICK-2099 JIRA ID: AICK-2100 JIRA ID: AICK-2104 JIRA ID: AICK-2109 Related PRs: #10921 #10927 #10928 #10934 #11048 #10931 #12544 ## Summary This PR lets three CK-Tile kernels host the gfx1250 TDM GEMM pipelines (`GemmPipelineAgBgCrCompTDMV1` / `V2`): - grouped GEMM; - GEMM multi-D; - batched contraction (and multi-ABD contraction, which wraps it). It also fixes two problems in the `comp_async` pipelines (`GemmPipelineAgBgCrCompAsync` and `...V2`), see "CompAsync fix" below: - an out-of-bounds K prefetch when K is not padded; - on gfx125, a CompAsyncV2 hot-loop wait that could leave async LDS loads in flight in blocks that reach past the A/B tensor (padded tiles and M/N tails). It changes only `include/` and `test/`. Enabling these pipelines in the TE/dispatcher bridge is a separate follow-up PR that depends on this one. For the non-CompAsync grouped GEMM, multi-D and contraction kernels, the compiled device code is byte-identical on gfx942 and gfx950. The CompAsync V1/V2 fixes change those pipelines on gfx950; their tail regressions and V1 performance measurements are listed below. ## Motivation gfx1250 bridge enablement is merged for grouped GEMM, multi-D and batched contraction (#10921, #10928, #10934). These kernels still cannot use the TDM pipelines, for three reasons: - **Tuple-only windows.** The TDM pipelines take A/B windows only as tuples. The grouped and contraction kernels call the pipeline with single windows, so a TDM instance does not compile. - **Hardware clipping, no padding.** TDM clips A/B/E in hardware and expects `kPadM/kPadN/kPadK == false`. The kernels need a path that builds unpadded tensor views. - **No D tensors in `TdmEpilogue`.** `TdmEpilogue` rejects D tensors, so multi-D GEMM has no epilogue to store E through TDM. ## Design note - **Detecting a TDM pipeline.** Only the TDM pipelines define `skipCheckValidLaunchParams`. `UniversalGemmKernel::kIsTdmPipeline` wraps the existing `has_skip_check_valid_launch_params` trait, and the grouped, multi-D and contraction kernels all use it. A non-TDM pipeline cannot take the TDM path by accident: it would also need a TDM epilogue (`is_tdm_epilogue_v`, which also accepts `TdmMultiDEpilogue`) to pass the new `static_assert`s. - **Shared TDM host checks.** `UniversalGemmKernel::IsTdmArgumentSupported` (k_batch == 1, K > 0, A/B/E aligned to their element size) and `IsTdmDeviceSupported` (gfx125 only) are shared by the three kernels, so they reject the same arguments with the same messages. - **New `TdmMultiDEpilogue`, separate from `TdmEpilogue`.** - It is a new type on purpose, so current `TdmEpilogue` users (grouped conv, mx flatmm, ...) keep rejecting D tensors. - It reuses `TdmEpilogue`'s LDS view, store distribution and TDM store. - It loads each D tile with the accumulator distribution and fuses it on the fp32 accumulator before the single cast to E. - It supports row-major E and row-major D only, with no ScaleM/ScaleN. - **Forced N pad on D loads.** TDM kernels run with `kPadN == false`, but the epilogue still reads D with regular tile loads. For TDM kernels with D tensors, `universal_gemm_kernel` forces the N pad on those loads, so a partial N tile reads zeros instead of running into the next row. - **Batched contraction host checks.** TDM uses one stride per merged dimension. The host checks that every merged G/M/N/K group collapses to a single stride, with K (and N for E) innermost; otherwise it throws `std::invalid_argument`. The batch stride comes from the fastest G dimension whose length is not 1. - **Arguments rejected on the TDM path:** padding (`kPadM/kPadN/kPadK` must be false, enforced with `static_assert` in the grouped and multi-D kernels), split-K (`k_batch != 1`), K <= 0, A/B/E pointers not aligned to their element size, persistent mode, cluster launch, and devices other than gfx125x. Batched contraction also rejects zero or negative strides in a merged dimension. They fail with `IsSupportedArgument(s)`, or with `static_assert` where the type decides. The TDM store always overwrites E, so split-K accumulation is not possible. - **Multi-ABD contraction.** A `static_assert` limits TDM to a single A and a single B tensor, because `TdmEpilogue` overwrites E. - **Contraction epilogue.** `RunGemmTdm` requires a TDM epilogue (`static_assert(is_tdm_epilogue_v<EpiloguePipeline>)`), so a CShuffle epilogue with a TDM pipeline fails to compile instead of producing wrong stores. ### CompAsync fix - **Bug.** With `kPadK == false`, the `comp_async` pipelines issue async prefetches for K tiles that do not exist: - the prologue loads tile 1 even when there is only one K tile (tail One); - with an even `num_loop` (tail Two), the last hot-loop iteration prefetches tile `num_loop`. On gfx1250 this reads past the A/B extent and gives wrong results on unpadded shapes. CI misses it because the gtests default to `PadK = true`. - **V1 fix:** - skip the tile-1 prologue prefetch for tail One; - for tail Two, peel the last hot-loop iteration and run it with a compile-time "no prefetch" flag. The steady-state loop body has no runtime branch, so gfx950 CompAsync performance is unchanged (see Testing). For tail Three the loop is not peeled; - add a `block_sync_lds_direct_load()` fence before the tail Two local read. Today the fence is redundant, but it protects the tail if the hot-loop fences are ever relaxed. - **V2 fix:** the same gating for the prologue and in-loop prefetches, plus a corrected tail handler. - **Pre-existing develop bug fixed in V2.** `CompAsyncV2` aliased the V1 base, so for an odd `num_loop` it dispatched tail Three, which V2 does not implement. K tiles were dropped and the output was wrong (for example K = 192, 320, 4160, 8256). The new V2 gtests fail 4/4 on develop headers and pass 4/4 with this PR. - **V2 async-load wait on gfx125 (also pre-existing in develop).** - The V2 hot loop waits with `s_wait_asynccnt(AB_Async_Load_Inst_Num)`, so only the loads of the newest window may still be in flight. - On gfx125 a lane whose element is out of bounds does not issue its async load (it zero-fills LDS instead). A wave can then issue fewer loads than `AB_Async_Load_Inst_Num`, the partial wait lets loads of the window about to be read stay in flight, and the block reads stale LDS. - This happens in padded tiles and also without padding, in the last block of an M (or N) tail, e.g. row-major A with M not a multiple of the tile. - Fix: before the loop each block checks whether the farthest A/B element it reads over the K loop is out of bounds (or a pad flag is set). Such blocks wait for all async loads; the other blocks keep the partial wait. The check is block uniform (scalar branch, about two SALU instructions per half iteration) and only exists on gfx125; other targets keep the partial wait at compile time, and the gfx950 V2 ISA is unchanged. - **Tests:** - new V2 tail-case gtests covering `num_loop` 1, 2, 3, 4, 5, 7 and 32, for XDL and WMMA, plus a padded tile and an unpadded M tail (`PaddedTailCoverage`, `UnpaddedMTailCoverage`). On gfx1250 the padded case fails 10/10 without the wait fix; - the CompAsync WMMA tests are added to `CK_TILE_GEMM_TEST_TARGETS`; - a new CMake helper `ck_tile_restrict_archs` limits the gfx125-only WMMA tests (TDM and CompAsync async LDS) to gfx125 targets, because the `_wmma` suffix only drops gfx9. - **Follow-up in #12544.** Once this PR merges, #12544 can drop its gfx1250 gate that forces non-MX `comp_async` to pad TTT. ## Changes ### New | File | Purpose | |---|---| | `include/ck_tile/ops/epilogue/tdm_multi_d_epilogue.hpp` | TDM epilogue with fused row-major D tensors | | `test/ck_tile/grouped_gemm/test_grouped_gemm_tdm_wmma.cpp`, `test_grouped_gemm_tdm_ut_cases.inc` | Grouped GEMM TDM V1/V2 gtests | | `test/ck_tile/gemm_multi_d/test_gemm_multi_d_tdm_wmma.cpp` | Multi-D TDM gtests | | `test/ck_tile/batched_contraction/test_batched_contraction_tdm_wmma.cpp` | Batched contraction TDM gtests (compared against the CPU reference with tolerances) | | `test/ck_tile/batched_contraction/test_batched_contraction_tdm_compile_fail.cpp`, `CMakeLists.txt` | Compile-fail tests for the TDM contraction `static_assert`s | ### Modified | File | Change | |---|---| | `ops/gemm/kernel/grouped_gemm_kernel.hpp` | Tuple-form `RunGemm` for TDM; no-padding `static_assert` for TDM; split-K, persistent, cluster-launch and D-tensor rejection; `_tdm` name suffix; `UsePersistentKernel` now uses the base kernel's detection | | `ops/gemm/kernel/gemm_multi_d_kernel.hpp` | Requires `TdmMultiDEpilogue` and no padding, row-major E and no cluster launch for TDM; gfx125 check | | `ops/gemm/kernel/universal_gemm_kernel.hpp` | Forced N pad on D loads for TDM kernels with D tensors; `kIsTdmPipeline`, shared `IsTdmArgumentSupported` / `IsTdmDeviceSupported` host checks | | `ops/epilogue/tdm_epilogue.hpp` | `is_tdm_epilogue_v` trait (true for `TdmEpilogue` and types derived from it) | | `ops/batched_contraction/kernel/batched_contraction_kernel.hpp` | `RunGemmTdm` (requires a TDM epilogue); host check that merged dimensions collapse to one positive stride; TDM batch stride; shared TDM host checks | | `ops/batched_contraction/kernel/batched_contraction_multi_abd_kernel.hpp` | `static_assert` limiting TDM to single A/B | | `ops/epilogue.hpp` | Includes the new epilogue | | `ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp`, `..._comp_async_v2.hpp` | CompAsync K-prefetch bound gating and tail fixes; V2 full async-load wait on gfx125 for blocks that reach past A/B | | `test/ck_tile/gemm/*comp_async*`, `test_gemm_pipeline_{kernel_types,util}.hpp`, `CMakeLists.txt` | CompAsync V2 tail-case tests (XDL + WMMA, unpadded, padded, unpadded M tail) with a shared tail-coverage loop; CompAsync WMMA tests in `CK_TILE_GEMM_TEST_TARGETS` | | `test/ck_tile/{CMakeLists.txt, grouped_gemm/*, gemm_multi_d/*}` | Registers the new tests; `ck_tile_restrict_archs`; TDM instance helpers in the test utils | ## Testing Results below are at head b38f77a unless noted. Every gfx1250 suite was run several times in a row (5 repeats for the TDM suites, 10 for the V2 tail cases) with no failures and no skips. The grouped TDM host validation rejects non-positive K in retained kernel arguments before a device query. `MakeKargs` continues to drop zero-sized groups, covered by the existing passing `EmptyGroup` cases. | Arch | Tests | Result | |---|---|---| | gfx1250 | Grouped GEMM TDM | 70/70 pass, including KTail, RaggedMN (M/N not multiples of the tile) and negative-K rejection | | gfx1250 | GEMM multi-D TDM | 360/360 pass | | gfx1250 | Batched contraction TDM | 80/80 pass, including ragged 200x136x100 and RejectZeroStrideMDim | | gfx1250 | Contraction TDM compile-fail | 4/4 fail to compile as expected (num D, pad N, multi-ABD, CShuffle epilogue) | | gfx1250 | CompAsync V1 WMMA suite (incl. unpadded-K tail cases) | 16/16 pass | | gfx1250 | CompAsync V2 tail cases (WMMA: unpadded, unpadded M tail, padded) | 3/3 pass, 10/10 repeats; the padded case fails 10/10 without the wait fix | | gfx1250 | TDM sub-tile gtests | 24/24 pass | | gfx1250 | Existing grouped GEMM / multi-D (cshuffle 231/231, default2d 294/294) / contraction tests | all pass | | gfx1250 | Combined with #12544 (comp_async, comp_async_v2, TDM sub-tile gtests) | 16/16, 2/2, 24/24 pass | | MI350X (gfx950) | CompAsync V1 / V2 tail cases (XDL) | V1 58/58, V2 4/4 pass. V2 fails 4/4 without the fix. At b38f77a the V2 tail cases (incl. padded and M tail) pass 5/5 repeats and the V2 device ISA is identical to the previous head. | | MI350X (gfx950) | CompAsync V1 perf vs develop, fp8, 10 shapes (even and odd `num_loop`) | median -0.22%, range -1.32% to +0.70% | | MI350X (gfx950) | grouped GEMM, multi-D, contraction suites | Same pass/skip counts before and after; non-CompAsync device ISA byte-identical (measured at a2ed839; later commits only change host checks and TDM-only code) | | MI300X (gfx942) | grouped GEMM, multi-D, contraction suites | Same pass/skip counts before and after (ctest 553 tests); device ISA byte-identical | | gfx90a (C++17 CI stage) | host build with `-Werror` | No new warnings (rebuilt at b38f77a) | A mutation check confirms the new stride test: with the stride check removed, RejectZeroStrideMDim fails 4/4. The gfx1250 CompAsync perf numbers from the previous revision (V2 fp16 4096^3 about +1%, other shapes within +/-0.6%) were measured before the loop peel and were not re-measured. All C++ files are formatted with clang-format-18. The V1 and V2 tail regressions share one tail-coverage loop between the XDL and WMMA suites. The V1 unpadded tail regression and covers `num_loop` = 1, 2, 3, 4, 5, 7 and 32 over four M/N shapes, with `k_batch = 1`. At head 26992d2 the gfx950/XDL regression passed 8/8 typed tests without skips (224 typed shape/loop combinations). ## Known limitations / reviewer notes - **Universal-GEMM TDM gtest not run here.** `test_ck_tile_gemm_pipeline_tdm_wmma` skipped all 112 of its cases because the gfx1250 part we tested on has no cluster launch. This PR does not change the universal TDM pipeline; it only changes how D tensors are padded. - **Pre-existing CompAsync issues, not changed by this PR:** CompAsyncV2 with a 256x256x64 block tile (32x32 warp tile) gives wrong results on gfx95x, and CompAsync/V2 faults when both N and K are padded. Both reproduce on develop and will be tracked separately. - **Host API change.** `BatchedContractionKernel::IsSupportedArguments` is no longer `constexpr`, because it now queries the device and checks pointer alignment. No in-tree caller uses it in a constant expression. - **Split-K.** Rejected on every TDM path. - **Multi-ABD `static_assert`.** It is only instantiated for gfx1250 TDM builds. - **Out of scope.** PRESHUFFLE_TDM is tracked in #10931; stream-K is excluded. - **Follow-up bridge PR.** It will enable COMPUTE_TDM_V1/V2 for grouped GEMM, multi-D and contraction in the TE/dispatcher bridge once this PR merges. --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JIRA ID : AICK-2098
JIRA ID : AICK-2099
JIRA ID : AICK-2100
JIRA ID : AICK-2137
Summary
Enables the regular (gemm_universal), grouped, multi-D, and multi-ABD GEMM Tile-Engine → Dispatcher bridges on gfx1250 (MI400, RDNA4 / WMMA), for fp16/bf16 (+ fp8 on the regular path).
This is the consolidation of the previously-stacked PRs #10904 / #10914 / #10915 / #10917 into one. Python + JSON only (no C/C++) — no clang-format gate.
The two partial variants (stream-K #10919, microscaling mx #10916) are not included — they need C++/kernel follow-ups (see below) and remain separate.
Do the merged bridges (Universal/grouped/multi-D) support gfx1250? — No.
As merged, none of these bridges run on gfx1250. They share a CDNA-only arch gate (
_SUPPORTED_ARCHES = gfx90a/gfx942/gfx950) and most pinned the MFMA32x32x16warp tile in their CI config, which is invalid on gfx1250's RDNA4 WMMA engine (fp16/bf16 WMMA is16x16x32, fp8 is16x16x64).Changes
Shared enablement (was #10904)
dispatcher/python/gemm_utils.py:gfx1250in_SUPPORTED_ARCHES+ gfx1250 wave combos; OCP fp8/bf8 codec (float8_e4m3fn/float8_e5m2) selected by arch (gfx950/gfx12xx → OCP, gfx942 → FNUZ).tile_engine/ops/gemm/gemm_validation_utils.py: gfx1250 warp-configuration combos.dispatcher/codegen/arch_specs_generated.py: gfx1250 fp8/bf8 WMMA[[16,16,64]](fp16/bf16[[16,16,32]]already present).tile_engine/ops/gemm/gemm_full_benchmark.py:SUPPORTED_DTYPES+= fp8, bf8.tile_engine/ops/gemm/configs/: gfx1250 regular CI configs (fp16/bf16 16x16x32; fp8 16x16x64).Variant configs + tests
grouped_gemm/configs/default_ci_config_gfx1250.json.gemm_multi_d/configs/default_ci_config_gfx1250.json+--archhelp.gemm_multi_abd/configs/default_ci_config_gfx1250.json(fp16-only, upstream design).dispatcher/tests/(gfx1250 arch, OCP fp8 codec, grouped/multi-D/multi-ABD configs). 63 tests pass.Validation (MI400 / gfx1250, ROCm 10.1)
Notes
arch_specs_generated.pyis auto-generated fromarch_specs.json, which has no gfx1250 entry on develop (the gfx1250 fp16/bf16 rows were hand-added to the generated file). This PR follows that precedent; a follow-up should add a full gfx1250 stanza toarch_specs.json.--verifyneeds a looser tol (~0.05–0.15) than the fp16 default (2e-2), since the reference is fp32 and fp8 carries E4M3 quantization..sobuildstatic_asserts gfx950-only viapreShuffleScaleBuffer_gfx950; ck_tile already shipspreShuffleScaleBuffer_gfx1250). Both need C++ follow-ups.Supersedes #10904, #10914, #10915, #10917.