[pull] develop from ROCm:develop - #261
Open
pull[bot] wants to merge 454 commits into
Open
Conversation
Fix per-layer conv2d int8 CPU verification reference path (#6656) case example_conv2d_fwd_xdl_perlayer_quantization_int8.exe 1 0 ## 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.
[CK] Fix divide-by-zero crash for grouped conv kernels (#6132) ## Motivation During run pytorch unit tests for conv3d: `test_dtypes_nn_functional_conv3d_cuda`, `test_fake_crossref_backward_amp_nn_functional_conv3d_cuda_float32` found divide-by-zero crash during CK kernel selection. Refs ROCM-20764 ## Technical Details Add assert for K0PerBlock equal 0, also covered other potential places related with k_batch calculation. ## Test Plan Run miopen command extracted from mentioned test: `MIOpenDriver convfp16 --spatial_dim 3 -I NCDHW -O NCDHW -f NCDHW -n 1 -c 1 -k 1 -g 1 --in_d 4 -H 4 -W 4 --fil_d 4 -y 4 -x 4 --pad_d 0 -p 0 -q 0 --conv_stride_d 2 -u 2 -v 2 --dilation_d 1 -l 1 -j 1 -m conv -F 4 -t 1` ## Test Result Passed ## Submission Checklist - [X] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Signed-off-by: Artem Kuzmitckii <artem.kuzmitckii@amd.com>
[CK] Fix out of bounds modifications caused by negative topk_ids in MoeSortingMultiPhaseKernel_P0_v1 (#6242) ## Motivation Fix sglang randomly crash by filter negative topk ids. ## Technical Details In sglang expert parallel mode, there may be idle batch (batch=0) fired, it will reuse batch=1 resource in cuda graph mode. But in topk op, it will set non used topk ids to -1, in idle batch case, all topk ids are set to -1. In `MoeSortingMultiPhaseKernel_P0_v1` negative expert id will cause overwrite somewhere and sglang may randomly crash. Except idle batch case, if the captured batch sizes are discrete, there may be -1 of expert id due to the similar logic. ## 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: zovonoir <jialzhu@amd.com>
[CK_TILE] fix(fmha): support >2GB KV cache in batch prefill via template dispatch (#6653)
## Motivation
The CK batch prefill kernel previously failed (silent overflow + page
faults) when the KV cache exceeded 2 GB, blocking long-context inference
workloads (e.g., 128K+ token contexts with paged KV).
Two distinct failure modes were addressed:
1. **>4GB SRD overflow (`page_size < kN0`):** The SRD
`buffer_load_dwordx4` path uses a 32-bit `voffset` register; for small
page sizes the rebased SRD spans the full KV pool and the offset wraps
past 2 GB, corrupting K/V loads.
2. **gfx950 page-table fault (`page_size >= kN0`):** On CDNA4 the
hardware validates the **full SRD `num_records` range** against
page-table permissions (CDNA3 only checks per-instruction `voffset`).
After per-tile SRD rebase, an un-trimmed `num_records` field extends
past the live page and faults on freed/protected memory.
## Technical Details
**Two-mode `tile_scatter_gather` selected by the `kUseGlobalLoad`
template parameter:**
| Case | `page_size` | KV cache size | Mode | Load path | Addressing |
|---|---|---|---|---|---|
| 1 | `>= kN0` (large pages) | any | SRD (`kUseGlobalLoad=false`) |
`buffer_load_dwordx4` | 32-bit `voffset`, bounded by per-page rebase |
| 2 | `< kN0` (small pages) | `<= 2 GB` | SRD (`kUseGlobalLoad=false`) |
`buffer_load_dwordx4` | 32-bit `voffset`, fits in INT32 byte range |
| 3 | `< kN0` (small pages) | `> 2 GB` | Global-load
(`kUseGlobalLoad=true`) | `async_load_tile_raw_flat` (K) +
`load_tile_flat` (V) | 64-bit |
**Dispatch:** the auto-gen API layer (`fmha_batch_prefill.py`) selects
the kernel instantiation at launch from `(page_block_size,
num_total_pages * batch_stride_k * kElementBytes)`, so the small-page
penalty is paid only when correctness requires it.
**gfx950 SRD `num_records` trimming:** in the K and V rebase lambdas of
`block_fmha_batch_prefill_pipeline_qr_ks_vs_async`,
`set_bottom_tensor_view_buffer_size(page_stride_k/v)` is called after
each rebase to constrain `num_records` to the live page. Required for
CDNA4 page-table validation; harmless on CDNA3.
**Pipeline sync for the global-load path:**
- V uses synchronous `load_tile_flat`; K uses
`async_load_tile_raw_flat`.
- `v_physical_pages_current` is double-buffered so the V flat load
doesn't race against the next iteration's K rebase computation.
**Arch guards:** `global_load_lds` intrinsics are gated to `__gfx94__` /
`__gfx950__` (CDNA3+). Other architectures hit a `dependent_false`
static_assert with a descriptive message.
**Device-side assertion convention:** SRD setters use
`__builtin_assume(cond)` (hint-only) rather than `<cassert>`'s
`assert()`. The latter introduces an `__assert_fail` call whose register
pressure scatters the K-SRD scalar register window across conditional
branches, corrupting `buffer_load_dwordx4` on gfx950.
## Test Plan
Tested on both MI308 (gfx942) and MI355 (gfx950) via the aiter wrapper
test suite. All coverage lives in **`op_tests/test_batch_prefill.py`**:
- **Functional matrix (96 cases)** — `test_batch_prefill`: `page_size ∈
{1, 16, 1024}` × `kv_layout ∈ {linear, vectorized}` × `dtype ∈ {bf16,
fp8 quant variants}` × `causal` × `soft_cap` × `LSE` × `batch_size ∈ {1,
4}` (parametrized to exercise per-sequence SRD rebase across batch
boundaries).
- **>2 GB coverage** — `test_batch_prefill_large_kvcache`: extended to
allocate a 5 GB+ KV cache pool and exercise both `kUseGlobalLoad=true`
(small-page) and `kUseGlobalLoad=false` (large-page rebase) paths.
Includes both single-batch and multi-batch (`batch_size=4`) cases to
exercise per-sequence SRD rebase across the >2 GB pool.
- Numerical reference: PyTorch SDPA, per-batch loop with `atol` / `rtol`
from the existing batch prefill test harness.
## Test Result
| Arch | `test_batch_prefill` | `test_batch_prefill_large_kvcache` (>2
GB) |
|------|----------------------|---------------------|
| MI308 (gfx942) | All passed | Passed |
| MI355 (gfx950) | All passed | Passed |
**Performance impact (gfx950, hot SRD path):**
- +2.67% kernel-time on `seqlen=1024 / page_sz=1024 / bf16 / sglang /
causal / soft_cap=30`, attributable in full to the two
`set_bottom_tensor_view_buffer_size` calls in the K/V rebase lambdas
(5-run median, signal/noise ≈ 9×).
- This cost is **mandatory for gfx950 correctness** on >2 GB workloads —
removing the setters re-introduces page-faults.
- gfx942: 0 regressions in the same range (all configs ≤ +0.97%).
## Submission Checklist
- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
[CK] Fix CI Failures for PR From Forks (#6701) ## Motivation Fork PRs fail CI when `RUN_AITER_TESTS` or `RUN_FA_TESTS` is enabled. The docker scripts run `git clone -b "$CK_*_BRANCH" https://github.com/ROCm/rocm-libraries.git`, but a fork's branch doesn't exist upstream: ``` fatal: Remote branch <fork-branch> not found in upstream origin ``` Example: [PR #6529 build #4](http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2FComposable%20Kernel/detail/PR-6529/4/pipeline). ## Technical Details **`Jenkinsfile`** — for PRs, use the upstream-visible PR ref instead of the head branch name: ```groovy CURRENT_BRANCH_NAME = env.CHANGE_ID ? "refs/pull/${env.CHANGE_ID}/head" : (env.CHANGE_BRANCH ? env.CHANGE_BRANCH : env.BRANCH_NAME) ``` **`Dockerfile.aiter` / `Dockerfile.fa`** — `git clone -b <ref>` only accepts branches (`refs/heads/*`) and tags (`refs/tags/*`), so it can't resolve `refs/pull/N/head`. Switch to `git fetch`, which accepts any refspec (and still works for plain branch names): ```sh mkdir rocm-libraries && cd rocm-libraries git init -q git remote add origin https://github.com/ROCm/rocm-libraries.git git fetch --depth 1 --filter=blob:none origin "$CK_*_BRANCH" git sparse-checkout init --cone git sparse-checkout set projects/composablekernel git checkout FETCH_HEAD ``` `git checkout FETCH_HEAD` lands in detached HEAD, which breaks the existing `git branch -m "$CK_*_BRANCH"` (and that name isn't a valid local branch anyway). Decouple the local branch name from the upstream ref: - Replace `git init` + `git branch -m` with `git init -b "$LOCAL_BRANCH"` (requires git ≥ 2.28, satisfied by base images) - `LOCAL_BRANCH="ck-import-${ROCM_LIBRARIES_SHA}"` in the rocm-libraries path; `LOCAL_BRANCH="$CK_*_BRANCH"` in the fallback - Downstream `git clone -b ... ../ck` uses `$LOCAL_BRANCH` ## Test Plan Manually trigger a build on this PR with `RUN_AITER_TESTS=true` and `RUN_FA_TESTS=true`; both docker images should build end-to-end. ## Test Result [jenkins / rocm-libraries-folder/Composable Kernel / PR-6701 / #3](http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2FComposable%20Kernel/detail/PR-6701/3/pipeline/) ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Improve the performance of qr_ks_vs_whole_k_prefetch pipeline (#6209) ## About qr_ks_vs_whole_k_prefetch pipeline This PR updates and enhances the qr_ks_vs_whole_k_prefetch pipeline to improve performance on both MI350 GPUs through better MFMA instruction usage, transposed V-loading support, and N0-loop implementation. The pipeline targets scenarios where the number of workgroups is low, enabling better CU occupancy by using smaller MTile sizes (kM0=64 vs 128) while prefetching entire K tiles. ## Changes: - Adds transposed V-loading support (qr_ks_vs_whole_k_prefetch_trload) to avoid using shuffle instructions on MI350 - Implements N0-loop based Gemm0 to reduce tile window movement overhead and eliminate `clear_tile` calls - Adds full support for hdim96/hdim160 without padding requirements - Updates MFMA instruction selection to ensure optimal choices for MI350 ## Performance results 1. For attention shapes which leads to kM0=64, `qr_ks_vs_async_whole_k_prefetch_trload` shows much better performance than `qr_ks_vs_async_trload` on the same case (execution time `41.02ms` by whole_k_prefetch_trload & `58.50ms` by async_load), and `qr_ks_vs_async_whole_k_prefetch_trload` also shows obviously better performance than the recently tuned `qr_ks_vs_async` on the same case (execution time `41.02ms` by whole_k_prefetch_trload 7 `47.60ms` by qr_ks_vs_async) 2. Also on MI300, for attention shapes which leads to kM0=64, `qr_ks_vs_async_whole_k_prefetch` shows much better performance than the `qr_ks_vs_async` (which is supposed to be very high-efficient) on the same case (execution time `64.50ms` by whole_k_prefetch & `80.20ms` by qr_ks_vs_async) 3. For attention shapes which leads to kM0=128, `qr_ks_vs_async_whole_k_prefetch_trload` show a little bit better performance than `qr_ks_vs_async` on mi350 (execution time `104.50ms` by whole_k_prefetch_trload & `106.50ms` by qr_ks_vs_async). And they shows completely on-par performance on MI300 ## Test/Verify 1. Use the ROCM xformers branch `test_whole_k_prefetch_n0loop` to test/verify qr_ks_vs_whole_k_prefetch pipeline since this pipeline can not be used by ck_tile fmha example so far 2. Use the following command-line for building/testing xformers >```bash > #> git clone -b test_whole_k_prefetch_n0loop https://github.com/ROCm/xformers > #> git submodule update --init --recursive > #> pip install --no-build-isolation -e ./ > #> pytest tests/test_mem_eff_attention.py::test_forward >``` 4. Any scripts which can run on xformers can be used to evaluate qr_ks_vs_whole_k_prefetch pipeline. Using the two environ variable to switch from using different pipelines > ```bash > #> export FMHA_DISABLE_SPECIAL_TREATMENT=1 #> to disable using FAV3 and qr_ks_vs_async_trload pipeline > #> export FMHA_ENABLE_ASYNC_PIPELINE=1 #> to disable using qr_ks_vs_async pipeline for comparing > ``` ## Discussion --------- Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: poyenc <1132573+poyenc@users.noreply.github.com> Co-authored-by: qianfengz <12429178+qianfengz@users.noreply.github.com> Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
[CK] restore fmha performance reporting and disable c++17 in CI. (#6741) ## Motivation This change restores monitoring of FMHA benchmarks performance in daily builds and removes the std=c++17 flag from CI builds on gfx90a. ## 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.
[CK Tile] Adding WMMA wrappers for dense builtins (#5801) ## Motivation This PR is part of the [WMMA/MFMA] unification work. It's the first of the series of PRs that add all the necessary MMA builtins as a `amdgcn_mma` structs. ## Technical Details This change adds new specializations for WMMA dense builtins. In total, we have now 9 RDNA4 builtins and 3 RDNA3 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. --------- Co-authored-by: Yung-sheng Tu <yung-sheng@streamhpc.com>
[CK][CK_TILE] Fix FMHA codegen group mode dispatch (#6764) ## Motivation FMHA codegen had incorrect dispatch behavior in group mode. Two root causes: 1. Wrong field names in dispatch conditions — Used batch-mode fields (seqlen_q, seqlen_k) instead of group-mode fields (max_seqlen_q, max_seqlen_k), causing wrong kernel selection at runtime on gfx950. 2. Missing kernel variants — Group mode was overly filtered out from smaller-tile specializations (bwd) and lacked spatial-padding pipeline variants on gfx950 (fwd). gfx942 don't support trload pipeline. ## Technical Details fmha_bwd.py: - max_seq_q_cond and extra_cond now emit t.max_seqlen_q / t.max_seqlen_k for group mode. - Relaxed kernel filtering: group mode no longer skips tiles with max_seq_q != 0. fmha_fwd.py: - get_bm0_cond emits a.max_seqlen_q for group mode tile-size dispatch. - Added two qr_async_trload pipeline variants with spatial padding for gfx950 group mode. ## Test Plan Triggering AITER CI job: ## Submission Checklist - [ x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
[CK_TILE] Add SageAttention v2 forward kernel with multi-granularity quantization (#6574) ## Summary Add a CK_TILE forward kernel implementing [SageAttention v2](https://arxiv.org/abs/2411.10958) — an attention algorithm that applies multi-granularity quantization to Q/K/V before computing attention, trading minimal accuracy loss for higher throughput on low-precision hardware. ### Quantization design | Tensor | Supported data types | Scale granularity options | |--------|---------------------|--------------------------| | Q | fp8 / int8 / int4 | per-tensor, per-block (128 tokens), per-warp (32 tokens), per-thread (4 tokens) | | K | fp8 / int8 / int4 | per-tensor, per-block (128 tokens), per-warp (64 tokens), per-thread (16 tokens) | | V | fp8 | per-channel (always) | | O | bf16 | — | Three precision combinations are supported: `fp8/bf16` (QKV fp8, O bf16), `i8/fp8/bf16` (QK int8, V fp8, O bf16), and `i4/fp8/bf16` (QK int4, V fp8, O bf16). ### Architecture support - **gfx9** (CDNA2/3, e.g. gfx90a, gfx942) — full tile set - **gfx950** (CDNA4) — restricted tile set (N-per-block capped at 64 for fp8-family dtypes) ### Implementation - Two pipeline variants: `QRKSVS` (synchronous) and `QRKSVS_ASYNC` (async copy) - Masking support: no mask, causal (top-left / bottom-right), and generic windowed - Batch and group (variable-length) modes - Head dimension: d=128, d_v=128 - Python codegen under `example/ck_tile/49_sageattention/codegen/` generates kernel instances per target/dtype/tile combination - Smoke tests included via `tile_example_sageattn_fwd` ### Test commands \`\`\`bash # fp8 QKV ./build/bin/tile_example_sageattn_fwd -v=1 -b=16 -h=8 -s=1024 -d=128 -kname=1 -prec=fp8bf16 -qscale=3 -init=3 # int8 QK, fp8 V ./build/bin/tile_example_sageattn_fwd -v=1 -b=16 -h=8 -s=1024 -d=128 -kname=1 -prec=i8fp8bf16 -qscale=3 -init=3 \`\`\` \`-qscale\` values: 1=per-tensor, 2=per-block, 3=per-warp, 4=per-thread
[CK] Dockerfile: auto-discover latest TheRock nightly tarball (#6972) ## Motivation Our docker containers with `--build-arg compiler_version=therock` should have the latest nightly build of TheRock in `/opt/rocm`. When I looked for `rocm_kpack` and other `kpack` artifacts, they were missing, and I realized we had pinned the version by date. Instead, we should look for the most recent linux-multiarch tarball. ## Summary - Auto-discover the latest TheRock nightly tarball at Docker build time instead of pinning a stale URL (previously hardcoded to a Feb 2026 nightly that predates kpack) - Logic is to `wget` the directory, and identify the latest tarball (alphabetically sorted by YYYYMMDD in filename). - Support manual override via `--build-arg TARBALL_URL=...` for pinning, and `--build-arg TARBALL_PATTERN=...` for selecting a specific arch variant - Fix sccache download URL: `/releases/latest/download/` was redirecting to v0.15.0 but the filename referenced v0.14.0, causing a 404 ## Test plan - [x] Verified tarball discovery logic resolves to `therock-dist-linux-multiarch-7.13.0a20260430.tar.gz` - [x] Built Docker image locally with `--build-arg compiler_version=therock` - [x] Confirmed sccache installs successfully with the fixed URL - [ ] Verify CI pipeline builds with the updated Dockerfile 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
[CK] Reduce per-file logging in cmake_dependency_analyzer (#6912) ## Motivation Current progress_callback function generates large volume of prints which creates noise in seeing actual CI failure logs. Only emit a progress line at the completion of each stage to avoid massive logs from the per-source-file extracting_dependencies callback. ## Technical Details Update the `progress` function to print only at the completion of each stage. https://github.com/ROCm/rocm-libraries/pull/6912/changes#diff-15971b83c7dfefb48fd788507a923017d93bbd9487ed6aeb414ad2c5e00be934R720 ## Test Plan to be tested in CI ## Test Result to be tested in CI ## Submission Checklist - [x ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
[CK] Fix OOB page table read in batch_prefill V prefetch (AICK-1171) (#6932)
## Summary
Fix a GPU memory access fault in `mha_batch_prefill` triggered when the
per-batch page table is tightly sized (no trailing slack).
**Affected configurations:**
- All FMHA batch prefill V2 kernels
(`block_fmha_batch_prefill_pipeline_qr_ks_vs_async`)
- Triggered by paged KV layouts where `kv_page_indices.numel() ==
ceil(seqlen_k / page_size)` exactly
- Manifests as: `Memory access fault by GPU node-X (Agent handle:
0x...)` followed by `Aborted (core dumped)`
- Silent corruption (no fault, wrong output) when the OOB read happens
to land in zero-initialized memory
### Root cause
`load_physical_pages` performs **lookahead reads** on the page table to
prefetch K/V tiles for the next iteration. When the page table for a
batch has exactly `N` entries, the V-tile prefetch indexes `page_idx[N]`
(one past the last valid entry), reading either uninitialized memory or
the next batch's slot. On gfx942 with a tightly-sized page table, the
read crosses into an unmapped page and triggers an HSA page fault.
The bug was masked in earlier testing because most test harnesses pad
`kv_page_indices` with trailing zeros — OOB reads then return `page_id =
0`, a valid in-cache page, producing silent numerical drift instead of a
fault.
### Fix design
Thread `max_page_table_idx = (seqlen_k - 1) / page_size` from the kernel
layer down to `load_physical_pages`, and clamp every page-table read
with `ck_tile::min()`. Applied to **all four code paths** in the V
prefetch:
| Branch | What it does | Clamp applied |
|--------|-------------|---------------|
| `kIsKcache` | K prefetch loop | `min(global_token_idx >>
kLog2PageSize, max_page_table_idx)` |
| V LINEAR (`page_size == 1`) | One token = one page |
`min(global_token_idx, max_page_table_idx)` |
| V crosses pages (`kVTileCrossesPages`) | Per-thread page lookup |
`min(global_token_idx >> kLog2PageSize, max_page_table_idx)` |
| V single page (lane0 broadcast) | `readfirstlane`-uniform lookup |
`min(... >> kLog2PageSize, max_page_table_idx)` |
### Key design decisions
**Mandatory parameter, not optional with a sentinel default.** An
optional `max_page_table_idx = INT32_MAX` default would let the bug
silently come back at any new callsite that forgets to pass it. Making
it mandatory forces every caller to opt in explicitly and surfaces
missed callsites at compile time.
**`seqlen_k == 0` clamps to 0** instead of underflowing `(0 - 1) /
page_size` to `-1`. The empty-batch case is rare but well-defined: clamp
every read to slot 0.
**Single computation in the kernel layer.**
`FmhaBatchPrefillWithPagedKVCacheKernel` computes `max_page_table_idx`
once per batch and forwards it through every QScale branch (PERTENSOR /
KV_BLOCKSCALE / default). All three `operator()` overloads of the
pipeline (rich, default forwarder, KV_BLOCKSCALE forwarder) take and
forward the parameter.
### Files changed
| File | Change |
|------|--------|
| `include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp` |
Compute `max_page_table_idx` per batch, forward to all 3 QScale branches
|
|
`include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp`
| Add `max_page_table_idx` to `load_physical_pages` and 3 `operator()`
overloads; clamp page-id reads in 4 code paths |
## Test plan
- [x] AICK-1171 reproducer verified on MI-308X (gfx942)
- [x] New pytest case `test_batch_prefill_aick1171_oob_page_table_read`
in aiter, parametrized over `total_blocks ∈ {160, 164, 168, 176, 208,
256}` (matches the `crash1_r8_*` bisect family)
- [x] Full FMHA batch prefill suite on gfx942 + gfx950
## Linked issue
AICK-1171.
[CK] fix CI git token. (#7046) ## Motivation Fix the CI breakage due to git PAT deprecation. ## 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.
[CK_TILE] Use Unified Workspace for FMHA BWD (#6152)
## Motivation
`dq_acc` is the intermediate accumulation buffer used in FMHA backward
pass for deterministic mode. The current implementation allocates it as
a **single rectangular tensor**:
```
shape = [shape_batch, nhead, nsplits, shape_seqlen_q, hdim_q]
```
where `nsplits = launcher.dq_acc_splits` (a single scalar), computed
from `max_seqlen_k` and shared across all batches.
### Problems
1. **Memory waste**: In group mode, each batch may have a different
`seqlen_k`, but `nsplits` is computed from `max_seqlen_k`, causing
batches with shorter `seqlen_k` to over-allocate in the split dimension.
2. **Interface coupling**: `fmha_bwd_args` exposes internal layout
details such as `stride_dq_acc`, `nhead_stride_dq_acc`,
`batch_stride_dq_acc`, and `split_stride_dq_acc`. The caller is
responsible for computing these strides, but this logic belongs inside
the kernel.
### Goals
1. Switch `dq_acc` buffer to a **compact layout**: batches are
concatenated contiguously, with each batch occupying `nhead * nsplits_i
* seqq_i * hdim_q` elements (nhead outermost).
2. **Remove all `*_stride_dq_acc` fields** from `fmha_bwd_args`,
replacing them with a single `workspace_ptr`; the kernel splits this
internally using a fixed layout.
4. `fmha_bwd_launcher` provides a **workspace management interface**:
the caller only needs to allocate GPU memory and call
`prepare_workspace()` — no layout computation required.
5. **Isolate kernel internals from the caller API**: the `dq_acc` layout
(nsplits, strides, buffer size) is determined entirely inside the
launcher/kernel. Future changes to block shape, pipeline type, or
persistent kernel strategy require no modifications to the caller's
`fmha_bwd_args` or workspace allocation logic.
## Technical Details
### Interface Design
#### New fields in `fmha_bwd_traits`
```cpp
struct fmha_bwd_traits
{
int seqlen_q;
int seqlen_k;
int batch;
int max_seqlen_q;
int max_seqlen_k;
int hdim_q;
int hdim_v;
int nhead_q;
int nhead_k;
std::string data_type;
bool is_group_mode;
mask_enum mask_type;
bias_enum bias_type;
bool has_dbias;
bool has_dropout;
bool is_store_randval;
bool is_deterministic;
// New: cumulative physical seqlen pointers for group mode (pass nullptr for batch mode).
// seqstart_qs[i+1] - seqstart_qs[i] = physical seqlen_q of batch i (including padding); length = batch+1
// seqstart_ks[i+1] - seqstart_ks[i] = physical seqlen_k of batch i (including padding); length = batch+1
const int* seqstart_qs = nullptr;
const int* seqstart_ks = nullptr;
};
```
#### `fmha_bwd_launcher` actual structure
```cpp
struct fmha_bwd_launcher
{
std::function<float(fmha_bwd_args, const ck_tile::stream_config&)> run{};
// Total workspace size in bytes (host_ws_size + device_ws_size), computed by init().
// Zero for kUseQrQtrDorPipeline (writes dq directly, no acc buffer needed).
size_t workspace_size = 0;
fmha_bwd_launcher(const fmha_bwd_traits&);
// Copies auxiliary data (nsplits[], offsets[]) via hipMemcpy to the head of the GPU workspace,
// and zeros the dq_acc buffer portion (tail of workspace) if required.
// The memory pointed to by device_ws must be >= workspace_size bytes.
std::function<void(void* device_ws)> prepare_workspace{};
template <typename... Args>
float operator()(Args&&... args) const { return run(std::forward<Args>(args)...); }
private:
size_t host_ws_size = 0; // CPU workspace size (nsplits[] + offsets[] arrays)
size_t device_ws_size = 0; // GPU-only data size (dq_acc buffer)
std::unique_ptr<char[]> ws_host; // host-side workspace buffer
public:
template <typename T0, typename T1, typename T2, typename Arch>
void init(const fmha_bwd_traits& traits);
};
```
The `init<>()` template method (invoked by codegen dispatch branches as
`this->init<...>(t)`) is responsible for:
1. Setting the `run` lambda
2. Calling `FmhaBwdDQDKDVKernel::GetWorkspaceHostSize(batch)` to obtain
`host_ws_size`
3. Allocating `ws_host` (host memory)
4. Calling `FmhaBwdDQDKDVKernel::PrepareWorkspaceHost(ws_host.get(),
...)` to fill nsplits/offsets; return value is `device_ws_size`
5. `workspace_size = host_ws_size + device_ws_size`
6. Setting the `prepare_workspace` lambda (captures `this`, calls
`PrepareWorkspaceDevice`)
When no kernel matches the given traits, both `run` and
`prepare_workspace` are initialized to default lambdas that print a
warning to `std::cerr` and return gracefully (no exception).
#### Workspace overall layout
The workspace is managed by `FmhaBwdWorkspaceManager` and consists of
two segments:
```
Offset 0 (CPU-prepared segment, host_ws_size bytes; also hipMemcpy'd to the head of GPU workspace):
index_t nsplits[batch or 1] — per-batch nsplits array
group mode: batch elements
batch mode / non-deterministic: 1 element
[group mode only] long_index_t dq_acc_offsets[batch+1]
— per-batch element offset (inclusive prefix sum)
offsets[0]=0, offsets[i+1] = offsets[i] + nhead*nsplits_i*seqq_i*hdim_q
Offset host_ws_size (device data segment, device_ws_size bytes):
AccDataType dq_acc[total_elements] — compact dq_acc buffer (zeroed if required)
total_elements = sum_i(nhead * nsplits_i * seqq_i * hdim_q)
layout within each batch: [nhead, nsplits_i, seqq_i, hdim_q]
note: seqq_i uses the physical length (including padding)
```
Alignment constant (`ALIGNMENT = 16`):
```
nsplits_size = align_up(sizeof(index_t) * N, 16) // N = batch (group) or 1 (batch/non-det)
offsets_size = align_up(sizeof(long_index_t) * (batch+1), 16) // group mode only
host_ws_size = nsplits_size + offsets_size
dq_acc_offset = host_ws_size // GetDqAccDataOffset(batch)
```
**Key benefits**:
- The kernel reads nsplits/offsets directly from the workspace head — no
device-side recomputation.
- `FmhaBwdConvertQGradKernel` is completely decoupled from the pipeline
block shape (`kN0`): nsplits is read from `nsplits_ptr`, `kN0` is no
longer a template parameter, and multiple dq_dk_dv tiles with different
`F_bn0` values now share a single convert_dq kernel instance (under
receipt 1/2, deterministic convert_dq kernel count drops from ~300 to
60).
- nsplits/offsets are computed on the host and transferred in one
`hipMemcpy`; the dq_acc buffer follows immediately, at the offset given
by `GetDqAccDataOffset`.
#### Workspace size by scenario
| Scenario | `workspace_size` | Notes |
|----------|-----------------|-------|
| **kUseQrQtrDorPipeline** (any mode) | `0` | Writes dq directly; no acc
buffer; `PrepareWorkspaceHost` returns 0 |
| **Non-deterministic + batch mode** | `> 0` | nsplits[1]=1; dq_acc used
for atomic add; `workspace_size = host_ws_size +
batch*nhead*seqlen_q*hdim_q*ebytes` |
| **Non-deterministic + group mode** | `> 0` | nsplits[1]=1; dq_acc
contiguous layout; `workspace_size = host_ws_size +
nhead*seqstart_qs[batch]*hdim_q*ebytes` |
| **Deterministic + group mode** | `> 0` | nsplits[batch],
offsets[batch+1], compact dq_acc; nsplits_i computed independently per
batch |
| **Deterministic + batch mode persistent** | `> 0` | nsplits[1]
(uniform across batches); dq_acc `batch*nhead*nsplits*seqlen_q*hdim_q` |
**NeedsZeroDqAcc** (determines whether `PrepareWorkspaceDevice` calls
`hipMemset`):
- Persistent kernel (deterministic batch mode) or non-deterministic:
**must zero** (atomic add requires zero initialization)
- Deterministic group mode + no mask: **no zeroing needed** (every tile
writes its full region)
- Deterministic + with mask: **must zero** (some blocks are skipped,
leaving uninitialized tiles that would contribute to the reduction)
#### Caller usage
```cpp
// 1. Create launcher (traits include seqstart_qs/ks pointers; workspace_size is computed during construction)
fmha_bwd_launcher launcher(fmha_traits);
// 2. Read launcher.workspace_size directly
const auto ws_size = launcher.workspace_size;
// 3. Allocate a single GPU workspace
ck_tile::DeviceMem ws_buf(ws_size);
// 4. Copy nsplits/offsets to GPU head and zero dq_acc if required
launcher.prepare_workspace(ws_buf.GetDeviceBuffer());
// 5. Build args with a single workspace pointer; the kernel splits it internally
fmha_bwd_args args{
...,
ws_size > 0 ? ws_buf.GetDeviceBuffer() : nullptr, // workspace_ptr
};
launcher(args, stream_config);
```
---
### Key Code Structure
#### FmhaBwdWorkspaceManager (`fmha_bwd_kernel.hpp`, new class)
```cpp
template <typename AccDataType, bool kIsGroupMode, bool kIsDeterministic>
struct FmhaBwdWorkspaceManager
{
static constexpr size_t ALIGNMENT = 16;
// CPU workspace (nsplits + offsets) sizes
static size_t GetDqAccSplitsSize(int batch); // align_up(sizeof(index_t)*N, 16)
static size_t GetDqAccOffsetsSize(int batch); // group mode only: align_up(sizeof(long_index_t)*(batch+1), 16)
static size_t GetWorkspaceHostSize(int batch); // = SplitsSize + OffsetsSize
// Starting offset of dq_acc data within the full workspace (= host_ws_size)
static size_t GetDqAccDataOffset(int batch); // = GetWorkspaceHostSize(batch)
// Fills nsplits/offsets in the CPU workspace; returns device_ws_size (dq_acc buffer bytes)
template <bool kUseQrQtrDorPipeline, index_t kN0>
static size_t PrepareWorkspaceHost(void* cpu_ws, index_t batch_size, index_t hdim_q,
index_t nhead_q, index_t seqlen_q, index_t seqlen_k,
const index_t* seqstart_qs, const index_t* seqstart_ks);
// hipMemcpy's cpu_ws to device_ws head; hipMemset's the dq_acc portion to 0 if required
template <bool kUseQrQtrDorPipeline, bool kHasMask>
static void PrepareWorkspaceDevice(void* device_ws, const void* host_ws,
size_t device_ws_size, size_t host_ws_size);
};
```
#### workspace_ptr parsing (inside the kernel)
The kernel parses three address regions from `kargs.workspace_ptr`:
**Group mode (`FmhaBwdDQDKDVKernel::MakeKargs`)**:
```cpp
const uint8_t* ws = reinterpret_cast<uint8_t*>(workspace_ptr);
// dq_acc_ptr (stored in FmhaBwdCommonKargs)
ws + WorkspaceManager::GetDqAccDataOffset(batch)
// dq_acc_batch_offset_ptr (FmhaBwdGroupModeKargs field)
reinterpret_cast<const long_index_t*>(ws + WorkspaceManager::GetDqAccOffsetsOffset(batch))
```
**Batch mode**:
```cpp
ws + WorkspaceManager::GetDqAccDataOffset(batch) // dq_acc_ptr
// No offsets pointer; batch offset is computed inside run_() from nsplits
```
**`FmhaBwdConvertQGradKernel`** follows the same pattern:
- Group mode: extracts `dq_acc_ptr`, `dq_acc_batch_offset_ptr`, and
`nsplits_ptr` (`GetDqAccSplitsOffset(batch)`) from workspace
- Batch mode: reads nsplits from `nsplits_ptr[0]`; batch offset computed
internally
### Addressing in `run_()` (group mode)
```cpp
// Per-batch processing:
const long_index_t batch_offset_dq_acc = kargs.dq_acc_batch_offset_ptr[i_batch];
// seqq_i (physical length) derived from seqstart_q_ptr
const index_t seqq_i = kargs.seqstart_q_ptr[i_batch+1] - kargs.seqstart_q_ptr[i_batch];
// nsplits_i read from nsplits_ptr (convert_dq kernel) or from GetDqAccSplits
const long_index_t split_stride_i = static_cast<long_index_t>(seqq_i) * kargs.hdim_q;
const long_index_t nhead_stride_i = static_cast<long_index_t>(nsplits_i) * split_stride_i;
// Final address:
dq_acc_base + batch_offset_dq_acc + i_nhead * nhead_stride_i + i_split * split_stride_i
```
#### nsplits computation (`PrepareWorkspaceHost`)
`PrepareWorkspaceHost` is a template method of `FmhaBwdWorkspaceManager`
that still takes `kN0` as a template parameter (from
`BlockFmhaShape::kN0` of the dq_dk_dv pipeline). However, this parameter
is **only used inside this host-side function** to compute nsplits — it
is no longer passed into the convert_dq kernel.
| Mode | nsplits computation |
|------|---------------------|
| kUseQrQtrDorPipeline | Writes dq directly; nsplits[0]=0; returns
device_ws_size=0 |
| Non-deterministic | nsplits[0]=1; dq_acc used for atomic add |
| Deterministic + group mode | `ceil((seqstart_ks[i+1]-seqstart_ks[i]) /
kN0)` computed per batch |
| Deterministic + batch mode persistent | Same logic as the original
`GetDqAccSplits` (`dqdqkdv_workers` based) |
### Removing kN0 dependency from `FmhaBwdConvertQGradKernel`
`FmhaBwdConvertQGradKernel` previously required `kN0` as a template
parameter (via `BlockFmhaBwdConvertQGradPipelineProblem`) for two
purposes:
1. In batch mode `operator()`: self-computing `nsplits = ceil(seqlen_k /
kN0)`
2. The `b{kM0}x{kN0}` component of the kernel name string
Both have been removed in this refactor:
- **Batch mode**: now reads `kargs.nsplits_ptr[0]` directly (guarded by
`if constexpr(kIsDeterministic)` to avoid accessing a non-existent field
in non-deterministic instances)
- **Kernel name**: simplified to `b{kM0}`, no longer includes `kN0`
- **Template parameters**: `BlockFmhaBwdConvertQGradPipelineProblem`
drops the `kN0_` parameter; `fmha_bwd_convert_dq_traits_` drops the
`kN0` parameter; `F_bn0`/`convert_dq_bn0` fields removed from codegen
Effect: all dq_dk_dv tiles sharing the same `(hdim, dtype, mode, pad,
deterministic)` combination — regardless of `F_bn0` value
(16/64/128/192/256) — now share a **single** convert_dq kernel instance.
---
## Test Plan
<!-- Explain any relevant testing done to verify this PR. -->
## 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.
[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels (#6914)
## Summary
Two related fixes to `ck_tile` FMHA so that StreamLLM-sink +
sliding-window
batch-prefill works correctly for fp8 KV / bf16 compute.
Review the commits in this order:
1. `fmha: emit sink kernels for fp8bf16 batch_prefill`
Extends `example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py` so
the fp8(KV) / bf16(QO) batch-prefill codegen also emits the
`mask=mask_enum::generic_with_sink` variant. Without this the runtime
could not dispatch to a sink-aware kernel for the fp8bf16 path.
2. `fmha: respect right-window in IsOutOfSinkBound`
The sink un-mask in `GenericAttentionMask::IsOutOfSinkBound` (local-mask
branch) used `(i_y + x) > 1` as the gate, which conditioned on the row
index instead of the column index. As a result, queries `1..sink-1`
could attend to *future* sink positions (violating causal /
right-window),
while query `0` fell back to the plain causal mask. The fix replaces the
guard with `i_x < i_y + x` so every query only sees sink columns up to
its own right-window boundary.
3. `fmha: clarify IsOutOfSinkBound predicate comment`
Doc-only follow-up that rewrites the comment above the predicate as a
clause-by-clause explanation (`i_x < sink`, `i_x < i_y + x`,
`y < y_total`, `i_y < x_total`).
## Test plan
- [x] Repro on aiter `op_tests/test_batch_prefill.py` (fp8 +
bf16_dequant
modes with `sink=4`, `win_left=1023`, `softcap=0.0`, `sal=True`)
now passes for all parametrized shapes.
- [x] Existing fp16/bf16 batch-prefill paths (no sink) unchanged —
codegen
diff only adds the `generic_with_sink` variant for fp8bf16; existing
kernel object lists unaffected.
## Submission Checklist
- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
---------
Co-authored-by: fengjunda.aml <fengjunda.aml@bytedance.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: root <root@smci350-rck-g03-f12-31.rck.dcgpu>
[CK] add swiglustep_and_mul activation to gridwise_moe_gemm (#6873) Title： feat(composablekernel): add swiglustep_and_mul activation to gridwise_moe_gemm Description： ## Motivation Step-3.5-Flash uses a clamped SwiGLU activation (`swiglu_limits[43]=7`, `swiglu_limits[44]=7`) for layers 43 and 44. Without this kernel path, those layers produce BOS token spam because unclamped gate/up values accumulate floating-point noise over 200+ decode steps, degrading output quality (cosine similarity drops from 0.999989 to ~0.998982). ## Changes Add `swiglustep_and_mul` as a new `Activation` enum branch in `gridwise_moe_gemm.hpp`, covering all 4 code paths: - Quantized (A×B scale) + IsInputGemm=true - Quantized (A×B scale) + IsInputGemm=false - Non-quantized + IsInputGemm=true - Non-quantized + IsInputGemm=false The activation computes: gate = silu(gate) gate = clamp(gate, max=7.0f) up = clamp(up, min=-7.0f, max=7.0f) output = gate * up Also handles the `MulRoutedWeight` case (topk weight multiplication) and `pk_i4_t` weight scaling (×16 dequant factor). ## Verification - Tested on gfx950 (MI350X, 8×GPU) - cosine similarity for layers 43/44: **0.999989** (vs 0.998982 before fix) - End-to-end Step-3.5-Flash inference: no BOS spam, output coherent - BF16 tp=2/tp=4 and FP8 tp=2/tp=4 all verified PASS - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
[CK_TILE] Enable V3 persistent kernel dispatch for FMHA forward on gfx950 (#6529)
[CK_TILE] Enable V3 persistent kernel dispatch for FMHA forward on
gfx950
## Motivation
Enable the existing V3 persistent kernel path for CK-Tile FMHA forward
on
gfx950 (MI350X/MI355X). The V3 kernel and codegen infrastructure already
exist but are disabled via hardcoded `F_is_v3_enabled=False`.
This change replaces the compile-time gate with a runtime environment
variable
`CK_FMHA_ENABLE_V3=1` (disabled by default, opt-in). When enabled:
- **Prefill** workloads (seqlen_q > 1) dispatch to V3 persistent
pipeline
- **Decode** workloads (seqlen_q == 1) always use V2 (memory-bound,
better suited)
The V3 persistent kernel uses grid-stride scheduling, XCD-interleave
tile
assignment for L2 locality, LPT reversal for causal masks, and gfx950
async
buffer loads.
## Technical Details
Single file: `example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py`
- Add `#include <cstdlib>` and `<string>` for `std::getenv`
- Replace `{F_is_v3_enabled}` template parameter with runtime env var
check
- Add `seqlen_q > 1` guard (decode always uses V2)
- Remove `.format()` call in `write_fwd_api()`
## Dependencies
Depends on ROCm/rocm-libraries#6501 — builds on
XCD-interleave and LPT scheduling infrastructure.
## 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 (V2): `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096
-d=128 -prec=bf16 -v=1 -warmup=1 -repeat=3`
- Command (V3): `CK_FMHA_ENABLE_V3=1 ./build/bin/tile_example_fmha_fwd
-b=2 -h=8 -s=4096 -d=128 -prec=bf16 -v=1 -warmup=1 -repeat=3`
- Command (decode, always V2): `./build/bin/tile_example_fmha_fwd -b=64
-h=32 -h_k=8 -s=1 -s_k=4096 -d=128 -prec=bf16 -mode=group -v=1 -warmup=1
-repeat=3`
## Test Result
Benchmark results (MI350X, gfx950, ROCm 7.0):
| Config | V2 (TFlops) | V3 (TFlops) | Speedup |
|--------|-------------|-------------|---------|
| Non-causal b=2 h=8 hk=2 s=4096 d=128 bf16 | 696.3 | 884.2 | **+27.0%**
|
| Causal b=2 h=8 hk=2 s=4096 d=128 bf16 | 371.3 | 494.9 | **+33.3%** |
| GQA b=2 h=32 hk=8 s=2048 d=128 bf16 | 671.3 | 831.7 | **+23.9%** |
| LLaMA-70B b=1 h=64 hk=8 s=4096 d=128 bf16 | 761.5 | 927.3 | **+21.8%**
|
| Causal GQA b=2 h=32 hk=8 s=2048 d=128 bf16 | 345.4 | 631.9 |
**+82.9%** |
| Long-seq b=1 h=16 s=16384 d=128 bf16 | 797.8 | 969.9 | **+21.6%** |
| Decode b=64 h=32 hk=8 s=1 s_k=4096 bf16 | 1828 GB/s | — (V2 path) |
unaffected |
Benchmark results (MI300X, gfx942, ROCm 6.4.1):
V3 has 0% effect on MI300X — V3 relies on gfx950 async buffer loads and
falls back to the V2 code path on gfx942. No regression on any config.
| Config | TFlops / GB/s | Time (ms) | Delta vs baseline |
|--------|-------------|-----------|-------------------|
| MHA bf16 b=2 h=8 s=4096 d=128 | 342.98 TFlops | 0.401 | +0.1% |
| MHA fp16 b=2 h=8 s=4096 d=128 | 411.18 TFlops | 0.334 | +4.9% |
| Causal MHA bf16 b=2 h=8 s=4096 d=128 | 232.61 TFlops | 0.296 | +2.4% |
| GQA 4:1 bf16 b=2 h=32 hk=8 s=2048 d=128 | 320.07 TFlops | 0.429 |
-1.4% |
| GQA 8:1 bf16 b=2 h=64 hk=8 s=2048 d=128 | 353.91 TFlops | 0.777 |
+1.7% |
| LLaMA-70B prefill b=1 h=64 hk=8 s=4096 d=128 bf16 | 381.53 TFlops |
1.441 | +1.2% |
| Long-seq bf16 b=1 h=16 s=16384 d=128 | 388.61 TFlops | 5.659 | +1.4% |
| Decode b=64 h=32 hk=8 s_k=4096 d=128 bf16 | 693.40 GB/s | 1.550 |
+0.3% |
All validation tests pass (`valid:y`) on both MI300X and MI350X.
Additional validation:
- `CK_FMHA_ENABLE_V3=0` correctly falls back to V2 (default behavior
unchanged)
- `CK_FMHA_ENABLE_V3=1` dispatches to V3 for prefill, V2 for decode
- Validation passes across fp16/bf16, batch/group mode,
causal/non-causal
- No regression on decode path
---------
Co-authored-by: Chao Zhou <chaozhou@fb.com>
Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
[CK] disable tile_engine by default, limit gfx1030 CI builds to develop only. (#7138) ## Motivation An attempt to reduce the build time and keep CI moving faster. Disable tile_engine by default since even the cmake step may take up to 30 minutes. Since we're down to a single gfx1030 CI node, use it only for develop builds. ## 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.
[CK] Filter out unsupported targets. (#6933) ## Motivation Filter out any unsupported targets, e.g., gfx900, gfx906, gfx90c, from the GPU_TARGETS or GPU_ARCHS lists. ## 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.
[CK_TILE] Fix typo in fmha_fwd_kernel K-dram unmerge tuple sizes (#7141)
## Summary
The qr_async_trload K-dram lambda's `else (XorLengthFold == 1)` branch
in `fmha_fwd_kernel.hpp` writes the outer-tile dim of its 3-tuple
unmerge/xor/merge as
```cpp
number<FmhaPipeline::kQKHeaddim / kDramTileK / FmhaPipeline::kAlignmentK>{}
```
which divides one extra time. For every fp16/bf16 hdim=128 configuration
the outer length collapses to **0**, e.g. `128 / 128 / 8 == 0`. The
3-tuple product no longer equals `kQKHeaddim`, so unmerge → xor → merge
stops round-tripping the head dimension.
This bug was masked by the async-load path: it only walks the descriptor
via stride and silently absorbs a length=0 outer dim. Any consumer that
actually traverses the descriptor (e.g. the TDM path on gfx1250)
immediately faults on the resulting `tuple<int, constant<0>>`.
The fix drops the extra `/ kAlignmentK` in all three call sites in the
same lambda so the outer dim becomes `kQKHeaddim / kDramTileK` and the
product is restored to `kQKHeaddim`. Strides are unaffected, so the
async path is bit-identical.
| Config (fp16/bf16) | hdim | kDramTileK | kAlignmentK | a (typo) | a
(fixed) | product (typo) | product (fixed) |
|---|---|---|---|---|---|---|---|
| hdim128, kKLoadOnce | 128 | 128 | 8 | 0 | 1 | **0** | **128** |
| hdim128, kK0=32 | 128 | 32 | 8 | 0 | 4 | **0** | **128** |
| hdim64, kKLoadOnce | 64 | 64 | 8 | 0 | 1 | **0** | **64** |
| hdim256, kK0=32 | 256 | 32 | 8 | 1 | 8 | **32** | **256** |
Bug introduced in 2cc0af6a815a (PR #2888 \"[CK_TILE] FMHA FWD bug
fix\"), where the original 2-tuple unmerge was generalized to a 3-tuple
and the typo slipped in.
## Test plan
- [x] Built `test_ck_tile_fmha_fwd` (umbrella, 5 gtest binaries) on
gfx950 native at develop b3bdc63a509 with `dev-gfx950` preset (clang 22,
ROCm 7.2.2). Compiles cleanly with `-Werror -Weverything`.
- [x] Ran `ctest -R test_ck_tile_fmha_fwd` on gfx950 native, baseline vs
patched: identical pass/fail (3 pass / 2 fail), identical failing case
set (114 gtest fails + 2 GPU memory access faults, all in pre-existing
fp16/bf16 group-mode `Alibi`/`Dropout` cases that reproduce on develop
without this patch). Total wall time 403s → 393s. Per-case latency drift
±8% (noise).
- [x] CI to verify on other gfx9 / gfx11 architectures.
[CK] Fix latest batch of staging compiler warnings (#7111) ## Motivation Suppress the new batch of clang lifetimebound and invalidation warnings with the latest 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.
[CK][CK TILE] Dispatcher kernel selection heuristic for grouped conv (#6327) ## Motivation The ML heuristic in dispatcher does not support grouped-conv operator yet. In this PR, the support for fwd, bdw-data, and bwd-weight grouped-conv kernels have been added. A tile_engine utility has also been added to compile and run any selected kernel configuration through dispatcher infrastructure. ## Technical Details 1. Tile engine utility is added to benchmark each shape with all the possible kernel+tile_size combinations here - [https://github.com/ROCm/rocm-libraries/blob/users/yraparti/ck/dispatcher-grouped-conv-heuristics/projects/composablekernel/tile_engine/ops/grouped_conv/grouped_conv_full_benchmark.py](url) 2. New LGBM regressor models for grouped conv are added to models directory. We have 3 separate models for fwd, bwd-data, and bwd-weights [https://github.com/ROCm/rocm-libraries/tree/users/yraparti/ck/dispatcher-grouped-conv-heuristics/projects/composablekernel/dispatcher/heuristics/models](url) 3. Implemented lazy GPU initialization (dispatcher/python) - **Issue**: ProcessPoolExecutor fork() + GPU context caused memory access faults - **Solution**: Mirror FMHA pattern - defer GPU initialization until first run() - **Changes**: - setup_multiple_grouped_conv_dispatchers() returns List[Path], not loaded libs - GpuGroupedConvRunner.__init__() no longer calls ctypes.CDLL - Added _ensure_initialized() method for lazy GPU loading - GPU context created only on first run() call - **Benefit**: Parallel compilation now works without GPU conflicts 4. Addressed few miscellaneous issues such as: - Fixed BF16->FP16 naming bug in the dispatcher wrapper - Added new tile sizes, and comp_v5 pipeline to the arch spec to expand the kernel selection - Added automatic padding support for unsupported shapes in dispatcher runner - Created a single source of truth between tile_engine and dispatcher about the architecture and tile_size details - Build a validation scripts to compare oracle_best vs ml_heuristic comparison ## Test Plan 1. Validated fwd, bwd-data, and bwd-weight kernels with both known and unseen data sets with up to 300 problems. 2. Ensured that test cases are added in both dispatcher and tile_engine to validate the heuristic. ## Test Result Results on Unseen shapes validated on gfx950 #### Forward Pass Model - **Training Data**: 48,845 measurements across 1,372 unique problem shapes - **Validation Set**: 300 unseen problems from model crawler - **Validation Performance** (vs. oracle): - Mean Efficiency: **93.05%** - Median Efficiency: **96.8%** - P10 Efficiency: **79.9%** #### Backward Data Gradient (bwd_data) Model - **Training Data**: 18,773 measurements across 891 unique problem shapes - **Validation Set**: 300 unseen problems from model crawler - **Validation Performance** (vs. oracle): - Mean Efficiency: **93.8%** - Median Efficiency: **96.5%** - P10 Efficiency: **82.9%** #### Backward Weight Gradient (bwd_weight) Model - **Training Data**: 34,900 measurements across 1,508 unique problem shapes - **Validation Set**: 300 unseen problems from model crawler - **Validation Performance** (vs. oracle): - Mean Efficiency: **96.1%** - Median Efficiency: **99.2%** - P10 Efficiency: **89.4%** ## Submission Checklist - [ x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Vidyasagar Ananthan <vidyasagar.ananthan@amd.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Jan Patrick Lehr <JanPatrick.Lehr@amd.com>
[ck_tile][fmha_bwd] Fix sink_host OOB in group mode reference runner (#7272)
## Summary
In `fmha_bwd_runner.hpp`, the `sink_host` `HostTensor` is allocated with
first
dimension `shape_batch` (= 1 in group mode), but the reference forward
loop
accesses `sink_host(wb, i_h)` with `wb ∈ [0, batch-1]`. For any `wb >=
1` this
is an out-of-bounds heap read, silently corrupting the reference forward
math
chain (`lse_host`, `o_host`) and turning the bwd-side `d_sink_head_acc`
reference into non-deterministic garbage.
`HostTensor::operator()` does not bounds check, so the OOB is not caught
at
runtime. This manifests as intermittent `tile_example_fmha_bwd` failures
(25–67% fail rate) when `-sink_grad=1` is combined with `-mode=1` (group
mode),
with bit-exact but spurious `max_err` values like 4.27 / 14.6.
## Fix
One-line: allocate `sink_host` with `batch` (the real per-batch dim)
instead of
`shape_batch`, mirroring how `sink_host` is accessed by the loop.
```diff
- sink_grad ? std::array<ck_tile::index_t, 2>{shape_batch, nhead}
+ sink_grad ? std::array<ck_tile::index_t, 2>{batch, nhead}
Repro
tile_example_fmha_bwd -b=2 -h=2 -s=516 -s_k=253 -prec=bf16 -d=72 \
-bias=n -dbias=0 -p_drop=0 -iperm=1 -operm=1 -deterministic=0 \
-v=3 -mode=1 -kname=1 -sink_grad=1
Verification
- 0/30 fail on the repro config after fix
- Baselines (before fix):
- sink=1, mask=n: 25% fail rate (p ≈ 1.8e-4)
- sink=1, mask=t: 67% fail rate (p ≈ 6e-15)
Attribution
Shape bug introduced together with sink_grad in #5504. Unrelated to
#6914
(which is a fwd-only fix on a different code path)
```
## Submission Checklist
- [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
---------
Signed-off-by: junlin12 <junlin12@amd.com>
Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com>
Skip numeric drop-out when PComputeWindow is a null_tile_window in Bl… (#7256) The BlockDropout implementation already provides very complete logic for generating random numbers and executing dropout for the P tensor after first attention Gemm with capability to support both Warp-Gemm 32x32 and 16x16 as well as to run on both wave32 and wave64 arch. But in some situation, we only need the block-layer process to generate random numbers, rather than simultaneously execute dropout in real-time on the vgpr tile. For example, xformers' `test_mem_eff_attention.py::test_dropout_ck` requires the host reference implementation of `attention forward with dropout` to use the same random numbers to compare & verify the device side implementation of `attention forward with dropout`, so a standalone kernel to generate random numbers only is required. This PR will enable xformers's random_val generating kernel (in file `ck_tiled_rand_uniform_kernel.h`) to depend on BlockDropout's `Run()` operator completely to generate random numbers for a `[MPerBlock, NPerBlock]` tile during the tile iteration, no need to replicate the logic of BlockDropout in the xformers kernel
Remove batch_prefill from FMHA_FWD_KNOWN_APIS (#6983) Remove `batch_prefill` from the `FMHA_FWD_KNOWN_APIS` list in `projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt`. **Change:** ```cmake # Before set(FMHA_FWD_KNOWN_APIS "fwd;fwd_splitkv;fwd_appendkv;pagedkv_prefill;batch_prefill") # After set(FMHA_FWD_KNOWN_APIS "fwd;fwd_splitkv;fwd_appendkv;pagedkv_prefill") ``` Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: asleepzzz <4926646+asleepzzz@users.noreply.github.com> Co-authored-by: asleepzzz <hanwen.chang@amd.com> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
[CK_TILE] Add async workspace prepare to FMHA BWD launcher (#7331) ## Motivation `aiter::mha_bwd` in group mode currently issues two synchronous `hipMemcpy` D2H copies to read `seqstart_q/k` for launcher construction. These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the stream, breaking CPU/GPU overlap on hot training paths. This PR adds a fully stream-async workspace preparation path on the FMHA BWD launcher so callers can pre-allocate the device workspace from upper-bound shapes and stage seqstart-dependent metadata via D2H/host-pack/H2D entirely on the user's stream. ## Technical Details - `FmhaBwdWorkspaceManager::GetWorkspaceDeviceSizeUpperBound` (`include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp`): computes the worst-case device dq_acc size from `(max_batch, hdim_q, nhead_q, max_seqlen_q, max_seqlen_k)` without dereferencing any seqstart array. Mirrors `PrepareWorkspaceHost`'s return value with worst-case bounds. - `fmha_bwd_launcher::prepare_workspace_async` (`example/ck_tile/01_fmha/fmha_bwd.hpp`): on the caller's stream, in order: 1. `hipMemsetAsync` of the dq_acc region (when `NeedsZeroDqAcc()`) 2. group mode: `hipMemcpyAsync` D2H of `seqstart_q/k` into a pinned host staging buffer 3. `hipLaunchHostFunc` runs `PrepareWorkspaceHost` on the pinned buffer 4. `hipMemcpyAsync` H2D of the packed metadata into `device_ws_ptr` The pinned staging buffer is held via `std::shared_ptr<void>` returned by a caller-provided `pinned_host_alloc` callback. Lifetime is extended past stream completion by a tail `hipLaunchHostFunc` scheduled in the launcher's destructor. - `ck_tile::pinned_host_releaser` (`include/ck_tile/host/pinned_host_releaser.hpp`): worker-thread utility for callers using bare `hipHostMalloc`. Defers `hipHostFree` off the HIP driver callback thread, which holds runtime locks and would deadlock against concurrent main-thread `hipFree`. PyTorch's `CachingHostAllocator` does not need this. - Example runner (`example/ck_tile/01_fmha/fmha_bwd_runner.hpp`): switched to the async path. ## Test Plan - `tile_example_fmha_bwd` (gfx950, dev preset `-Werror -Weverything`): - batch + nondet / batch + det / group + nondet / group + det - group + det 4-batch varlen (`-b=4 -h=8 -s=4096,3072,2048,1024 -d=128`) - FA (`flash-attention`) integration on ROCm 7.1.1 + PyTorch 2.9.1: - `tests/test_flash_attn_ck.py::test_flash_attn_varlen_deterministic` - `tests/test_flash_attn_ck.py::test_flash_attn_bwd_varlen_seqq_zero` ## Test Result - All CK runner cases `valid:y`. - FA pytest: **1952 passed in 44.82s**. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Added custom FMHA codegen receipt for TransformerEngine (#6867) ## Motivation TE uses AITER to build static MHA libraries, which ultimately rely on CK kernels. We use the `600` receipt which generates more kernels than TE truly needs. This bespoke receipt allows us to minimize the kernel count, compile time, and memory footprint of our MHA library. ## Technical Details Extended the receipt mechanism to include a custom `700` receipt for TE's needs ## Test Plan Test by building TE using the same receipt profile ## Test Result Build validated in TE using a custom feature branches of AITER/CK to temporarily apply the patch ## Submission Checklist - [ ] 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> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
[CK] Add rocm_ck directory structure with feature flag (#7090) ## Summary Adds initial rocm_ck directory structure, #7119. - Establishes production `rocm_ck/` directory at `composablekernel/rocm_ck/`, peer to `tile_engine/` and `dispatcher/` - Adds `CK_ENABLE_ROCM_CK` option (default OFF) as a CK-internal feature flag — no superbuild or TheRock changes needed - Creates `rocm_ck` INTERFACE library, `ck_tile_headers` target, GTest integration with builder-style convenience targets (`smoke-rocm-ck`, `check-rocm-ck`) - Adds Jenkins `RUN_ROCM_CK_TESTS` parameter for CI, following the `RUN_BUILDER_TESTS` pattern - README explains the constexpr schema model: host-device separation via constexpr data rather than template parameters, enabling multi-arch distribution through kpack archives ## Test plan - [x] `cmake -DCK_ENABLE_ROCM_CK=ON` configures without errors - [x] `ninja check-rocm-ck` passes (4 host-only index type tests) - [x] Default build (`CK_ENABLE_ROCM_CK=OFF`) is unaffected — no rocm_ck targets present - [x] Jenkins `RUN_ROCM_CK_TESTS=true` enables the flag and runs `check-rocm-ck` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com>
[CK] Fix smart build false positives from merged commits (#7289) ## Motivation Current smart-build infrastructure triggers full build for almost every PR which is draining our CI infrastructure. Need to update the test selection logic based on diffs from the current workspace instead of entire repo. ## Technical Details Use three-dot syntax and scope BUILD_INFRA_PATTERN to composablekernel. Changes: - Switch from two-dot (..) to three-dot (...) in git diff - Three-dot shows only PR-specific changes - Excludes commits merged from develop (prevents false positives) - Scope BUILD_INFRA_PATTERN to projects/composablekernel/ paths only - Avoids triggering on other projects (hipblas, hipdnn, etc.) - Only composablekernel build infra changes trigger full build - Update both ci_safety_check.sh and validate_pr.sh ## Test Plan Test with PR 7112 and 7223 ## Test Result Impact: - PR 7112: Was 620 files (false positive) → Now 6 files (correct) - PR 7223: Was full build (false positive) → Now selective build (correct) ## Submission Checklist - [ x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
feat(ck-tile): multi-D GEMM TE to dispatcher bridge
ISSUE ID: #8997
## Motivation
The TileEngine → Dispatcher bridge had no path for the **gemm_multi_d**
op, which
fuses one or more extra D operands into the GEMM epilogue
(`E = elementwise_op(A@B, D0, D1, ...)`). This is a real Old-TE
capability used for
fused bias/residual-style epilogues with no dispatcher equivalent, so
this PR adds a
complete bridge so the dispatcher can generate, build, and launch
multi_d at parity
with the legacy Tile Engine version.
The capability set matches the Old-TE `gemm_multi_d_instance_builder.py`
exactly:
`fp16`, the 4-char layouts `{rcrr, rrrr, ccrr, crrr}` (A/B vary, C and D
row-major),
the element-wise ops `{MultiDAdd, MultiDMultiply, PassThrough}`, and a
swept number of
D tensors (1 and 2). It follows the registry-bypass bridge pattern used
by the grouped
(#9000) and stream-K (#9028) bridges.
## Test Plan
- Run the CPU-only unit tests (no GPU required):
`python3 -m pytest dispatcher/tests/test_multi_d_bridge.py -v`
- On-GPU numeric verify over the full capability matrix
(fp16 × {rcrr, rrrr, ccrr, crrr} × {MultiDAdd, MultiDMultiply} × {num_d
1, 2} = 16
combos) at M=N=K=1024 against an fp32 reference, gate 2e-2.
- Confirm the CI config builds real kernels and the sweep covers all ops
× D counts.
## Test Result
- CPU-only unit tests pass (10 passed).
- On-GPU numeric verify: 16/16 combos pass at M=N=K=1024, worst-case
`max_rel = 6.16e-4` (~30x under the 2e-2 gate). Col-major `ccrr` /
`crrr` have real
on-GPU numeric evidence.
- CI config now builds real kernels (was zero); the sweep expands evenly
across
`{MultiDAdd, MultiDMultiply} × {num_d 1, 2}` per layout.
- clang-format (18.1.8) clean on `multi_d_gemm_ctypes_lib.cpp`.
- Serialized A/B perf-parity vs Old-TE (MI300X / gfx942, fp16, 4 layouts
× 2 ops ×
num_d=1 = 8 stems × 5 shapes = 40 rows, interleaved, fair
50/100/flush/rotating
both sides): **at parity, bridge consistently faster** — median gap
+9.44%, 100%
within ±15% (range [+3.76%, +14.99%]; positive = bridge faster, from the
registry-bypass direct launch avoiding the Old-TE profiler's per-call
overhead).
num_d=1 is the fair slice since the Old-TE `gemm_multi_d` benchmark is
single-D.
See the parity comment for details.
---
**Related PRs / references (TileEngine → Dispatcher GEMM bridge
series):** #8997 (regular GEMM fp16/bf16 all-layout), #9000 (grouped
GEMM), #9028 (stream-K), #8887 (fp8/bf8/int8). This PR is a sibling in
the same bridge effort tracked across those PRs.
---------
Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
fix(ck): Stream-K Tile Engine GPU Query Fix ## Motivation The Stream-K tile engine validation utilities call rocminfo at build time to detect the GPU architecture, which prevents building on CPU-only nodes. The GPU target is already known from CMake's SUPPORTED_GPU_TARGETS, so runtime hardware detection is unnecessary during code generation. ## Technical Details The gemm_streamk_validation_utils.py file used subprocess.check_output(["rocminfo"]) to query GPU hardware at CMake configure time. This call originated from get_gpu_name_by_id() and used during tile configuration validation. On CPU-only build nodes, this fails because rocminfo either doesn't exist or returns no GPU devices. The main changes are as follows: - Removed runtime GPU detection infrastructure: deleted get_gpu_name_by_id(), set_gpu_targets(), get_configured_gpu_targets(), the _configured_gpu_targets module variable, and the GPU_NAME_PATTERN regex. - Added gpu_target as an explicit parameter. is_tile_config_valid(), validate_gemm(), validate_warp_tile_combination(), validate_warp_configuration(), and validate_lds_capacity() now accept gpu_target as a required parameter instead of querying hardware internally. The corresponding changes were also made in gemm_streamk_instance_builder.py and CMakeLists.txt - Added WARP_SUPPORTED_COMBINATIONS for per-GPU warp config validation, and LDS_SIZE_MAP / DEFAULT_LDS_SIZE for GPU-aware LDS capacity checks. ## Test Plan The benchmarks were compiled and run on a GPU as well as a CPU only node to verify correctness. ## Test Result All tests passed ## Related JIRA ID : AICK-1635 ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Maksim (Max) Podkorytov <Maksim.Podkorytov@amd.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
fix(ck): correct GU-fusion B-up vector loading ## Motivation Updating latest CK in Aiter will cause ATOM Qwen tests failure. Bisect CK commit narrow down to the changes in ROCm/rocm-libraries#4798. JIRA ID AICK-1709 ## Technical Details The fix is to pass b_thread_vec_up instead of b_thread_vec in thread_buf_to_vec_loader for LoadBUp. ## Test Plan default CK CI test manual testing for ATOM Qwen tests. ## Test Result default CK CI test pass ATOM Qwen tests pass ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Users/mkulikow/ck/data prefetch in mxgemm pipeline JIRA ID : AICK-1670 ## Motivation Added example for data cache prefetch in mx gemm pipeline while also fixing some bugs in data cache prefetch pipeline ## Technical Details Add a standalone flatmm example (mx_flatmm_data_cache_prefetch) that runs MX GEMM through the compute TDM v1 pipeline (GemmPipelineAgBgCrCompTDMV1) with hardware data cache prefetch on gfx1250. Prefetch destination is selectable per operand (A/B) between L1, L2 or None via the DataCachePrefetchKind trait, exposed through -prefetch_a_l1 / -prefetch_b_l1 CLI flags, with an optional -compare mode against a no-prefetch run. Guarded behind gfx125 in the 18_flatmm CMakeLists. Also fix data cache prefetch being silently disabled in the scaled operator() of GemmPipelineAgBgCrCompTDMV1. The scaled path defaulted data_cache_prefetch_a/b to false and only set them under UseClusterLaunch, so with cluster launch off the runtime guards folded away every prefetch and no global_prefetch_b8 was emitted despite an L1/L2 policy. Default them to true (matching the non-scaled operator()); emission stays gated by the compile-time UseDataCachePrefetch policy check, so None still emits nothing. ## Test Plan Checked on simulators: test name: tile_example_mx_flatmm_mxgemm_data_cache_prefetch ## Test Result fp4 for 512x512x4096: no prefetch/L1 prefetch: 71,334 / 54,889 ( 23.1% increase ) ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
fix(ck-tile): Prefer using amd-smi over rocm-smi ## Motivation CK tooling invokes `rocm-smi` directly in many places for GPU discovery and info. `amd-smi` is the modern replacement and is preferred going forward, but call sites should not need to know which tool is present. This PR centralizes GPU SMI access behind a single set of wrappers that prefer `amd-smi` and transparently fall back to `rocm-smi`, so every call site gets consistent behavior with no duplicated detection logic. ## Technical Details - Added `tile_engine/ops/common/smi_utils.py` as the single source of truth for SMI access: parsers for both tools plus wrappers `detect_gpu_ids()`, `count_gpus()`, `show_gpu_info()`, `check_gpu_available()`, and `show_version()`. - Tool order is chosen by `_smi_order()`: `amd-smi` first, `rocm-smi` fallback if `amd-smi` is missing or fails. `CK_SMI_TOOL=rocm-smi` forces rocm-smi first (mainly for testing). - Added `tile_engine/ops/common/smi_cli.py`, an argparse CLI (`list-ids`, `count`, `show-info`, `check`, `show-version`) so shell scripts can reuse the same Python logic. - Added thin `ck_smi_*` delegates in `script/tools/common.sh` that call `smi_cli.py` (no parsing logic in Bash). - Migrated call sites off direct `rocm-smi`: `gemm_full_benchmark.py` now uses `detect_gpu_ids()`; `generate_test_dataset.sh`, `ck-status`, `ck-start`, `ck-docker`, `ck-exec`, `ck-shell`, `ck-rocprof.md`, and the GEMM `README.md` now use the wrappers. ## Test Plan - `python3 -m unittest tile_engine.ops.common.test_smi_utils -v` live tests comparing `rocm-smi` vs `amd-smi` normalized fields (GPU IDs, product, gfx, driver) and verifying each wrapper against live output. - `CK_SMI_TOOL=rocm-smi python3 -m unittest tile_engine.ops.common.test_smi_utils -v` verifies the rocm-smi override path. - `bash script/tools/test_ck_smi_helpers.sh` pure-bash live comparison of the two tools' fields (no Python dependency). - Ran on a GPU host. ## Test Result All 11 unittest cases pass; the bash comparison reports all fields matching. Both the default (amd-smi first) and `CK_SMI_TOOL=rocm-smi` paths return identical GPU IDs. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. JIRA ID : AICK-1649 --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
fix(ck-tile): repair #9308 merge truncations in gemm_utils + unified_gemm_codegen (develop broken)
ISSUE ID: #9308
Fixes the develop breakage introduced by the #9308 multi-D merge (commit
7fcb5f36). When I was fixing merge conflict, some line was deleted and
CI did not test it ( bridge test was not added CI that time). These two
issue has been fixed with this PR.
## Summary
The multi-D merge **#9308** (commit `7fcb5f36`) left **two independent
truncations** on `develop`, both of which make the dispatcher Python
fail to parse:
1. `dispatcher/python/gemm_utils.py` — `import gemm_utils` raises
`SyntaxError: '(' was never closed`.
2. `dispatcher/codegen/unified_gemm_codegen.py` — `SyntaxError:
unterminated string literal` (multi_d / multi_abd codegen can't run).
Both verified against the GitHub `develop` blob. Every consumer of the
multi-D / multi-ABD Python paths is broken today.
## Root cause
- **gemm_utils.py:** the multi-D `run()` landed **inside
`GpuMultiABDRunner`** and was truncated at an unterminated `return
MultiDGemmResult(`; `GpuMultiDGemmRunner` was left with only `__init__`.
- **unified_gemm_codegen.py:** `_multi_d_single_include()` returns an
f-string of C++ `#define`/exports whose closing `"""` was dropped.
## Why CI stayed green
The dispatcher Python tests that import these modules were **not
registered in `dispatcher/tests/CMakeLists.txt`**, so they never ran in
the gate.
## Changes
- **gemm_utils.py:** move multi-D `run()` (+
`kernel_name`/`num_d_tensors`) into `GpuMultiDGemmRunner` and complete
the `MultiDGemmResult(...)` return; remove the stray block from
`GpuMultiABDRunner`.
- **unified_gemm_codegen.py:** close the truncated f-string in
`_multi_d_single_include`.
- **Tests + CI:** add `TestModuleImportsAndRunnerShape` to
`test_gemm_utils.py` (import canary + multi-D runner-shape assertions +
`ast.parse` canary over `unified_gemm_codegen.py`), and **register
`test_gemm_utils.py` in CMake (`dispatcher_test_gemm_utils`)** so
ctest/CI runs it.
## Test plan
- [x] `import gemm_utils` succeeds (was SyntaxError)
- [x] `ast.parse` of `unified_gemm_codegen.py` succeeds (was
SyntaxError)
- [x] `python3 -m unittest discover -p test_gemm_utils.py` -> 21 passed
- [ ] CI runs `dispatcher_test_gemm_utils` (newly wired)
## Follow-up (separate)
Enable the dispatcher Python test suite broadly in the PR CI gate to
fully close the coverage gap.
---------
Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
[CK_TILE] Retune RDNA FMHA D128 tile selection after pipeline update ## Motivation The FMHA forward pipeline changed in [5961a2e](ROCm/rocm-libraries@5961a2e), shifting the optimal D128 tile from `64x64` to `128x64`. This PR fixes the performance regression on RDNA for sequence lengths ≤ 4096 introduced in #6498. ## Technical Details - Remove the `64x64` D128 tile and its sequence-length constraint from the GFX11 and GFX12 FMHA forward kernel tables. - Use `128x64` for D128 at all sequence lengths. - GFX115 targets inherit the updated GFX11 tile selection. - This only changes kernel selection and does not modify the pipeline's numerical behavior. ## Test Plan - Compared the baseline and updated tile selection on `gfx1100`, `gfx1151`, and `gfx1201`. - Tested BF16 BSHD forward with `B=1`, `H=24`, `hdim=128`, batch/group APIs, and sequence lengths 1024 and 4096. ## Test Result The table shows candidate throughput relative to baseline. | GPU | Mode | L=1024 | L=4096 | |---|---|---:|---:| | gfx1100 | batch | +10.99% | +2.90% | | gfx1100 | group | +10.69% | -2.22% | | gfx1151 | batch | +10.57% | +0.59% | | gfx1151 | group | +10.71% | -0.23% | | gfx1201 | batch | +12.75% | +11.42% | | gfx1201 | group | +12.88% | +11.41% | On `gfx1100` and `gfx1151`, L=4096 already uses `128x64` in both versions, so the differences represent same-kernel measurement variation. On `gfx1201`, both tested sequence lengths change from `64x64` to `128x64` and improve by approximately 11–13%. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
feat(ck): add SPIR-V target support Add initial support for the amdgcnspirv target in Composable Kernel, enabling target-agnostic SPIR-V compilation that is JIT-compiled to native code at runtime JIRA ID - ROCM-27738 ## Motivation Enable Composable Kernel to compile with the SPIR-V target, producing target-agnostic SPIR-V binaries that are JIT-compiled to native GPU ISA at runtime via comgr. ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan - Built with GPU_TARGETS=gfx90a;amdgcnspirv - Compared wall time (native vs SPIR-V), cold/hot 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. --------- Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> Co-authored-by: illsilin_amdeng <Illia.Silin@amd.com>
fix(ck): [CK] LCOPMILER-2487: Remove erroneous `__restrict__` qualifier ## Motivation <!-- Explain the purpose of this PR and the goals it aims to achieve. --> Memory referenced via `__restrict__`-qualified pointers may not be modified in any way other than through said pointer so long as that pointer is alive. This is ambiguated in the context of multiple threads, but insofar as it is modeled by `noalias` in LLVM, accessing memory through a `__restrict__` pointer in one thread after having it modified by another thread violates this contract. JIRA ID: https://amd-hub.atlassian.net/browse/LCOMPILER-2487 ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> This is analogous to #9629. The semantics for LLVM's `noalias` between threads was clarified in llvm/llvm-project#211507 to also prohibit modifications through the same pointer from other threads. Unless `__restrict__` adopts a weaker guarantee in the future, `p_shared_block` is in violation of this contract. ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> Build and run: `test_gemm_splitk` ## Test Result <!-- Briefly summarize test outcomes. --> Before: ``` [----------] Global test environment tear-down [==========] 32 tests from 8 test suites ran. (77718 ms total) [ PASSED ] 29 tests. [ FAILED ] 3 tests, listed below: [ FAILED ] TestGemmSplitK_MK_NK/0.SmallM, where TypeParam = std::tuple<_Float16,_Float16,_Float16> [ FAILED ] TestGemmSplitK_MK_NK/0.MidLargeM, where TypeParam = std::tuple<_Float16,_Float16,_Float16> [ FAILED ] TestGemmSplitK_MK_NK/0.Regular, where TypeParam = std::tuple<_Float16,_Float16,_Float16> 3 FAILED TESTS ``` With this patch, all tests pass. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
fix(ck-tile): prefer amd-smi over rocm-smi in remaining TE GEMM benchmarks ## Motivation PR #9731 centralized GPU SMI access behind `tile_engine/ops/common/smi_utils.py`, which prefers `amd-smi` and transparently falls back to `rocm-smi`, and migrated `gemm_full_benchmark.py` off direct `rocm-smi` calls. Three Tile Engine benchmark scripts were left calling `rocm-smi --showid` directly (with an `amd-smi list` fallback), so the migration is incomplete. Although wrapped in try/except today, these direct calls will break once `rocm-smi` is removed from future ROCm releases. This completes the migration for the remaining scripts, as requested by the tech lead and tracked in JIRA **ROCM-28734**. ## Technical Details Migrated the identical hand-rolled `detect_devices()` (env → `rocm-smi --showid` → `amd-smi list` → hardcoded `["0"]`) to a single `smi_utils.detect_gpu_ids()` call, mirroring the already-landed `gemm_full_benchmark.py`: - `tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_full_benchmark.py` - `tile_engine/ops/gemm/gemm_multi_d_full_benchmark.py` - `tile_engine/ops/gemm/streamk_gemm_full_benchmark.py` Each script adds `common/` to `sys.path`, imports `detect_gpu_ids`, and drops the now-unused `import re`. Net -67 lines. ## Test Plan / Result Validated on an MI350 (gfx950) host, exercising the exact `sys.path` arithmetic + `smi_utils` import + `detect_gpu_ids()` for both the shallow (`gemm/`) and deep (`block_scale_gemm/gemm_bquant/`) script locations: - Default (amd-smi first): `detect_devices()` returns the live GPU ids; `count_gpus()`/`check_gpu_available()` consistent. - `CK_SMI_TOOL=rocm-smi` (forced fallback): identical result. Both `python3 -m py_compile` cleanly. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. https://amd-hub.atlassian.net/jira/software/c/projects/AICK/boards/4399?selectedIssue=ROCM-28734 JIRA ID : ROCM-28734 Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
feat(ck-tile): add AQuant and ABQuant grouped GEMM dispatcher with ctypes bridge ## Motivation This PR extends the CK tile dispatcher with full quantization coverage for grouped GEMM operations. Building on the BQuant dispatcher introduced in #9166, this PR adds dispatcher bridges for AQuant (A-side quantization) and ABQuant (both sides quantized), completing the three-mode quantization dispatch surface for grouped GEMM on gfx950 (MI350X). ## Technical Details Stacked on #9166 (BQuant grouped GEMM dispatcher). Adds AQuant (C = dequant(A, AQ) @ B) and ABQuant (C = dequant(A, AQ) @ dequant(B, BQ)) dispatcher bridges following the same three-layer pattern (ctypes C API → Python dispatcher → codegen). Each mode adds: *_ctypes_lib.cpp — C API exposing the kernel to Python via ctypes unified_*_codegen.py — generates kernel instantiations across the full dtype/layout/quant-group matrix *_utils.py — Python dispatcher handling argument validation, dtype selection, and parallel hipcc builds via ThreadPoolExecutor Example script and CPU-only unit tests codegen_common.py is extended with shared epilogue selection helpers reused across all three quant modes. CMakeLists.txt updated with build targets for both new ctypes shared libraries. Supported dtypes: fp8/bf8 activations, fp8i4/bf8i4 packed int4, preshuffle layouts, and MX microscaling variants — targeting gfx950 (MI350X). ## Test Plan CPU-only unit tests covering: Kernel name generation across all dtype/layout/quant-group combinations Config serialization and dimension helpers Epilogue selection logic for AQuant and ABQuant variants End-to-end example scripts (14_grouped_gemm_aquant.py, 15_grouped_gemm_abquant.py) with CPU reference verification for correctness checks. ## Test Result CPU unit tests pass for both AQuant and ABQuant dispatcher utilities. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.com>
feat(ck): Added Gelu with Tanh approx to XDL 2-stage MoE epilogue ## Motivation Enable the tanh-approximation GELU activation `(gelu_tanh, 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))))` in the Composable Kernel XDL 2-stage MoE path. The MoE gridwise kernel epilogue currently supports only `silu/gelu/swiglustep/swiglu_oai`; this adds `gelu_tanh_and_mul` so models whose MoE experts use the GELU tanh approximation (e.g. [Gemma-family MoE](https://huggingface.co/google/gemma-4-26B-A4B/blob/main/config.json)) can use this path. JIRA ID : ROCM-27619 ## Technical Details - `gridwise_gemm_xdl_cshuffle_common.hpp`: add `Activation::gelu_tanh_and_mul = 4` to the activation enum. - `gridwise_moe_gemm.hpp`, `gridwise_moe_gemm_blockscale.hpp`, `gridwise_moe_mx_gemm_<>.hpp`: wire `gelu_tanh_and_mul` into epilogue paths, delegating to the existing `ck::tensor_operation::element_wise::FastGelu` helper (the single source of truth for the tanh-GELU math, `FastGelu(gate) * up`). Also added `static_assert` for validation of supported activations - The activation is applied in fp32 in the epilogue and is orthogonal to the GEMM compute (MFMA/tile/pipeline untouched) and to quantization (existing per-token dequant reused). Only the non-blockscale gridwise kernel is changed. - Then I plan to port these changes to AITER after ROCm/aiter#3886 to avoid merge conflicts ## Test Plan Use `ActOP = 4` in the example `moe_gemm1_xdl_fp8`, rebuild example and launch ctest ## Test Result ``` ctest -R "^example_moe_gemm1_xdl_fp8$" -V' Constructing a list of tests Done constructing a list of tests Updating test list for fixtures Added 0 tests to meet fixture requirements Checking test dependency graph... Checking test dependency graph end test 257 Start 257: example_moe_gemm1_xdl_fp8 257: Test command: example_moe_gemm1_xdl_fp8 257: Working Directory: example/65_gemm_multiply_multiply 257: Test timeout computed to be: 1500 257: a0_t_k: dim 2, lengths {16384, 6144}, strides {6144, 1} 257: b0_e_n_k: dim 3, lengths {8, 6144, 8192}, strides {50331648, 1, 6144} 257: d1_e_n: dim 2, lengths {8, 8192}, strides {8192, 1} 257: d2_e_n: dim 2, lengths {32768, 4096}, strides {1, 0} 257: d0_t_n: dim 2, lengths {16384, 4096}, strides {1, 16384} 257: d2_e_n: dim 2, lengths {32768, 4096}, strides {1, 0} 257: e_t_n: dim 3, lengths {16384, 2, 4096}, strides {8192, 4096, 1} 1/1 Test #257: example_moe_gemm1_xdl_fp8 ........ Passed 83.51 sec The following tests passed: example_moe_gemm1_xdl_fp8 100% tests passed, 0 tests failed out of 1 Label Time Summary: SMOKE_TEST = 83.51 sec*proc (1 test) Total Test time (real) = 83.59 sec ``` ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
fix(ck): downstream library pipeline docker name override ## Motivation Jira ID AICK-1834 The CI job triggers multiple downstream library pipelines that share a common variable for Docker image names. The failure occurred because the Docker image name was overwritten by another pipeline before the docker pull step executed. ## Technical Details Add def to make the name variable "retimage" local) ## Test Plan See if this failure happens again in daily scheduled Aiter CI pipeline after the fix.
fix(CK): Forward port convolutions update from standalone repo ## Motivation Forward-port grouped-conv-fwd perf-regression fix (im2col GEMM-size 2GB check) to rocm-libraries mainline ## Technical Details Port this commit 248f155#diff-d921922979d6cfc604769c3bc231107ae574375875ecce04bcb1065d43a5ce2e Introduced by a mainline change for grouped-conv large-tensor / global-load-store CK enablement, which added a LargeTensors template flag to GridwiseGemmMultiD_xdl_cshuffle_v3 and reworked its CheckValidity, but did not exempt the convolution im2col-view path from the pre-existing logical-GEMM-size 2GB check ## Test Plan Regular CI ## Test Result CI should pass ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. JIRA ID : AICK-1818
Update CODEOWNERS ## Motivation Add Jia Luo as a code owners, remove Thomas Ning from the list. JIRA ID : AICK-1882 ## Technical Details N/A ## Test Plan N/A ## Test Result N/A ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
fix(ck): route scattered page_size=1 paged-KV batch-prefill to GLOBAL_LOAD_LDS page_size=1 paged-KV batch-prefill faults on gfx950 with the BUFFER_LOAD gather even for in-bounds addresses. Route page_block_size == 1 to the 64-bit-safe GLOBAL_LOAD_LDS path; #9214's fast path for 1 < page_block_size < kN0 is unchanged. Validated on MiniMax-M3-MXFP4 TP=4 (gfx950): 3000/3000 requests, zero faults. Root cause of the negative BUFFER_LOAD offset is still open, tracked separately. Override rationale: the only red check is the sles16 RPM install lane, which is broken repo-wide and unrelated to this change. Tracking: ROCm/TheRock#7161.
feat(ck) [CK] Wavelet gemm pipeline for conv fwd
## Motivation
In the current CShuffleV3 conv fwd kernel, the in-kernel conv-to-GEMM
transform generates significant INT32 VALU pressure per MFMA
instruction. On VALU-heavy shapes (e.g., G=1, 3×3, C=256), these index
computation ops compete with MFMA for VALU issue slots, creating a
bottleneck that cannot be resolved by pipeline prefetching alone.
This PR adds a wave-specialized ("wavelet") convolutions forward kernel
that splits workgroup threads into two roles:
- **Load waves**: conv-to-GEMM address computation + global memory loads
+ LDS writes (all VALU/VMEM)
- **Math waves**: LDS reads + MFMA + CShuffle epilogue (no index
computation)
By physically separating the two instruction classes onto different
waves, VALU and MFMA execute on different hardware functional units
without contention.
## Technical Details
**Wave pipeline (modified):**
- `gridwise_gemm_waveletmodel.hpp` — load/math wave pipeline structs
with `sched_group_barrier` scheduling hints to front-load VMEM reads
before address-advance VALU
**Two wave ratios:**
- **(4,4)**: 256 load + 256 math = 512 threads (8 waves). Best on large
shapes.
- **(4,2)**: 256 load + 128 math = 384 threads (6 waves). Best on small
shapes (fewer sync barriers, denser MFMA per math wave).
JIRA ID : ROCM-21620
---------
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Add build # to workspace path. ## Motivation Since we recently lost VMs running on gfx950 machines, building CI jobs in parallel on the same machine for the same branch (e.g. develop) causes those jobs to delete each other's workspace folders, which is not ideal. In order to prevent this from happening, we can add the build # to the workspace path. This way we can be sure each workspace path is unique. JIRA ID : AICK-1920 ## 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.
fix(ck-tile): prefer amd-smi for dispatcher GPU-arch detection
## Summary
Extends the amd-smi-first migration (#9731, #10204) from the Tile Engine
to the
dispatcher's GPU-architecture detection.
- Adds `detect_gpu_arch()` to the shared
`tile_engine/ops/common/smi_utils.py`
wrapper: amd-smi `static` (`TARGET_GRAPHICS_VERSION`) first, rocm-smi
`--showproductname` (`GFX Version`) fallback — same amd-smi-first policy
as the
rest of the wrapper.
- `dispatcher_common.detect_gpu_arch()`: prefers amd-smi via the
wrapper, then
rocminfo, then the caller-supplied default. Also removes a **duplicate**
definition of this function.
- `gemm_utils._get_arch()`: tries amd-smi via the wrapper before
rocminfo,
preserving the raise-on-unresolved (no silent default) behavior.
## Motivation
Per the directive behind #9731, device-property access should prefer
`amd-smi` (rocm-smi is being phased out). The Tile Engine was already
migrated
(#9731, #10204); the dispatcher still detected the GPU arch by shelling
out to
`rocminfo` directly. Tracked in JIRA **ROCM-28734**.
## Design note
The dispatcher reaches the sibling `tile_engine/ops/common/smi_utils.py`
through
a guarded import and falls back to `rocminfo` if the wrapper is
unavailable, so
behavior is unchanged on hosts without amd-smi. (Open to moving the
shared
wrapper to a neutral location if reviewers prefer avoiding the
dispatcher →
tile_engine reach.)
## Test plan (MI350X / gfx950)
- [x] New mock-based unit tests for `detect_gpu_arch()` (amd-smi
preference,
rocm-smi fallback, default) + a live check — 4/4 pass.
- [x] `dispatcher_common.detect_gpu_arch()` and `gemm_utils._get_arch()`
both
return `gfx950` via amd-smi on the live host.
- [x] `test_gemm_utils` 21/21 still pass (no arch-detection regression).
---------
Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
fix(ck): fix test_gemm_mx correctness failures on gfx1250 (LDS read drain)
JIRA ID : AICK-1890
## Summary
`test_gemm_mx` fails the correctness check on gfx1250 (A0). Five tests
produce sparse wrong values, non-deterministically, at roughly a 50% hit
rate per run:
```
TestGemmMX_MK_NK/0.Large f8 x f8 -> f16
TestGemmMX_MK_NK/1.Large f8 x f8 -> bf16
TestGemmMX_MK_NK/2.Large f4 x f4 -> f16
TestGemmMX_MK_NK/3.Large f6 x f6 -> f16
TestGemmMX_MK_NK/4.Large bf6 x bf6 -> bf16
```
The MX v3 pipeline double-buffers LDS. Each iteration reads one buffer
with `ds_load` while the hardware async copy
`global_load_async_to_lds_b128` fills the other. The barrier at the top
of `LoopFunc` is what keeps a wave from overwriting a buffer another
wave is still reading.
On gfx1250 that barrier is `block_sync_lds_async_load()`, which waits on
**ASYNCcnt**. LDS reads are tracked by **DScnt**, a different counter,
so the barrier does not wait for them. The compiler normally emits the
LDS wait itself, but in this hot loop it computes the weakest wait its
per-wave dependency analysis requires and sinks it *past* the barrier.
In the generated ISA a barrier retires with **28 `ds_load`s still
outstanding** (`s_wait_loadcnt_dscnt 0x11c`), immediately followed by an
async write into that same buffer.
Per-wave dependency analysis cannot see the cross-wave contract — that
the barrier exists so *other* waves may overwrite the buffer — so the
ordering has to be explicit in the source.
## The change
One line, in the existing gfx1250 arm of this one pipeline:
```cpp
#if defined(__gfx125__)
llvm_amdgcn_s_wait_dscnt(0);
block_sync_lds_async_load();
#else
```
A full drain the compiler cannot move past the barrier.
## Test plan
All runs on gfx1250 (ASIC rev 0x0).
- [x] `test_gemm_mx` full suite: **5 failing tests on every run → 0
failures, 3/3 runs**
- [x] Isolated instance, 20 repetitions per build, measured back to
back: **11/20 failures → 0/20**
- [x] `example_gemm_mx_fp8` (configured identically to the failing
instance): **10/10 incorrect → 0/10**
- [x] ISA verified: the added drain appears before the barrier at every
async site; `ds_load` count unchanged, so this is an ordering fix and
not a data-flow change
- [x] Non-gfx1250 paths untouched (change is inside `#if
defined(__gfx125__)`)
Reproducer, for anyone verifying:
```
./bin/example_gemm_mx_fp8 1 2 1 0 5120 5120 4096 4096 4096 5120 1 20 50
```
This fails on **every** run before the fix, which makes it a much better
regression gate than the test suite's ~50% hit rate.
## Performance
~1.6% throughput on the affected kernel (351.5 → 345.9 TFlops at
5120x5120x4096). The baseline is computing wrong answers, so this is the
cost of correctness rather than a regression against a working build.
[CK] [ci] upgrade to rocm7.14 and disable faulty daily pytorch/FA tests ## Motivation Upgrade to the latest ROCm 7.14 release as the default compiler in CI and temporarily disable the daily tests that had been broken due to issues outside of CK control. We can re-enable the tests once the underlying problems are resolved. JIRA ID: AICK-1955 ## 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.
feat(CK): Enable CK backend for Inductor on GFX1250
JIRA ID: AICK-1464
## Summary
Extends the `ck4inductor` enumerators so PyTorch Inductor's max-autotune
CK backend has
WMMA candidates on gfx1250, ships the headers those kernels need, adds a
header-resolution diagnostic, and fixes the wheel's version label.
Python-only; no CK C++ changes. gfx9 behaviour is unchanged — the new
enumerators are
additive and the existing XDL paths are byte-identical.
This is the CK half of the work. The corresponding PyTorch-side changes
(arch-aware
instance pruning, the opt-in `CKWMMA` backend tokens, and the CK pin
bump) ship
separately and depend on this landing first.
### WMMA instance enumeration
Adds a WMMA enumerator alongside the existing XDL one for all three op
families:
| family | device op |
|---|---|
| universal GEMM | `DeviceGemmMultiD_Wmma_CShuffleV3` |
| batched GEMM | `DeviceBatchedGemmMultiD_Wmma_CShuffleV3` |
| grouped conv fwd | `DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle_V3`
|
### Packaging
Two `package-data` additions in `pyproject.toml`:
- **CK-Tile headers** (`ck_tile/**/*.hpp`, `ck_tile/**/*.inc`). CK-Tile
kernels
previously compiled only if a separate ROCm CK install happened to
supply these; the
wheel now carries them.
- **Batched WMMA instance sources**
(`batched_gemm/**/*wmma_universal*.cpp`). Unlike
`gemm_universal`, the `batched_gemm` folder keeps its template-argument
tuples only in
the instance `.cpp` files, with no `.hpp` aggregator — so without this
glob the batched
WMMA enumerator returns zero instances from an installed wheel while
working fine from
a source tree.
### Header-resolution diagnostic
Adds `check_headers()` and `include_roots()`. These report which of
`ck/ck.hpp`,
`ck/config.h` and `ck_tile/core.hpp` are reachable, and from which root.
Note
`ck/config.h` is CMake-generated and supplied by `$ROCM_HOME` rather
than the wheel;
when it is missing, the failure otherwise surfaces as an opaque compile
error at
autotune time, long after the real cause.
### Wheel version
The wheel was labelled `7.1.1.dev912+gdec6e1`: an old release baseline
plus a
repo-wide commit count.
The label is now a manually maintained `BASE_VERSION` plus the commit
hash.
## Test plan
- [x] `pytest python/test/test_gen_instances.py` — 13/13. New coverage
asserts the WMMA
pools are non-empty, all warp-16x16, dtype in {F16, BF16} and Ds-free;
that the XDL
pools contain no WMMA ops; and the shape of `include_roots()` /
`check_headers()`
- [x] XDL enumeration unchanged: identical counts and sorted-alias
hashes across all six
enumerators
- [x] Out-of-tree packaging gate: wheel built, installed outside the
source tree, and
enumerated from the installed package — this is what catches the batched
`.cpp`
glob, since an in-tree run passes either way
- [x] Wheel version correct in a source checkout, a shallow tagless
clone, and a tree
with no `.git`; PEP 440-valid and satisfies `>=7.14`
- [x] gfx1250 hardware: full PyTorch `TestCKBackend` suite green (59
passed, 8 skipped,
0 failed), with WMMA kernels winning autotune for GEMM, batched GEMM and
conv
- [x] gfx942 / gfx950 regression run: every test keeps its name and
result
docs(hipcub, hipthreads, rocprim, rocthrust, ck, rpp) added rocm theme info to the conf.py files
## Motivation
New buttons are needed at the top of the component documentation landing
pages. This PR adds them to hipcub, hipthreads, rocthrust, rocprim, ck,
and rpp
JIRA ID : AIROCDOC-4239
Co-authored-by: {spolifroni-amd} <{sandra.polifroni@amd.com}>
[GFX1250][CK_TILE] Coalesce MX scale16 scale load ## Problem On gfx1250 the MX `scale16` scale load was strided across lanes while `scale32` was coalesced. `scale16` used an identity host pre-shuffle plus a K-fastest `[packs_m, MThreadPerXdl, packs_k]` scale descriptor, so each lane's per-K-iteration offset was `lane * num_scale_k` (a separate cache line per lane, every K iteration). ## Fix Put `scale16` on the same **M/N-fastest** `[packs_m, packs_k, MThreadPerXdl]` descriptor as `scale32`, and fold both host pre-shuffles into **one formula** in `preShuffleScaleBuffer_gfx1250`. Each lane's offset becomes `lane * 1` (unit-dword) instead of `lane * num_scale_k`, so consecutive lanes hit consecutive addresses. This is a bijective transpose of both the data and the descriptor, so the scale values are unchanged. `scale32` is byte-for-byte unchanged. ### The layout is parameterized by `WarpTile::M`, not `ScaleBlockSize` This is the key point from the review of #8202. The single pre-shuffle formula keys off `MThreadPerXdl` = the WMMA `WarpTile` M (A scales) / N (B scales) — the number of lanes holding a distinct scale row per warp — **not** off `ScaleBlockSize`: - a 32×32 WMMA has 32 distinct scale lanes → `MThreadPerXdl = 32`; - a 16×16 WMMA has 16 (the other wavefront lanes replicate the same scale rows) → `MThreadPerXdl = 16`. Both `scale16` and `scale32` are instantiated with **both** 32×32 and 16×16 WMMA tiles, so the lane count — not the block size — selects the layout. Deriving it from `ScaleBlockSize` mislays the `scale16` configs that use a 32×32 `WarpTile`. `ScaleBlockSize` only changes how many `int32` K-packs a lane holds: `scale16` halves the K span per scale, so a lane holds twice as many packs and the **device** reads two adjacent `int32` as one `int64` — a device-side read concern. Host packing stays `int32` (`PackSize = 4`) for both block sizes. The callers pass `M_Warp_Tile` (A) / `N_Warp_Tile` (B) into the pre-shuffle from the pipeline and grouped-gemm MX tests. ## Validation **Correctness (measured).** `test_ck_tile_mx_gemm_pipeline_tdm_wmma` on gfx1250, full suite = 148 subtests over 37 configs: **147 / 148 pass** against the fp32 reference — every `scale16` config (`WarpTile` 16 and 32, TDM V1 and V2) and every `scale32` config. The lone failure, `/14.SmallM` (Col-major-A FP4×FP8 `scale32`), is **pre-existing and unrelated**: it reproduces identically on `develop` with this PR reverted. **Coalescing (by construction).** The per-lane scale offset going from `lane * num_scale_k` to `lane * 1` follows directly from the unified `[packs_m, packs_k, MThreadPerXdl]` descriptor (reviewer-verifiable from the diff); the host pre-shuffle was checked to reproduce this device descriptor layout index-for-index for both block sizes. **Performance (not measured).** Runtime throughput / bandwidth was not benchmarked here; this PR is validated for correctness and for the descriptor-level coalescing property, not for end-to-end kernel performance. --------- Co-authored-by: Andriy Roshchenko <107577548+andriy-ca@users.noreply.github.com>
fix(ci): aiter pipeline failure on develop CI 2092 ## Motivation Fix Aiter pipeline failure on develop CI 2092 (http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2fComposable%20Kernel/detail/develop/2092/pipeline/290) Jira ID AICK-1947 ## Technical Details The AITER stage aborts before any test runs: newly updated rocm/pytorch:latest aborts HIP_VISIBLE_DEVICES=-1 and kills the process, so the FlyDSL AOT pre-compile dies during import. Changing HIP_VISIBLE_DEVICES=99 as a workaround. ## Test Plan Aiter CI pass this stage and can detect real issue. ## 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.
fix(ck): prevent int32 overflow in tensor descriptor element space size
JIRA ID : ROCM-29071
Fixes ROCM-29071 (synced with SWSPLAT-48860).
## Summary
Fixes a signed 32-bit integer overflow (CWE-190 → CWE-787) in Composable
Kernel's runtime tensor-descriptor construction, reported via bug bounty
as **ROCM-29071**.
`calculate_element_space_size_impl()` (the active path under
`CK_WORKAROUND_SWDEV_275126`) and the fallback lambda in
`make_naive_tensor_descriptor()` both computed:
```cpp
auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i];
```
For a **runtime** (dynamic) descriptor, `lengths[i]` and `strides[i]`
are `index_t` (int32), so the multiply is performed in **32-bit before**
being widened into the `long_index_t` (int64) accumulator. For tensors
where a dimension product exceeds `INT32_MAX`, the product wraps.
**Example (K=C=65537, grouped conv bwd weight):**
- `(65537-1) * 65537 = 4,295,032,832` → wraps to `131,072` (int32)
- `GetElementSpaceSize()` returns `131,073` instead of `4,295,098,369`
- Workspace allocated ≈ 524 KB, but the kernel writes the full ≈ 17 GB
region (sized correctly via `accumulate_n<long_index_t>`) →
out-of-bounds GPU write.
## Fix
Widen **only the runtime operands** to `long_index_t` before the
multiply, via a small helper:
```cpp
template <typename T>
__host__ __device__ constexpr auto widen_runtime_index_to_long(T v)
{
if constexpr(is_number_v<T> || is_long_number_v<T>)
return v; // compile-time operand: leave untouched
else
return static_cast<long_index_t>(v); // runtime operand: widen before multiply
}
```
```cpp
auto acc_new = acc_old + widen_runtime_index_to_long(lengths[i] - Number<1>{})
* widen_runtime_index_to_long(strides[i]);
```
### Why not a blanket `static_cast<long_index_t>` on both operands?
An unconditional cast (the first revision of this PR) also widened the
**compile-time** (`Number<>`) operands, turning the element space size
of a **fully-static** descriptor from a compile-time `integral_constant`
into a runtime `long_index_t`. That flips
`TensorDescriptor::IsKnownAtCompileTime()` to `false`, and every
static-descriptor consumer gated on it (e.g.
`threadwise_tensor_slice_transfer`, contraction/gemm instances) fails to
instantiate — the gfx950 / gfx1201 / Windows `math-libs` build breaks
seen in the previous CI run.
`widen_runtime_index_to_long()` widens the runtime path (fixing the
overflow) while leaving compile-time operands as `Number<>`, so static
descriptors keep a compile-time-constant element space size — identical
to the pre-fix type behavior.
## Affected path
`device_grouped_conv_bwd_weight_xdl_cshuffle.hpp` → `GetWorkSpaceSize()`
→ `GetWorkspaceSizeBytes()` →
`make_naive_tensor_descriptor(...).GetElementSpaceSize()`. Any
conv-bwd-weight (XDL cshuffle) call with `K·C > INT32_MAX`.
## Notes for reviewers
- Both sites are patched so the fix holds regardless of
`CK_WORKAROUND_SWDEV_275126`.
- Mirrors the intent of the already-correct ck_tile pattern
(`ck_tile/core/tensor/tensor_descriptor.hpp`,
`detail::calculate_element_space_size_impl`), adapted to preserve CK's
compile-time `LongNumber<>` element-space-size for static descriptors
(ck_tile instead clamps to a runtime `index_t`).
## Test plan
- [x] Host type-check: a fully-static descriptor keeps a
compile-time-constant element space size (`IsKnownAtCompileTime()` stays
true); the runtime path computes `65536 * 65537 = 4,295,032,832` with no
int32 wrap.
- [x] GPU regression in `test/grouped_convnd_bwd_weight/` at `K·C >
INT32_MAX`: reported workspace size matches the 64-bit
`c_space_size_bytes`; kernel run returns `hipSuccess`.
- [ ] Audit `GetWorkspaceSizeBytes()` / `c_space_size_bytes` callers
agree end-to-end.
---------
Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com>
Co-authored-by: muozturk <Osman.Ozturk@amd.com>
Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>
fix(ck_tile): gfx1250 test failure: test_ck_tile_gemm_pipeline_comp_async_wmma ## Motivation JIRA ID AICK 1886 ## Technical Details misaligned parameter
feat(ck-tile): batched GEMM + batched-contraction TE to dispatcher bridges ISSUE ID: #8997 ## Summary This PR combines two sibling TileEngine → Dispatcher bridge ops into a single PR: - **batched GEMM** (previously #9306) - **batched_contraction** (previously #9328) Both follow the same **direct-launch, registry-bypass** pattern (as the stream-K bridge #9028), because their launch ABIs carry variable-length / batch-specific arguments the single-pointer registry backend cannot express. The two ops touch disjoint files except `dispatcher/tests/CMakeLists.txt`, where both GPU-correctness test registrations are kept. ## Motivation The dispatcher had no path for batched GEMM (same GEMM across many independent problems with per-batch strides) or batched_contraction (generalized batched tensor contraction `E[G.., M.., N..] = sum_K A[G.., M.., K..] * B[G.., N.., K..]` with multi-dim G/M/N/K index groups). Both are real Old-TE ops; this bridge lets Python callers generate, build, and launch them at parity with the legacy Tile Engine — without writing C++. ## Test Plan / Result - CPU-only unit tests for both bridges + gemm_utils: **74 passed** (`pytest dispatcher/tests/test_batched_bridge.py dispatcher/tests/test_batched_contraction_bridge.py dispatcher/tests/test_gemm_utils.py`). - Batched GEMM: end-to-end name-parity + correctness across batch counts 1/2/4/8 (`max_rel ~5e-4`), non-packed strides, split-K; full `default_config` codegen 6672 kernels / 0 failures. - Batched contraction: on-GPU verify (gfx950) across dtype × layout × shape × multi-dim × pipeline + D-tensor epilogue, all PASS. - clang-format-18 clean on both ctypes libs. ### Perf parity vs Old-TE - Batched GEMM (MI300X, fp16 rcr, batch=8): at parity / slightly ahead — median gap +4.00%, 100% within ±15%. - Batched contraction (MI350X, fp16 rcr): at parity — median gap -0.95%, 100% within ±15%. ## Scope / known limitations - Batched contraction: `rcr` only, `k_batch==1` only (split-K is a shared Old-TE kernel defect — hard-rejected, never silently-wrong), non-tile-multiple M/N/K rejected by `IsSupportedArguments`. --- Supersedes and closes #9306 (batched GEMM) and #9328 (batched_contraction). **Related PRs (TileEngine → Dispatcher GEMM bridge series):** #8997 (regular GEMM), #9000 (grouped), #9028 (stream-K), #8887 (fp8/bf8/int8), #9305 (multi-ABD), #9307 (preshuffle), #9308 (multi-D), #10439 (block-scale quant, 5 ops). --------- Co-authored-by: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Thrupti Raj Lakshmana Gowda <thruptiraj.lakshmanagowda@amd.com>
[CK] [MHA] Fix for MHA with softmax-sink ## Motivation MHA softmax-sink-learnable support was needed in the Maxtext->TransformerEngine->Aiter framework, but testing revealed several bugs in the implementation. While I was fixing it, I found that tests is not sensitive to changes in the code. JIRA ID : AIJAX-322 ## Technical Details **Problem 1)** The external MHA interface contract for the backward (BWD) path uses the **_sink_** buffer with an incorrect shape [B, H]. This breaks the symmetry with the forward (FWD) path, where the corresponding _sink_ buffer has shape [H]. As a result, the backward computation produces incorrect sink gradients. The [H] shape is the only correct representation because this tensor corresponds to a trainable parameter with one shared value per attention head, not a unique value for every element in the batch. Therefore, introducing the batch dimension is semantically incorrect. In addition, [H] is the shape used by other frameworks. Because of this mismatch, framework integrations currently fall back to Triton implementations for Sink attention instead of using Composable Kernel (CK). _Fix:_ The first commit fixes the BWD. **Problem 2**) After applying the first commit, the existing test suite continues to pass even though a critical part of the kernel behavior has changed. Since the tests themselves were not modified, this indicates that they are effectively insensitive to the bug and unable to detect incorrect behavior. ``./bin/test_ck_tile_fmha_bwd_bf16 --gtest_filter='*Sink*'`` The root cause is the initialization range used by the tests. The generated input values are too large, which causes the Sink outputs to become extremely small and converge toward zero. As a consequence, the numerical differences introduced by the bug are masked, allowing the tests to pass regardless of whether the implementation is correct. Failed 3 out of 15, but should all 15 <img width="1088" height="228" alt="image" src="https://github.com/user-attachments/assets/23b38cc1-0f81-44da-8b46-0f2cf2eb7b9e" /> _Fix:_ The second commit adjusts the input initialization range to produce numerically meaningful outputs and improve test sensitivity. <img width="1154" height="521" alt="image" src="https://github.com/user-attachments/assets/678d5c1c-7e97-4d12-9117-6f4e814db8e0" /> _Fix:_ The 3-rd commit fixed test code, now all 15 passing. <img width="717" height="338" alt="image" src="https://github.com/user-attachments/assets/d20ad772-fb77-4cfc-9f7b-a9dbaa62e741" /> **Problem 3**) There is an issue with the Sliding Window Attention (SWA) window progression when used together with Sink attention. The bug is observable in external framework tests, but the current CK test suite does not contain any coverage capable of exposing this behavior. As a result, the implementation can regress without any CK tests failing. Test Coverage Improvement The fourth commit adds a set of test configurations specifically designed to exercise the affected SWA + Sink code path. These new tests reproduce the issue and fail with the current implementation, demonstrating the gap in the existing test coverage and providing a reliable way to validate the fix. <img width="1256" height="340" alt="image" src="https://github.com/user-attachments/assets/c9f050e1-a973-4aff-9d38-17ed21387a6f" /> Fix The fifth commit corrects the SWA window movement logic when Sink attention is enabled. With this change applied, the newly added tests pass, and the behavior matches the expectations observed in external framework implementations. ### **Related PRs** This is PR in CK ROCm/rocm-libraries#10519 this is PR in TransformerEngine ROCm/TransformerEngine#678 ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )