Repository navigation
feat(ck-tile): enable multi-ABD GEMM bridge on gfx1250 (MI400) - #11048
Merged
Merged
Conversation
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.
Adds the gfx1250 (MI400) enablement for the multi-ABD GEMM TE->dispatcher bridge (#9305), split out of #10921 so it can land independently once its MI400 validation is green. - gemm_multi_abd/configs/default_ci_config_gfx1250.json: WMMA-shaped CI config (warp_tile 16x16x32) so kernels actually instantiate on gfx1250 WMMA instead of the CDNA MFMA 32x32x16 tile. - test_multi_abd_bridge.py: gfx1250 CPU test class asserting the WMMA warp tile and a gfx1250-supported wave triple. Depends on #10921 (shared gfx1250 arch enablement in gemm_utils.py / gemm_validation_utils.py). Known: some multi-ABD MI400 correctness cases still failing under investigation before merge.
ozturkosu
changed the base branch from
users/muozturk/ck/gemm-bridges-gfx1250-enable
to
develop
August 19, 2026 22:51
Contributor
Author
|
Warning Merge order: #10921 must merge first. This PR targets
Additionally, multi-ABD MI400 correctness is not yet green — do not merge until both #10921 lands and the failing gfx1250 cases are resolved. |
✅ 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. |
ozturkosu
marked this pull request as draft
August 22, 2026 21:16
gemm_multi_abd shipped no configs/example_problems.json, so the benchmark fell back to DEFAULT_PROBLEMS which includes a non-tile-divisible 257^3 shape. The pad=false a2_b2_d2 kernels can't run it, so SelectedKernel::launch() threw and the ctypes bridge returned status -2 (16/64 measurements failing on MI400). Every other variant (gemm_universal, bquant, ...) ships tile-divisible example problems; add the same for multi_abd. Verified on gfx1250 (MI400): 80/80 measurements verify vs CPU (incl. the non-square 1024x512x256 case), 0 failures.
Contributor
Author
|
Diagnosed the 16/64 multi-ABD failures — it's a test-config gap, not a kernel bug. The a2_b2_d2 kernels were failing with Fixed by adding Re-verified on gfx1250 (MI400, Helios The non-square 1024x512x256 case passing rules out a multi-tensor stride/marshalling bug — the a2_b2_d2 path is correct for every divisible shape. Note: if we want a2_b2_d2 to also handle arbitrary non-divisible shapes, that's a separate change (enable padded kernels, or have the bridge check |
Contributor
Author
|
quick update on the multi-ABD MI400 PR. I dug into why a few cases were failing and it turned out not to be a real kernel problem. The test was just feeding it one odd input size that this kind of kernel isn't meant to handle, and multi-ABD was missing the small config file that every other op already ships to keep it on valid sizes. I added that file and re-ran everything on the MI400 box — all cases pass now, verified against the CPU reference with zero failures. Fix is already pushed to the PR. |
ozturkosu
marked this pull request as ready for review
August 22, 2026 23:17
#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.
Contributor
Author
|
Interim fix for #11161 pushed here (54473b9). The full Verification of the gate: the previously-failing config now expands to 0 kernels on gfx1250 ( A full post-fix |
Contributor
Author
Post-gate
|
Contributor
Author
|
| op | supported matrix | kernels | result rows | fail |
|---|---|---|---|---|
gemm_multi_abd |
fp16 × rcrr (only) | 1,640 | 4,920 | 0 |
Method: full default_config (whole tile/trait grid) → compile each kernel for gfx1250 → run on GPU → verify vs fp32 CPU reference, 3 tile-divisible shapes (768³/1536³/3072³). Run against this branch as-is (interim gate on), so the compv3+intrawave+8-warp class is excluded — that's the #11161 defect, fixed for real by #11175 and validated separately there.
Every kernel in multi_abd's supported (fp16/rcrr) surface is correct on gfx1250.
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.
Contributor
Author
On-device validation —
|
| 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).
#11161 note: the raw sweep initially showed 552 failures, all in the compv3 + intrawave + 8-warp (2x4x1 / 4x2x1) class — the known #11161 defect (a wave64 assumption in the comp_v3 pipeline, wrong on gfx1250 wave32). An interim gate excludes that class here; the real fix, #11175 (Use8WarpSchedule), is already on develop, so once this branch is on develop that class validates directly.
Net: 0 correctness defects in multi-ABD's supported surface on gfx1250.
yraparti
approved these changes
Sep 1, 2026
ozturkosu
enabled auto-merge (squash)
September 1, 2026 21:49
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Sep 2, 2026
feat(ck-tile): enable multi-ABD GEMM bridge on gfx1250 (MI400) (#11048) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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).
Contributor
TheRock Submodule Bump ActivityNewest first
|
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-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 head56e7bc3):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.jsoninstantiates zero runnable kernels there.Supported surface (gfx1250)
The
a2_b2_d2variant is intentionally narrow — the codegen, ctypes lib, and runner are all fp16-only and rcrr-only (expand_sweeprejects 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_configsweep of the fp16/rcrr surface, 3 tile-divisible shapes (768³ / 1536³ / 3072³):gemm_multi_abdEvery 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 thecomp_v3pipeline, wrong on gfx1250 wave32). An interim gate excludes that class so the supported surface reports clean; the real fix, #11175 (Use8WarpSchedule), is already ondevelop.Notes / dependencies
developvia the sibling series).