mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-10 16:28:38 +00:00
01bd52bdb5fdf55f4047cce5e4abc1ebbf302f5f
1042 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Hosang Yoon
|
e7e8801dc3 |
[rocm-libraries] ROCm/rocm-libraries#7586 (commit c18f2c7)
[CK_TILE] Use gfx11 float buffer atomics in FMHA Bwd ## Motivation FlashAttention CK backward on gfx11 can hit out-of-bounds/tail writes in the dQ accumulator atomic-add path when sequence rows are padded at the tile level but not marked invalid in the DQDKDV main tensor view. With the generic global atomic fallback, an incorrectly-valid tail element can issue an actual pointer-based `atomicAdd`. With the buffer atomic path, the write is issued through a buffer resource with bounds information and follows the same backend already used by gfx9/gfx12. This fixes the gfx11 FMHA BWD failure without changing the gfx11 default for unrelated CK Tile kernels. ## Technical Details This PR enables the existing CK Tile AMD buffer float atomic-add path only for generated FMHA BWD gfx11 translation units. gfx11 normally uses the generic global atomic fallback for floating-point `buffer_view::atomic_add`. That fallback performs the atomic through a raw computed pointer and depends on the software validity predicate to avoid invalid elements. In FMHA BWD dQ accumulation, padded tail rows can reach this path, so using the buffer atomic backend is safer: it uses a buffer resource with base pointer, bounds information, and an element offset, matching the backend already used by gfx9/gfx12. Enabling `CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT` globally for gfx11 is too broad and can break unrelated gfx11 CK builds such as GEMM. Instead, `config.hpp` now preserves an explicitly pre-defined `CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT`, while keeping the existing default disabled for gfx11. ## Test Plan Validated the change with the FlashAttention CK full test suite with backward pass enabled on gfx11. pytest -q -s tests/test_flash_attn_ck.py ## Test Result FlashAttention CK gfx11 test result: 260680 passed, 152076 skipped ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com> |
||
|
Andriy Roshchenko
|
d5c9215064 |
[rocm-libraries] ROCm/rocm-libraries#7359 (commit dd62f9f)
[CK_TILE][GFX1250] Enable MX GEMM FLATMM with ASYNC ## Motivation Enables MX GEMM FLATMM pipeline on gfx1250. The pipeline uses an async load instruction for tensor A, which complements the existing MX GEMM FLATMM pipeline with TDM load. At this time, only FLATMM MX pipelines are enabled on gfx1250. ## Technical Details The existing gfx950 implementation was extended to support gfx1250 architecture. All three MX FP data types are supported across the two ASICs. It should be noted that while the TDM pipeline uses an emulated 32x32x128 warp-tile instruction, the present submission relies on the built-in 16x16x128 instruction, called 4 times per warp. ## Test Plan Existing `test/ck_tile/flatmm` tests were extended to cover new gfx1250 functionality. To help facilitate the testing in development, `example/ck_tile/18_flatmm/script/smoke_test_mx.sh` script was introduced to verify various combinations of supported data types and pipeline versions. ## Test Result The present submission is expected to work on both gfx950 and gfx1250 hardware for all reasonable sizes and all MX FP8/FP6/FP4 data types. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. - [x] Relies on #6978 and should only be merged after the changes are merged to the `develop`. |
||
|
Illia Silin
|
8bd8094012 |
[rocm-libraries] ROCm/rocm-libraries#7833 (commit 8a444cd)
[CK] Replace deprecated load_module function in python (#7833) ## Motivation Recent pytorch builds with python 3.15 failed in CK due to deprecation of load_module function. This should fix the issue. ## 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. |
||
|
ltqin
|
bf07a0150e |
[rocm-libraries] ROCm/rocm-libraries#7723 (commit 4ed6c51)
[CK Tile] Enable LSE output for fp8bf16 V3 FMHA kernels (#7723) ### Motivation The V3 pipeline (qr_async_trload_v3) for fp8bf16 FMHA kernels did not support LSE (Log-Sum-Exp) output. This PR enables LSE output support for fp8bf16 V3 FMHA kernels, allowing users to retrieve attention statistics alongside attention outputs. ### Technical Details - StandardAttention: lse = softmax_scale * m + log(l) - LogitsSoftCap: lse = (m / log2(e)) + log(l) ### Test Plan Run FMHA forward example with fp8bf16 precision and LSE output enabled: - Test 1: Basic LSE functionality ./build/bin/tile_example_fmha_fwd -v=1 -b=1 -h=8 -s=1024 -d=128 -prec=fp8bf16 -init=3 -qscale=1 -lse=1 - Test 2: LSE with LogitsSoftCap (CMakeList should remove Logits filter) ./build/bin/tile_example_fmha_fwd -v=1 -b=1 -h=8 -s=1024 -d=128 -prec=fp8bf16 -init=3 -qscale=1 -lse=1 -logits_soft_cap=30.0 |
||
|
Zoltán Lakatos
|
58e2ab1fc7 |
[rocm-libraries] ROCm/rocm-libraries#6761 (commit d19f6f1)
[CK] Large tensor gemm workaround (#6761) ## Motivation Customer qeruested large tensor gemm support for 8bit and 4bit data types. Currently CK triggers “This GEMM not supported” error. The root cause appears to be the 2 GB limit on the input/output matrix, triggered by buffer offset constraints when testing a larger shape such as M = 699,904 (which is an exact multiple of MPerBlock = 256). ## Technical Details Quick workaround to have support ASAP. Split the tensors into inputs / outputs smaller than 2GB limit. Iterate on host and call all subproblems without device code change. Support is restricted to rowise layout in A, Ds and E All changes were implemented in DeviceGemm structures to avoid secondory affect on grouped convolutions. Got lots of AI generated comments. Addressed the ones that seemed relevant on the functionality. ## Test Plan Within CK the following examples can be used with modified input sizes: example_gemm_multiply_multiply_xdl_fp8 example_gemm_mx_fp4 Tested with Aiter tuning on provided shapes. ## Test Result All gemms run and provide correct results. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Zoltán Lakatos <zoltan.lakatos@streamhpc.com> Co-authored-by: Márton Bidlek <marton.bidlek@streamhpc.com> Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com> |
||
|
Illia Silin
|
c24e528481 |
[rocm-libraries] ROCm/rocm-libraries#7760 (commit a61bc76)
[CK] suppress compiler warnings while building pytorch. (#7760) ## Motivation Recently added compiler flags that are required to suppress false warnings by latest staging compiler are not recognized by older compiler versions and are triggering an avalanche of warnings. Previous attempt to suppress them by using -Wno-unknown-warning-option flag didn't help, because that flag wasn't recognized either and just added more warnings. I've verified that current approach by checking the clang version actually works as intended and makes the warnings go away. ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
JH-Leon-KIM-AMD
|
00e1d82ae7 |
[rocm-libraries] ROCm/rocm-libraries#7732 (commit b0e29d9)
[CK] Fix grouped conv bwd data stride>1 silent miscompute (ALMIOPEN-1959) (#7732) ## Motivation Fix silent miscompute in the grouped convolution backward-data kernel (`DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1`) when stride > dilation (ALMIOPEN-1959). PR #6208 introduced a flat-descriptor fast path that dropped all but the first sub-GEMM, producing zeroed slices of `dx` on the (G=1, stride>1, 2D, NumDTensor=0) intersection. Restore correctness without giving up the perf gains PR #6208 delivered on stride=1 shapes. ## Technical Details - Tighten the flat-descriptor fast-path gate to require `arg.gemms_count_ == 1` (i.e. a single sub-GEMM per dispatch — its original purpose). For stride > 1, the implicit GEMM is split into `gemms_count_` sub-GEMMs whose output cells tile `dx` disjointly; routing them through the flat path required dropping all but the first, which was the source of the bug. - Stride > 1 now falls through to the existing grouped CShuffle path, which packs all sub-GEMMs into one descriptor array and walks them on-device in a single kernel launch. This is the pre-PR-6208 production path; correctness is established and per-dispatch launch count is minimised. - Add regression coverage for the (G=1, stride>1, 2D, NumDTensor=0) intersection in `test/grouped_convnd_bwd_data/test_grouped_convnd_bwd_data.cpp` with `gemms_count` ∈ {4, 9, 36}. Pre-existing cases did not hit this intersection (all stride>1 cases used G=2; all G=1 cases used stride=1), which is why PR #6208's regression slipped past CI. ## Test Plan - `ctest -L SMOKE_TEST -R 'grouped_convnd_bwd_data'` on gfx942 (smoke tier — runs on every PR via `smart_build_and_test.sh`). - End-to-end verify (`verify=1`) via `example_grouped_conv_bwd_data_xdl_fp16` on stride 1/2/3/6 shapes including the original ALMIOPEN-1959 case and a cross-bucket (`gemms_count=36`) case spanning two `MaxGroupedGemmGroupsNum=32` buckets. - ckProfiler A/B sweep on MI300X (gfx942) toggling the flat-path gate via an environment variable: full kernel-family enumeration, winning kernel + its avg_time reported under each gate. 33/41 shapes completed before the sweep was stopped; the remaining 8 were the largest i2v/synthetic shapes where ckProfiler exceeded its 300s per-shape enumeration budget (not relevant to the verdict). ## Test Result ### Correctness | Test | Result | |---|:---:| | `test_grouped_convnd_bwd_data` (12 type parameterizations × Test2D, includes 3 new regression shapes) | **12/12 PASSED** in 14.18 s | | `test_grouped_convnd_bwd_data_interface` (API checks) | **PASSED** in 0.28 s | | ALMIOPEN-1959 stride=2 (`verify=1`) | **PASSED** | | stride=1 K3 (`verify=1`) | **PASSED** | | stride=3 K3 `gemms_count=9` (`verify=1`) | **PASSED** | | stride=6 K6 `gemms_count=36` cross-bucket (`verify=1`) | **PASSED** | ### Performance (ckProfiler A/B on gfx942 / MI300X) Comparing the **post-fix gate** (flat path only when `gemms_count_==1`, column "B") vs the **inner-loop variant** that keeps the flat path on stride>1 (column "A") across 25 stride>1 shapes where production picks a `_v1` instance (so the gate actually fires): | Stride | Shapes | A wins | Tie | B wins | Notes | |:------:|:------:|:------:|:---:|:------:|---| | 1 (sanity, gate moot) | 3 | 0 | 3 | 0 | gate doesn't differentiate — A == B as expected | | > 1 (gate fires) | 25 | **0** | 11 | **14** | B wins +6% to +32%; A never wins | Highlights from the firing-gate cases: | Shape (G=1, stride=2 unless noted) | A ms | B ms | B vs A | |---|---:|---:|---:| | ALMIOPEN-1959 (N=16, K=256, C=128, 5×5, 40×175) | 0.183 | 0.171 | **B +6%** | | Retinanet-L61 (N=32, K=C=256, 3×3, 25×25) | 0.054 | 0.045 | **B +17%** | | i2v-010 (N=1, K=C=384, 3×3, 277×209) | 0.174 | 0.125 | **B +28%** | | Synthetic 50×50 K3 N=32 K=C=256 | 0.131 | 0.088 | **B +32%** | Why B wins everywhere the gate fires: for `gemms_count = N`, the flat path needs N kernel launches (one per sub-GEMM), while the grouped path loops over the same N sub-GEMMs on-device in 1 launch. The (N−1) × launch-tax is a structural disadvantage A can't recover from. ### Diff | File | Lines | |---|---:| | `include/.../device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp` | +14 / −8 (one extra condition + expanded dispatch comment) | | `test/.../test_grouped_convnd_bwd_data.cpp` | +9 / −0 (3 new shapes) | ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Aviral Goel
|
0df3523ef1 |
[rocm-libraries] ROCm/rocm-libraries#6807 (commit ddda8ac)
[CK_TILE] Add save_matrix_txt() and extract HostTensor I/O to free functions (#6807) ## Summary - Extract `loadtxt`, `savetxt`, and `save_matrix_txt` from `HostTensor` member functions into standalone free functions in `host_tensor_io.hpp` (Single Responsibility Principle) - Add `save_matrix_txt()` for writing 2D tensors to space-separated text files with configurable output limit (default 256x256, pass 0 to dump all) - Supports float, int, and int8_t output formats via a `dtype` parameter - Validate dtype early and throw on unsupported values in all three functions - Update callers in `15_fused_moe/main.cpp` to use free function syntax |
||
|
Anton Gorenko
|
66d6714376 |
[rocm-libraries] ROCm/rocm-libraries#5388 (commit 45583bd)
[CK_TILE][FMHA] Improve precision of mxfp4 FMHA with fp6 for matrix P (#5388) ## Motivation Improve precision of mxfp4 without performance penalties. ## Technical Details Since performance of scale MFMAs is the same when neither A nor B is fp8/bf8, it is possible to use fp6 x fp4 instead of fp4 x fp4 for the second GEMM, while types of Q, K, V stay the same. This allows to improve overall precision significantly because fp6 has 32 non-negative values used for P quantization compared to just 8 values for fp4. It was found that there is a compiler bug with `__builtin_amdgcn_cvt_scalef32_2xpk16_fp6_f32` (described in LCOMPILER-561) but a workaround seems to fix all failing instances. ## Test Plan ``` ninja test_ck_tile_fmha_fwd_mxfp4 && bin/test_ck_tile_fmha_fwd_mxfp4 ``` ## Test Result The tests must pass. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Yi DING
|
6a9c03f692 |
[rocm-libraries] ROCm/rocm-libraries#7450 (commit 402dbad)
[CK_TILE] Use Persistent Scheduling for FMHA BWD Group Deterministic (#7450) ## Motivation FMHA BWD group-mode deterministic currently uses a non-persistent scheduler: each `(batch, head, K-row)` work-item is launched as its own block, with no work-stealing across CUs. On uneven workloads (varlen, GQA, many heads with few K-rows) this leaves CUs idle and forces a larger dq_acc workspace than necessary. This PR ports the persistent + deterministic scheduling already used in batch mode to group mode: a fixed-grid kernel that pre-computes per-CU work ranges on the host and uses sparse dq_acc slot indexing so multiple K-rows handled by the same CU share one accumulator slot via intra-CU atomic adds. Stacked on #7331; merge that first. ## Technical Details Single file changed: `ops/fmha/kernel/fmha_bwd_kernel.hpp`. A new `kUsePersistent` path is added to the group-mode deterministic kernel, mirroring the batch-mode persistent scheduler. The host pre-computes a fixed per-CU partition of the total `(batch, head, K-row)` work and packs it into `cu_states[]` so the GPU consumes it in a single launch. Host preparation happens in four steps: 1. Build per-batch `seqstart` prefix sums. 2. Fill per-batch `(sq_w, nc)` with a placeholder `nsplits` (bumped in step 3). 3. Two-pointer scan over CUs to fill `cu_states[c]` (`isplit`, `head_start`, `c_start`, `w_lo`, `w_hi`), accumulating `nsplits[b]` as `max(cs->isplit + 1)`. 4. Compute compact per-batch dq_acc offsets from the finalized `nsplits`. `isplit` is the sparse dq_acc slot index — one CU's multi-K-row writes share slot `ceil(wc_start / denom)`, enabling intra-CU atomic accumulation instead of one slot per K-row. `denom = max(sq_w, target_w)`, splitting two regimes: - `target_w >= sq_w` (large work): `denom = target_w`, intra-CU atomic optimization engaged. - `target_w < sq_w` (sub-K-row sharding, multiple CUs sharing one K-row): `denom = sq_w` collapses to per-K-row indexing (`= c_start`), keeping `isplit ∈ [0, nc-1]` and matching the `nsplits_max = ceil(s_k/kN0) = nc` upper bound that #7331's `GetWorkspaceDeviceSizeUpperBound` assumes for group+det. `isplit` is additionally clamped to `nc-1` to absorb empty CUs (rounded-up `wc_start` past the last K-row); they don't write dq_acc on GPU so the slot value is harmless. `nsplits[b]` is accumulated dynamically in step 3 rather than via a closed form so it tightly matches the actual sparse slots used; step 4 (offsets) follows step 3 since offsets now depend on the dynamic `nsplits`. Group mode also allows batches with `seqlen_q == 0`. The persistent scheduler skips them on the dQ path (no work) but dK/dV are still zero-filled. ## Test Plan Built `tile_example_fmha_bwd` with receipt 5 (fp16, no-bias, no-dropout, `dpad == dvpad`, group + batch) on gfx950 (MI355X). - 8-case smoke (shapes that exercise the sub-K-row regime). - 44-case sweep covering: mask 0/1/2, GQA, var seqlen, `d != d_v`, extreme small seqlen / `nc=1`, CU >> work, huge batch, batch-mode regression. - 12-case perf comparison vs the non-persistent baseline (warmup=10, repeat=50). ## Test Result - All 8 + 44 cases `valid:y`. - Perf: ±5% noise, average -0.4% across the 12 cases — neutral. - Batch-mode deterministic / non-deterministic regression unchanged. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Michal Kulikowski
|
8de4cb72fb |
[rocm-libraries] direct push (commit 49b73ad)
[CK][CK_TILE] POC for Instruction Cache prefetch. Signed-off-by: Michal Kulikowski <Michal.Kulikowski@amd.com> |
||
|
JP-Fernando
|
74bc86240b |
[rocm-libraries] ROCm/rocm-libraries#5647 (commit 490437a)
[CK Tile] Add gemm universal preshuffle to MX GEMM (#5647) ## Motivation Add gemm universal preshuffle support to existing MX GEMM pipeline. The straightforward way to do this is to port the `mx_flatmm` pipeline to the existing `gemm_mx` framework. ## Technical Details The `mx_flatmm` pipeline was not deleted, to allow for back-compatibility. ## Test Plan Add `preshuffle` option to example: `tile_example_mx_gemm`. Add new configurations with enabled preshuffle to the existing `test/ck_tile/gemm_mx` tests. ## Test Result Example and tests were successful on `gf950` architecture in the `Alola` cluster. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Fernando Jiménez <fernando.jimenez@streamhpc.com> |
||
|
Bartłomiej Kocot
|
ebb97044f4 |
[rocm-libraries] ROCm/rocm-libraries#7664 (commit de5d6b1)
Revert "[CK] Enable grouped conv bwd data to match non-grouped perf" (#7664) ## Motivation Incorrect results has been introduced for some conv bwd cases. ## Technical Details This reverts commit 33424f65346d6330d0fd94b5a4e6f843f24e52c3. ## Test Plan CI ## Test Result Pending ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. ALMIOPEN-1959 |
||
|
Illia Silin
|
e02c566795 |
[rocm-libraries] ROCm/rocm-libraries#7612 (commit 5427d24)
[CK] upgrade CI to rocm7.13 as default compiler (#7612) ## Motivation Upgrade the default docker and compiler version in CI to rocm7.13. In order to pass all the checks I had to also clean up a lot of non-ascii characters in the source code comments and modify a couple of tests that were affected by a new compiler logic. ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Aviral Goel <aviral.goel@amd.com> |
||
|
kensclin
|
fc2862d712 |
[rocm-libraries] ROCm/rocm-libraries#6846 (commit 377def4)
[CK_TILE] Add fmha forward hdim 256 support (#6846) ## Motivation Enable Composable Kernel FMHA forward kernel for **hdim=256 BF16** on AMD gfx950 (MI350X). Prior to this change the (256, 256) head-dim configuration either failed to compile, was filtered out by the compatibility rules, or produced incorrect kernel output due to an LDS layout accounting bug. ## Technical Details Four files changed, all to enable hdim=256 BF16 on gfx950. - **`fmha_fwd.py`** — Allow `(256, 256)` in gfx950 compatibility rule; set `(256,256)` BF16 tile to `M0=128, N0=64` (the LDS-feasible shape on gfx950); emit minimal valid instance set for d=256 to bound compile time. - **`fmha_fwd_kernel.hpp`** — Gate Prefill launch path off for d=256 (`PrefillCase = kM0 > 64 && kQKHeaddim < 256`); the double-buffer Prefill variant overflows the 160 KB LDS budget. - **`trload_policy.hpp`** — **Critical correctness fix**: the LDS layout accounting in `GetSmemSize` was wrong (`max(Q, K+S+V)` instead of `max(Q, K) + V + S`), under-allocating LDS and silently corrupting d=256 output (~2% wrong values). - **`trload.hpp`** — Thread `LoadOnce=true` through all d=256 K-LDS descriptors so the compiler picks the matching XOR swizzle period; recompute the S-tile LDS offset to match the corrected `GetSmemSize` formula. ## Test Plan Built and ran `tile_example_fmha_fwd` on gfx950 (MI350X) with the canonical d=256 BF16 configurations: ```bash cd build && ninja tile_example_fmha_fwd ./bin/tile_example_fmha_fwd -prec=bf16 -d=256 -d_v=256 -b=1 -h=32 -h_k=2 -s=1024 -s_k=1024 -bias=n -mask=t -lse=0 -p_drop=0 -warmup=3 -repeat=10 -kname=1 -v=1 ./bin/tile_example_fmha_fwd -prec=bf16 -d=256 -d_v=256 -b=8 -h=32 -h_k=2 -s=16384 -s_k=16384 -bias=n -mask=t -lse=0 -p_drop=0 -warmup=3 -repeat=10 -kname=1 -v=1 ``` ## Test Result ```bash -b=1 -s=1024 [bf16|batch|bhsd] b:1, h:32/2, s:1024/1024, d:256/256, scale_s:0.0625, bias:n, p_drop:0, lse:0, qscale:n, mask:t(-1:0), v:r, fmha_fwd_d256_bf16_batch_b128x64x32x256x32x256_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_psddv_nlogits_nbias_mc_nlse_ndropout_nskip_nqscale_ntrload_nsink, 0.058 ms, 298.42 TFlops, 618.68 GB/s, valid:y -b=4 -s=16384 [bf16|batch|bhsd] b:8, h:32/2, s:16384/16384, d:256/256, scale_s:0.0625, bias:n, p_drop:0, lse:0, qscale:n, mask:t(-1:0), v:r, fmha_fwd_d256_bf16_batch_b128x64x32x256x32x256_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_psddv_nlogits_nbias_mc_nlse_ndropout_nskip_nqscale_ntrload_nsink, 42.797 ms, 822.18 TFlops, 106.63 GB/s, valid:y ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: poyenc <1132573+poyenc@users.noreply.github.com> |
||
|
JP-Fernando
|
e7798e9560 |
[rocm-libraries] ROCm/rocm-libraries#7112 (commit a6e5eac)
Add asynchronous XOR shuffle support to the Async GEMM pipeline and the MX GEMM pipeline (#7112) ## Motivation The goal of this work is to apply XOR shuffle (swizzle) to the current `comp_async` GEMM pipeline and the `gemm_mx` pipeline. XOR swizzling has been helpful to avoid LDS bank conflicts, as data are redistributed across LDS banks, such that simultaneous threads accessing different rows land on different LDS banks. ## Technical Details A similar approach to the work in the existing eight-waves pipeline was followed. Currently, XOR swizzle support is available for FP8 and BF8 types. FP4 support is also available for MX GEMM. Should the types not match, or should the async vector width be of an unsupported size, then the pipeline falls through to the previously existing ('unswizzled') path. ## Test Plan Execute `test_ck_tile_gemm_pipeline_comp_async` for the Async GEMM pipeline. Execute `test_ck_tile_mx_gemm_fp8` and `test_ck_tile_mx_gemm_fp4` for the MX GEMM pipeline. ## Test Result The tests passed successfully in the `Alola` cluster with MI350 hardware. ## Submission Checklist - [X] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Fernando Jiménez <fernando.jimenez@streamhpc.com> Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> |
||
|
Kiefer van Teutem
|
b5f8bef97f |
[rocm-libraries] ROCm/rocm-libraries#6088 (commit 6ac353c)
[CK Tile][MFMA/WMMA unification] Add support for packed datatypes (tiny types) (#6088) ## Motivation This MR makes all the changes required for the unified architecture to be able to deal with packed datatypes i.e. int4, fp4, fp6, and bf6. The crux is that layout parameters should be interpreted as describing the pure mathematical matrix fragments, while the ext_vectors and tile distribution encodings describe everything in terms of packed datatype units. This matches how packed types are dealt with in ck_tile and should play nicely with the load and store tile ops once we integrate the unified framework into CK tile. The bf6 datatype was added to CK tile in the form of pk_bf6x16_t and pk_bf6x32_t, which did not exist before. The ext_vector implementations of pk_fp6x16_t and pk_bf6x16_t (vec size 1 and 2) were extended to make the subscripting operator work as expected. The layout test was adapted to be compatible with all packed datatypes, and all new intrinsics were added to the test. This MR adds ALL intrinsics across ALL architectures which use packed datatypes, as well as ALL scale intrinsics: mfma_scale_f32_16x16x128_f8f6f4 gfx950 (F8xF8, BF8xBF8, F4xF4, F6xF6, BF6xBF6) mfma_scale_f32_32x32x64_f8f6f4 gfx950 (F8xF8, BF8xBF8, F4xF4, F6xF6, BF6xBF6) wmma_i32_16x16x16_iu4_w32 wmma_i32_16x16x16_iu4_w32_gfx12 wmma_i32_16x16x32_iu4_w32_gfx12 ## Testing All intrinsics were tested on all architectures. |
||
|
Enrico Degregori
|
9565ca21ec |
[rocm-libraries] ROCm/rocm-libraries#5552 (commit 369c7a2)
[CK Tile] Eight Waves pipeline for MX GEMM (#5552) ## Motivation Integrate Eight Waves pipeline in MX GEMM ## Technical Details - EightWaves pipeline: - Add pipeline, policy and block gemm (internally using existing implementation used by GEMM and ABQuant) - Extend support of EightWaves policy for FP4 (packed types) - Async pipeline: - Fix pipeline with packed scales (requires MRepeat and NRepeat to be contiguous) - block gemm specific for MX GEMM is defined because distribution encodings have changed - CShuffle: - Add new functionality to support MRepeat and NRepeat contiguous (defined by `TilesPacked`) - Examples: - Refactor examples to easily switch different configurations (similar to GEMM universal) - Scales values generated consistently with other microscale implementations in CK Tile - Add configuration for EightWaves pipeline - Tests: - Unify existing FP8 and FP4 tests - Add tests for EightWaves pipeline - Scales values generated consistently with other microscale implementations in CK Tile Note: FP6 support for MX GEMM was added later and the support for the Eight Waves pipeline will be done in following PR ## Test Plan Add new pipeline to tests: `test_ck_tile_mx_gemm_async` for both FP4 and FP8 ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Po Yen Chen
|
a11f53564f |
[rocm-libraries] ROCm/rocm-libraries#7530 (commit 378e049)
[CK] Fix FMHA sink dispatch when init_sink_value is set (#7530) ## Summary - Fix `traits.has_sink` in `fmha_fwd_runner.hpp` to also check `init_sink_value != 0`, so the GPU kernel dispatches with sink support when `-init_sink=1` is passed. - Gate `run_sink_mask_tests` (StreamLLM) and `run_sink_init_tests` (GPT-OSS) behind opt-in flags `-m` and `-g` in `smoke_test_fwd.sh`. These tests require sink=true kernel instances which are excluded by the `BUILD_TESTING` CMake filter (`*_nsink*`), causing unconditional "not supported yet" failures (48 tests in CI). The opt-in flag approach was borrowed from PR #6057. ## Why gate tests instead of compiling sink=true kernels? The `BUILD_TESTING` filter in `CMakeLists.txt` uses `*_nsink*` glob patterns for the `fwd` and `fwd_splitkv` APIs, excluding sink=true kernel instances from compilation. We chose opt-in flags over widening the filter because: - **Compile time**: Enabling sink=true kernels doubles the kernel variants for `fwd` and `fwd_splitkv` APIs. The filter exists specifically to reduce CI build times. - **Incremental enablement**: Sink support (StreamLLM / GPT-OSS) is still maturing. Gating lets teams opt in explicitly (`smoke_test_fwd.sh -g`) while keeping the default CI path fast. - **Precedent**: splitkv (`-s`) and appendkv (`-a`) tests already follow this opt-in pattern. ## Test plan - [ ] Run `smoke_test_fwd.sh -g` with sink=true kernels compiled and verify sink-enabled kernels are dispatched - [ ] Verify `smoke_test_fwd.sh` still passes without `-m` / `-g` flags - [ ] Confirm CI no longer fails on sink tests (they are now opt-in) |
||
|
Johannes Graner
|
3727d5220a |
[rocm-libraries] ROCm/rocm-libraries#5652 (commit 7dc7d1d)
[CK Conv] Wavelet gemm pipeline for bwd_weight convolution (#5652) ## Motivation In the current CShuffleV3 backward weight 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") convolution backward weight 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 **Core kernel (new files):** - `gridwise_gemm_xdl_waveletmodel_cshuffle_conv_v3.hpp` — wave-specialized gridwise GEMM for conv bwd weight (2-way split: load + math) - `device_grouped_conv_bwd_weight_xdl_waveletmodel_cshuffle_v3.hpp` — device op following CShuffleV3 patterns; `BlockSize = TileMathThreadGroupSize` for MFMA wave assignment, `LaunchBlockSize = TileLoad + TileMath` for kernel launch **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). **Instance coverage (F16 and BF16 symmetric):** | Ratio | Tiles | Layouts | ConvSpecs | |-------|-------|---------|-----------| | (4,4) | M128×N128, M64×N64, M128×N64, M64×N128 | 2D NHWGC, 3D NDHWGC | Default, Filter1x1Stride1Pad0 | | (4,2) | M64×N64, M128×N64, M64×N128 | 2D NHWGC | Default, Filter1x1Stride1Pad0 | **Existing wavelet model fixes:** - `BlockSize` corrected from `math::max(TileLoad, TileMath)` to `TileMathThreadGroupSize` in the flat-GEMM wavelet device op and gridwise kernel ## Test Plan - `test_grouped_convnd_bwd_weight` GTest: 34 hardcoded test cases covering 1D/2D/3D, F16/BF16, G=1/2/16, various spatial sizes - Performance benchmark: all 37 RetinaNet bwd_weight shapes on gfx950 ```bash ninja -C build test_grouped_convnd_bwd_weight ./build/bin/test_grouped_convnd_bwd_weight ``` ## Test Result **Correctness:** 34/34 GTest cases passed (F16/BF16 × 1D/2D/3D × Default/Filter1x1Stride1Pad0 × various G/N/K/C combinations). **Performance:** Wavelet is the fastest overall instance on 12/37 RetinaNet shapes — all G=1, 3×3 convolutions with C=256 (the VALU-heavy target shapes): | Shape | Uplift vs best baseline | |-------|------------------------| | K=36, 7×7 | 1.91x | | K=36, 100×100 | 1.60x | | K=36, 13×13 | 1.43x | | K=36, 25×25 | 1.38x | | K=36, 50×50 | 1.38x | | K=256, 100×100 | 1.24x | | K=256, 13×13, s=2 | 1.20x | | K=256, 25×25, s=2 | 1.20x | | K=256, 7×7 | 1.17x | | K=256, 13×13 | 1.13x | | K=2376, 50×50 | 1.05x | | K=2376, 100×100 | 1.06x | Where wavelet does not win (25/37): 1×1 convolutions (explicit kernel does host-side transform), grouped convolutions with small per-group channels, and shapes where standard CShuffleV3 already amortizes VALU overhead. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: jakpiase <jakpia21@gmail.com> |
||
|
JH-Leon-KIM-AMD
|
9a5d1ea791 |
[rocm-libraries] ROCm/rocm-libraries#6208 (commit 33424f6)
[CK] Enable grouped conv bwd data to match non-grouped perf via NoShuffle + packed descriptors (#6208) ## Motivation Improve performance of grouped convolution backward-data kernels to match non-grouped kernel performance for G=1 cases. ## Technical Details - Add NoShuffle epilogue path (direct VGPR→Global writes) by setting `CDEBlockTransferScalarPerVector_NPerBlock = 1` - Add nongrouped-match instances with optimized BBlockTransfer parameters for better thread utilization - Add packed (flat) descriptor path for G=1 2D convolutions, using simpler tensor descriptors with fewer transform layers to reduce address computation overhead in the GEMM main loop - Cherry-pick PR #6090 for fair benchmarking (cache flush, include dX zeroing cost) ## Test Plan - Benchmark grouped vs non-grouped kernels on MI300X (589 shapes, BF16) - Verify correctness with existing conv bwd data tests ## Test Result | Metric | Before | After | |--------|--------|-------| | Mean ratio (grouped/nongrouped) | 1.159 | **1.028** | | Median ratio | 1.142 | **1.026** | | Cases within 2% | 26 (4.4%) | **186 (31.8%)** | | Cases >20% slower | 188 (32%) | **2 (0.3%)** | NoShuffle + nongrouped-match instances achieve **~2.8% average gap** with non-grouped kernels (down from ~16%). ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: root <root@ctr-cx64-mi300x-4.amd.com> Co-authored-by: root <root@ctr-cx71-mi300x-01.amd.com> Co-authored-by: root <root@ctr-cx63-mi300x-21.amd.com> Co-authored-by: Bartłomiej Kocot <barkocot@amd.com> Co-authored-by: root <root@gt-ccs-aus-h17-18.cs-aus.dcgpu> Co-authored-by: Cursor <cursoragent@cursor.com> |
||
|
Illia Silin
|
717f2efef7 |
[rocm-libraries] ROCm/rocm-libraries#6978 (commit e58096d)
[CK] add composable kernel support on gfx1250 (#6978) ## Motivation Add composable kernel support on gfx1250. ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Qun Lin <qlin@amd.com> Co-authored-by: jialuo12_amdeng <jia.luo@amd.com> Co-authored-by: Andriy Roshchenko <andriy.roshchenko@amd.com> Co-authored-by: hsivasun_amdeng <haresh.sivasuntharampillai@amd.com> |
||
|
Illia Silin
|
ac18460782 |
[rocm-libraries] ROCm/rocm-libraries#7384 (commit 10e9d70)
[CK] Suppress new staging compiler errors (#7384) ## Motivation This should make new builds with staging compiler pass. ## 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. |
||
|
Meekail Zain
|
0937b002d8 |
[rocm-libraries] ROCm/rocm-libraries#6867 (commit 3cb0219)
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> |
||
|
Yi DING
|
af7118e342 |
[rocm-libraries] ROCm/rocm-libraries#7331 (commit 5692db0)
[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. |
||
|
Copilot
|
c39bff93d0 |
[rocm-libraries] ROCm/rocm-libraries#6983 (commit f4e9a84)
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> |
||
|
Linjun-AMD
|
5003f7ef8a |
[rocm-libraries] ROCm/rocm-libraries#7272 (commit d02f3c0)
[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> |
||
|
Illia Silin
|
22b9feb40f |
[rocm-libraries] ROCm/rocm-libraries#7111 (commit 651947f)
[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. |
||
|
Chao
|
3d8c21e838 |
[rocm-libraries] ROCm/rocm-libraries#6529 (commit 93a6097)
[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 https://github.com/ROCm/rocm-libraries/pull/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> |
||
|
Linjun-AMD
|
5a2a362c46 |
[rocm-libraries] ROCm/rocm-libraries#6914 (commit b791478)
[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> |
||
|
Yi DING
|
2c677e8471 |
[rocm-libraries] ROCm/rocm-libraries#6152 (commit 36b016a)
[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. |
||
|
ltqin
|
501e7ef12a |
[rocm-libraries] ROCm/rocm-libraries#6574 (commit b3db057)
[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 |
||
|
ArthurLiu
|
67aa854621 |
[rocm-libraries] ROCm/rocm-libraries#6764 (commit 8c20d70)
[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. |
||
|
Wojciech Laskowski
|
640bd560ec |
[rocm-libraries] ROCm/rocm-libraries#5801 (commit 27f6d15)
[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> |
||
|
Jeff Huang
|
08097fa515 |
[rocm-libraries] ROCm/rocm-libraries#6653 (commit 1df887e)
[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. |
||
|
KateJu
|
7ae21043fd |
[rocm-libraries] ROCm/rocm-libraries#6656 (commit 1c958f8)
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. |
||
|
Illia Silin
|
d16061f578 |
[rocm-libraries] ROCm/rocm-libraries#6550 (commit c396de9)
[CK] Fix/suppress clang lifetimebound warnings with staging compiler. (#6550) ## Motivation New changes from upstream llvm-project cause an avalanche of warnings in CK. Gonna disable them by ignoring the lifetime-safety-intra-tu-suggestions flag until a better permanent solution is found. ## 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. |
||
|
Linjun-AMD
|
dfc1305685 |
[rocm-libraries] ROCm/rocm-libraries#6479 (commit 0705c2d)
CK][fmha] Add StreamLLM sink support to batch_prefill pipeline (#6479) ## Motivation The existing paged-KV attention pipelines (pagedkv, splitkv) support StreamLLM-style sink tokens — a fixed set of initial tokens kept in attention alongside the sliding window. The `batch_prefill` pipeline (chunked-prefill with VLLM-style block tables) previously hardcoded `kHasSink = false`, making it incompatible with sink-based attention patterns in LLM serving scenarios. This PR extends `batch_prefill` to support `kHasSink` and wires it into `fmha_fwd_runner` for validation against the existing CPU reference. ## Technical Details **Pipeline** (`block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp`): - When `kHasSink`, the K/V loop splits into a sink phase [0, sink_seq_end) and a window phase [seqlen_k_start, seqlen_k_end), mirroring pagedkv. - K advance at the sink→window transition jumps `seqlen_k_start - sink_seq_end + kN0` to bridge the gap. - V scatter-gather offsets are re-initialized at the transition to fix a window mismatch bug: V was lagging kN0 behind K after the large jump, loading from the wrong sequence position. - Bias window, dropout seq_offset, and mask type (LogitsSinkMask) updated for sink-awareness. **Traits / codegen** (`tile_fmha_traits.hpp`, `fmha_fwd.hpp`, `fmha_batch_prefill.py`): - `TileFmhaBatchPrefillTraits` gains `kHasSink_` (was hardcoded `false`). - Codegen adds `F_sink` field; skips batch-mode kernels (group mode required). - CMake test filter broadened from 9 → 33 instances covering fp16/bf16 × mask/nmask × lse/nlse × sink/nsink. **Runner** (`fmha_fwd_runner.hpp`, `CMakeLists.txt`): - `fmha_batch_prefill()` dispatched from `run_fwd` when: group mode + paged KV + num_splits == 1. - K/V strides corrected for runner's [num_pages, nhead_k, page_block_size, hdim] layout. - `page_block_size % 128` check relaxed: batch_prefill supports ps=16. - CPU reference paged-KV reordering guards extended with `CK_TILE_FMHA_FWD_BATCH_PREFILL_API`. ## Test Plan Build with `-DFMHA_FWD_ENABLE_APIS="fwd;batch_prefill"`, run `tile_example_fmha_fwd` in group mode with page_block_size=16. Test matrix: - Mask: no-mask, causal, sliding window - Sink: nsink, sink=1..128 - dtype: fp16, bf16 - LSE output: on/off - seqlen ∈ {512,1024,2048,4096} × window ∈ {32,256,512,1024} - GQA, chunked prefill, large batch×seqlen - page_block_size: 16, 32 ## Test Result 171 test cases, all valid:y: - nmask + nsink: ✓ - causal + nsink: ✓ - causal + sink=8: ✓ - sliding window + sink=8 (d=128, d=256): ✓ - bf16, LSE output, GQA: ✓ ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Hosang Yoon
|
720cc88a31 |
[rocm-libraries] ROCm/rocm-libraries#6253 (commit 61934c6)
[CK_TILE] Enable canonical-NaN BF16 conversion for FMHA on RDNA (#6253) ## Motivation - On gfx11/gfx12, the existing float -> bf16 conversion path in FMHA forward adds noticeable overhead and causes a meaningful performance gap versus fp16. The asm-based path (mode 3) does not improve this on RDNA and can perform even worse. - In particular, on gfx12, bf16 FMHA forward can be up to ~20% slower than the corresponding fp16 path. - This PR reduces that gap by switching FMHA forward to a different BF16 conversion strategy based on Triton’s canonical-NaN round-to-nearest-even behavior. ## Technical Details - Add a new `standard_cnan` BF16 conversion mode to CK Tile. - Implement a canonical-NaN RTN `float -> bf16` conversion path based on the Triton implementation. - Enable this conversion mode by default for FMHA forward builds targeting gfx11/gfx12. - Retune gfx11/gfx12 FMHA forward kernel selection thresholds for some `hdim=128` cases to keep kernel selection aligned with the updated conversion behavior. ## Test Plan ./build/bin/tile_example_fmha_fwd -prec=bf16 -mode={0/1} -b=1 -h=16 -d={hdim} -s={seqlen} -s_k={seqlen} -lse=0 -iperm={0/1} -operm={0/1} ## Test Result - all tests passed when running `test_ck_tile_fmha` - BF16 FMHA forward performance improves by up to ~5% on gfx11. - BF16 FMHA forward performance improves by up to ~10% on gfx12. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Hosang Yoon
|
7e4e291771 |
[rocm-libraries] ROCm/rocm-libraries#6450 (commit b75fed1)
[CK_TILE] Skip padded k/n fragment work in qr_hpad FMHA fwd (#6450) ## Motivation `qr_hpad` currently executes work for padded head-dim fragments even when only a subset of the values are valid. This adds unnecessary computation for head dimensions that require padding, such as `hdim=72` and `hdim=80`, and hurts FMHA forward performance. The goal of this PR is to make the padded-head-dim path skip invalid work based on the actual valid fragment count, while preserving the existing behavior for the non-padded path. ## Technical Details This PR improves the `qr_hpad` FMHA forward path in three parts: - Skip padded `k`/`n` fragments in the GEMM/pipeline path when only part of the fragment is valid. - Add partial GEMM0 tail handling for `qr_hpad` so the kernel uses the valid fragment range instead of always computing over the padded extent. - Retune the gfx11 `qr_hpad` kernel configuration after enabling the partial-fragment path. To keep the existing path stable, the implementation adds overloads for the updated GEMM/pipeline interfaces. This allows existing full-tile callers to keep using the previous form, while the `qr_hpad` path can pass valid fragment counts when needed. ## Test Plan ./build/bin/tile_example_fmha_fwd -prec=bf16 -mode={0/1} -b=1 -h=16 -d={72/80} -s={seqlen} -s_k={seqlen} -lse=0 -iperm={0/1} -operm={0/1} ## Test Result - On gfx11 and gfx12, for head dimensions that require padding, `tile_example_fmha_fwd` shows about 20-30% performance improvement at `hdim=72/80`. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Po Yen Chen
|
c0d9614af6 |
[rocm-libraries] ROCm/rocm-libraries#6305 (commit 19e10a0)
[CK] Remove obsolete benchmark_fwd_v3.sh script and README reference (#6305) The tile_example_fmha_fwd_v3 target no longer exists in this project, making this benchmark script non-functional. |
||
|
Max Podkorytov
|
3aee45e115 |
[rocm-libraries] ROCm/rocm-libraries#5383 (commit b660b8c)
[CK_TILE] Add CShuffleLds microbenchmark suite (#5383) ## Summary Microbenchmarks isolating LDS store/load operations in CShuffleEpilogue for bank conflict analysis. ## Motivation CShuffleEpilogue performs LDS store (MFMA registers → LDS) and load (LDS → registers for coalesced global writes). This suite isolates each operation to: - Identify which operation causes bank conflicts - Measure pure LDS bandwidth per access pattern - Validate access patterns across MFMA tile sizes and wave layouts ## Components - **Microkernels** (`tile_load_store_microkernels.hpp`): `StoreTile<Setup>`, `LoadTile<Setup>` - **Setup Adapters** (`benchmark_cshuffle_lds.hpp`): Wire CShuffleEpilogue to microkernels - **Template** (`benchmark_template.cpp.in`): Generated benchmarks with timing ## Build ```bash cmake -G Ninja -B build -S . \ -DGPU_TARGETS=gfx950 \ -DBUILD_CK_EXAMPLES=ON \ -DBUILD_CK_TILE_CSHUFFLE_LDS_BENCHMARKS=ON ninja -C build bench_lds_fp8_16x16x128_2x2_fp8 ``` ## New CMake Options | Option | Default | Description | |--------|---------|-------------| | `BUILD_CK_TILE_CSHUFFLE_LDS_BENCHMARKS` | OFF | LDS microbenchmarks | | `BUILD_CK_TILE_FMHA_TESTS` | ON | FMHA tests | | `BUILD_CK_TILE_ENGINE` | ON | Tile engine | | `BUILD_CK_TILE_ENGINE_TESTS` | ON | Tile engine tests | | `BUILD_CK_EXAMPLES` | ON | Examples | | `BUILD_CK_TUTORIALS` | ON | Tutorials | | `BUILD_CK_DEVICE_INSTANCES` | ON | Device instances | | `BUILD_CK_PROFILER` | ON | Profiler | Setting guards to OFF reduces cmake configure from ~150s to ~5s. --------- Made-with: Claude Code, Opus 4.5 |
||
|
msaffari-amd
|
cf517ec050 |
[rocm-libraries] ROCm/rocm-libraries#5863 (commit 31d9247)
[CK_TILE] Separate PermuteN epilogue from CShuffle epilogue into standalone file (#5863) ## Motivation The PermuteN epilogue was previously embedded within cshuffle_epilogue.hpp, despite having fundamentally different behaviour. Coupling these two independent strategies in one file introduced unnecessary complexity, SFINAE guards, and a dual operator() overload selected at compile time via TiledMMAPermuteN_ template parameter. This PR separates PermuteN into its own standalone file(pertmuten_epilogue.hpp), simplifying both implementations and making the codebase easier to maintain and extend independently. ## Technical Details **New file: permuten_epilogue.hpp:** contains PermuteNEpilogueProblem and PermuteNEpilogue, extracted from the permuteN code path in cshuffle_epilogue.hpp. **Cleanup of cshuffle_epilogue.hpp:** - Removed the TiledMMAPermuteN_ template parameter from [CShuffleEpilogueProblem] - Removed the SFINAE-guarded permuteN operator() overload - Removed the EnablePermuateN_ SFINAE alias - CShuffle now only contains CShuffle logic; EightWave support (independent feature) is retained **Consumer migration :** All consumer files now use compile-time epilogue selection via [std::conditional_t] `using GemmEpilogue = std::conditional_t< TiledMMAPermuteN, PermuteNEpilogue<PermuteNEpilogueProblem<...>>, CShuffleEpilogue<CShuffleEpilogueProblem<...>>>;` **Files modified:** - flatmm_basic.cpp, moe_flatmm.cpp, a16w4_moe_flatmm.cpp, mixed_prec_flatmm.cpp, mx_flatmm_instance.hpp — flatmm examples - run_gemm_quant_example.inc — block-scale GEMM example - gemm_weight_preshuffle_invoker.hpp — weight preshuffle invoker - test_gemm_quant_fixtures.hpp, test_gemm_persistent_async_input.cpp, test_gemm_pipeline_util.hpp — test utilities - universal_gemm_invoker.hpp — universal GEMM invoker - epilogue.hpp — add header updated to include permuten_epilogue.hpp ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com> |
||
|
Kiefer van Teutem
|
b7156a8bbf |
[rocm-libraries] ROCm/rocm-libraries#5515 (commit 40f66c3)
[CK Tile] Add Tile Distribution Encoding Calculator (#5515) ## Motivation We want to be able to calculate TileDistributionEncodings describing register mappings for any MmaOp. This is necessary for further integration with CK Tile. This MR adds a new struct TileDistrEncCalc, which takes an amdgcn_mma type (MmaOp) and provides ABC warp distribution encodings for mapping matrix fragment coordinates to register coordinates (lane, vector item) and vice versa. It is able to take CTranpose, Swizzle, and NumAccessA / NumAccessB template parameters for tweaking the tile distributions. Swizzle modification will be implemented later. The current implementation can deal with all intrinsic types and block-hiding. This MR also adds some additional static asserts and derived params within amdgcn_mma_base, to enforce consistency and help calculate Tile Distributions for block-hiding intrinsics. An Example was added that uses the Tile Distr Enc Calc to calc and print register layouts for Tile Distributions for some of our amdgcn_mma structs. It also makes sure that the CTranspose modifier works as intended. Some additional gfx9 intrinsics were added to test block-hiding layouts for the different types of C-block-hiding layouts. The sparse intrinsic wrappers were updated according to Chris's recent changes in another branch (https://github.com/ROCm/rocm-libraries/pull/5508), which moved the compression step outside of the intrinsic itself. This is necessary to make sure that the Calculator can deal with this new interpretation of the sparse intrinsics. I directly copied the new amdgcn structs from Chris's branch and changed nothing else to avoid more complex merges in the future. Note that this means I did not update a bunch of related sparse code since that would be a lot, and therefore I disabled test_amdgcn_sparse_mma for now. The amdgcn_mma_layout test was refactored a bit: - The old register mapping utility was removed and its use was replaced by the new TileDistrEncCalc - More tests were added to test layouts for different types of block-hiding and sparse intrinsics - The Selector method was removed and the tests were split up over target architectures, with each target arch having a direct list of amdgcn structs to be tested. This ensures that we force specific tests on specific architectures and makes sure that the selector doesn't quietly do some workarounds like creating compound intrinsics. ## Test Results Layout tests based on calculated tile distribution encodings pass on all architectures. Calculator works for all currently added amdgcn structs, which includes different types of block-hiding and sparse intrinsics. Printed layouts from new example verified by eye. CTranspose modifier tested for large set of intrinsics. |
||
|
Aviral Goel
|
81a5826132 |
[rocm-libraries] ROCm/rocm-libraries#6323 (commit a668483)
CK: Extract shared boilerplate from 47 gemm_quant test files (#6323) Depends on #6303 ## Summary Extract shared test boilerplate (includes, type aliases, test fixture macros) from 47 `test_gemm_quant_*` files into a single `test_gemm_quant_common.hpp` header. Each test file is reduced from ~50 lines of boilerplate to ~5 lines. | Metric | Value | |--------|-------| | Files changed | 48 | | Insertions | +413 | | Deletions | −1,106 | | **Net lines removed** | **−693** | ### What changed | Before | After | |--------|-------| | 47 test files, each with ~50 lines of identical includes, type aliases, and fixture macros | 1 shared header (`test_gemm_quant_common.hpp`) + 47 thin files (~5 lines each: include + params) | ### Readability assessment A code realist review confirmed this change **improves readability**: the 47 test files had identical boilerplate obscuring the only meaningful content — the `GemmConfig` type alias and test dimensions. After the refactoring, each file's unique configuration is immediately visible, and adding a new test variant requires specifying only the varying parameters instead of copying 50 lines. ### Cumulative cleanup series stats | PR | Description | Net lines | |----|-------------|-----------| | #6300 | Remove 61 dead `#if 0` blocks | −2,648 | | #6302 | Remove 41 commented-out dead code blocks | −2,861 | | #6303 | Remove 4 orphaned files | −3,886 | | This PR | Extract gemm_quant test boilerplate | −693 | | **Total** | | **−10,088** | |
||
|
Aviral Goel
|
c7eb33078c |
[rocm-libraries] ROCm/rocm-libraries#6302 (commit 8d419e8)
CK: Remove 41 commented-out dead code blocks (~200 lines) (#6302) Depends on #6300 ## Summary Remove 41 commented-out code blocks across 33 files in Composable Kernel, totaling ~200 lines. Identified using an automated dead code scanning skill (`ck-dead-code`) with a calibrated two-stage pipeline: 1. **Pre-filter**: Keyword-based scan found 1,338 `//`-commented blocks. Calibrated heuristics (trained on 50-sample expert classification) reduced to 89 high-confidence candidates — 93% noise reduction. 2. **Expert triage**: LLM expert classified each block in context as CODE_REMOVE, CODE_KEEP, or NOT_CODE. | Classification | Count | |---------------|-------| | Removed (this PR) | 41 | | Kept (debug helpers, alt configs, reference impls) | 32 | | Not code (false positives) | 16 | Removed blocks include: superseded implementations, old test data, abandoned stubs, unreachable code, and buggy dead code. |
||
|
Hosang Yoon
|
00faecdfcf |
[rocm-libraries] ROCm/rocm-libraries#6156 (commit 367565a)
[CK_TILE] Optimize FMHA head-dim padded path on gfx11/gfx12 (#6156) ## Motivation On gfx11/gfx12, FMHA forward kernels that require head-dim padding show a large performance drop compared to the exact-head-dim path. In practice, padded cases such as `HDIM=72` and `HDIM=80` were falling too far off the fast path. This PR improves padded-head-dim FMHA performance on gfx11/gfx12 while keeping the behavior for other GPUs unchanged. ## Technical Details - Add/scope a dedicated padded-head-dim (`qr_hpad`) FMHA forward path for gfx11/gfx12. - For `receipt=0`, keep support conservative and only enable the padded fast path for vector-safe cases (`head_dim % 8 == 0`), matching the existing assumption used on other GPUs. - Move `v_prefetch` later only for the head-dim-padded path on gfx11/gfx12. This reduces live ranges and removes the register-spill behavior seen in the earlier scheduling. - Enable the buffer-load OOB check offset trick for the padded path on gfx11/gfx12. ## Test Plan ./build/bin/tile_example_fmha_fwd -prec=bf16 -mode={0/1} -b=1 -h=16 -d={72/80} -s={seqlen} -s_k={seqlen} -lse=0 -iperm={0/1} -operm={0/1} ## Test Result Observed padded-head-dim performance improvements for HDIM=72/80: - gfx11: about ~3.5x - gfx1151: about ~2.0x - gfx12: about ~1.3x ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Copilot
|
31fd3fe9c1 |
[rocm-libraries] ROCm/rocm-libraries#6003 (commit b4c8c52)
[CK_TILE] Refine FMHA Readme (#6003) Updates the FMHA README to document fp8 precision support more accurately, replacing the outdated "experimental" section and incomplete CLI arg descriptions. ## Changes - **`-prec` arg**: expanded supported values from `fp16/bf16/fp8/bf8` → `fp32/fp16/bf16/fp8/fp8bf16/fp8fp32/mxfp8/mxfp4` - **`-qscale` arg**: replaced single-line `1: per-tensor quantization` with all four modes: `pt/1`, `bs/2`, `kvbs/3`, `mx/4` - **FP8 support section**: replaced "FP8 experimental support" paragraph with: - Supported targets: gfx942/gfx950 + ROCm 6.0+ - Table distinguishing `fp8` / `fp8bf16` / `fp8fp32` by Q/K/V input type and output type - Table for all `-qscale` modes with descriptions - Note that `-vlayout=r` (`seqlen*hdim` for V) is the only supported layout for fp8 types <!-- START COPILOT ORIGINAL PROMPT --> <details> <summary>Original prompt</summary> Please open a PR against base branch `develop` in repository `ROCm/rocm-libraries` applying the following documentation updates within the composable kernel path. ## Scope Update the file: - `projects/composablekernel/example/ck_tile/01_fmha/README.md` ## Changes to apply Apply the combined edits described in the diffs below (two consecutive patches). Ensure the final file content includes **both** sets of changes. ### Patch 1 - In the CLI args section: - Update `-qscale` description lines to include: - `pt or 1, per-tensor scale` - `bs or 2, block scale` - `kvbs or 3, Q per-tensor, K/V per-page block scale` - `mx or 4, microscaling (exclusively for mxfp8/mxfp4)` - Update `-prec` supported data types from `fp16/bf16/fp8/bf8` to `fp32/fp16/bf16/fp8/fp8bf16/fp8fp32/mxfp8/mxfp4`. - Replace the existing "FP8 experimental support" section with an "FP8 support" section stating: - FP8 FMHA kernels supported on gfx942/gfx950 with ROCm 6.0+ - Precision selectable via `-prec=fp8` (or `fp8bf16`, `fp8fp32`) for `tile_example_fmha_fwd` - Add a table describing `-qscale` modes: - `n` or `0`: No quantization scale (default) - `pt` or `1`: Per-tensor quantization scale - `bs` or `2`: Per-block quantization scale - `kvbs` or `3`: Q per-tensor + K/V per-page block scale - `mx` or `4`: Microscaling (MX format), exclusively for `mxfp8` and `mxfp4` - Add/keep note that currently only `-vlayout=r` (`seqlen*hdim` for V matrix) is supported for fp8 data types. ### Patch 2 Further refine the "FP8 support" paragraph to explain the difference between `fp8`, `fp8bf16`, and `fp8fp32` via a table: | `-prec` value | Q/K/V input type | Output type | Description | |---|---|---|---| | `fp8` | fp8 | fp8 | Fully fp8: both inputs and output are in fp8 | | `fp8bf16` | fp8 | bf16 | Mixed precision: fp8 inputs, bf16 output — useful when the consumer expects a wider-range output format | | `fp8fp32` | fp8 | fp32 | Mixed precision: fp8 inputs, fp32 output — highest-precision output, suitable for debugging or further fp32 processing | Keep the rest of the `-qscale` table and the `-vlayout=r` limitation note. ## Notes - PR title must be: `[CK_TILE] Add fp8 in FMHA readme` - Ensure markdown formatting renders correctly (tables, code formatting). - Only modify the file listed above. The following is the prior conversation context from the user's chat exploration (may be truncated): User: 能幫我上這個pr嗎 在composable kernel裡的路徑 diff --git a/example/ck_tile/01_fmha/README.md b/example/ck_tile/01_fmha/README.md index 0b526f4e9fc..1627435863b 100644 --- a/example/ck_tile/01_fmha/README.md +++ b/example/ck_tile/01_fmha/README.md @@ -62,14 +62,17 @@ args: -d_v head dim for v, -1 means equal to d (default:-1) -scale_s scale factor of S. 0 means equal to 1/sqrt(hdim). (default:0) -qscale n or 0, no scaling (default:n) - 1: per-tensor quantization. + pt or 1, per-tensor scale + bs or 2, block scale + kvbs or 3, Q per-tensor, K/V per-page block scale + mx or 4, microscaling (exclusively for mxfp8/mxfp4) -iperm permute input (default:1) if true, will be b*h*s*d, else b*s*h*d -operm permute output (default:1) -bias n or 0, no bias (default:n) e(lementwise) or 1, elementwise bias with 1*1*s*s. e:1, 1*h*s*s. e:2, b*h*s*s a(libi) or 2, alibi with 1*h. a:1, b*h - -prec data type. fp16/bf16/fp8/bf8 (default:fp16) + -prec data type. fp32/fp16/bf16/fp8/fp8bf16/fp8fp32/mxfp8/mxfp4 (default:fp16) -mask 0: no mask, 1: top-left(same as 't'), 2:bottom-right(same as 'b') (default:0) 't', top-left causal mask, 'b', bottom-r causal mask 't:l,r', top-left sliding window attn(swa) with FA style left right size @@ -161,7 +164,17 @@ We support sequence padding and variable-length processing in both batch and gro Both approaches optimize memory access patterns while supporting flexible sequence length requirements commonly found in transformer inference scenarios. -## FP8 experimental support -As described in [this blog](https://blog.hippoml.com/8bit-hippoattention-up-to-3x-faster-compared-to-flashattentionv2-8f9def90b482), we have an experimental support for fp8 fmha kernels, you can evaluate the performance by setting the arg `-prec=fp8` to the `tile_example_fmha_fwd`, on a gfx942 machine and ROCm 6.0+. +## FP8 support +FP8 FMHA kernels are supported on gfx942/gfx950 machines with ROCm 6.0+. You can select fp8 precision by setting the arg `-prec=fp8` (or `fp8bf16`, `fp8fp32`) to the `tile_example_fmha_fwd`. -Currently we only support `-vlayout=r`( `seqlen*hdim` for V matrix) for fp8 and fp8bf16 now. Full feature support will come later. +The following quantization scale modes are available via `-qscale`: + +| `-qscale` value | Description | +|---|---| +| `n` or `0` | No quantization sca... </details> <!-- START COPILOT CODING AGENT SUFFIX --> *This pull request was created from Copilot chat.* > --------- 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> |
||
|
Po Yen Chen
|
bf736dfa74 |
[rocm-libraries] ROCm/rocm-libraries#6051 (commit f0838b2)
[CK] Add FP8 per-tensor quantization support for FMHA V3 pipeline (#6051) ## Motivation The existing FMHA V3 pipeline only supports fp16/bf16 data types. This PR extends V3 to handle FP8 inputs with per-tensor descaling on gfx950, enabling higher throughput for FP8 inference workloads using the assembly-optimized V3 code path. ## Technical Details **Warp GEMM:** - Add FP8 32x32x32 warp gemm with C-transposed distribution (`WarpGemmMfma_f32_32x32x32_fp8_fp8_CTransposed`) and dispatcher entries **V3 Kernel (`fmha_fwd_v3_kernel.hpp`):** - Add per-tensor descale support for Q, K, V tensors, passing descale pointers through to pipeline kargs **V3 Pipeline (`block_fmha_fwd_v3_pipeline.hpp`):** - Add FP8 data path with dtype-aware type selection - Add asm volatile P matrix conversion from f32 to fp8 - Add FP8-aware instruction scheduling in `CoreLoopScheduler` **V3 Pipeline Policy (`block_fmha_fwd_v3_pipeline_default_policy.hpp`):** - Add FP8 QK warp gemm selection (SwizzleB variant for V tile distribution compatibility) **Codegen (`fmha_fwd.py`):** - Add gfx950 FP8BF16 V3 tile size (256x64x128x128x64x128) - Add FP8BF16 V3 pipeline variants (mask: no/causal, qscale: no/pertensor) - Extend `can_dispatch_v3` condition for fp8bf16 + pertensor **Misc:** - Add LLVM scheduler `TRANS` mask to `LLVMSchedGroupMask` enum (`arch.hpp`) - Fix `mask_info` default initialization for `no_mask` case (`mask.hpp`) V3 dispatch for FP8 is disabled by default (`F_is_v3_enabled=false`) pending further validation. ## Performance: fmha_fwd V3 FP8 (avg runs 2-6, stock ROCm 7.1.1, gfx950) | Problem | Regular (TFlops) | Varlen (TFlops) | |---|---:|---:| | batch=1 heads=6/1 seqlen=1024 causal | 48.9 | 47.6 | | batch=1 heads=6/1 seqlen=2048 causal | 119.8 | 117.4 | | batch=1 heads=6/1 seqlen=4096 causal | 263.7 | 259.2 | | batch=1 heads=6/1 seqlen=8192 causal | 548.9 | 543.6 | | batch=1 heads=6/1 seqlen=16384 causal | 1043.0 | 1063.7 | | batch=1 heads=6/1 seqlen=32768 causal | 1237.2 | 1279.6 | | batch=1 heads=6/1 seqlen=65536 causal | 1315.4 | 1382.7 | | batch=1 heads=6/1 seqlen=131072 causal | 1326.3 | 1402.2 | | batch=1 heads=16/1 seqlen=65536 causal | 1298.7 | 1388.4 | | batch=1 heads=40/40 seqlen=37200 non-causal | 1248.9 | 1326.1 | ## Test Plan Tested with aiter's `test_mha_fp8.py` test suite (176 cases) covering batch sizes (1-2), sequence lengths (113-4096), head counts (5/8/32/40), GQA ratios (1:1, 1:8), and causal/non-causal modes. Verified all cases dispatch to the V3 pipeline by enabling `F_is_v3_enabled` and confirming kernel names contain `qr_async_trload_v3`. ## Test Result 176/176 tests passed with V3 enabled. All cases correctly dispatched to V3 pipeline with `pertensor` quantization. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |
||
|
Hosang Yoon
|
e18f4d8e6d |
[rocm-libraries] ROCm/rocm-libraries#6038 (commit d7041a2)
[CK_TILE] Restrict FMHA codegen to the kernel subset used by FlashAttention (#6038) ## Motivation Currently, the CK FlashAttention integration generates a broader FMHA kernel set than the FlashAttention wrappers can actually dispatch, which increases compile time without improving runtime coverage. ## Technical Details The FlashAttention CK wrappers do not use all logits/LSE variants emitted by the default FMHA codegen. The direct `fmha_fwd` path always uses softcap-disabled, LSE-enabled kernels, and the `fmha_fwd_splitkv` path only uses softcap-disabled kernels. This change trims codegen to that subset and stops generating the unused logits/LSE variants. This reduces the generated forward kernel set without changing `fmha_fwd_appendkv` or `fmha_bwd`. The reduced kernel set was validated by building and running the [FlashAttention](https://github.com/Dao-AILab/flash-attention) CK backend. Across targets, the total generated FMHA kernel count is reduced by: - `gfx942`: 29.3% - `gfx1100`: 33.7% - `gfx1201`: 31.3% ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> pytest test/test_flash_attn_ck.py from https://github.com/Dao-AILab/flash-attention ## Test Result all tests passed <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. |