[CK Tile] Add sink token gradient support in FMHA backward pass by LJ-underdog · Pull Request #5504 · ROCm/rocm-libraries

LJ-underdog · 2026-03-17T03:17:05Z

Motivation

Adds sink token support to the FMHA backward kernel (dot_do_o pipeline):

Technical Details

Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType
Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs
Compute per-head sink gradient via atomic accumulation in the pipeline
Update example runner with reference validation for sink gradient

Test Plan

Add new test case

Test Result

WIP

Submission Checklist

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

Adds sink token support to the FMHA backward kernel (dot_do_o pipeline): - Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType - Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs - Compute per-head sink gradient via atomic accumulation in the pipeline - Update example runner with reference validation for sink gradient

Replace the fixed scalar sink_val with per-head random values sampled from uniform distribution [30, 60] to improve test coverage.

Copilot

Pull request overview

Adds sink token support for the FMHA backward dot(dO, O) pipeline, including optional per-head sink gradient accumulation and example-side reference validation.

Changes:

Extend the dot(dO, O) pipeline problem/kernel args to carry LSE and sink-related pointers/metadata.
Compute per-head sink gradient in-kernel via atomic accumulation (optional path).
Update the example runner + reference path to validate sink gradients; adjust FMHA bwd codegen + codegen CMake args.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp	Adds LSEDataType to pipeline problem typing for sink-related computations.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp	Adds optional sink-grad path to dot(dO,O) computation and atomic accumulation.
projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp	Threads sink/lse pointers and nhead into kargs; wires sink value + sink grad pointer into pipeline.
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd_runner.hpp	Adds sink score generation, device buffers, and reference validation for sink gradients.
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd.hpp	Extends example args struct and kargs creation to pass sink pointers.
projects/composablekernel/example/ck_tile/01_fmha/example_fmha_bwd.cpp	Enables sink_grad validation by default in the example runner invocation.
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py	Updates codegen template instantiation to include LSEDataType.
projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt	Changes codegen optdim set (now only 128) and adds commented compile options.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

- Always accumulate sink gradient in float regardless of DDataType to avoid precision loss (fp16/bf16) and ensure atomicAdd compatibility across all architectures. - Replace per-thread atomicAdd with a warp-level shuffle reduction followed by a single atomicAdd per warp, reducing global memory contention by warp_size (e.g. 256 atomics -> 4 for a 256-thread block). - Allocate sink_buf and d_sink_buf conditionally on sink_grad to avoid unnecessary device memory allocation and H2D transfers. - Eliminate intermediate sink_buf_dev and d_sink_ptr_dev variables, passing GetDeviceBuffer() directly into fmha_bwd_args.

Copilot

Pull request overview

Adds sink-token gradient support to the FMHA backward “dot(dO, O)” pipeline and threads the required inputs/outputs through the kernel + example validation.

Changes:

Extends the bwd dot-do-o pipeline problem/kernel to carry LSE and sink pointers, and accumulates per-head sink gradients via atomics.
Updates the example runner to generate sink inputs and validate d_sink against a reference.
Adjusts codegen/CMake configuration for FMHA bwd example builds.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp	Adds `LSEDataType` to the pipeline problem types to support LSE loads for sink grad.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp	Implements optional sink-gradient accumulation path inside dot(dO,O) pipeline.
projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp	Wires `lse_ptr`, `sink_ptr`, `d_sink_ptr`, and `nhead` into kargs and invokes pipeline with sink params.
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd_runner.hpp	Adds sink input allocation/init and validation of `d_sink` against reference.
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd.hpp	Extends `fmha_bwd_args` and kargs creation to pass sink + LSE into dot-do-o kernel.
projects/composablekernel/example/ck_tile/01_fmha/example_fmha_bwd.cpp	Enables sink-grad validation by default (currently hard-coded).
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py	Updates generated pipeline instantiation to include `LSEDataType`.
projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt	Restricts `--optdim` to 128 (and adds commented compile options).

Comments suppressed due to low confidence (2)

projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd_runner.hpp:1

sink_buf/d_sink_buf are constructed with size 0 when sink_grad == false, but GetDeviceBuffer() is still passed into fmha_bwd_args. If DeviceMem(0) returns a non-null sentinel pointer, the kernel will treat sink as enabled and read/atomicAdd through an invalid pointer. Pass nullptr explicitly when sink_grad is false (e.g., sink_grad ? sink_buf.GetDeviceBuffer() : nullptr and same for d_sink_buf).

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.

projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt:1

The PR description focuses on sink-token gradient support, but this change restricts codegen --optdim from multiple dims to only 128, which can reduce supported configurations/perf coverage for the example binaries. If this restriction is required for this PR (e.g., compile time constraints), please note it in the PR description or keep the previous list.

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

sink scores and their gradients live in the same log-space as LSE, so they should share LSEDataType rather than a hardcoded float. - fmha_bwd_args: annotate sink_ptr/d_sink_ptr as LSEDataType in comments; keep void* for API consistency with lse_ptr and other typeless pointers. - FmhaBwdOGradDotOCommonKargs: type sink_ptr/d_sink_ptr as const LSEDataType*/LSEDataType* so the kernel body needs no casts. - MakeKargs: cast void* -> LSEDataType* at the single API boundary. - Kernel body: sink_value and atomic_sink_grad_ptr use LSEDataType; -inf sentinel uses numeric<LSEDataType>::infinity(). - Pipeline: sink_value and atomic_sink_grad_ptr parameters use LSEDataType; type_convert<float> guards float arithmetic in exp computation. - Runner: sink_host and d_sink_host changed from AccDataType to LSEDataType.

- example_fmha_bwd.cpp: expose sink_grad as a command-line argument (-sink_grad=0/1, default 0) instead of a hardcoded constant. - CMakeLists.txt: restore full --optdim list (32,64,96,128,256) and remove stray debug compile options. - script/smoke_test_bwd_sink.sh: new smoke test script for backward sink gradient validation, covering fp16/bf16 x hdim(64/128/256) x mode x bias x dropout, plus non-standard hdims (40/48/72/96). - script/smoke_test_fwd_sink.sh: refactor existing script to match the standard smoke test structure: add run_exe wrapper with fail tracking, fix GPU_arch default, add CK_WARMUP/CK_REPEAT exports, use COMMON_ARGS, expand coverage to bf16 and hdim 64/128/256, and organise cases into run_sink_mask_tests / run_sink_init_tests functions.

Consolidate sink test coverage into the primary smoke test scripts instead of maintaining separate scripts: - smoke_test_bwd.sh: add sink gradient test loops (same coverage as main tests but with -sink_grad=1) and non-standard hdim cases. - smoke_test_fwd.sh: add run_sink_mask_tests() and run_sink_init_tests() functions covering sink-specific mask patterns and init_sink path. - run_full_test.sh: remove now-redundant separate sink script calls. - Delete smoke_test_fwd_sink.sh and smoke_test_bwd_sink.sh.

Replace ck_tile::exp with exp2 in the sink backward path to match the pattern used by other bwd pipelines (e.g. block_fmha_bwd_dq_dk_dv_*). exp2 maps directly to the v_exp_f32 hardware instruction on AMD GPUs, avoiding the extra log2e multiply that exp requires internally. Pre-multiply sink_value by log2e at the kernel call site so that exp2(sink_value - log2e*lse) == exp(raw_sink - lse) preserving numerical equivalence. -inf is unchanged by the scaling.

- Rename ls_e_dram_block_window_tmp to lse_dram_block_window_tmp for consistency with the rest of the codebase - Add static_assert(std::is_same_v<LSEDataType, float>) to guard the reinterpret_cast<float*> in the sink atomicAdd path - Allocate sink_host with a dummy shape when sink_grad is disabled, matching the pattern used by bias_host/alibi_slope_host

…ariant docs Rename nhead_stride_d/batch_stride_d -> nhead_stride_lsed/batch_stride_lsed in FmhaBwdOGradDotOKernel kargs to make explicit that LSE and D always share the same layout, matching the naming convention already used by the main FmhaBwdKernel (nhead_stride_lsed/batch_stride_lsed). This eliminates the implicit assumption that was previously undocumented and potentially confusing when layouts differ. Also rename the local batch_offset_d -> batch_offset_lsed for consistency.

- block_fmha_bwd_dot_do_o.hpp: add comment at atomicAdd site noting that d_sink accumulation is non-deterministic across runs even when deterministic=1, due to floating-point atomicAdd ordering. - fmha_bwd_runner.hpp: only print sink info in the log line when sink_grad=true; previously printed unconditionally, misleadingly showing "sink:(rand[30,60], const)" when sink was disabled.

LJ-underdog · 2026-03-31T03:08:56Z

Review notes (post item-by-item triage)

Three minor items after verifying each finding against the codebase:

block_fmha_bwd_dot_do_o.hpp:153 — The atomicAdd into d_sink_ptr makes the result non-deterministic across runs, even when -deterministic=1. Since d_sink is a single scalar per head the practical impact is negligible, but a one-line comment at the atomicAdd call site noting this would be helpful.

fmha_bwd_runner.hpp:439 — The log line ", sink:(rand[30,60], " << (sink_grad ? "grad" : "const") << ")" prints unconditionally. When sink_grad=false, sink is completely disabled (sink_ptr is nullptr), but the output says sink:(rand[30,60], const) which is misleading.

smoke_test_fwd.sh — run_sink_init_tests — The third run_exe call inside the for mask in 0 1 loop doesn't use $mask. It runs the identical command twice (once per loop iteration), wasting CI time. It should either use -mask=$mask or be moved outside the mask loop.

…k + large seqlen_k The sink modification block in the BWD reference ran unconditionally even when sink_grad=false. The block computed exp(lse) which overflows float32 (FLT_MAX ~3.4e38, threshold exp(88.72)) when the attention LSE is large. With alibi bias and a bottom-right causal mask, LSE can exceed 88.72 when: - h=5 (non-power-of-2) introduces head slope=0.5 (largest among 5 heads) - seqlen_k > seqlen_q, extending the visible key range per query row - e.g. h=5, s_q=128, s_k=256: LSE for q>=45 reaches ~89, exp(89) overflows to +inf, causing p_scale = inf/inf = NaN, corrupting all gradients. Two fixes applied to fmha_bwd_runner.hpp: 1. Guard the sink block with if(sink_grad) so it is skipped entirely when the kernel does not use a sink pointer (sink_grad=false). 2. Replace exp(lse_old)/exp(lse_new) with numerically stable log-domain arithmetic: diff = sink_val - lse_old lse_new = lse_old + log(1 + exp(diff)) p_scale = 1 / (1 + exp(diff)) This avoids exp(lse) entirely and remains correct for all LSE values. All 1008 fp16 and bf16 fmha_bwd gtests pass (672 passed, 336 skipped due to missing kernel instances).

poyenc

All review comments addressed. LGTM.

…ked rows) The previous fix used diff = sink_val - lse_old to compute lse_new, but when lse_old = -inf (fully-masked rows, e.g. s_q > s_k with bottom-right causal mask), diff = +inf and lse_new = -inf + log(1 + inf) = NaN. This NaN lse was written to lse_buf and sent to the GPU kernel, causing the kernel to produce NaN gradients while the reference correctly outputs 0. Fix: use the max-based log-sum-exp formula which is stable for all inputs: hi = max(lse_old, sink_val) lo = min(lse_old, sink_val) lse_new = hi + log(1 + exp(lo - hi)) p_scale = exp(lse_old - lse_new) When lse_old = -inf: hi = sink_val, lo = -inf, exp(-inf - sink_val) = 0, lse_new = sink_val (finite), p_scale = exp(-inf) = 0. No NaN. When lse_old >> sink_val: hi = lse_old, lo - hi = diff < 0, reduces to the same stable formula as before.

[CK Tile] Add sink token gradient support in FMHA backward pass (#5504) ## Motivation Adds sink token support to the FMHA backward kernel (dot_do_o pipeline): ## Technical Details - Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType - Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs - Compute per-head sink gradient via atomic accumulation in the pipeline - Update example runner with reference validation for sink gradient ## Test Plan Add new test case ## Test Result WIP ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

This reverts commit 7481fd6.

* CK mha bwd: add sink attention score gradient support * test: add varlen sink bwd tests to test_mha_sink_bwd * Update op_tests/test_mha_sink_bwd.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * style: apply black formatting to test_mha_sink_bwd * test: move sink bwd tests into test_mha.py and test_mha_varlen.py * style: apply black formatting to sink bwd tests in test_mha and test_mha_varlen * fix: adapt mha bwd to updated CK fmha_bwd API and zero dq_accum Three fixes required after the CK submodule was updated to the sink_bwd_cherry_pick branch: 1. fmha_bwd_traits no longer carries seqlen/batch/nhead fields. Remove the now-stale seqlen_q, seqlen_k, batch, max_seqlen_*, nhead_q, nhead_k arguments from the traits initializer lists in mha_bwd.cu, mha_bwd_kernels.cu, and mha_varlen_bwd_kernels.cu. 2. nhead_stride_dq_acc / batch_stride_dq_acc are int64_t in mha_bwd_args but ck_tile::index_t (int) in fmha_bwd_args. Add explicit static_cast<ck_tile::index_t> to silence the narrowing-conversion errors. 3. fmha_bwd_launcher was removed from the new CK API. Replace launcher.dq_acc_splits with the equivalent expression ceil(seqlen_k / 16) for deterministic mode and 1 otherwise, matching the logic documented in fmha_bwd_runner.hpp. Replace launcher.needs_zero_dq_acc with unconditional torch::zeros: the dq_dk_dv kernel always writes dq_acc via atomicAdd (even in non-deterministic mode), so an uninitialized accumulator silently corrupts dQ for hdim >= 128 where the convert_dq kernel is active. All 22 sink-bwd tests pass after this change. * update ck to ROCm/rocm-libraries#5504 * Revert "update ck to ROCm/rocm-libraries#5504" This reverts commit 7481fd6. * update ck commit Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * update bwd args Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * [CK] update mha bwd traits args and fix sink_ptr comments * [CK] fix mha_bwd_args initializer in benchmark_mha_bwd.cpp for sink_ptr/d_sink_ptr --------- Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>

## Motivation Adds sink token support to the FMHA backward kernel (dot_do_o pipeline): ## Technical Details - Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType - Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs - Compute per-head sink gradient via atomic accumulation in the pipeline - Update example runner with reference validation for sink gradient ## Test Plan Add new test case ## Test Result WIP ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…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>

[CK Tile] Add sink token gradient support in FMHA backward pass (#5504) ## Motivation Adds sink token support to the FMHA backward kernel (dot_do_o pipeline): ## Technical Details - Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType - Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs - Compute per-head sink gradient via atomic accumulation in the pipeline - Update example runner with reference validation for sink gradient ## Test Plan Add new test case ## Test Result WIP ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

## Motivation 2 upstream CK Tile PRs were pushed and broke `rocm_ck`, so it needed to be adapted to the changes. This PR updates the `rocm_ck` bridge and tests. - **PR #5504 — `[CK Tile] Add sink token gradient support in FMHA backward pass`** `OGradDotO` `lse_ptr` / `sink_ptr` / `d_sink_ptr` / `p_undrop` / `seqlen_q` / `hdim_v` / `nhead` added to common kargs, `*_stride_d` renamed to `*_stride_lsed` (LSE and D has the same layout, so it covers both, mode dependent kargs split, shifting the `LSEDataType` parameter. - **PR #6152 — `[CK_TILE] Use Unified Workspace for FMHA BWD`** `dq_acc` is not provided by `acc_buf` field anymore. It is now a device tensor together with `nsplits_ptr` and, in group/varlen mode, `dq_acc_batch_offset_ptr` (per-batch element offset into the `dq_acc` buffer). ## Technical Details Fix plan is described in: #7865 | File | Description | | ---- | ----------- | | args.hpp | Increase `kMaxTensors` to 20 and update `Args` size/static_asserts. | | tests/test_args.cpp | Update ABI/size expectations for `Args` and capacity constants. | | tests/test_signature.cpp | Update capacity-limit expectation (`kMaxTensors`). | | ops/fmha_bwd/dqdkdv_spec.hpp | Add deterministic workspace slots (`NSPLITS`, `DQ_ACC_BATCH_OFFSET`) and update `requiredTensors()`. | | include/rocm_ck/ops/fmha_bwd/dqdkdv_api.hpp | Extend debug validation to include new slots and skip group-only slots in batch mode. | | include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp | Update DqDkDv device bridge to match CK Tile deterministic `Kargs` changes. | | tests/test_fmha_bwd_validate_args.cpp | Update death test to populate newly-required deterministic workspace slot. | | include/rocm_ck/ops/fmha_bwd/convert_dq_spec.hpp | Update ConvertDQ slot layout to include workspace-derived `NSPLITS`/offsets and revised `requiredTensors()`. | | include/rocm_ck/ops/fmha_bwd/convert_dq_dev.hpp | Update ConvertDQ device bridge to match CK Tile `Kargs` changes (nsplits ptr, nhead). | | tests/test_fmha_bwd_convert_dq.cpp | Update required-tensor-count expectations for the new slot layout. | | include/rocm_ck/ops/fmha_bwd/ograd_dot_o_dev.hpp | Update OGradDotO device bridge for CK Tile `Kargs` signature changes (LSE/sink fields). | </details> ## Test Plan ## Test Result - [x] `ctest -L rocm_ck --output-on-failure`: 64/64 pass - [x] `ctest -L compile_fail --output-on-failure`: 45/45 pass - [x] `ninja kpack_archive` produces non-zero .hsaco files for every entry in `KERNEL_VARIANTS` against `GPU_TARGETS=gfx942` - [x] `kernels.kpack` archive is produced and contains entries for all variants. - [x] `pack.py` integrity check (no duplicates, every CMake-listed variant present in manifest, every manifest entry has matching `.hsaco`) passes. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- ## Resolves Closes #7865 Closes #7879 Closes #7880 Closes #7881 ### Additional fixes folded in (beyond the CK Tile interface-drift sync) To reach a fully clean `ninja kpack_archive` *and* host-example build, four follow-up commits were added on top of the Arg-structure sync: - **dqdkdv `rand_val_ptr`** — `const_cast<void*>(t_randval.ptr)`; pre-existing const-discard that broke the `*_dropout` variants (`TensorArg::ptr` is `const void*`, CK Tile's `rand_val_ptr` is `void*`). - **ConvertDQ Kargs** — two-path init (`{}` placeholder + named `nsplits_ptr` under `if constexpr(K.is_deterministic)`), mirroring CK Tile's own `MakeKargs`, so a non-deterministic ConvertDQ instantiation also compiles. - **dqdkdv wave64 guard** — fall back to `__GFX9__` because clang ≥23 dropped the `__AMDGCN_WAVEFRONT_SIZE` predefine, which otherwise breaks every dqdkdv variant on rocm7.13+. - **host example** — `variant.spec.mode` (the flattened spec has no nested `signature` member); unblocks the `kpack_rocm_ck_fmha_bwd` executable. **Verified** on `rocm7.13` / clang 23 / `gfx942`: all 40 `KERNEL_VARIANTS` compile with 0 errors, `kernels.kpack` (933 KB) is produced, and the host loader links. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Adam Osewski <Adam.Osewski@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…#5504) ## Motivation Adds sink token support to the FMHA backward kernel (dot_do_o pipeline): ## Technical Details - Extend BlockFmhaBwdOGradDotOPipelineProblem with LSEDataType - Add sink_ptr/d_sink_ptr/lse_ptr/nhead to FmhaBwdOGradDotOCommonKargs - Compute per-head sink gradient via atomic accumulation in the pipeline - Update example runner with reference validation for sink gradient ## Test Plan Add new test case ## Test Result WIP ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Second-round /ck-code-review fixes for the FMHA BWD port: - convert_dq_api validateArgs: fix out-of-bounds read of tensor_names[] (6 entries) for GROUP mode, which needs 8 (requiredTensors==8). Use a per-mode, correctly-sized name table; also fixes the wrong BATCH-mode slot-2 name (was "SEQSTART_Q", is NSPLITS). - dqdkdv_dev: drop two phantom positional initializers (stride_dq_acc, nhead_stride_dq_acc) from the FmhaBwdCommonKargs aggregate-init. The bundled CK Tile struct has 30 fields, not 32 -- dq_acc strides are derived in-kernel. 32-initializer list was ill-formed once instantiated. - ograd_dot_o_dev: reorder the FmhaBwdOGradDotOCommonKargs aggregate-init to the current (post-#5504) layout -- lse/sink/d_sink pointers belong after d_ptr, and nhead before stride_do. The old layout placed p_undrop (float) into the lse_ptr (const void*) slot -> ill-formed init. - Cardinal-rule (ASCII-only) fixes: replace em-dashes with -- in ck_type_map.hpp, dqdkdv_spec.hpp, and two FMHA unit tests. - ck_type_map.hpp: use <ck_tile/core.hpp> angle include (W9 straggler). Verified: rocm_ck host suite (401 unit + 58 compile-fail) passes under -Werror (Clang 23); device bridges compile clean against bundled CK Tile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

github-actions Bot added the project: composablekernel label Mar 17, 2026

LJ-underdog force-pushed the sink_bwd branch from e06dd97 to 1e04a99 Compare March 17, 2026 03:22

assistant-librarian Bot added the organization: ROCm label Mar 17, 2026

[CK Tile] Use per-head random sink values in FMHA backward validation

e5a9df3

Replace the fixed scalar sink_val with per-head random values sampled from uniform distribution [30, 60] to improve test coverage.

LJ-underdog requested a review from Copilot March 17, 2026 05:26

Copilot started reviewing on behalf of LJ-underdog March 17, 2026 05:31 View session

Copilot AI reviewed Mar 17, 2026

View reviewed changes

LJ-underdog requested a review from Copilot March 17, 2026 06:10

Copilot started reviewing on behalf of LJ-underdog March 17, 2026 06:15 View session

Copilot AI reviewed Mar 17, 2026

View reviewed changes

LJ-underdog requested a review from DDEle March 17, 2026 06:30

LJ-underdog added 6 commits March 17, 2026 02:00

[CK Tile] Add smoke_test_bwd_sink to run_full_test.sh

c3102f8

Merge branch 'develop' into sink_bwd

301916a

poyenc reviewed Mar 18, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp

poyenc reviewed Mar 18, 2026

View reviewed changes

Comment thread projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd_runner.hpp Outdated

poyenc reviewed Mar 18, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp

poyenc reviewed Mar 18, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp Outdated

LJ-underdog added 6 commits March 18, 2026 00:24

Merge branch 'develop' into sink_bwd

87e5b19

[CK Tile] Apply clang-format to sink bwd files

60313b9

Merge branch 'develop' into sink_bwd

b13346c

Merge branch 'develop' into sink_bwd

a07053c

LJ-underdog marked this pull request as ready for review March 30, 2026 07:04

LJ-underdog added 2 commits March 31, 2026 11:09

Merge branch 'develop' into sink_bwd

6fde36f

[CK_TILE] Apply clang-format to fmha_bwd_runner.hpp

ff9d45d

LJ-underdog force-pushed the sink_bwd branch from 12d361b to 0ffb795 Compare March 31, 2026 07:00

LJ-underdog force-pushed the sink_bwd branch from 0ffb795 to c9044ef Compare March 31, 2026 07:12

poyenc approved these changes Mar 31, 2026

View reviewed changes

LJ-underdog added 2 commits March 31, 2026 22:04

fmha: apply clang-format to fmha_bwd_runner.hpp

814e85a

LJ-underdog merged commit 47f86c7 into develop Apr 2, 2026
52 of 54 checks passed

LJ-underdog deleted the sink_bwd branch April 2, 2026 03:17

LJ-underdog added a commit to ROCm/aiter that referenced this pull request Apr 2, 2026

update ck to ROCm/rocm-libraries#5504

7481fd6

LJ-underdog added a commit to ROCm/aiter that referenced this pull request Apr 2, 2026

Revert "update ck to ROCm/rocm-libraries#5504"

852a009

This reverts commit 7481fd6.

LJ-underdog mentioned this pull request May 11, 2026

[ck_tile][fmha_bwd] Fix sink_host OOB in group mode reference runner #7272

Merged

1 task

This was referenced May 29, 2026

[KPACK-FIX] experimental/rocm_ck kpack example does not compile against current CK Tile #7865

Closed

[KPACK-FIX] ograd_dot_o_dev.hpp regression from PR #5504 (missing LSEDataType template arg) #7880

Closed

bidlekm mentioned this pull request Jun 1, 2026

Synchronizing Arg structure based on the CKTile changes #7928

Merged

6 tasks

aosewski mentioned this pull request Jun 3, 2026

Porting existing FMHA infra from users/shumway/ck/exp-kpack to develop #7891

Merged

Navigation Menu

Uh oh!

[CK Tile] Add sink token gradient support in FMHA backward pass - #5504

[CK Tile] Add sink token gradient support in FMHA backward pass#5504
LJ-underdog merged 24 commits into
developfrom
sink_bwd

LJ-underdog commented Mar 17, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

LJ-underdog commented Mar 31, 2026 •

edited

Loading

Review notes (post item-by-item triage)

Uh oh!

poyenc left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

LJ-underdog commented Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

LJ-underdog commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review notes (post item-by-item triage)

Uh oh!

poyenc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

LJ-underdog commented Mar 17, 2026 •

edited

Loading

LJ-underdog commented Mar 31, 2026 •

edited

Loading