refactor(ck): mx gemm kernel unification

[EnricoDeg]

Motivation

CK tile currently has two separate MX GEMM kernels for gfx950 and gfx1250. This pull request refactors and modernizes the MX GEMM kernel and example to use new scale tensor handling, improved kernel argument structures, and updated pipeline and kernel APIs. The changes simplify the interface and improve type safety.

JIRA ID ROCM-26313

Technical Details

• Add support for gfx950 in MX GEMM kernel for gfx1250 and remove unused kernel
• Unify comp async pipeline for GEMM and MX GEMM
• Unify eight waves pipeline for GEMM and MX GEMM
• Move preshuffle MX GEMM pipeline to gemm ops and remove gemm_mx ops
• Unify testing framework for MX GEMM
• Add gfx950 tests for grouped MX GEMM

Test Plan

• test_mx_gemm_async.cpp for MX GEMM on gfx950
• test_mx_grouped_gemm_comp_async.cpp for grouped MX GEMM on gfx950

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[aosewski]

Padding test removal (N=160) may be masking a real PadN correctness bug

The HEAD commit "Remove padding tests to avoid numerical issues" narrows the MNPadding test in test/ck_tile/gemm_mx/test_mx_gemm_async_rcr.cpp from Ns{96,160,224} → Ns{96,224} while keeping Ms{96,160,224}.

That M/N asymmetry is the concerning part: N=160 against N_Tile=128 is one full block + a 32 remainder, which satisfies the kernel's own documented PadN precondition (include/ck_tile/ops/gemm/kernel/mx_gemm_kernel.hpp:124-129 — PadN is supported "when the GEMM has at least one full block along that dimension"). So N=160 is inside the supported envelope, yet it reproducibly mis-compares. Removing the test makes a supported PadN configuration ship silently wrong rather than fixing it — and any real caller with N not a multiple of N_Tile would hit it.

I'd ask for one of two resolutions before merge:

1. Preferred — root-cause and restore. Investigate the N-pad scale / CShuffle epilogue path for the partial-N-tile case and re-add N=160. (Worth ruling out the test-side tolerance issue too: max_accumulated_value is computed with *std::max_element(...), i.e. the max signed value, which under-sizes atol for all-negative result tiles — but that isn't N-specific, so it's unlikely to be the whole story here.)

2. Acceptable fallback — reject it explicitly. If partial-N-tile padding is genuinely unsupported for this pipeline, make it a deterministic, documented rejection in MxGemmKernel::IsSupportedArgument rather than a silently-untested gap, e.g.:

   // PadN with a partial trailing N-tile is not numerically correct on the
   // <pipeline> path (<short reason>); reject so callers get a clear skip
   // instead of silently wrong results.
   if(kPadN && (kargs.N % TilePartitioner::NPerBlock != 0))
   {
       if(log)
           CK_TILE_ERROR("MX GEMM: PadN with a partial trailing N tile is not supported.");
       return false;
   }

   That way the catching test can stay in the suite (it just becomes a clean skip), and downstream users can't trip the broken path.

This may also be related to the Jenkins Math CI failure on this branch — worth confirming that log isn't the same N-padding case.

[EnricoDeg]

Padding test removal (N=160) may be masking a real PadN correctness bug

The HEAD commit "Remove padding tests to avoid numerical issues" narrows the MNPadding test in test/ck_tile/gemm_mx/test_mx_gemm_async_rcr.cpp from Ns{96,160,224} → Ns{96,224} while keeping Ms{96,160,224}.

That M/N asymmetry is the concerning part: N=160 against N_Tile=128 is one full block + a 32 remainder, which satisfies the kernel's own documented PadN precondition (include/ck_tile/ops/gemm/kernel/mx_gemm_kernel.hpp:124-129 — PadN is supported "when the GEMM has at least one full block along that dimension"). So N=160 is inside the supported envelope, yet it reproducibly mis-compares. Removing the test makes a supported PadN configuration ship silently wrong rather than fixing it — and any real caller with N not a multiple of N_Tile would hit it.

I'd ask for one of two resolutions before merge:

1. **Preferred — root-cause and restore.** Investigate the N-pad scale / CShuffle epilogue path for the partial-N-tile case and re-add `N=160`. (Worth ruling out the test-side tolerance issue too: `max_accumulated_value` is computed with `*std::max_element(...)`, i.e. the max _signed_ value, which under-sizes `atol` for all-negative result tiles — but that isn't N-specific, so it's unlikely to be the whole story here.)

2. **Acceptable fallback — reject it explicitly.** If partial-N-tile padding is genuinely unsupported for this pipeline, make it a _deterministic, documented_ rejection in `MxGemmKernel::IsSupportedArgument` rather than a silently-untested gap, e.g.:
   ```c++
   // PadN with a partial trailing N-tile is not numerically correct on the
   // <pipeline> path (<short reason>); reject so callers get a clear skip
   // instead of silently wrong results.
   if(kPadN && (kargs.N % TilePartitioner::NPerBlock != 0))
   {
       if(log)
           CK_TILE_ERROR("MX GEMM: PadN with a partial trailing N tile is not supported.");
       return false;
   }
   ```
   
   
       
         
       
   
         
       
   
       
     
   That way the catching test can stay in the suite (it just becomes a clean skip), and downstream users can't trip the broken path.

This may also be related to the Jenkins Math CI failure on this branch — worth confirming that log isn't the same N-padding case.

I re-enabled all padding tests but they are skipped for 8 waves pipeline. I will investigate further in follow up issue if there are problems. The test coverage is already better than develop because there, only FP8 is tested with padding and only for CompAsync pipeline.

This feature should not be a blocker for Aiter integration so I suggest to move forward

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/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.

[therock-pr-bot]

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