feat(ck-tile): comp_tdm_v3 bf16/fp16 GEMM on gfx1250 by jefyang1 · Pull Request #15005 · ROCm/rocm-libraries

jefyang1 · 2026-11-03T21:12:30Z

JIRA ID : AICK-2966

Motivation

comp_tdm_v3 (from #14466) is a fast gfx1250 GEMM pipeline, but it accepts only fp8/bf8 A/B in the RCR layout. GEMM workloads on gfx1250 also need 16-bit inputs, in every layout and at shapes that are not tile multiples. This PR opens comp_tdm_v3 to bf16 and fp16 A/B in every A/B layout, makes it correct for any M x N x K with either epilogue, and adds split-K with the CShuffle epilogue. Everything is opt-in: kernels that do not use comp_tdm_v3 compile the same.

Technical Details

bf16 / fp16 A/B in every layout. comp_tdm_v3 accepts 16-bit A/B as RCR, RRR, CRR and CCR (warp tile 16x16x32). 8-bit A/B stays RCR-only, and C stays row-major.
Any M x N x K. TDM clips every A/B box to the real extent, so A/B never limit the shape. The TDM epilogue clips C in hardware and accepts every shape; the CShuffle epilogue gets the kernel's usual C checks (see Default behaviour changes).
Split-K with the CShuffle epilogue. Each K range runs as its own workgroup layer and adds its partial C atomically; each split's TDM loads are clipped to its own K range, including a partial last K tile. The TDM epilogue overwrites C, so it still rejects k_batch > 1.
Example 54, tile_example_gemm_comp_tdm_v3_bf16_bench. A correctness check and benchmark for any shape: -layout, -epilogue (tdm or cshuffle), -tile (256x256x64 or 128x128x64), -vector_c, -k_batch; modes verify-int (exact), verify-rand, verify-sparse (exact at any split-K order), peak and sustained. Exit code 0 pass, 1 wrong result, 2 rejected, 3 error.

Scope of `comp_tdm_v3` after this PR

	A/B dtypes	A/B layouts	Epilogue	Split-K
8-bit	fp8, bf8	RCR	TDM, CShuffle	CShuffle, C vector of at least 2
16-bit	bf16, fp16	RCR, RRR, CRR, CCR	TDM, CShuffle	CShuffle, C vector of at least 2

C is row-major in every case. Warp tiles: 32x32x128 (8-bit), 16x16x32 (16-bit).

Default behaviour changes

UniversalGemmKernel::IsSupportedArgument, TDM pipeline with a non-TDM epilogue. Such a kernel skipped the C checks, so a CShuffle comp_tdm_v3 instance accepted an unaligned N without kPadN, or an N that its C vector size does not divide, and a partial vector could write past the row end into the next row. It now applies the same C checks as the other pipelines. TDM-epilogue kernels are unchanged: the TDM store clips C in hardware.
Split-K for a TDM pipeline. IsTdmArgumentSupported rejected k_batch != 1 for every TDM pipeline. The rejection now applies only with a TDM epilogue; with CShuffle the epilogue's own split-K rule decides (EpilogueSupportsSplitK, which also covers develop's SupportsSplitK = false on the TDM epilogue).

No other pipeline or kernel changes.

Test Plan

On gfx1250:

comp_tdm_v3 gtests:
- fp8, fp8 long K, and bf16 in every A/B layout (one, two and three K tiles and 512 x 512 x 1024), bf16 4096 cubed;
- unaligned M/N/K: M, N or K unaligned, all three unaligned with odd strides, below one tile, M = 1 and N = 1, for the padded and unpadded CShuffle instances, the TDM-epilogue instances, early refill (RefillPoint::AfterArrive) and fp8. Every shape must either be rejected by the C rule or computed correctly;
- split-K: which k_batch the kernel accepts (including one beyond the K steps, which must be rejected), and k_batch 2, 3 and 4 over a K with a partial last tile, on sparse inputs (exact at any split order) and random inputs; the TDM-epilogue instances must reject split-K.
The TDM/MX gtest suites that share IsSupportedArgument, for regressions.
Example 54 over a 153-shape production GEMM list, with the exact integer check.

Test Result

comp_tdm_v3 gtests, after the rebase onto develop: fp8 66/66, fp8 long K 16/16, bf16 84/84, bf16 4096 cubed 18/18, unaligned M/N/K (including split-K) 115/115.
TDM/MX suites: comp_tdm_pc 91/91 and the TDM epilogue 1/1 after the rebase; before it, every TDM/MX gtest suite passed, on one gfx1250 host with all six commits and on a second with the first five.
Example 54: every shape of the 153-shape list passes the exact integer check, on two gfx1250 hosts. A split-K sparse check (k_batch 16) is exact on every sampled point.

Notes

Split-K adds bf16 partial sums atomically, so the low bits vary from run to run, and an odd N cannot split (CShuffle needs a C vector of at least 2). An fp32-workspace reduction would give deterministic results and cover odd N; it is a follow-up.
Performance data and follow-ups are tracked in AICK-2966.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

…; bf16 gtests comp_tdm_v3 required A row-major, B column-major because the other A/B layouts spill VGPRs in the K loop at the 16-wave tile. That holds for 8-bit A/B with the 32x32x128 warp tile, not for 16-bit A/B: with the 16x16x32 warp tile, the only bf16/fp16 WMMA warp tile on gfx1250, every A/B layout compiles without spills (16 and 4 waves, K tiles 64 and 128; CCR spills only at 256x256x128). The pipeline now accepts every A/B layout for 16-bit A/B and keeps RCR for 8-bit A/B. C stays row-major. The comp_tdm_v3 gtest helper takes the K tile, the warp tile and the A/B layouts (defaults: the fp8 setup, so the fp8 instances are unchanged), packs A and B in their layouts and checks the result on several host threads. New gtests, bf16 x bf16 -> bf16, default policy: - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma (SMOKE_TEST): 14 instances over RCR, RRR, CRR and CCR at 16 and 4 waves; one, two and three K tiles and 512x512x1024, integer inputs (bit-exact) and random inputs; split-K rejected. - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma (REGRESSION_TEST): 4096x4096x4096 on six instances. - the scratch-use check also runs on the bf16 binary. Gate: on gfx1250 hardware, bf16 70/70 and 4096^3 12/12 passed; the fp8 gtests 55/55 and 16/16 long-K still pass; scratch check: 14 bf16 kernels, no scratch, no VGPR spills. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…without a TDM epilogue; unaligned M/N/K gtests for comp_tdm_v3 UniversalGemmKernel::IsSupportedArgument returns the TDM checks for every TDM pipeline and skips the C checks: a TDM epilogue clips C in hardware. But comp_tdm_v3 writes C through the CShuffle epilogue, whose vector stores need them. With an unpadded comp_tdm_v3 kernel, a C row-major N that is not a multiple of NPerBlock was accepted, and the last partial vector of each row wrote zeros past the row end into the start of the next row (bf16, every layout and tile: N = 1, 45, 300, 517 failed). The C checks move into IsCTensorSupported, which IsSupportedArgument calls as before and the TDM branch now also calls when the epilogue is not a TDM epilogue: an N (M) that is not a multiple of the tile needs kPadN (kPadM), and N (M) must be a multiple of the C vector size. TDM pipelines with a TDM epilogue are unchanged. The comp_tdm_v3 gtest helper takes kPadM/kPadN and a fixed C vector size (defaults unchanged). New gtest for M, N and K that are not multiples of the tile (SMOKE_TEST; test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma here, renamed to ..._bf16_unaligned_mnk_wmma in the next commit): M, N or K unaligned, odd strides, below one tile, M = 1, N = 1; for each shape the kernel must accept or reject it by the rule above, and every accepted shape must be correct. Unpadded instances in every layout reject an unaligned N and pass the rest; padded instances (every layout, C vector 8, and RCR/CCR with C vectors 2 and 1) accept it. The 4096^3 test adds an unaligned 4111x4096x4105. Gate: on gfx1250 hardware, bf16 unaligned M/N/K 50/50, bf16 70/70, bf16 4096^3 18/18, fp8 55/55 and long-K 16/16 passed. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ue instances - The unaligned-shape gtest takes CK's name for it, unaligned_mnk (as test_gemm_unaligned_mnk.cpp): test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma. "Ragged" stays only where it describes TDM clipping, as elsewhere in CK. - The gtest helper can build comp_tdm_v3 with TdmEpilogue instead of CShuffleEpilogue. TdmEpilogue stages C in LDS and writes it with one TDM store that clips to the real M x N, so the unpadded kernel serves every N with no narrow C vector. Same VGPRs, spills and LDS as with CShuffle (compile probe, 20 bf16 configs over all layouts). New suite: unpadded TdmEpilogue instances in every layout must accept and pass every unaligned shape. Gate: on gfx1250 hardware, unaligned M/N/K 68/68 (18 TdmEpilogue), bf16 70/70, 4096^3 18/18, fp8 55/55, long-K 16/16; the 18 configured TDM/MX gtests pass or skip on their own hardware gates (cluster launch and one MX TDM variant need a newer gfx1250 revision). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…nchmark for any shape tile_example_gemm_comp_tdm_v3_bf16_bench runs comp_tdm_v3 with bf16 A/B and bf16 C for any M x N x K, default 4096^3. Instances are chosen at run time: -layout rcr|rrr|crr|ccr, -epilogue tdm|cshuffle, -tile 256x256x64 (16 waves) or 128x128x64 (4 waves), warp tile 16x16x32, and for CShuffle -vector_c auto|8|2|1. The default epilogue is TdmEpilogue, which clips C in hardware and so serves every shape. The CShuffle instances are padded and store C in vectors of 8, 2 or 1; auto picks the widest that divides N, so CShuffle also serves every shape. Modes and output follow example 53: verify-int (bit-exact), verify-rand, peak and sustained, one JSON line per result. The exit code is the result (0 pass, 1 wrong, 2 rejected or bad arguments, 3 error), so a shape list can be checked with one process per shape. Gate: on gfx1250 hardware, 4096^3 RCR passes verify-int and verify-rand with both epilogues; every layout, both epilogues and both tiles on seven shapes (333x517x1000, 256x300x128, 256x264x128, 1x256x128, 256x1x128, 37x48x40, 4111x4096x4105, plus 512x1026x512 for CShuffle): TdmEpilogue and CShuffle with -vector_c=auto pass all (C vectors 8, 2 and 1 as N allows); a forced C vector that does not divide N is rejected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…, bf16 and fp8 RefillPoint::AfterArrive has its own K-tile bookkeeping and was only tested with K a multiple of the K tile (the fp8 gtests) or not at all on unaligned shapes (the bf16 unaligned test ran the default policy only). The unaligned M/N/K test now also runs it: - bf16, AfterArrive with TdmEpilogue, RCR 256x256x64 and 128x128x64 and CRR 256x256x64: every shape runs; - fp8 (RCR, warp tile 32x32x128, K tile 256), the default policy and AfterArrive at 16 and 4 waves with unpadded CShuffle as fp8 ships (an unaligned M or K runs, an unaligned N is rejected), and AfterArrive with TdmEpilogue at 16 waves (every shape runs). The shapes include a partial last K tile (two tiles plus 8, one tile plus 1, K = 1000 and K = 40), integer inputs bit-exact and random inputs within tolerance. V3Bf16Config takes the policy (default unchanged). The test now covers both data types, so it is renamed test_ck_tile_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma. Gate: on gfx1250 hardware, unaligned M/N/K 84/84, bf16 70/70, bf16 4096^3 18/18, fp8 55/55 and long-K 16/16 passed. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

UniversalGemmKernel rejected k_batch != 1 for every TDM pipeline. That rejection is about the TdmEpilogue, which overwrites C with a plain TDM store. comp_tdm_v3 with the CShuffle epilogue adds each K split into C atomically, and each split's TDM loads are clipped to its own K range (including a partial last K tile), so split-K is now checked only with a TDM epilogue. The CShuffle epilogue's own rule still applies: bf16 C needs a C vector of at least 2. Tests (comp_tdm_v3 gtests): - SplitKSupport (fp8 and bf16): which k_batch the kernel accepts, including a k_batch beyond the warp-tile K steps, which must be rejected. - SplitK: k_batch 2, 3 and 4 over five K tiles plus a partial one, on sparse inputs (exact at any split order: every partial and running sum is an integer that bf16 holds exactly) and random inputs (bounded by each split's bf16 rounding). - Unaligned M/N/K with k_batch 3: the padded CShuffle instances run every shape they accept; the C-vector-1 and TdmEpilogue instances reject split-K. Example 54 gains -k_batch (C is zeroed before every split-K launch, inside the timed loop) and -mode=verify-sparse. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

therock-pr-bot · 2026-11-03T21:21:21Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-11-03T21:21:25Z

🎉 All checks passed! This PR is ready for review.

feat(ck-tile): comp_tdm_v3 bf16/fp16 GEMM on gfx1250 JIRA ID : AICK-2966 ## Motivation `comp_tdm_v3` (from #14466) is a fast gfx1250 GEMM pipeline, but it accepts only fp8/bf8 A/B in the RCR layout. GEMM workloads on gfx1250 also need 16-bit inputs, in every layout and at shapes that are not tile multiples. This PR opens `comp_tdm_v3` to bf16 and fp16 A/B in every A/B layout, makes it correct for any M x N x K with either epilogue, and adds split-K with the CShuffle epilogue. Everything is opt-in: kernels that do not use `comp_tdm_v3` compile the same. ## Technical Details - **bf16 / fp16 A/B in every layout.** `comp_tdm_v3` accepts 16-bit A/B as RCR, RRR, CRR and CCR (warp tile 16x16x32). 8-bit A/B stays RCR-only, and C stays row-major. - **Any M x N x K.** TDM clips every A/B box to the real extent, so A/B never limit the shape. The TDM epilogue clips C in hardware and accepts every shape; the CShuffle epilogue gets the kernel's usual C checks (see Default behaviour changes). - **Split-K with the CShuffle epilogue.** Each K range runs as its own workgroup layer and adds its partial C atomically; each split's TDM loads are clipped to its own K range, including a partial last K tile. The TDM epilogue overwrites C, so it still rejects `k_batch > 1`. - **Example 54, `tile_example_gemm_comp_tdm_v3_bf16_bench`.** A correctness check and benchmark for any shape: `-layout`, `-epilogue` (tdm or cshuffle), `-tile` (256x256x64 or 128x128x64), `-vector_c`, `-k_batch`; modes `verify-int` (exact), `verify-rand`, `verify-sparse` (exact at any split-K order), `peak` and `sustained`. Exit code 0 pass, 1 wrong result, 2 rejected, 3 error. ### Scope of `comp_tdm_v3` after this PR | | A/B dtypes | A/B layouts | Epilogue | Split-K | |---|---|---|---|---| | 8-bit | fp8, bf8 | RCR | TDM, CShuffle | CShuffle, C vector of at least 2 | | 16-bit | bf16, fp16 | RCR, RRR, CRR, CCR | TDM, CShuffle | CShuffle, C vector of at least 2 | C is row-major in every case. Warp tiles: 32x32x128 (8-bit), 16x16x32 (16-bit). ### Default behaviour changes - **`UniversalGemmKernel::IsSupportedArgument`, TDM pipeline with a non-TDM epilogue.** Such a kernel skipped the C checks, so a CShuffle `comp_tdm_v3` instance accepted an unaligned N without `kPadN`, or an N that its C vector size does not divide, and a partial vector could write past the row end into the next row. It now applies the same C checks as the other pipelines. TDM-epilogue kernels are unchanged: the TDM store clips C in hardware. - **Split-K for a TDM pipeline.** `IsTdmArgumentSupported` rejected `k_batch != 1` for every TDM pipeline. The rejection now applies only with a TDM epilogue; with CShuffle the epilogue's own split-K rule decides (`EpilogueSupportsSplitK`, which also covers develop's `SupportsSplitK = false` on the TDM epilogue). No other pipeline or kernel changes. ## Test Plan On gfx1250: - `comp_tdm_v3` gtests: - fp8, fp8 long K, and bf16 in every A/B layout (one, two and three K tiles and 512 x 512 x 1024), bf16 4096 cubed; - unaligned M/N/K: M, N or K unaligned, all three unaligned with odd strides, below one tile, M = 1 and N = 1, for the padded and unpadded CShuffle instances, the TDM-epilogue instances, early refill (`RefillPoint::AfterArrive`) and fp8. Every shape must either be rejected by the C rule or computed correctly; - split-K: which `k_batch` the kernel accepts (including one beyond the K steps, which must be rejected), and `k_batch` 2, 3 and 4 over a K with a partial last tile, on sparse inputs (exact at any split order) and random inputs; the TDM-epilogue instances must reject split-K. - The TDM/MX gtest suites that share `IsSupportedArgument`, for regressions. - Example 54 over a 153-shape production GEMM list, with the exact integer check. ## Test Result - `comp_tdm_v3` gtests, after the rebase onto develop: fp8 66/66, fp8 long K 16/16, bf16 84/84, bf16 4096 cubed 18/18, unaligned M/N/K (including split-K) 115/115. - TDM/MX suites: `comp_tdm_pc` 91/91 and the TDM epilogue 1/1 after the rebase; before it, every TDM/MX gtest suite passed, on one gfx1250 host with all six commits and on a second with the first five. - Example 54: every shape of the 153-shape list passes the exact integer check, on two gfx1250 hosts. A split-K sparse check (k_batch 16) is exact on every sampled point. ## Notes - Split-K adds bf16 partial sums atomically, so the low bits vary from run to run, and an odd N cannot split (CShuffle needs a C vector of at least 2). An fp32-workspace reduction would give deterministic results and cover odd N; it is a follow-up. - Performance data and follow-ups are tracked in AICK-2966. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

… pipelines Since #15005, UniversalGemmKernel::IsSupportedArgument runs IsCTensorSupported for a TDM pipeline with a non-TDM epilogue: a ragged contiguous C dimension needs kPadN (kPadM), and that dimension must be a multiple of the C vector size. StreamKKernel::IsSupportedArgument calls it after its TDM block, so the Stream-K copy of the C-tail check is redundant for the unpadded TDM configs (the gtest and the dispatcher never pad a TDM pipeline), and its comment, which says the universal kernel skips that check, no longer holds. Drop the copy and keep the gfx125 device check. A padded C goes through the same padded C window as the other Stream-K pipelines.

JIRA ID : AICK-2966 ## Motivation `comp_tdm_v3` (from #14466) is a fast gfx1250 GEMM pipeline, but it accepts only fp8/bf8 A/B in the RCR layout. GEMM workloads on gfx1250 also need 16-bit inputs, in every layout and at shapes that are not tile multiples. This PR opens `comp_tdm_v3` to bf16 and fp16 A/B in every A/B layout, makes it correct for any M x N x K with either epilogue, and adds split-K with the CShuffle epilogue. Everything is opt-in: kernels that do not use `comp_tdm_v3` compile the same. ## Technical Details - **bf16 / fp16 A/B in every layout.** `comp_tdm_v3` accepts 16-bit A/B as RCR, RRR, CRR and CCR (warp tile 16x16x32). 8-bit A/B stays RCR-only, and C stays row-major. - **Any M x N x K.** TDM clips every A/B box to the real extent, so A/B never limit the shape. The TDM epilogue clips C in hardware and accepts every shape; the CShuffle epilogue gets the kernel's usual C checks (see Default behaviour changes). - **Split-K with the CShuffle epilogue.** Each K range runs as its own workgroup layer and adds its partial C atomically; each split's TDM loads are clipped to its own K range, including a partial last K tile. The TDM epilogue overwrites C, so it still rejects `k_batch > 1`. - **Example 54, `tile_example_gemm_comp_tdm_v3_bf16_bench`.** A correctness check and benchmark for any shape: `-layout`, `-epilogue` (tdm or cshuffle), `-tile` (256x256x64 or 128x128x64), `-vector_c`, `-k_batch`; modes `verify-int` (exact), `verify-rand`, `verify-sparse` (exact at any split-K order), `peak` and `sustained`. Exit code 0 pass, 1 wrong result, 2 rejected, 3 error. ### Scope of `comp_tdm_v3` after this PR | | A/B dtypes | A/B layouts | Epilogue | Split-K | |---|---|---|---|---| | 8-bit | fp8, bf8 | RCR | TDM, CShuffle | CShuffle, C vector of at least 2 | | 16-bit | bf16, fp16 | RCR, RRR, CRR, CCR | TDM, CShuffle | CShuffle, C vector of at least 2 | C is row-major in every case. Warp tiles: 32x32x128 (8-bit), 16x16x32 (16-bit). ### Default behaviour changes - **`UniversalGemmKernel::IsSupportedArgument`, TDM pipeline with a non-TDM epilogue.** Such a kernel skipped the C checks, so a CShuffle `comp_tdm_v3` instance accepted an unaligned N without `kPadN`, or an N that its C vector size does not divide, and a partial vector could write past the row end into the next row. It now applies the same C checks as the other pipelines. TDM-epilogue kernels are unchanged: the TDM store clips C in hardware. - **Split-K for a TDM pipeline.** `IsTdmArgumentSupported` rejected `k_batch != 1` for every TDM pipeline. The rejection now applies only with a TDM epilogue; with CShuffle the epilogue's own split-K rule decides (`EpilogueSupportsSplitK`, which also covers develop's `SupportsSplitK = false` on the TDM epilogue). No other pipeline or kernel changes. ## Test Plan On gfx1250: - `comp_tdm_v3` gtests: - fp8, fp8 long K, and bf16 in every A/B layout (one, two and three K tiles and 512 x 512 x 1024), bf16 4096 cubed; - unaligned M/N/K: M, N or K unaligned, all three unaligned with odd strides, below one tile, M = 1 and N = 1, for the padded and unpadded CShuffle instances, the TDM-epilogue instances, early refill (`RefillPoint::AfterArrive`) and fp8. Every shape must either be rejected by the C rule or computed correctly; - split-K: which `k_batch` the kernel accepts (including one beyond the K steps, which must be rejected), and `k_batch` 2, 3 and 4 over a K with a partial last tile, on sparse inputs (exact at any split order) and random inputs; the TDM-epilogue instances must reject split-K. - The TDM/MX gtest suites that share `IsSupportedArgument`, for regressions. - Example 54 over a 153-shape production GEMM list, with the exact integer check. ## Test Result - `comp_tdm_v3` gtests, after the rebase onto develop: fp8 66/66, fp8 long K 16/16, bf16 84/84, bf16 4096 cubed 18/18, unaligned M/N/K (including split-K) 115/115. - TDM/MX suites: `comp_tdm_pc` 91/91 and the TDM epilogue 1/1 after the rebase; before it, every TDM/MX gtest suite passed, on one gfx1250 host with all six commits and on a second with the first five. - Example 54: every shape of the 153-shape list passes the exact integer check, on two gfx1250 hosts. A split-K sparse check (k_batch 16) is exact on every sampled point. ## Notes - Split-K adds bf16 partial sums atomically, so the low bits vary from run to run, and an odd N cannot split (CShuffle needs a C vector of at least 2). An fp32-workspace reduction would give deterministic results and cover odd N; it is a follow-up. - Performance data and follow-ups are tracked in AICK-2966. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

assistant-librarian · 2026-11-06T12:55:37Z

TheRock Submodule Bump Activity

Newest first

2026-11-06 — Included in TheRock via ROCm/TheRock#9759.

JIRA ID : AICK-2991 ## Motivation `comp_tdm_v3` (#14466, #15005) is a gfx1250 GEMM pipeline for fp8/bf8 and bf16/fp16 A/B. MX (block-scaled) FP8 GEMM on gfx1250 runs only on the earlier TDM pipelines (`comp_tdm`, `comp_tdm_v2`). This PR adds MX FP8 to `comp_tdm_v3`: fp8 x fp8 with one E8M0 scale per 32 K elements, with the TDM or the CShuffle epilogue and split-K. It also adds an MX FP8 example and lets the plain FP8 example 53 choose its tile, epilogue and split-K at run time, as example 54 does. Everything is opt-in: kernels that do not use `comp_tdm_v3` with scales compile the same. ## Technical Details - **MX FP8 on `comp_tdm_v3`.** The pipeline takes `MxGemmKernel`'s A/B scale block windows. Each K tile's packed E8M0 scales are loaded from global memory into one of two register buffers (buffer s holds the scales of the K tile in LDS stage s), one K tile ahead, in both the split-barrier and the joint-barrier K loop. The WMMAs use the scaled 16x16x128 form: the row-streamed A path (`BlockGemmRows`) and the block GEMM's scaled `operator()` take the scales of each row and column. Without scale windows the pipeline is unchanged. - **Split-K.** With the CShuffle epilogue each K range runs as its own workgroup layer and adds its partial C atomically; its TDM loads and scales are clipped to its own K range. The TDM epilogue overwrites C, so it still rejects `k_batch > 1`. - **Split-K only on `comp_tdm_v3`, for every data type.** A TDM pipeline runs `k_batch > 1` only if it declares `kSupportsTdmSplitK` (which `comp_tdm_v3` does) and the epilogue supports split-K. `UniversalGemmKernel` detects the trait (`kPipelineSupportsTdmSplitK`) and `MxGemmKernel` uses the same flag, so plain and MX GEMM follow one rule. Since #15005, `UniversalGemmKernel` let every TDM pipeline run split-K with CShuffle, although only `comp_tdm_v3` is tested with it. - **K tails.** #14780 lets `MxGemmKernel` take K that is not a multiple of the K tile (FP8/FP4, 32-element scales, row-major A, column-major B) on the TDM pipelines. `comp_tdm_v3` MX FP8 matches that configuration, so it takes such K with the TDM epilogue; with CShuffle, `IsSupportedArgument` rejects K tails. - **Example 55, `tile_example_gemm_comp_tdm_v3_mx_fp8_bench`.** MX FP8 correctness check and benchmark for any M x N x K with K a multiple of 128: `-c_dtype` (fp16 or bf16 C), `-epilogue` (tdm or cshuffle), `-tile` (16-, 32-, 64- and 256-row tiles), `-k_batch`; modes `verify-sparse` (exact where C holds the result exactly, at any split order; scales random per row or column and K block), `verify-rand` (bounded), `peak` and `sustained`. Exit code 0 pass, 1 wrong result, 2 rejected, 3 error. - **Example 53, `tile_example_gemm_comp_tdm_v3_fp8_bench`.** Takes `-tile` (the 256x256x256 default and 4-wave tiles for small M), `-epilogue` and `-k_batch`, and lets `Kernel::IsSupportedArgument` decide which M, N, K and `k_batch` an instance takes (it required multiples of 256). With `k_batch > 1` both verify modes use a split-K error bound. ### Scope of `comp_tdm_v3` after this PR | | A/B dtypes | Scales | A/B layouts | Epilogue | Split-K | |---|---|---|---|---|---| | 8-bit | fp8, bf8 | none | RCR | TDM, CShuffle | CShuffle | | 8-bit, MX | fp8 x fp8 | E8M0, 32 K elements | RCR | TDM, CShuffle | CShuffle | | 16-bit | bf16, fp16 | none | RCR, RRR, CRR, CCR | TDM, CShuffle | CShuffle | C is row-major (fp16 or bf16) in every case. MX uses the 16x16x128 warp tile. ### Default behaviour changes - **Example 53's default instance** (256x256x256, CShuffle, `k_batch` 1) keeps its bench name and resources, but its CShuffle epilogue is now padded in M and N (as in example 54), and M, N, K need no longer be multiples of 256. - **Plain GEMM split-K on `comp_tdm`, `comp_tdm_v2` and the producer/consumer TDM pipeline** with the CShuffle epilogue: `IsSupportedArgument` rejects `k_batch > 1` again (accepted since #15005, never tested). No test, example or dispatcher configuration runs these pipelines with `k_batch > 1`. No other pipeline or kernel changes behaviour: `comp_tdm` and `comp_tdm_v2` MX still reject `k_batch > 1` with every epilogue. ## Test Plan On gfx1250: - `test_ck_tile_mx_gemm_tdm_v3` (new): MX FP8 on `comp_tdm_v3`, fp16 and bf16 C, the TDM epilogue (`k_batch` 1, the derived MX tolerance) and the CShuffle epilogue (`k_batch` 1 and 2, the split-K tolerance). - `test_ck_tile_mx_gemm_tdm_k_tail`: adds `comp_tdm_v3` rows with the TDM epilogue (fp16 and bf16 C, 128- and 256-deep K tiles) to the K-tail cases. - `test_ck_tile_mx_gemm_splitk_support`: `comp_tdm_v3` with CShuffle accepts `k_batch` 2; `comp_tdm` and `comp_tdm_v2` with CShuffle reject it.- Regressions: the MX TDM suites that share `MxGemmKernel` (`tdm_gates`, `pipeline_tdm_wmma`, `pipeline_tdm_wmma_smoke`, `tdm_bf16c`, `tdm_fp6`, `tdm_fp6_kloop_shapes`, `e8m0_gfx1250`, `host_helpers`) and the `comp_tdm_v3` gemm suites. - Example 55: `verify-sparse` and `verify-rand` over every tile and epilogue, `k_batch` 1, 2 and 4 with CShuffle, at M = 16 and several M = 512 and 4096 shapes. - Example 53: `verify-int` over every tile and epilogue, ragged M/N, and split-K 2 and 4. ## Test Result On gfx1250, at this PR's head merged with develop (with #15096) and with split-K limited to `comp_tdm_v3`: - `test_ck_tile_mx_gemm_tdm_v3` 20/20; `test_ck_tile_mx_gemm_tdm_k_tail` 60/60 (with the `comp_tdm_v3` rows and develop's FP6 rows); `test_ck_tile_mx_gemm_splitk_support` 7/7. - MX TDM suites: `tdm_gates` 35/35, `pipeline_tdm_wmma_smoke` 111/111, `tdm_bf16c` 12/12, `tdm_fp6` 34/34, `tdm_fp6_kloop_shapes` 14/14, `e8m0_gfx1250` 8/8, `host_helpers` 23/23, `grouped_gemm_mx_tdm` 8/8. - `comp_tdm_v3` gemm suites: fp8 66/66, bf16 84/84, bf16 large 18/18, unaligned M/N/K (including split-K) 115/115. - TDM suites that share `UniversalGemmKernel`: `grouped_gemm_tdm` 70/70, `batched_contraction_tdm` 80/80. At this PR's head before the develop merge (unchanged by it): - `pipeline_tdm_wmma` 147/148: the one failure, `TestCkTileMxGemmPipelineCompTDMWmma/14.SmallM` (`comp_tdm`, column-major fp4 A, M = 1, `IsSupportedArgument` rejects it), fails the same way on develop without this PR. `comp_tdm_v3` fp8 long K 16/16. - Example 55: 126 of 126 checks pass: every tile with the TDM epilogue and with CShuffle at `k_batch` 1, 2 and 4 at 16 x 49152 x 32768 (fp16 and bf16 C); TDM and CShuffle `k_batch` 2 at 512 x 7168 x 3072 and 4096 x 4096 x 8192 (bf16 C); random-data checks for every tile and epilogue. - Example 53: 90 of 90 checks pass, every tile and epilogue at 16 x 49152 x 32768, 512 x 7168 x 3072 and a ragged 144 x 112 x 4096, split-K 2 and 4, and random-data checks. ## Notes - Split-K adds 16-bit partial sums atomically, so the low bits vary from run to run. - Performance data and follow-ups are tracked in AICK-2991. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

…iles (#15199) JIRA ID : AICK-3064 ## Motivation Since #15005, the dispatcher CI stage fails on gfx942 and gfx950 in `test_preshuffle_gpu_correctness.py`. The preshuffle kernels are reported as unsupported (status -3, "N is not a multiple of vector load size for C tensor!"). #15005's code is correct. The failure comes from how the dispatcher compiles its kernels, which #15005 only exposed. ## Technical Details **Root cause.** The dispatcher's Python bridges JIT-compile each kernel with `--offload-arch`, but without `CK_CMAKE_GPU_TARGET_IDS`. Without that define, ck_tile's target detection (`arch.hpp`) resolves the host compilation pass to the host target, while the device pass sees the real GPU. So the host-side argument checks run against a different warp GEMM than the device code does: - On gfx942 and gfx950, the preshuffle kernel's `DefaultGemm2DEpilogue::GetVectorSizeC()` evaluates to 0 on the host. - `IsSupportedArgument` then computes `N % GetVectorSizeC()`, i.e. `N % 0`, which is undefined behaviour. - Before #15005, the compiler happened to fold this into "accepted". After #15005 moved the C-tensor checks into `IsCTensorSupported`, it folds into "rejected". - The device code is identical before and after #15005; only the host code changed. The CMake-built dispatcher libraries for the quant kernels already passed this define by hand, which is why they were not affected. **Fix.** - `dispatcher_common.gpu_target_id_define(arch)` returns `-DCK_CMAKE_GPU_TARGET_IDS=0x<id>`, spelled the way the dispatcher CMake spells it (gfx942 -> `0x942`). `arch_feature_defines()` includes it, so every bridge that already uses `arch_feature_defines()` compiles its host pass for the real target. - `bindings/ctypes/CMakeLists.txt` sets the define once in `set_single_target_arch_features()` for every single-target library, including the MX library, instead of per library. With the define, the host side computes `GetVectorSizeC()` = 1 on gfx942 and gfx950, and the 128x128x64 preshuffle case is accepted. ## Test Plan - `tests/test_dispatcher_common.py`: new `TestGpuTargetIdDefine` tests that the define matches the CMake spelling for every dispatcher arch, that the arch name is normalized, that unknown arches get no define, and that every dispatcher arch passes it to the bridges. The existing gfx942 and gfx950 `arch_feature_defines` expectations are updated. - Related dispatcher suites: `test_mx_gemm_cmake` (CMake/Python flag parity), `test_quant_review_regressions`, `test_blockscale_quant_arch_gates`, `test_streamk_gfx1250`, `test_batched_bridge`, `test_mx_gemm_bridge`, `test_codegen_common`. - `test_preshuffle_gpu_correctness.py --gfx gfx1250` on gfx1250 hardware, built with the same steps the CI stage uses. - gfx942 and gfx950: the host-side check for the failing kernel was compiled with the dispatcher's flags and evaluated. ## Test Result - Dispatcher unit suites above: 693 passed. - gfx1250 `test_preshuffle_gpu_correctness.py`: PASS. fp16 and bf16 ran 94 cases each, and the 8 partial-tile shapes that must be rejected were rejected. - gfx942 and gfx950: with the fix the host side computes `GetVectorSizeC()` = 1 and accepts the 128x128x64 case. Without the fix it computes 0. The CI stage will run the GPU test on both. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

jefyang1 and others added 6 commits November 3, 2026 17:15

jefyang1 requested review from a team as code owners November 3, 2026 21:12

github-actions Bot added documentation project: composablekernel labels Nov 3, 2026

assistant-librarian Bot added the organization: ROCm label Nov 3, 2026

illsilin approved these changes Nov 4, 2026

View reviewed changes

jefyang1 merged commit 00e31d2 into develop Nov 4, 2026
155 of 161 checks passed

jefyang1 deleted the users/jefyang/ck/aick-2966 branch November 4, 2026 16:45

ozturkosu mentioned this pull request Nov 5, 2026

feat(ck-tile): enable gfx1250 Stream-K on the comp_tdm_v3 pipeline #15118

Draft

This was referenced Nov 5, 2026

feat(ck-tile): comp_tdm_v3 MX FP8 GEMM on gfx1250 #15153

Merged

fix(ck-tile dispatcher): pass CK_CMAKE_GPU_TARGET_IDS to the JIT compiles #15199

Merged

Navigation Menu

Repository navigation

feat(ck-tile): comp_tdm_v3 bf16/fp16 GEMM on gfx1250 - #15005

feat(ck-tile): comp_tdm_v3 bf16/fp16 GEMM on gfx1250#15005
jefyang1 merged 6 commits into
developfrom
users/jefyang/ck/aick-2966

jefyang1 commented Nov 3, 2026

Uh oh!

therock-pr-bot Bot commented Nov 3, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Nov 3, 2026

Uh oh!

Uh oh!

assistant-librarian Bot commented Nov 6, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jefyang1 commented Nov 3, 2026

Motivation

Technical Details

Scope of comp_tdm_v3 after this PR

Default behaviour changes

Test Plan

Test Result

Notes

Submission Checklist

Uh oh!

therock-pr-bot Bot commented Nov 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Nov 3, 2026

Uh oh!

Uh oh!

assistant-librarian Bot commented Nov 6, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Scope of `comp_tdm_v3` after this PR

therock-pr-bot Bot commented Nov 3, 2026 •

edited

Loading