Commit Graph

1646 Commits

Author SHA1 Message Date
Enrico Degregori
1b4fbd95fd [rocm-libraries] ROCm/rocm-libraries#6089 (commit c876d18)
[CK Tile] Extend type support EightWave pipeline

## Motivation

EightWave pipeline was designed for 8 bit types. This PR extend support
for any FP type

## Technical Details

 - Generalize policy to support any FP type
- Change LDS layout to fix bank conflicts. This removes all bank
conflicts in the pipeline (checked for all supported types). Remaining
bank conflicts are related to Cshuffle epilogue.

## Test Plan

Added GEMM tests with new supported types. Note that FP6 is also
supported for MX GEMM but the PR was reverted so no tests were added for
it.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-05 23:54:40 +00:00
Yung-sheng Tu
e826b2eb7e [rocm-libraries] ROCm/rocm-libraries#6768 (commit 43ca43f)
=?UTF-8?q?[CK=20TILE]=20Unification=20Work=20=E2=80=93=20?=
 =?UTF-8?q?Add=20MFMA=20specialisations=20for=20`tf32=5Ft`=20(#6768)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

This PR adds two specialisations related to `tf32_t`.

## Technical Details

This change treats `tf32_t` as a concrete type rather than an empty
`struct`. It also adds two new specialisations for MFMA dense builtins
and resolves existing circular include issues.

## Test Plan

All the new wrappers were added to the test suite in
test_amdgcn_mma_layout.inc.

## Test Result

Test should pass.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-05 12:27:41 +00:00
Sami Remes
ad4e2e7624 [rocm-libraries] ROCm/rocm-libraries#7199 (commit 23f7320)
[CK_TILE] [QuantGEMM] Fix SplitK tail handling and other
 improvements (#7199)

This pull request introduces improved and more robust split-K support
for quantized GEMM. The main changes add runtime validation, utility
functions for split-K batch calculations, pointer offset handling for
split-K in grouped kernels, and enhanced support for various tensor
layouts. The changes also improve error handling and provide more
flexibility for runtime tail handling in split-K pipelines.

**Split-K Support and Validation Enhancements:**

* Added runtime validation to ensure `k_batch` is a positive integer and
that split-K configurations do not produce empty final batches or
mismatched pipeline tails, with detailed error messages and logging for
misconfiguration.
[[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1184-R1211)
[[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1161-R1250)
* Introduced utility functions `get_splitk_batch_k_read` and
`get_splitk_last_batch_k` to compute per-batch K read sizes and handle
split rounding, ensuring correct and consistent split-K batch
partitioning.
[[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R206-R234)
[[2]](diffhunk://#diff-635b89bdffa96b2b42f1632520cde36701d7d631e864185591f6b32f7645cf47L104-R107)
[[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L388-R417)
[[4]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1161-R1250)
* Changed the default value of `k_batch` in `QuantGemmHostArgs` to 1 (no
split-K) for safer default behavior.

**Pointer Offsets and Grouped Kernel Handling:**

* Updated `QuantGroupedGemmKernel` to apply split-K per-batch offsets to
all input pointers, mirroring the behavior of non-grouped kernels and
ensuring correctness for split-K launches.
* Modified AQ tensor view handling to correctly reflect the remaining
K-groups from the split-K batch's offset position, improving accuracy
for split-K in grouped kernels.

**Pipeline and Layout Flexibility:**

* Added support for runtime selection of split-K tail handling via a new
template parameter `RuntimeSplitKTail_`, with new helper methods to
dispatch GEMM pipelines accordingly.
[[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R273)
[[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1496-R1567)
[[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1427)
[[4]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1447-R1629)
[[5]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L1459-R1641)
* Improved handling for tensor layout cases, including preshuffled B and
both row-major and column-major AQ layouts, ensuring correct pointer
arithmetic and compatibility checks.
[[1]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R438-R454)
[[2]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871L464-R516)
[[3]](diffhunk://#diff-d000149a681cd42bfb9947872c603e556cea26cbd7fd4f8f60afc6595d975871R1184-R1211)
2026-06-05 11:41:49 +00:00
Enrico Degregori
7b9245f18c [rocm-libraries] ROCm/rocm-libraries#5854 (commit 8e2d46d)
[CK Tile] Async support preshuffle GEMM

## Motivation

Add async support to existing preshuffle GEMM pipeline

## Technical Details

Notes:
the implementation avoids previous strategy of duplicating pipelines for
async support and instead add a switch `Async` to the ops Problem to
enable async pipeline. Then, integrate the async pipeline in the
existing one. This allows to avoid code duplication and facilitate the
integration of buffer load to lds in existing pipelines. In my opinion,
it should be used also for other pipelines which don't support buffer
load to lds yet and it would also be a good idea to refactor the
existing async GEMM pipelines with the same approach.

Summary:

 - integrate buffer load to lds in existing pipeline
- add optimal tensor descriptors for vmem loading and lds reading. They
are currently optimized for 16x16 wave tiles but they also work for
32x32 wave tiles. Optimizations for 32x32 wave tile requires different
lds layout and it will be done in a follow-up issue
 - Add async config to examples
 - Add test (gfx950 only)

## Test Plan

New test for gfx950 `test_ck_tile_gemm_pipeline_wp_async`

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-05 07:17:09 +00:00
Aviral Goel
267ca67001 [rocm-libraries] ROCm/rocm-libraries#8028 (commit c1cb112)
[CK_Tile] Add wmma_bf16f32_16x16x32_bf16 via
 fused-downconvert override (#8028)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Adds `__builtin_amdgcn_wmma_bf16f32_16x16x32_bf16` (fp32 accumulate →
bf16 output) to the CK Tile WMMA warp-gemm path. **API only** — the unit
test is split into a stacked PR (#8035) so this API change can be
reviewed in isolation.

## Changes (4 files)

- **16-bit trait:** `wmma_intrinsic_downconvert` (calls the bf16f32
builtin — fp32 C in, bf16 C out) plus `COutDataType = bf16_t` /
`COutVecType`.
- **`WarpGemmAttributeWmmaImpl` / `WarpGemmAttributeWmma`:**
`mac_downconvert(c_fp32, a, b)` (kTransC-aware) returning the bf16
C-output vector.
- **`WarpGemmImpl`:** `mac_downconvert` tail handler producing a bf16
C-output tile from the fp32 accumulator tile, reusing
`CWarpDstrEncoding` (output layout identical to the f32 C tile).

Verified on gfx1250 (via the stacked test PR #8035): the test passes;
the existing WMMA warp-gemm test is unaffected (additive change only).
2026-06-05 05:01:31 +00:00
Enrico Degregori
bdd7a8333d [rocm-libraries] ROCm/rocm-libraries#6672 (commit bda3f97)
[CK Tile] PermuteN support MX GEMM

## Motivation

Add PermuteN support to preshuffle MX GEMM

## Technical Details

 - Modify `shuffle_b_permuteN` to support MX preshuffled layout
- Add `preShuffleScalePermuteN` with same functionality of
`preShuffleScale` but layout consistent with PermuteN
 - Include MX pre-processing functions in the library

## Test Plan

Add test configuration for permuteN with preshuffle (both FP4 and FP8)

## Submission Checklist

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

Co-authored-by: Cong Ma <congma13@amd.com>
2026-06-05 03:04:43 +00:00
Illia Silin
aef7b42883 [rocm-libraries] ROCm/rocm-libraries#7816 (commit f6324af)
[CK] Fix latest build issues with staging compiler.

## Motivation

Fixing new warnings with staging compiler.

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-04 17:41:09 +00:00
John Afaganis
96c39b331e [rocm-libraries] ROCm/rocm-libraries#7829 (commit 13af7da)
[ck] Enforce ASCII-only C/C++ sources for hipRTC
 compatibility (#7829)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

CK source files must be compilable via **hipRTC (HIP runtime
compilation)**, whose preprocessor does not accept non-ASCII bytes
anywhere in a translation unit — **including in comments**. Bytes that
are harmless under `hipcc` (em-dashes, smart quotes, multiplication
signs, Greek letters, box-drawing glyphs, etc.) cause hipRTC to fail at
preprocessing time. These regularly leak in via LLM-assisted authoring
or copy/paste from formatted documents and silently break hipRTC paths
that are not exercised by the default `hipcc`-based build matrix.

This PR (a) cleans every existing violation (53 files) and (b) adds a
pre-checkin gate so new violations are rejected before merge.

## File extensions covered

Both the cleanup scan and the new Jenkins enforcement stage use the same
predicate:

```
*.h  *.hpp  *.cpp  *.h.in  *.hpp.in  *.cpp.in  *.inc  *.cl
```

(excluding `*/build/*` and `*/include/rapidjson/*`). This is a strict
superset of the existing `Clang Format` stage's predicate — `*.inc` is
added so test-fixture include files are also gated. The local pre-commit
hook's `c++/inc` type filter covers the same set.

## Why no enforcement today

CK is opted out of the rocm-libraries root `.pre-commit-config.yaml`, so
the existing `pre-commit` workflow doesn't touch CK. The local CK
`.pre-commit-config.yaml` only runs for developers who installed hooks.
The **authoritative gate is therefore the new Jenkins stage** in this
PR; the local hook is convenience.

## Commit layout (bisect-friendly)

1. `79798aa6261` — **`[ck] Convert reflect/ rendering to ASCII for
hipRTC compatibility`**
Behavior change, isolated. `TreeFormatter` swaps `├─ / └─ / │ ` for `|-
/ +- / | ` (3-col width preserved so alignment is unchanged).
`conv_description.hpp` swaps `×` for `x` as the dimension separator.
`test_conv_description.cpp` expected strings updated in lockstep so the
snapshot test stays green. This is the only commit in the series with
observable runtime impact.

2. `738fdb0d81c` — **`[ck] Strip non-ASCII bytes from C++ sources for
hipRTC compatibility`**
Mechanical text cleanup across 53 files. Replacements happen in comments
or in `std::cout` strings that are not asserted on by any test. None of
the 174 `.inc` files in the tree required edits, but they were in the
scan's predicate so the enforcement stage's predicate is a superset of
what was scanned. Full replacement table in the commit message.

3. `1d7cd8ba235` — **`[ck] Enforce ASCII-only C/C++ sources for hipRTC
compatibility`**
- New `projects/composablekernel/script/check_ascii_only.sh` (modeled on
`check_copyright_year.sh`).
- New entry in `projects/composablekernel/.pre-commit-config.yaml` under
the local-hooks block (`types_or: [c++, inc]`).
- New `ASCII Only Check` parallel stage in
`projects/composablekernel/Jenkinsfile`'s `Static checks` block,
mirroring the existing `Clang Format` stage but with `*.inc` added to
the find predicate. Always-on, no `RUN_CPPCHECK` gate.

The tree is buildable at every commit boundary. Commit 1 leaves 50 known
violations; commit 2 leaves 0; commit 3 wires the gate.

## Demo

Script output on a synthesized violation:

```
$ printf '// em-dash test \xe2\x80\x94 here\n' > /tmp/bad.cpp
$ projects/composablekernel/script/check_ascii_only.sh /tmp/bad.cpp
ERROR: /tmp/bad.cpp contains non-ASCII bytes:
1:// em-dash test — here
  Fix: replace with ASCII (em-dash -> --, smart quotes -> ", arrows -> ->, etc.)
$ echo $?
1
```

Full repo scan after the cleanup commits (note the `-name '*.inc'`
clause):

```
$ cd projects/composablekernel && find . -type f \( -name '*.h' -o -name '*.hpp' -o -name '*.cpp' \
    -o -name '*.h.in' -o -name '*.hpp.in' -o -name '*.cpp.in' -o -name '*.inc' -o -name '*.cl' \) \
    -not -path '*/build/*' -not -path '*/include/rapidjson/*' -print0 \
  | xargs -0 -P 8 -n 64 script/check_ascii_only.sh
$ echo $?
0
```

## Test plan

- [ ] Jenkins PR build: confirm new `Static checks -> ASCII Only Check`
stage runs green over the full predicate (incl. `*.inc`) and existing
`Clang Format` stage is unaffected.
- [ ] `test_conv_description` passes against the ASCII tree-formatter
output (touched in commit 1).
- [ ] Local: `pre-commit run ascii-only-checker --all-files` runs
cleanly after installing CK pre-commit hooks via
`script/install_precommit.sh`.
- [ ] Manually inject a non-ASCII byte in any `.cpp/.hpp/.inc` file,
push: confirm Jenkins fails the new stage with a clear error.
- [ ] Spot-check a representative subset of touched files under hipRTC
compilation to confirm no remaining hipRTC-blocking content (optional,
since the static byte check is a sufficient condition for hipRTC
preprocessor acceptance on this dimension).

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-06-04 15:00:17 +00:00
Copilot
4fcd73a98e [rocm-libraries] ROCm/rocm-libraries#7974 (commit 9df2c76)
composablekernel: remove stray *.hpp.bk backup artifacts
 (#7974)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Four `*.hpp.bk` files were accidentally committed to
`projects/composablekernel/`, likely as leftovers from a prior merge or
conflict resolution. Each is an older snapshot of its `.hpp` counterpart
— the canonical `.hpp` files are newer and contain the correct current
content.

## Deleted files

| File | vs. `.hpp` counterpart |
|---|---|
| `ck_tile/core/tensor/tile_window.hpp.bk` | Older version: uses legacy
`bool isL1Cache`/`PrefetchL1` template params; missing
`DataCachePrefetchKind`-based prefetch API and `data_cache_prefetch.hpp`
include |
| `ck_tile/core/tensor/load_tile_transpose.hpp.bk` | Older version:
missing `#if defined(__gfx950__)` guard and `Quad` struct (~90 lines)
for gfx1250 architecture |
| `ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp.bk` | Older version:
missing `WmmaTag`, `IsScale16` template param, and several newer
dispatcher specializations |
|
`ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp.bk`
| Older version: `KPackA`/`KPackB` (since renamed `KPack`); uses
`static_ford` (since refactored to nested `static_for`) |

## Verification

- No other `.bk` files exist in `projects/composablekernel/`.
- No build scripts, CMake files, includes, or documentation reference
these `.bk` files.
- No `.hpp` files were modified.
2026-06-04 03:06:43 +00:00
apophis
42c82b093e [rocm-libraries] ROCm/rocm-libraries#7786 (commit 7842dfd)
[CK TILE][Windows] add `msvc::no_unique_address` support for
 Windows (#7786)

## Motivation

While building Flash Attention 2 with CK backend, this warning will spam
in every kernel:
```
DEBUG [1/1837] hipcc.exe ...
DEBUG In file included from H:\ROCm\flash-attention\build\fmha_fwd_d32_bf16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_alibi_mask_lse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cu:6:
DEBUG In file included from H:\ROCm\flash-attention\csrc\composable_kernel\example\ck_tile\01_fmha\fmha_fwd.hpp:6:
DEBUG In file included from H:\ROCm\flash-attention\csrc\composable_kernel\include\ck_tile/core.hpp:111:
DEBUG H:\ROCm\flash-attention\csrc\composable_kernel\include\ck_tile/core/tensor/tile_scatter_gather.hpp:1246:7: warning: unknown attribute 'no_unique_address' ignored [-Wunknown-attributes]
DEBUG  1246 |     [[no_unique_address]] std::conditional_t<kUseGlobalLoad_, PageIdxArray, gl_field_empty_t>
DEBUG       |       ^~~~~~~~~~~~~~~~~
DEBUG H:\ROCm\flash-attention\csrc\composable_kernel\include\ck_tile/core/tensor/tile_scatter_gather.hpp:1254:7: warning: unknown attribute 'no_unique_address' ignored [-Wunknown-attributes]
DEBUG  1254 |     [[no_unique_address]] std::conditional_t<kUseGlobalLoad_, index_t, gl_field_empty_t>
DEBUG       |       ^~~~~~~~~~~~~~~~~
DEBUG 2 warnings generated when compiling for host.
...
```

## Technical Details

`[[no_unique_address]]` is not working on Windows LLVM, should use
`[[msvc::no_unique_address]]`.

## Test Plan

Build FA2 with CK backend.

## Test Result

No warnings, no errors.

## Submission Checklist

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

Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
2026-06-04 02:28:12 +00:00
Aviral Goel
e01603bc31 [rocm-libraries] ROCm/rocm-libraries#7725 (commit eef7e12)
[GFX1250][CK_TILE] Add scale16 warp gemm unit tests

## Summary
- Add scale16 WMMA intrinsic overloads and int64_t forwarding to warp
gemm layers for gfx1250
- Add comprehensive wave-level unit tests for scale16 warp gemm
(16x16x128 and 32x32x128 tile sizes)
- Test all fp8/bf8 type combinations and TransposeC variants
- Fix WarpGemm wrapper for non-uniform scale16 configurations

Stacked on #7724 (FillUniformScaleDistribution / MX GEMM scale init).
Pipeline enablement follows in the next PR.
2026-06-03 22:05:29 +00:00
chris-tsiaousis-hpc
db05d61136 [rocm-libraries] ROCm/rocm-libraries#6212 (commit ccee58d)
=?UTF-8?q?[CK=20TILE]=20Unification=20Work=20=E2=80=93=20?=
 =?UTF-8?q?More=20accurate=20tests=20for=20MmaPipelines=20(#6212)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

This PR solves several issues:

#### More accurate tests for MmaPipelines

The current tests for the MmaPipelines (test_amdgcn_sparse_mma,
test_amdgcn_wavewise_mma) use explicit input fragment vectors filled
with 1s, and only check the output of a single lane. We should have
tests that actually use the MmaPipelines with non-trivial input matrices
and verify the complete output.
Some other aspects of the current MmaPipelines tests that I noticed and
deserve some attention:

1. There is sometimes iteration over K outside of the pipeline, which is
then included in WaveTileK or FragK, which is not correct. We should
remove it, move K iteration inside of the pipeline, or be more clear
about this outer-K loop size and how it propagates downwards.
2. There is very tight coupling between the kernel, gtest code, and
test_pipeline helper, requiring a lot of information and functions to be
passed back and forth.
3. The test_pipeline helper is doing a bunch of register-related logic
on the host (related to point 1)
4. Without this register logic the only thing it does is check the
device, call the kernel, and check the output, but with a lot of
boilerplate.

#### Test helper for detecting target arch at HOST runtime

There is a really apparent issue we faced while writing tests:

Scenario:
1. Compile a test that supports both gfx950 and gfx1201 for gfx950
2. Run the test on a server that only has gfx1201 GPU

Actual:
Segmentation fault

Expected:
The test can correctly detect from HOST runtime that the DEVICE
target_id was different and skips the test.

Notes:

The only way of detecting the COMPILER_TARGET_ID in the existing "arch"
framework is launching a kernel and calling `get_compiler_target()` (so,
from a DEVICE code). This will create a segmentation fault if the
current arch differs from the target arch. To cope with this issue, we
propose to export the compiler target(s) (note they can be many) through
`projects/composablekernel/test/ck_tile/core/arch/CMakeLists.txt` and
define a test helper to deal with such cases.

#### Add composition support to Transforms

We have a small number of Transforms which act on MmaOp input and output
data, before and after the MmaOp call respectively. These are currently
implemented to work on an MmaTile level, but in theory they are also
supposed to work at a WaveTile level, i.e. after composition of multiple
MmaTiles to create larger effective MNK dimensions. Currently the
composed MmaTiles look like 2D C-style arrays of the individual MmaTile
level register vectors (see WaveWiseMmaPipeline). The transforms should
be able to take these and perform the proper transforms to the whole
WaveTile at once. This might allow for better performing
transformations.

Note: This PR handles the SparseTransform case and if we don't end up
doing scale as a transformation, there isn't really much left to do. If
we end up having only the sparse transform as a non-trivial transform,
then we could also consider removing the Transform framework.
2026-06-03 14:35:18 +00:00
Ville Pietilä
88f8d24c34 [rocm-libraries] ROCm/rocm-libraries#7936 (commit 3dc91e6)
[CK Tile] Fix V6 pipeline applicability and split-image
 initialization (#7936)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

After adding code generation via CK Tile Dispatcher, some fwd and bwd
weight tests for CK Tile convolutions are failing. This PR introduced
correct applicability checks and fixes the split-image parameter
initialization such that non-applicable instances are not invoked during
test execution and split-image instances are correctly initialized.

## Technical Details

Investigation revealed two distinct problems

1. For bwd weight, the compute V3 uses prefetch of 3 distinct tiles,
which works incorrectly when the number of K-slices addressed by the
workgroup is 1. This occurs when a large split-K value is used for a
problem that results in a small Gemm-K value.
2. For fwd direction, the current CK Profiler/test infrastructure
doesn't initialize the split-image parameters for instance where
split-image is enable. Uninitialized split-image values result in
non-deterministic behavior where the tests might randomly fail.

Fixed problem 1. by adding a check in `IsSupportedArgument` that marks
the instance invalid if the `num_loops = ceil(GemmK / (k_batch *
KPerBlock)) < 4` for V6 pipeline kernel instances. The check is
compile-time eliminated for other kernels.

Fixed problem 2. by adding initialization of split-image parameters when
split-image is enabled. The default initialization corresponds to full
image with no split, i.e., the number of splits is 1 and it has the size
of the full image.

Added unit tests for the added logic.

## Test Plan

Running the following test suites cover the logic added in this PR
- test_grouped_convnd_fwd_tile
- test_ck_tile_grouped_conv_fwd
- test_grouped_convnd_bwd_weight_tile
- test_ck_tile_grouped_conv_bwd_weight

All test suites above are included in the automated test runs.

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-03 08:40:03 +00:00
Anton Gorenko
7ecbf82708 [rocm-libraries] ROCm/rocm-libraries#7500 (commit f5cd4fd)
[CK_TILE][FMHA] Optimize long-context decoding on gfx11/12
 (#7500)

## Motivation

Relevant issue: ROCM-22065

FMHA has less-than-optimal performance of long-context decoding (i.e.
when seqlen_q = 1) on gfx11/12.
This PR optimizes the splitkv pipeline and configs for such scenarios.

## Technical Details

Optimizations applied in this PR:
1. use tiles with smaller M0 (16 vs 64), these tiles are used when
seqlen_q <= 16
2. adapt qr_nwarp_sshuffle pipeline for gfx11, it allows to use more
warps even for M0 = 16 (the qr pipeline parallelizes work between warps
in M dim so with M0 = 16 it allows to use only 1 warp)
3. enable kMergeNumHeadGroupsSeqLenQ (an optimization that merges one
group of heads in GQA) for all hdim values, not only 128
4. increase the number of splits (multiply by the number of head groups)
if (3) is used
5. increase the number of splits for RDNAs (`multiProcessorCount` is the
number of WGPs on RDNAs, not CUs, so it should be doubled to have
meaning similar to CDNAs)

Performance on gfx1151:

| Case | develop (GB/s) | This PR (GB/s) |
|:-------|-------:|-------:|
| [fp16\|group\|bshd] b:1, h:32/32, s:1/45056, d:64/64 | 127.58 | 183.11
|
| [fp16\|group\|bhsd] b:1, h:32/32, s:1/45056, d:64/64 | 153.64 | 215.02
|
| [fp16\|group\|bshd] b:1, h:16/8, s:1/77184, d:128/128 | 120.51 |
225.76 |
| [fp16\|group\|bhsd] b:1, h:16/8, s:1/77184, d:128/128 | 130.62 |
223.84 |
| [fp16\|group\|bshd] b:1, h:32/32, s:1/9600, d:128/128 | 82.65 | 138.44
|
| [fp16\|group\|bhsd] b:1, h:32/32, s:1/9600, d:128/128 | 105.75 |
220.45 |
| [fp16\|group\|bshd] b:1, h:8/1, s:1/401024, d:256/256 | 16.27 | 187.89
|
| [fp16\|group\|bhsd] b:1, h:8/1, s:1/401024, d:256/256 | 16.28 | 188.19
|

## Test Plan

An additional test case is added to the exiting test. It uses seqlen_q =
1, GQA, no mask to trigger the changes
```
ninja test_ck_tile_fmha_fwd_fp16 && bin/test_ck_tile_fmha_fwd_fp16 --gtest_filter="*SplitKV*
ninja test_ck_tile_fmha_fwd_bf16 && bin/test_ck_tile_fmha_fwd_bf16 --gtest_filter="*SplitKV*
```

Manual testing can be done with these commands:
```
bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=32 -h_k=32 -d=64  -s=1 -s_k=$((352 * 128))  -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1
bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=16 -h_k=8  -d=128 -s=1 -s_k=$((603 * 128))  -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1
bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=32 -h_k=32 -d=128 -s=1 -s_k=$((75 * 128))   -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1
bin/tile_example_fmha_fwd -prec=fp16 -mode=1 -page_block_size=128 -b=1 -h=8  -h_k=1  -d=256 -s=1 -s_k=$((3133 * 128)) -lse=1 -mask=0 -num_splits=0 -kname=1 -v=1
```

## Test Result

All the tests must pass.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-03 06:16:10 +00:00
Yi DING
01bd52bdb5 [rocm-libraries] ROCm/rocm-libraries#7925 (commit a8f0845)
[CK] Fix gfx950 AITER Sync Regressions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Fixes three gfx950 regressions in the AITER downstream CI that surfaced
after the internal/gfx1250 re-sync (ROCm/rocm-libraries#6978):

> **Companion aiter PR:** ROCm/aiter#3392 — host-side adaptations
(`Kernel::BlockSize()` `constexpr` drops, blockscale `KBatch=1` clamp)
plus the CK submodule bump used to validate these fixes together.

- **FlyDSL MoE AOT cache miss** — the AITER MoE tests run with
`check_aot_cache=True` and fail on any FlyDSL JIT cache miss, but the CI
never pre-compiles the FlyDSL MoE kernels, so gfx950 always misses.
Pre-compile them at the start of the AITER test stage.
- **`buffer.load.lds.v4i32` link error** — ROCm/rocm-libraries#6978
reintroduced a clang-version guard mapping
`llvm.amdgcn.raw.buffer.load.lds` to a `.v4i32`-suffixed name. That name
exists in no LLVM (the rsrc operand is a fixed, non-overloaded `<4 x
i32>`, so the intrinsic is never type-mangled), so gfx950 4-DWORD
direct-to-LDS (e.g. fp4 MoE bpreshuffle) fails to link with `lld:
undefined symbol: llvm.amdgcn.raw.buffer.load.lds.v4i32`. Use the
canonical plain name unconditionally.
- **mixed-precision flatmm warp-GEMM call** — ROCm/rocm-libraries#6978
generalized the scaled `WarpGemmImpl::operator()` from a fixed `<index_t
opselA, index_t opselB>` signature to a variadic `<typename... Params>`
one and updated the `mx_flatmm` pipeline to pass the op-selectors as
`OpSelA<>`/`OpSelB<>` types, but missed the mixed-precision flatmm
pipeline (`F8xMXF4`/`F16xMXF4`), which still passed raw integer
op-selectors. These no longer bind to `typename... Params` (`error: no
matching member function for call to 'operator()'`), breaking
compilation of the fp8/bf16 × fp4 cktile MoE gemm1 instances on gfx950
(aiter `test_moe_2stage`). Wrap the op-selectors in
`OpSelA<>`/`OpSelB<>`.

## Changes

- `Jenkinsfile`: pre-compile the FlyDSL MoE AOT cache (`python3
aiter/aot/flydsl/moe.py`) before the AITER tests.
- `include/ck/utility/amd_buffer_addressing_builtins.hpp` and
`include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp`: drop the
`__clang_major__` guard and always use
`__asm("llvm.amdgcn.raw.buffer.load.lds")`. The plain name is the
canonical one for all sizes including the gfx950 16-byte form, as the
upstream LLVM gfx950 tests confirm.
-
`include/ck_tile/ops/flatmm/pipeline/mixed_prec_flatmm_pipeline_agmem_bgmem_creg_v1.hpp`:
wrap the warp-GEMM op-selectors in `OpSelA<>`/`OpSelB<>` at the five
call sites, matching the `mx_flatmm` pipeline.

## Test plan

Validated via CI.
2026-06-03 02:09:05 +00:00
Illia Silin
5720589311 [rocm-libraries] ROCm/rocm-libraries#7960 (commit ddac5cf)
[CK] Upgrade to new gfx1250 compiler and fix build issues
 (#7960)

## Motivation

The docker image we've been using to build for gfx1250 is a few months
old, so we need to upgrade. Some of the changes in the latest compiler
version require changes in the code. TDM is temporarily disabled due to
changes in the lds load/store intrinsics.

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-03 01:58:59 +00:00
Maksim (Max) Podkorytov
d574cc4757 [rocm-libraries] ROCm/rocm-libraries#6696 (commit 9627b91)
Replace nested static_for lambdas with compile-time search
 helper (#6696)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

- Add `sequence_find_value` and `find_in_tuple_of_sequences`
compile-time search helpers with O(1) template depth
- Replace nested `static_for` lambdas in
`TensorDescriptor::GetTransformAndItsUpperDimension` and
`InitializeElementSize`
- Apply same optimizations to `TensorAdaptor`

Supersedes #4287. Conflict-resolved rebase of
ROCm/composable_kernel#3600 onto current develop.

## Motivation

The `TensorDescriptor` and `TensorAdaptor` classes had excessive
template instantiation from:
1. Nested `static_for` loops with lambdas creating unique closure types
at every call site
2. `generate_tuple` with lambdas causing per-type instantiation overhead

The new helpers use constexpr array lookup and pack expansion instead of
recursive template patterns, achieving O(1) template depth.

## Results (`example_grouped_conv_fwd_xdl_fp16`, n=10, interleaved,
`-j1`, `-ftime-trace`)

| TU | Baseline (mean) | New (mean) | Delta | Wilcoxon p | Mann-Whitney
p |

|----|-----------------|------------|-------|-----------|---------------|
| `grouped_conv_fwd_xdl_fp16` (host) | 14,886 ms | 13,353 ms |
**-10.3%** | **0.002** | **0.0002** |
| `grouped_conv_fwd_xdl_fp16` (device) | 27,762 ms | 25,629 ms |
**-7.7%** | **0.002** | **0.0002** |
| **Total (all TUs)** | **57,732 ms** | **54,030 ms** | **-6.4%** | | |

Unrelated TUs (`device_memory`, `host_tensor`, `convolution_parameter`)
show no significant difference (p > 0.3), serving as negative controls.

### Methodology

- 10 interleaved runs (baseline₁, new₁, baseline₂, new₂, ...) on the
same node to eliminate ordering/warmup bias
- Wilcoxon signed-rank test (paired, non-parametric) and Mann-Whitney U
test (unpaired)
- Built with patched clang (LLVM 22) on ctr2-alola-compile-11, `-j1` for
accurate per-TU timing
- Raw data available in Slurm job 275230 results

## Test plan

- [x] 11 unit tests added (5 for `sequence_find_value`, 6 for
`find_in_tuple_of_sequences`)
- [x] Compile-time benchmark with statistical significance (p < 0.01)
- [ ] Full CI

Tracking issue: #4229
2026-06-02 23:15:10 +00:00
Aviral Goel
99ab4c4ef7 [rocm-libraries] ROCm/rocm-libraries#7830 (commit 590fe58)
[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32)
 (#7830)

Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250,
enabling bf16-input bf16-output WMMA at the warp GEMM level.

- Add WmmaTraits specialization for <gfx125_t, bf16, bf16, bf16,
16,16,32>
- Add WarpGemmAttributeWmmaImpl typedef and WarpGemmWmma alias
- Add Dispatcher entry for bf16->bf16 16x16x32
- Add warp_gemm test with reference GEMM validation

## Motivation

<!-- Explain the purpose of this PR and the goals it aims to achieve.
-->

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-02 13:54:16 +00:00
Sami Remes
919096fde8 [rocm-libraries] ROCm/rocm-libraries#7935 (commit 5c96097)
[CK] Allow skipping split-K C-buffer zero-init in
 xdl_cshuffle blockscale GEMM (#7935)

Add a `skip_zero_init` flag (default false) to the Problem/Argument of
the xdl_cshuffle block-scale GEMM device ops (multiple_d ab_scale and
blockscale b-preshuffle). When the flag is set, the device invoker skips
the internal hipMemsetAsync that zeroes p_c_grid before the KBatch > 1
split-K atomic-accumulation path. The flag is declared on the gridwise
Problem struct (inherited by Argument), so it is visible on both the
rotating-cache (arg_) and the normal (arg) launch paths in each device
op.

Why: callers that already pre-zero the output buffer otherwise pay for a
redundant device-wide memset before split-K atomic accumulation. Gating
the memset behind an opt-in flag lets such callers avoid the duplicate
work. Because the flag defaults to false, every existing call site is
unaffected and the observable behavior is unchanged.

## Motivation

<!-- Explain the purpose of this PR and the goals it aims to achieve.
-->

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

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

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-02 13:08:46 +00:00
Johannes Graner
b7c8fb164f [rocm-libraries] ROCm/rocm-libraries#7937 (commit abe276d)
[CK Tile] Add conv Wavelet GEMM pipeline and bwd_weight
 instances (#7937)

## Motivation

CK Tile had no pipeline competitive with old CK's wavelet on the
RetinaNet K=36 C=256 3x3 conv bwd_weight class. This adds a
wave-specialized "wavelet" GEMM pipeline so CK Tile has a competitive
kernel for spatial small-K shapes.

## Technical Details

- New wavelet GEMM pipeline (`gemm_pipeline_ag_bg_cr_wavelet.hpp`):
workgroup split into math waves (LDS read + MFMA) and load waves (DRAM
read + LDS write).
- VGPR role-split: `operator()` has two top-level mutually-exclusive
`is_math` branches so the allocator overlays both roles onto the same
physical VGPRs, cutting arch VGPR ~33-40% and raising occupancy.
Correctness depends on identical `block_sync_lds` counts on both arms
plus a matching load-wave barrier stub in the epilogue
(`cshuffle_epilogue.hpp`).
- Kernel dispatch (`grouped_convolution_backward_weight_kernel.hpp`):
`kIsWavelet` path, `LaunchBlockSize`, load-wave barrier stub.

Uplift: wavelet is the fastest CK Tile pipeline on the RetinaNet K=36
C=256 3x3 family, beating the best non-wavelet CK Tile kernel by 10-27%
(googlenet K=320 by 16-23%); the role-split roughly halves the parity
gap vs old CK on the 13x13 fp16 shape.

## Test Plan

- `ckProfiler grouped_conv_bwd_weight`, NHWGC layout, fp16/bf16,
`split_k=all`, CPU verify on RetinaNet K=36 shapes (7x7, 13x13) and a
broad 2D sweep.
- Correctness: `-v=1` across `split_k` in {-1,1,2,4,8,16,32,64}
(barrier-parity / deadlock check).
- `test_grouped_convnd_bwd_weight` over the tests `.conf` wavelet
instances.

## Test Result

- All wavelet instances CPU-verify correct across the split-K sweep; no
hangs (dual-arm barrier sequence matches).
- Wavelet wins the RetinaNet K=36 C=256 3x3 family (10-27% over best
non-wavelet CK Tile) and googlenet K=320 (16-23%); at parity-or-better
vs old CK on the majority of spatial shapes.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-06-02 08:51:17 +00:00
Chao
c56c6750d0 [rocm-libraries] ROCm/rocm-libraries#6498 (commit 5961a2e)
[CK_TILE] Fix conditional rescale numerical instability in
 FMHA forward (#6498)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

[CK_TILE] Fix conditional rescale numerical instability in FMHA forward

## Motivation

Fix numerical instability in the conditional O-accumulator rescaling
optimization
for CK-Tile FMHA forward (FlashAttention-4, Algorithm 6, Eq. 6).

The conditional rescale optimization skips the expensive O-accumulator
rescale when
the running row-max shift is within a threshold (tau = log2(256) = 8.0).
The original
implementation had a bug: attention weights P were computed in the
`m_new` reference
frame before the skip/rescale decision. In the skip branch, `m` was
reverted to
`m_old`, but P remained in the `m_new` frame, causing incorrect softmax
normalization.

This fix introduces a `p_row_correction` factor: in the skip branch, P
is multiplied
by `exp2(m_new - m_old)` to bring it back to the `m_old` reference
frame.

- **Correctness:** Fixes broken inference on long sequences where
running-max drift
causes exp2 overflow (observed as degraded image quality on MI350X Flux2
generation)
- **Performance:** Neutral to +4% depending on workload shape

## Technical Details

6 pipeline header files (same pattern in each):
- `block_fmha_pipeline_qr_ks_vs.hpp`
- `block_fmha_pipeline_qr_ks_vs_async.hpp`
- `block_fmha_pipeline_qr_ks_vs_async_trload.hpp`
- `block_fmha_pipeline_qr_ks_vs_fp8.hpp`
- `block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp`
- `block_fmha_pipeline_qs_ks_vs.hpp`

In each file:
- Lower threshold from 10.0 to 8.0 (tau = log2(256))
- Add `p_row_correction` distributed tensor initialized to 1.0
- Rescale branch: standard rescale of O_acc and l; correction = 1.0
- Skip branch: compute correction = exp2(-acc_scale_log2), update l,
revert m, store correction
- New `p_spans` sweep applies per-row correction to `p_compute` before
P*V GEMM
- Move P-to-PDataType cast to after correction sweep

## Dependencies

None — this PR is standalone.

## Test Plan

- GPU validation on MI300X (gfx942, ROCm 6.4.1):
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=bf16 -v=1 -warmup=1 -repeat=3`
- GPU validation on MI350X (gfx950, ROCm 7.0):
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=bf16 -v=1 -warmup=1 -repeat=3`
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=fp16 -v=1 -warmup=1 -repeat=3`

## Test Result

Accuracy vs FP32 reference (MI350X, gfx950):

| Shape | max_diff | mean_diff |
|-------|----------|-----------|
| B=1 H=24 M=4096 K=128 bf16 | 9.1e-4 | 4.6e-5 |
| B=4 H=32 M=4096 K=128 bf16 | 9.9e-4 | 4.6e-5 |
| B=1 H=24 M=4096 K=128 fp16 | 1.2e-4 | 9.0e-6 |

Performance (MI350X, gfx950, ROCm 7.0):

| Shape | FA4 (TFlops) | Always-rescale (TFlops) | Delta |
|-------|-------------|------------------------|-------|
| B=1 H=24 M=4096 K=128 bf16 | 425.9 | 428.5 | neutral |
| B=2 H=8 M=2048 K=256 bf16 | 513.9 | 509.0 | +1.0% |
| B=1 H=64 M=2048 K=64 bf16 | 481.7 | 464.3 | +3.7% |

Benchmark results (MI300X, gfx942, ROCm 6.4.1):

No regression on MI300X. This correctness fix is performance-neutral.

| Config | TFlops / GB/s | Time (ms) |
|--------|-------------|-----------|
| MHA bf16 b=2 h=8 s=4096 d=128 | 342.49 TFlops | 0.401 |
| MHA fp16 b=2 h=8 s=4096 d=128 | 391.70 TFlops | 0.351 |
| Causal MHA bf16 b=2 h=8 s=4096 d=128 | 227.07 TFlops | 0.303 |
| GQA 4:1 bf16 b=2 h=32 hk=8 s=2048 d=128 | 324.69 TFlops | 0.423 |
| GQA 8:1 bf16 b=2 h=64 hk=8 s=2048 d=128 | 348.09 TFlops | 0.790 |
| LLaMA-70B prefill b=1 h=64 hk=8 s=4096 d=128 bf16 | 376.71 TFlops |
1.459 |
| Long-seq bf16 b=1 h=16 s=16384 d=128 | 383.42 TFlops | 5.735 |
| Decode b=64 h=32 hk=8 s_k=4096 d=128 bf16 | 691.64 GB/s | 1.554 |

All validation tests pass (`valid:y`) on both MI300X and MI350X.

Additional validation:
- Uniform scores: softmax output matches FP32 reference (max_diff <
1e-3)
- Large seqlen (4096+): no overflow or NaN in O-accumulator
- Spike pattern: correct handling of sudden row-max jumps
- Multiple spikes: correction applied correctly across multiple
skip/rescale transitions
- Deterministic: identical outputs across repeated runs
- No performance regression on standard workloads
2026-05-30 10:34:06 +00:00
Tianyuan Wu
22a99f97e8 [rocm-libraries] ROCm/rocm-libraries#7677 (commit 308af93)
[CK_Tile] Add scale16 Support for F4 WMMA in CK_Tile

## Motivation
This PR adds CK Tile support for the scale16 F4 WMMA path on gfx1250 and
improves warp GEMM unit test coverage/structure for gfx1250-specific
cases.

## Technical Details

- Scale16 support in warp GEMM dispatch and WMMA trait plumbing: added
IsScale16 plumbing to warp GEMM dispatcher path
- Warp GEMM test restructuring for gfx1250: added Warp GEMM gfx1250
coverage to verify all F4 WMMA paths

## Test Plan
Run ./test_ck_tile_wg_32x16x128_fp4.

## Test Result
```
./test_ck_tile_wg_32x16x128_fp4
[----------] Global test environment tear-down
[==========] 3 tests from 1 test suite ran. (1751 ms total)
[  PASSED  ] 3 tests.
```

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-30 01:28:48 +00:00
Hosang Yoon
e7e8801dc3 [rocm-libraries] ROCm/rocm-libraries#7586 (commit c18f2c7)
[CK_TILE] Use gfx11 float buffer atomics in FMHA Bwd

## Motivation

FlashAttention CK backward on gfx11 can hit out-of-bounds/tail writes in
the dQ accumulator atomic-add path when sequence rows are padded at the
tile level but not marked invalid in the DQDKDV main tensor view.

With the generic global atomic fallback, an incorrectly-valid tail
element can issue an actual pointer-based `atomicAdd`. With the buffer
atomic path, the write is issued through a buffer resource with bounds
information and follows the same backend already used by gfx9/gfx12.

This fixes the gfx11 FMHA BWD failure without changing the gfx11 default
for unrelated CK Tile kernels.

## Technical Details

This PR enables the existing CK Tile AMD buffer float atomic-add path
only for generated FMHA BWD gfx11 translation units.

gfx11 normally uses the generic global atomic fallback for
floating-point `buffer_view::atomic_add`. That fallback performs the
atomic through a raw computed pointer and depends on the software
validity predicate to avoid invalid elements. In FMHA BWD dQ
accumulation, padded tail rows can reach this path, so using the buffer
atomic backend is safer: it uses a buffer resource with base pointer,
bounds information, and an element offset, matching the backend already
used by gfx9/gfx12.

Enabling `CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT` globally for gfx11 is
too broad and can break unrelated gfx11 CK builds such as GEMM. Instead,
`config.hpp` now preserves an explicitly pre-defined
`CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT`, while keeping the existing
default disabled for gfx11.

## Test Plan

Validated the change with the FlashAttention CK full test suite with
backward pass enabled on gfx11.
pytest -q -s tests/test_flash_attn_ck.py

## Test Result

FlashAttention CK gfx11 test result:
260680 passed, 152076 skipped

## Submission Checklist

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

Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
2026-05-30 00:10:26 +00:00
Emily Martins
95c916369c [rocm-libraries] ROCm/rocm-libraries#7584 (commit 060bad5)
[CK_TILE] Fix Stream-K k_size calculation

## Motivation

In a recent benchmarking task for CK Tile Stream-K algorithm, we
identified that certain instances segfault. This change works to fix the
bug and adds necessary regression tests.

## Technical Details

The StreamK kernel constructs tensor views using a `k_size` parameter
that determines how much of the K dimension to process in each
iteration. Previously, this was calculated as:
 ```cpp
index_t k_size = num_loop_sk * TilePartitioner::KPerBlock;
```
This calculation assumes all macro tiles along K are exactly `KPerBlock` in size. However, when `K % KPerBlock != 0`, the final macro tile along K has a remainder size of `K % KPerBlock`, not a full `KPerBlock` (see the figure below):
<img width="961" height="488" alt="image" src="https://github.com/user-attachments/assets/3e1cceed-5dcd-4980-8b02-cee24eecf262" />
With the old code, a workgroup working with the `MPerBlock x (K % KPerBlock)` tile in A and B risk accessing illegal memory.

Hence, this change ensures that when `K % KPerBlock != 0`, workgroups processing iterations that include the final macro-tile along K calculate the correct `k_size` based on the remainder rather than assuming a full `KPerBlock`.

## Test Plan
I added the following tests:
1. Unit tests added for the Stream-K Tile Partitioner:
- `StreamKTilePartitionerBaseGetKSize/NoRemainderTiles` - validates full tiles
- `StreamKTilePartitionerBaseGetKSize/RemainderTiles` - validates remainder handling
2. Regression tests that test a case where `K % KPerBlock != 0`

## Test Result

Tests passed locally on gfx90a, gfx942, and gfx950.

## Submission Checklist

- [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-29 21:36:49 +00:00
Aviral Goel
15c904b460 [rocm-libraries] ROCm/rocm-libraries#7724 (commit 4cb149a)
ck_tile: add FillUniformScaleDistribution and fix MX GEMM
 scale init (#7724)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

### Problem
MX GEMM pipeline tests were passing vacuously: scale bytes were drawn
from a fixed range (40–60) which, for e8m0, maps to scales ≈ 10⁻²⁷ — far
below FP16 min denorm. Both GPU and CPU produced all-zero outputs, so
numerical checks passed without exercising the GEMM.

### Changes

**`include/ck_tile/host/fill.hpp`** — new
`FillUniformScaleDistribution<ScaleType>` functor
- Accepts human-readable float bounds and maps them to the raw byte
range of any ExMy scale type (e8m0, e4m3, e5m3) by re-centering the IEEE
754 exponent into the type's bias space
- Sampling is uniform over raw bytes → uniform over representable values
- Fixes left-shift UB: uses multiplication instead of `<< mant_bits` to
avoid shifting negative signed integers (C++17 UB)
- Adds `assert(min_r <= max_r)` to catch inverted-range UB when both
bounds exceed the type's representable range
- Provides default member values (0.125f, 2.0f) and `std::optional` seed
consistent with sibling fillers
- `/** */` Doxygen style with `@note` on snapping asymmetry

**`test/ck_tile/gemm_mx/test_mx_gemm_pipeline_util.hpp`** — fix scale
initialization
- Replace manual byte-range distribution with
`FillUniformScaleDistribution<>{0.125f, 2.0f}`
- Use distinct seeds for scale_a (11941) and scale_b (11943) to avoid
correlated scale tensors that were causing 60 test failures for
fp4+e5m3/e4m3 combinations

**`test/ck_tile/utility/test_fill.cpp`** — new unit tests for
`FillUniformScaleDistribution`
- 16 typed tests across e8m0, e4m3, e5m3: validity, range,
reproducibility, coverage, snapping, stress, nullopt seed, and range
overload
- Test helper `expected_raw_range` mirrors implementation clamping
exactly
2026-05-29 18:45:13 +00:00
Andriy Roshchenko
d5c9215064 [rocm-libraries] ROCm/rocm-libraries#7359 (commit dd62f9f)
[CK_TILE][GFX1250] Enable MX GEMM FLATMM with ASYNC

## Motivation

Enables MX GEMM FLATMM pipeline on gfx1250. The pipeline uses an async
load instruction for tensor A, which complements the existing MX GEMM
FLATMM pipeline with TDM load. At this time, only FLATMM MX pipelines
are enabled on gfx1250.

## Technical Details

The existing gfx950 implementation was extended to support gfx1250
architecture. All three MX FP data types are supported across the two
ASICs.
It should be noted that while the TDM pipeline uses an emulated
32x32x128 warp-tile instruction, the present submission relies on the
built-in 16x16x128 instruction, called 4 times per warp.

## Test Plan

Existing `test/ck_tile/flatmm` tests were extended to cover new gfx1250
functionality.

To help facilitate the testing in development,
`example/ck_tile/18_flatmm/script/smoke_test_mx.sh` script was
introduced to verify various combinations of supported data types and
pipeline versions.

## Test Result

The present submission is expected to work on both gfx950 and gfx1250
hardware for all reasonable sizes and all MX FP8/FP6/FP4 data types.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
- [x] Relies on #6978 and should only be merged after the changes are
merged to the `develop`.
2026-05-29 17:02:45 +00:00
Bartłomiej Kocot
5d912538d3 [rocm-libraries] ROCm/rocm-libraries#7847 (commit b995ef2)
[CK] Remove IsPackedTensor function

## Motivation

Fix codegen hipRTC

## Technical Details

Remove not needed function. Since MakeArgument supports long_index_t
strides.

## Test Plan

Codegen tests.

## Test Result

Passed.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-29 14:00:06 +00:00
Ville Pietilä
78d657c4f7 [rocm-libraries] ROCm/rocm-libraries#7284 (commit e7d25b2)
[CK_TILE] Integrate CK Tile Dispatcher code generation into
 CK Tile Profiler (#7284)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

CK Tile is going to be delivered to hipDNN via CK Dispatcher. Currently
the CK Tile Profiler using CK Builder for generating the profiled
instances from the configuration files that identify the instances that
old CK exposes. We need to replace this instance generation with the CK
Tile Dispatcher codegen.

## Technical Details
The old CK Profiler config files are converted to JSON files that the CK
Tile Dispatcher can digest. The conversion script for configurations is
stored to source control in case we need to update the JSON
configurations later. The dispatcher generates instance libraries per
conv direction (fwd, bwd data, and bwd weight) that are linked to the CK
Profiler executable. I also implemented codegne for the stream-K and
depthwise conv instances. The proposed solution replaces the CK Builder
codegen with the CK Tile Dispatcher codegen.

There are two new methods that are exposed via the dispatcher backend

- `is_supported` - required to enabled the profiler workflow where we
check the applicability of the kernel instance before running it.
- `get_instance_string` - this mainly for verification. This provide the
CK Builder instance string for verifying that the old CK Builder based
profiler and the new CK Tile Dispatcher based profiler have the same
instances.

The rules that limit the generated instances are now collected to a
single location under the dispacther. The CK Builder codegen uses these,
which ensures that the two codegen pipelines are in sync. The next step
(different PR) is to remove the CK Builder codegen pipeline altogether.

## Test Plan

Verified that the old CK Builder based profiler and the new CK Tile
Dispatcher based profiler have the same instances, that is, the
Dispatcher based codgen can generate the same instances as the old CK
Builder.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-28 21:03:37 +00:00
ltqin
bf07a0150e [rocm-libraries] ROCm/rocm-libraries#7723 (commit 4ed6c51)
[CK Tile] Enable LSE output for fp8bf16 V3 FMHA kernels
 (#7723)

###  Motivation
The V3 pipeline (qr_async_trload_v3) for fp8bf16 FMHA kernels did not
support LSE (Log-Sum-Exp) output. This PR enables LSE output support for
fp8bf16 V3 FMHA kernels, allowing users to retrieve attention statistics
alongside attention outputs.
### Technical Details
    - StandardAttention: lse = softmax_scale * m + log(l)
- LogitsSoftCap: lse = (m / log2(e)) + log(l)

### Test Plan
Run FMHA forward example with fp8bf16 precision and LSE output enabled:
- Test 1: Basic LSE functionality
./build/bin/tile_example_fmha_fwd -v=1 -b=1 -h=8 -s=1024 -d=128
-prec=fp8bf16 -init=3 -qscale=1 -lse=1
- Test 2: LSE with LogitsSoftCap (CMakeList should remove Logits filter)
./build/bin/tile_example_fmha_fwd -v=1 -b=1 -h=8 -s=1024 -d=128
-prec=fp8bf16 -init=3 -qscale=1 -lse=1 -logits_soft_cap=30.0
2026-05-28 15:58:54 +00:00
Zoltán Lakatos
58e2ab1fc7 [rocm-libraries] ROCm/rocm-libraries#6761 (commit d19f6f1)
[CK] Large tensor gemm workaround (#6761)

## Motivation

Customer qeruested large tensor gemm support for 8bit and 4bit data
types. Currently CK triggers “This GEMM not supported” error. The root
cause appears to be the 2 GB limit on the input/output matrix, triggered
by buffer offset constraints when testing a larger shape such as M =
699,904 (which is an exact multiple of MPerBlock = 256).

## Technical Details

Quick workaround to have support ASAP. Split the tensors into inputs /
outputs smaller than 2GB limit. Iterate on host and call all subproblems
without device code change.
Support is restricted to rowise layout in A, Ds and E

All changes were implemented in DeviceGemm structures to avoid secondory
affect on grouped convolutions.

Got lots of AI generated comments. Addressed the ones that seemed
relevant on the functionality.

## Test Plan

Within CK the following examples can be used with modified input sizes:
example_gemm_multiply_multiply_xdl_fp8
example_gemm_mx_fp4
Tested with Aiter tuning on provided shapes.

## Test Result

All gemms run and provide correct results.

## Submission Checklist

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

---------

Co-authored-by: Zoltán Lakatos <zoltan.lakatos@streamhpc.com>
Co-authored-by: Márton Bidlek <marton.bidlek@streamhpc.com>
Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>
2026-05-27 18:55:15 +00:00
Michael Halkenhäuser
5dc8fbd1a8 [rocm-libraries] ROCm/rocm-libraries#6900 (commit 28608c2)
[CK] Fix and expand CK's commit records in version.h (#6900)

## Motivation

In `version.h` of a CK installation `CK_COMMIT_ID` would be empty for
out-of-source builds.
Additionally, if it worked, it would show the parent repo's
(`rocm-libraries`) commit.

## Technical Details

Dropped "required" constraint so "unknown" string becomes a graceful
option.

Changed process of determining the CK commit, now uses
`WORKING_DIRECTORY`.
Thus, `CK_COMMIT_ID` holds only the last CK-relevant commit.

Added `CK_PARENT_COMMIT_ID` which holds the parent's, e.g.
`rocm-libraries`, commit.
This can be the same as `CK_COMMIT_ID`, or not even applicable,
depending on the scenario.

## Test Plan

Ran CMake configuration and installation of CK to verify happy path.

## Test Result

Commit SHA's showed the expected values depending on the repo state.

## Submission Checklist

- [ x ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-27 10:17:02 -07:00
Aviral Goel
4aecc8de5b [rocm-libraries] ROCm/rocm-libraries#7442 (commit b7d57ef)
[CK] CompV4: remove redundant barrier (+5.7% gfx942, +1% gfx950) (#7442)

## Summary

- Remove one redundant `block_sync_lds()` from the pong phase of the
CompV4 GEMM pipeline hot loop
- The pong phase had 2 barriers while ping had 1 — the second pong
barrier (after LDS writes, before global loads) was unnecessary because
the sync at the top of the next ping iteration already ensures LDS
coherence
- Removing this barrier allows global loads to overlap with LDS write
drain, restoring the latency hiding the ping-pong design was built to
provide
- Abstracting away Ping Pong phases into generic lambda avoids making
such mistake again.

## Benchmark

### gfx942 (MI300X), 86 fp16 GEMM shapes

| Metric | Value |
|---|---|
| Improved (>1%) | **80** |
| Neutral (±1%) | **4** |
| Regressed | **2** |
| Average gain | **+5.7%** |
| Best gain | +18.0% (4096x256x16384) |
| Worst regression | -2.9% (12288x3072x4096) |

### gfx950 (MI355X), 86 fp16 GEMM shapes

| Metric | Value |
|---|---|
| Improved (>1%) | **32** |
| Neutral (±1%) | **54** |
| Regressed | **0** |
| Best gain | +9.0% (4096x2048x28672) |

### Top gains by workload

| Shape (MxNxK) | Source | gfx942 BL | gfx942 Opt | gfx942 Gain | gfx950
BL | gfx950 Opt | gfx950 Gain |
|---|---|---|---|---|---|---|---|
| 4096x256x16384 | bloom_fc2 | 38.3 | 45.2 | **+18.0%** | 75.6 | 77.0 |
+1.9% |
| 4096x512x22016 | llama2_7b | 77.8 | 90.8 | **+16.7%** | 152.4 | 154.9
| +1.7% |
| 256x1536x7168 | deepseek | 14.4 | 16.7 | **+16.0%** | 27.2 | 28.0 |
+2.8% |
| 4096x1024x22016 | llama2_7b | 156.2 | 180.8 | **+15.7%** | 304.8 |
311.6 | +2.2% |
| 4096x1024x16384 | bloom_fc2 | 154.6 | 178.5 | **+15.4%** | 303.1 |
309.5 | +2.1% |
| 4096x4096x22016 | llama2_7b | 371.0 | 412.3 | **+11.1%** | 819.8 |
823.6 | +0.5% |
| 4096x2048x28672 | llama3_8b | 235.5 | 259.5 | **+10.2%** | 530.0 |
577.7 | **+9.0%** |
| 250880x256x4096 | bloom_logits | 289.0 | 335.9 | **+16.2%** | 595.5 |
599.1 | +0.6% |
| 8192x8192x8192 | square | 411.8 | 432.9 | **+5.1%** | 825.1 | 825.8 |
+0.1% |
| 7168x4096x8192 | llama70b | 362.9 | 374.7 | **+3.3%** | 775.8 | 782.5
| +0.9% |

## Hardware counter analysis (rocprof-compute, 8192x8192x8192, gfx942)

| Metric | Baseline | Optimized | Delta |
|---|---|---|---|
| s_barrier per ping+pong | 5 | 4 | **-1** |
| MFMA Utilization | 47.8% | 55.5% | **+7.7pp** |
| IPC | 0.17 | 0.21 | **+23.5%** |
| MFMA F16 % of peak | 30.6% | 33.5% | **+2.8pp** |
| VALU (instructions) | 41.67M | 41.67M | identical |
| MFMA (instructions) | 65.91M | 65.91M | identical |
| Spill/Stack Read | 8.27M | 8.27M | identical |

All instruction counts are identical — the optimization removed one
synchronization point, not any compute instructions.

## Correctness

- gfx942: GPU verification (`-v=2`) passed on 4 shapes (8192x8192x8192,
4096x4096x4096, 22016x4096x4096, 4096x512x28672)
- gfx950: GPU verification (`-v=2`) passed on all 86 shapes
2026-05-27 12:23:43 -04:00
Illia Silin
c24e528481 [rocm-libraries] ROCm/rocm-libraries#7760 (commit a61bc76)
[CK] suppress compiler warnings while building pytorch. (#7760)

## Motivation

Recently added compiler flags that are required to suppress false
warnings by latest staging compiler are not recognized by older compiler
versions and are triggering an avalanche of warnings. Previous attempt
to suppress them by using -Wno-unknown-warning-option flag didn't help,
because that flag wasn't recognized either and just added more warnings.
I've verified that current approach by checking the clang version
actually works as intended and makes the warnings go away.

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-27 06:56:58 -07:00
Bartłomiej Kocot
60df79085d [rocm-libraries] ROCm/rocm-libraries#7631 (commit d591a7c)
[CK] Grouped Convolution Global Load/Store support (#7631)

## Motivation

Grouped Convolution Global Load/Store support to cover large tensor
cases.

## Technical Details

Utilize global load for grouped convolution forwad kernels. Update
Indexes to use int64.

## Test Plan

- test utils
- test conv kernels in next pr with instances

## Test Result

CI pending

## Submission Checklist

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

AICK-1255

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
2026-05-27 08:21:54 +00:00
JH-Leon-KIM-AMD
00e1d82ae7 [rocm-libraries] ROCm/rocm-libraries#7732 (commit b0e29d9)
[CK] Fix grouped conv bwd data stride>1 silent miscompute (ALMIOPEN-1959) (#7732)

## Motivation

Fix silent miscompute in the grouped convolution backward-data kernel
(`DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1`) when stride >
dilation (ALMIOPEN-1959). PR #6208 introduced a flat-descriptor fast
path that dropped all but the first sub-GEMM, producing zeroed slices of
`dx` on
the (G=1, stride>1, 2D, NumDTensor=0) intersection. Restore correctness
without giving up the perf gains PR #6208 delivered on stride=1 shapes.

## Technical Details

- Tighten the flat-descriptor fast-path gate to require
`arg.gemms_count_ == 1` (i.e. a single sub-GEMM per dispatch — its
original purpose). For stride > 1, the implicit GEMM is split into
`gemms_count_` sub-GEMMs whose output cells tile `dx` disjointly;
routing them through the flat path required dropping all but the first,
which was the source of the bug.
- Stride > 1 now falls through to the existing grouped CShuffle path,
which packs all sub-GEMMs into one descriptor array and walks them
on-device in a single kernel launch. This is the pre-PR-6208 production
path; correctness is established and per-dispatch launch count is
minimised.
- Add regression coverage for the (G=1, stride>1, 2D, NumDTensor=0)
intersection in
`test/grouped_convnd_bwd_data/test_grouped_convnd_bwd_data.cpp` with
`gemms_count` ∈ {4, 9, 36}. Pre-existing cases did not hit this
intersection (all stride>1 cases used G=2; all G=1 cases used stride=1),
which is why PR #6208's regression slipped past CI.

## Test Plan

- `ctest -L SMOKE_TEST -R 'grouped_convnd_bwd_data'` on gfx942 (smoke
tier — runs on every PR via `smart_build_and_test.sh`).
- End-to-end verify (`verify=1`) via
`example_grouped_conv_bwd_data_xdl_fp16` on stride 1/2/3/6 shapes
including the original ALMIOPEN-1959 case and a cross-bucket
(`gemms_count=36`) case spanning two `MaxGroupedGemmGroupsNum=32`
buckets.
- ckProfiler A/B sweep on MI300X (gfx942) toggling the flat-path gate
via an environment variable: full kernel-family enumeration, winning
kernel + its avg_time reported under each gate. 33/41 shapes completed
before the sweep was stopped; the remaining 8 were the largest
i2v/synthetic shapes where ckProfiler exceeded its 300s per-shape
enumeration budget (not relevant to the verdict).

## Test Result

### Correctness

| Test | Result |
|---|:---:|
| `test_grouped_convnd_bwd_data` (12 type parameterizations × Test2D,
includes 3 new regression shapes) | **12/12 PASSED** in 14.18 s |
| `test_grouped_convnd_bwd_data_interface` (API checks) | **PASSED** in
0.28 s |
| ALMIOPEN-1959 stride=2 (`verify=1`) | **PASSED** |
| stride=1 K3 (`verify=1`) | **PASSED** |
| stride=3 K3 `gemms_count=9` (`verify=1`) | **PASSED** |
| stride=6 K6 `gemms_count=36` cross-bucket (`verify=1`) | **PASSED** |

### Performance (ckProfiler A/B on gfx942 / MI300X)

Comparing the **post-fix gate** (flat path only when `gemms_count_==1`,
column "B") vs the **inner-loop variant** that keeps the flat path on
stride>1 (column "A") across 25 stride>1 shapes where production picks
a `_v1` instance (so the gate actually fires):

| Stride | Shapes | A wins | Tie | B wins | Notes |
|:------:|:------:|:------:|:---:|:------:|---|
| 1 (sanity, gate moot) | 3 | 0 | 3 | 0 | gate doesn't differentiate — A
== B as expected |
| > 1 (gate fires) | 25 | **0** | 11 | **14** | B wins +6% to +32%; A
never wins |

Highlights from the firing-gate cases:

| Shape (G=1, stride=2 unless noted) | A ms | B ms | B vs A |
|---|---:|---:|---:|
| ALMIOPEN-1959 (N=16, K=256, C=128, 5×5, 40×175) | 0.183 | 0.171 | **B
+6%** |
| Retinanet-L61 (N=32, K=C=256, 3×3, 25×25) | 0.054 | 0.045 | **B +17%**
|
| i2v-010 (N=1, K=C=384, 3×3, 277×209) | 0.174 | 0.125 | **B +28%** |
| Synthetic 50×50 K3 N=32 K=C=256 | 0.131 | 0.088 | **B +32%** |

Why B wins everywhere the gate fires: for `gemms_count = N`, the flat
path needs N kernel launches (one per sub-GEMM), while the grouped path
loops over the same N sub-GEMMs on-device in 1 launch. The (N−1) ×
launch-tax is a structural disadvantage A can't recover from.

### Diff

| File | Lines |
|---|---:|
|
`include/.../device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp`
| +14 / −8 (one extra condition + expanded dispatch comment) |
| `test/.../test_grouped_convnd_bwd_data.cpp` | +9 / −0 (3 new shapes) |

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-27 09:59:14 +03:00
assistant-librarian[bot]
6181eb2adf [rocm-libraries] ROCm/rocm-libraries#4279 (commit 5b3f4b7)
[CK_TILE] Stream-K XCD remapping (#4279)

## Proposed changes

This PR adds support for XCD remapping as detailed in this
[document](https://amdcloud.sharepoint.com/:w:/r/sites/ComposableKernels/Shared%20Documents/Stream-K/Design%20Docs/XCD%20Mapping.docx?d=w2df1b0737dc54614970d99a2e26022d1&csf=1&web=1&e=mLVN4A).
On gfx942, workgroups are typically scheduled round-robin across XCDs,
which can lead to poor locality. We will use a remapping to assign
workgroups to contiguous tiles in the XCDs improving the locality and
the cache hit rate. This is done through a function that computes this
contiguous mapping from this
[PR](https://github.com/ROCm/composable_kernel/pull/3161), which we have
added to the StreamKTilePartitioner. This will require minimal changes
to the Stream-K algorithm, only requiring a remap at the time the
workgroups are partitioned. Through this approach we can improve the
data locality by improving cache hits therefore closing performance gaps
that are seen with the default scheduling. There have been unit tests
added to verify the function in isolation. This is an optimization that
is not specialized to just Stream-K GEMM and can be applied across GEMM.

Note: This only applies to the gfx942 as they introduce the XCDs.

Please put an `x` into the boxes that apply. You can also fill these out
after creating the PR. If you're not sure, please don't hesitate to ask.

- [x] I have added tests relevant to the introduced functionality, and
the unit tests are passing locally
- [ ] I have added the test to REGRESSION_TESTS list defined at the top
of CMakeLists.txt in tests/CMakeLists.txt, **IF** the test takes more
than 30 seconds to run.
- [x] I have added inline documentation which enables the maintainers
with understanding the motivation
- [ ] I have removed the stale documentation which is no longer relevant
after this pull request
- [ ] (If this change is user-facing) I have added release notes which
provide the end users with a brief summary of the improvement from this
pull request
- [x] I have run `clang-format` on all changed files
- [x] Any dependent changes have been merged

---
🔁 Imported from
[ROCm/composable_kernel#3652](https://github.com/ROCm/composable_kernel/pull/3652)
🧑‍💻 Originally authored by @arai713

---------

Co-authored-by: Astha <astha.rai713@gmail.com>
Co-authored-by: systems-assistant[bot] <systems-assistant[bot]@users.noreply.github.com>
Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>
Co-authored-by: arai713 <67439843+arai713@users.noreply.github.com>
2026-05-26 09:43:03 -07:00
Aviral Goel
0df3523ef1 [rocm-libraries] ROCm/rocm-libraries#6807 (commit ddda8ac)
[CK_TILE] Add save_matrix_txt() and extract HostTensor I/O to free functions (#6807)

## Summary
- Extract `loadtxt`, `savetxt`, and `save_matrix_txt` from `HostTensor`
member functions into standalone free functions in `host_tensor_io.hpp`
(Single Responsibility Principle)
- Add `save_matrix_txt()` for writing 2D tensors to space-separated text
files with configurable output limit (default 256x256, pass 0 to dump
all)
- Supports float, int, and int8_t output formats via a `dtype` parameter
- Validate dtype early and throw on unsupported values in all three
functions
- Update callers in `15_fused_moe/main.cpp` to use free function syntax
2026-05-26 11:07:18 -04:00
Anton Gorenko
66d6714376 [rocm-libraries] ROCm/rocm-libraries#5388 (commit 45583bd)
[CK_TILE][FMHA] Improve precision of mxfp4 FMHA with fp6 for matrix P (#5388)

## Motivation

Improve precision of mxfp4 without performance penalties.

## Technical Details

Since performance of scale MFMAs is the same when neither A nor B is
fp8/bf8, it is possible to use fp6 x fp4 instead of fp4 x fp4 for the
second GEMM, while types of Q, K, V stay the same.
This allows to improve overall precision significantly because fp6 has
32 non-negative values used for P quantization compared to just 8 values
for fp4.

It was found that there is a compiler bug with
`__builtin_amdgcn_cvt_scalef32_2xpk16_fp6_f32` (described in
LCOMPILER-561) but a workaround seems to fix all failing instances.

## Test Plan

```
ninja test_ck_tile_fmha_fwd_mxfp4 && bin/test_ck_tile_fmha_fwd_mxfp4
```

## Test Result

The tests must pass.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-26 06:55:17 -07:00
Yung-sheng Tu
760f9e1d0a [rocm-libraries] ROCm/rocm-libraries#7104 (commit 0fab8d8)
[CK TILE] Unification Work – Add MFMA specialisations for `fp64_t` (#7104)

## Motivation

This PR adds two specialisations related to `fp64_t`.

## Technical Details

This adds two new specialisations for MFMA dense builtins, and adjusts
ABLayout and CLayout to L{K1BM} and L{M1BN}.

## Test Plan

All the new wrappers were added to the test suite in
test_amdgcn_mma_layout.inc.

## Test Result

Test should pass.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-26 10:49:36 +00:00
Yi DING
6a9c03f692 [rocm-libraries] ROCm/rocm-libraries#7450 (commit 402dbad)
[CK_TILE] Use Persistent Scheduling for FMHA BWD Group Deterministic (#7450)

## Motivation

FMHA BWD group-mode deterministic currently uses a non-persistent
scheduler: each `(batch, head, K-row)` work-item is launched as its own
block, with no work-stealing across CUs. On uneven workloads (varlen,
GQA, many heads with
few K-rows) this leaves CUs idle and forces a larger dq_acc workspace
than necessary.

This PR ports the persistent + deterministic scheduling already used in
batch mode to group mode: a fixed-grid kernel that pre-computes per-CU
work ranges on the host and uses sparse dq_acc slot indexing so multiple
K-rows handled
by the same CU share one accumulator slot via intra-CU atomic adds.

Stacked on #7331; merge that first.

## Technical Details

Single file changed: `ops/fmha/kernel/fmha_bwd_kernel.hpp`.

A new `kUsePersistent` path is added to the group-mode deterministic
kernel, mirroring the batch-mode persistent scheduler. The host
pre-computes a fixed per-CU partition of the total `(batch, head,
K-row)` work and packs it into
`cu_states[]` so the GPU consumes it in a single launch. Host
preparation happens in four steps:

1. Build per-batch `seqstart` prefix sums.
2. Fill per-batch `(sq_w, nc)` with a placeholder `nsplits` (bumped in
step 3).
3. Two-pointer scan over CUs to fill `cu_states[c]` (`isplit`,
`head_start`, `c_start`, `w_lo`, `w_hi`), accumulating `nsplits[b]` as
`max(cs->isplit + 1)`.
4. Compute compact per-batch dq_acc offsets from the finalized
`nsplits`.

`isplit` is the sparse dq_acc slot index — one CU's multi-K-row writes
share slot `ceil(wc_start / denom)`, enabling intra-CU atomic
accumulation instead of one slot per K-row.

`denom = max(sq_w, target_w)`, splitting two regimes:

- `target_w >= sq_w` (large work): `denom = target_w`, intra-CU atomic
optimization engaged.
- `target_w < sq_w` (sub-K-row sharding, multiple CUs sharing one
K-row): `denom = sq_w` collapses to per-K-row indexing (`= c_start`),
keeping `isplit ∈ [0, nc-1]` and matching the `nsplits_max =
ceil(s_k/kN0) = nc` upper bound that #7331's
`GetWorkspaceDeviceSizeUpperBound` assumes for group+det.

`isplit` is additionally clamped to `nc-1` to absorb empty CUs
(rounded-up `wc_start` past the last K-row); they don't write dq_acc on
GPU so the slot value is harmless.

`nsplits[b]` is accumulated dynamically in step 3 rather than via a
closed form so it tightly matches the actual sparse slots used; step 4
(offsets) follows step 3 since offsets now depend on the dynamic
`nsplits`.

Group mode also allows batches with `seqlen_q == 0`. The persistent
scheduler skips them on the dQ path (no work) but dK/dV are still
zero-filled.

## Test Plan

Built `tile_example_fmha_bwd` with receipt 5 (fp16, no-bias, no-dropout,
`dpad == dvpad`, group + batch) on gfx950 (MI355X).

- 8-case smoke (shapes that exercise the sub-K-row regime).
- 44-case sweep covering: mask 0/1/2, GQA, var seqlen, `d != d_v`,
extreme
  small seqlen / `nc=1`, CU >> work, huge batch, batch-mode regression.
- 12-case perf comparison vs the non-persistent baseline (warmup=10,
  repeat=50).

## Test Result

- All 8 + 44 cases `valid:y`.
- Perf: ±5% noise, average -0.4% across the 12 cases — neutral.
- Batch-mode deterministic / non-deterministic regression unchanged.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-26 10:01:54 +08:00
Michal Kulikowski
8de4cb72fb [rocm-libraries] direct push (commit 49b73ad)
[CK][CK_TILE] POC for Instruction Cache prefetch.

Signed-off-by: Michal Kulikowski <Michal.Kulikowski@amd.com>
2026-05-25 11:26:26 +02:00
chris-tsiaousis-hpc
c7fac341de [rocm-libraries] ROCm/rocm-libraries#4871 (commit 7d4c040)
[CK] Decouple EpilogueArgs from GridwiseGemm implementation (#4871)

This is duplicate of #4537. I could not re-open it since te target
branch got deleted and could not change the target branch since it was
closed... :)

## Motivation

Right now, all the Epilogues structs are declared inside the base
gridwise struct. They should be independent of it and the specialization
of the selected Epilogue Type should be declared within the the kernel
function.

## Technical Details

All Epilogue structs depend on template parameters that are known to the
base Gridwise Gemm struct. In this PR, we export them to be used
independently by any struct that might need to extract them. This
approach will serve the decoupling purposes for the Epilogues, but also
enable future constructs to use and expand this approach. See
30e2a4c01b64bdea68857c7badd9d7cffbf1adb9.

Right now an issue that arises is that when implementing a new Epilogue
Type, the developer is not forced to decide where this struct should/can
be used or not. To fix this I propose defining an `enum struct
EpilogueType` that will be used to fetch the Epilogue specialization
through a helper struct. See a943ac8d130e12d6843715b322181186e54ba15c.
Note that all the instantiation details will stay in this helper struct.
Also note the static assertion in the else statement.

## Test Plan

Test with existing CI, as nothing is added/removed.

## Test Result

All relevant existing CI tests should pass.

## Submission Checklist

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

---------

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>
2026-05-22 18:39:01 +00:00
JP-Fernando
74bc86240b [rocm-libraries] ROCm/rocm-libraries#5647 (commit 490437a)
[CK Tile] Add gemm universal preshuffle to MX GEMM  (#5647)

## Motivation

Add gemm universal preshuffle support to existing MX GEMM pipeline.

The straightforward way to do this is to port the `mx_flatmm` pipeline
to the existing `gemm_mx` framework.

## Technical Details

The `mx_flatmm` pipeline was not deleted, to allow for
back-compatibility.

## Test Plan

Add `preshuffle` option to example: `tile_example_mx_gemm`.

Add new configurations with enabled preshuffle to the existing
`test/ck_tile/gemm_mx` tests.

## Test Result

Example and tests were successful on `gf950` architecture in the `Alola`
cluster.

## Submission Checklist

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

---------

Co-authored-by: Fernando Jiménez <fernando.jimenez@streamhpc.com>
2026-05-22 16:07:53 +02:00
Bartłomiej Kocot
ebb97044f4 [rocm-libraries] ROCm/rocm-libraries#7664 (commit de5d6b1)
Revert "[CK] Enable grouped conv bwd data to match non-grouped perf" (#7664)

## Motivation

Incorrect results has been introduced for some conv bwd cases.

## Technical Details

This reverts commit 33424f65346d6330d0fd94b5a4e6f843f24e52c3.

## Test Plan

CI

## Test Result

Pending

## Submission Checklist

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

ALMIOPEN-1959
2026-05-22 12:28:49 +00:00
Wojciech Laskowski
3ea9ce7e37 [rocm-libraries] ROCm/rocm-libraries#6567 (commit 753c7a8)
[CK Tile] Adding WMMA wrappers for sparse builtins (#6567)

## Motivation

This PR is part of the [WMMA/MFMA] unification work. It's the third of
the series of PRs (after
https://github.com/ROCm/rocm-libraries/pull/5801 and
https://github.com/ROCm/rocm-libraries/pull/6014) that add all the
necessary MMA builtins as amdgcn_mma structs. This PR focuses on sparse
WMMA intrinsics.

## Technical Details

This change adds new specializations for WMMA sparse builtins. In total,
we add 8 WMMA builtins.

## Test Plan

All the new wrappers were added to the test suite in
`test_amdgcn_mma_layout.inc`.

## Test Result

Test pass locally, waiting for the CI.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-22 13:34:33 +02:00
Illia Silin
e02c566795 [rocm-libraries] ROCm/rocm-libraries#7612 (commit 5427d24)
[CK] upgrade CI to rocm7.13 as default compiler (#7612)

## Motivation

Upgrade the default docker and compiler version in CI to rocm7.13.
In order to pass all the checks I had to also clean up a lot of
non-ascii characters in the source code comments and modify a couple of
tests that were affected by a new compiler logic.

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

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

---------

Co-authored-by: Aviral Goel <aviral.goel@amd.com>
2026-05-22 02:43:50 +00:00
kensclin
fc2862d712 [rocm-libraries] ROCm/rocm-libraries#6846 (commit 377def4)
[CK_TILE] Add fmha forward hdim 256 support (#6846)

## Motivation

Enable Composable Kernel FMHA forward kernel for **hdim=256 BF16** on
AMD gfx950 (MI350X). Prior to this change the (256, 256) head-dim
configuration either failed to compile, was filtered out by the
compatibility rules, or produced incorrect kernel output due to an LDS
layout accounting bug.

## Technical Details

  Four files changed, all to enable hdim=256 BF16 on gfx950.

- **`fmha_fwd.py`** — Allow `(256, 256)` in gfx950 compatibility rule;
set `(256,256)` BF16 tile to `M0=128, N0=64` (the LDS-feasible shape on
gfx950); emit minimal valid instance set for d=256 to bound compile
time.

- **`fmha_fwd_kernel.hpp`** — Gate Prefill launch path off for d=256
(`PrefillCase = kM0 > 64 && kQKHeaddim < 256`); the double-buffer
Prefill variant overflows the 160 KB LDS budget.

- **`trload_policy.hpp`** — **Critical correctness fix**: the LDS layout
accounting in `GetSmemSize` was wrong (`max(Q, K+S+V)` instead of
`max(Q, K) + V + S`), under-allocating LDS and silently corrupting d=256
output (~2% wrong values).

- **`trload.hpp`** — Thread `LoadOnce=true` through all d=256 K-LDS
descriptors so the compiler picks the matching XOR swizzle period;
recompute the S-tile LDS offset to match the corrected `GetSmemSize`
formula.

## Test Plan

Built and ran `tile_example_fmha_fwd` on gfx950 (MI350X) with the
canonical d=256 BF16 configurations:

  ```bash
  cd build && ninja tile_example_fmha_fwd
./bin/tile_example_fmha_fwd -prec=bf16 -d=256 -d_v=256 -b=1 -h=32 -h_k=2
-s=1024 -s_k=1024 -bias=n -mask=t -lse=0 -p_drop=0 -warmup=3 -repeat=10
-kname=1 -v=1
./bin/tile_example_fmha_fwd -prec=bf16 -d=256 -d_v=256 -b=8 -h=32 -h_k=2
-s=16384 -s_k=16384 -bias=n -mask=t -lse=0 -p_drop=0 -warmup=3
-repeat=10 -kname=1 -v=1
  ```

## Test Result

  ```bash
-b=1 -s=1024
[bf16|batch|bhsd] b:1, h:32/2, s:1024/1024, d:256/256, scale_s:0.0625,
bias:n, p_drop:0, lse:0, qscale:n, mask:t(-1:0), v:r,
fmha_fwd_d256_bf16_batch_b128x64x32x256x32x256_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_psddv_nlogits_nbias_mc_nlse_ndropout_nskip_nqscale_ntrload_nsink,
0.058 ms, 298.42 TFlops, 618.68 GB/s, valid:y

-b=4 -s=16384
[bf16|batch|bhsd] b:8, h:32/2, s:16384/16384, d:256/256, scale_s:0.0625,
bias:n, p_drop:0, lse:0, qscale:n, mask:t(-1:0), v:r,
fmha_fwd_d256_bf16_batch_b128x64x32x256x32x256_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_psddv_nlogits_nbias_mc_nlse_ndropout_nskip_nqscale_ntrload_nsink,
42.797 ms, 822.18 TFlops, 106.63 GB/s, valid:y
  ```

## Submission Checklist

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

---------

Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: poyenc <1132573+poyenc@users.noreply.github.com>
2026-05-22 01:57:41 +08:00
JiaLuo-CAN
5ff7497fa7 [rocm-libraries] ROCm/rocm-libraries#7537 (commit 07123f4)
[CK Tile] Fix Grouped Gemm quant mixed precision (#7537)

<Migrate from Internal repo PR>
test_ck_tile_grouped_gemm_quant_tensor would fail for mixed FP8/BF8
cases:
std::tuple<Row, Col, Row, FP8, F32, BF8, F32, F32, F16, TensorQuant,
False, True, False>,
std::tuple<Row, Col, Row, BF8, F32, FP8, F32, F32, F16, TensorQuant,
False, True, False>

GFX1250 would fail with incorrect results, GFX950 would fail when
compiling BF8+FP8 and give incorrect results for FP8+BF8.
The issue is due to the wrong ComputeDataType selection.
The fix is to consider original ADataType and BDataType even when
ComputeDataType is not void. For compiling error on gfx950, the bf8,
fp8, 16x16x32 warp Gemm is added.
2026-05-21 08:36:23 -07:00
JP-Fernando
e7798e9560 [rocm-libraries] ROCm/rocm-libraries#7112 (commit a6e5eac)
Add asynchronous XOR shuffle support to the Async GEMM pipeline and the MX GEMM pipeline (#7112)

## Motivation

The goal of this work is to apply XOR shuffle (swizzle) to the current
`comp_async` GEMM pipeline and the `gemm_mx` pipeline.
XOR swizzling has been helpful to avoid LDS bank conflicts, as data are
redistributed across LDS banks, such that simultaneous threads accessing
different rows land on different LDS banks.

## Technical Details

A similar approach to the work in the existing eight-waves pipeline was
followed.
Currently, XOR swizzle support is available for FP8 and BF8 types.
FP4 support is also available for MX GEMM.
Should the types not match, or should the async vector width be of an
unsupported size, then the pipeline falls through to the previously
existing ('unswizzled') path.

## Test Plan

Execute `test_ck_tile_gemm_pipeline_comp_async` for the Async GEMM
pipeline.
Execute `test_ck_tile_mx_gemm_fp8` and `test_ck_tile_mx_gemm_fp4` for
the MX GEMM pipeline.

## Test Result

The tests passed successfully in the `Alola` cluster with MI350
hardware.

## Submission Checklist

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

---------

Co-authored-by: Fernando Jiménez <fernando.jimenez@streamhpc.com>
Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
2026-05-21 09:36:41 +02:00
Wojciech Laskowski
275629fe34 [rocm-libraries] ROCm/rocm-libraries#6014 (commit 2f8259d)
[CK Tile] Adding MFMA wrappers for dense builtins (#6014)

## Motivation

This PR is part of the [WMMA/MFMA] unification work. It's the second of
the series of PRs (after #5801) that add all the necessary MMA builtins
as `amdgcn_mma` structs. This PR focuses on dense MFMA intrinsics.

## Technical Details

This change adds new specializations for WMMA dense builtins. In total,
we add 55 MFMA builtins.

## Test Plan

All the new wrappers were added to the test suite in
`test_amdgcn_mma_layout.inc`.

## Test Result

Test pass locally, waiting for the CI.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
2026-05-21 09:05:19 +02:00