Synchronizing Arg structure based on the CKTile changes by bidlekm · Pull Request #7928 · ROCm/rocm-libraries

bidlekm · 2026-06-01T09:06:04Z

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 [CK Tile] Add sink token gradient support in FMHA backward pass #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 [CK_TILE] Use Unified Workspace for FMHA BWD #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).

Test Plan

Test Result

ctest -L rocm_ck --output-on-failure: 64/64 pass
ctest -L compile_fail --output-on-failure: 45/45 pass
ninja kpack_archive produces non-zero .hsaco files for every entry in KERNEL_VARIANTS against GPU_TARGETS=gfx942
kernels.kpack archive is produced and contains entries for all variants.
pack.py integrity check (no duplicates, every CMake-listed variant present in manifest, every manifest entry has matching .hsaco) passes.

Submission Checklist

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.

…ompile_fix

Copilot

Pull request overview

This PR updates the rocm_ck::Args ABI/capacity and FMHA-BWD argument slot layouts to stay in sync with recent CK Tile kernarg (Kargs) changes (notably new workspace-derived pointers and updated positional field ordering), and updates unit tests accordingly.

Changes:

Increase kMaxTensors (and derived Args size) to accommodate new FMHA-BWD deterministic workspace slots.
Extend FMHA-BWD DqDkDv and ConvertDQ slot definitions to include workspace-derived pointers (e.g., NSPLITS, DQ_ACC_BATCH_OFFSET), and adjust validation/tests.
Update device bridges (*_dev.hpp) to match CK Tile Kargs layout changes.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/args.hpp	Increase `kMaxTensors` to 20 and update `Args` size/static_asserts.
projects/composablekernel/experimental/rocm_ck/tests/test_args.cpp	Update ABI/size expectations for `Args` and capacity constants.
projects/composablekernel/experimental/rocm_ck/tests/test_signature.cpp	Update capacity-limit expectation (`kMaxTensors`).
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_spec.hpp	Add deterministic workspace slots (`NSPLITS`, `DQ_ACC_BATCH_OFFSET`) and update `requiredTensors()`.
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_api.hpp	Extend debug validation to include new slots and skip group-only slots in batch mode.
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp	Update DqDkDv device bridge to match CK Tile deterministic `Kargs` changes.
projects/composablekernel/experimental/rocm_ck/tests/test_fmha_bwd_validate_args.cpp	Update death test to populate newly-required deterministic workspace slot.
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/convert_dq_spec.hpp	Update ConvertDQ slot layout to include workspace-derived `NSPLITS`/offsets and revised `requiredTensors()`.
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/convert_dq_dev.hpp	Update ConvertDQ device bridge to match CK Tile `Kargs` changes (nsplits ptr, nhead).
projects/composablekernel/experimental/rocm_ck/tests/test_fmha_bwd_convert_dq.cpp	Update required-tensor-count expectations for the new slot layout.
projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/ograd_dot_o_dev.hpp	Update OGradDotO device bridge for CK Tile `Kargs` signature changes (LSE/sink fields).

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

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

aledudek

main.cpp:906 example host doesn't compile.
variant.spec.signature.mode - FmhaBwdDQDKDVSpec has no signature member

use variant.spec.mode (the correct usages at lines ~1232/1393). Pre-existing, but it blocks the example build.

const_cast t_randval.ptr (const void*) when assigning the non-const rand_val_ptr Kargs field. Without it the dropout variant (fmha_bwd_dqdkdv_fp16_d128_batch_dropout) fails to compile. Also clarify the conditional-base placeholder comments and update the DQ_ACC stride doc to reflect the workspace-derived layout (CK Tile #6152), with NSPLITS / DQ_ACC_BATCH_OFFSET slot notes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Zero-init the conditional deterministic base and assign nsplits_ptr by name under if constexpr(K.is_deterministic), in both group and batch paths. An EmptyKargs base cannot accept {nsplits_ptr}, so this lets a future non-deterministic ConvertDQ instantiation compile. Kargs are bit-identical for the existing (deterministic) built variants. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Explain why nullptr lse/sink/d_sink and nhead=0 are safe: LSE is only loaded when atomic_sink_grad_ptr != nullptr and the sink-score read is guarded by sink_ptr != nullptr, so these pointers are address-computed but never dereferenced. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

clang >=23 dropped the __AMDGCN_WAVEFRONT_SIZE (and transitional __AMDGCN_WAVEFRONT_SIZE__) predefines, breaking the wave64 static_assert in FmhaBwdDQDKDVTypes -- every dqdkdv kpack variant failed to compile against rocm7.13 with "use of undeclared identifier '__AMDGCN_WAVEFRONT_SIZE'". Fall back to the __GFX9__ arch macro (gfx9/CDNA is wave64 by construction) when no wavefront-size predefine is available. clang <=22 still uses the original predefine; wave32 targets (gfx10/11/12) hit the #else and are rejected as before. Verified: ninja kpack_archive now compiles all 40 variants (gfx942, rocm7.13/clang23) and produces kernels.kpack. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

FmhaBwdDQDKDVSpec is the flattened validated descriptor and has no nested `signature` struct; main.cpp:906 used variant.spec.signature.mode which broke the host kpack_rocm_ck_fmha_bwd executable build (the `all` target). Use variant.spec.mode, matching every other access in the file. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…cm-libraries into users/bidlekm/cktile_compile_fix

Resolve dqdkdv_dev.hpp conflict: keep the base's multidim getTileConfig refactor (drops the d128-only static_assert) combined with the portable wavefront-size guard (clang >=23 dropped __AMDGCN_WAVEFRONT_SIZE; fall back to __GFX9__). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

bidlekm added 4 commits June 1, 2026 08:55

fixing convert_dq

aa0bff1

fixing dqdkdv

788bd8c

fixing ograd_dot_o

e94d46b

clang format

90e37f6

bidlekm requested review from aledudek, aosewski and chris-tsiaousis-hpc June 1, 2026 09:06

bidlekm self-assigned this Jun 1, 2026

bidlekm requested a review from a team as a code owner June 1, 2026 09:06

github-actions Bot added the project: composablekernel label Jun 1, 2026

bidlekm changed the title ~~Users/bidlekm/cktile compile fix~~ Synchronizing Arg structure based on the CKTile changes Jun 1, 2026

bidlekm added organization: ROCm organization: streamhpc contributors from streamhpc labels Jun 1, 2026

Merge branch 'users/shumway/ck/exp-kpack' into users/bidlekm/cktile_c…

78b3ea6

…ompile_fix

bidlekm requested a review from Copilot June 1, 2026 09:29

Copilot started reviewing on behalf of bidlekm June 1, 2026 09:29 View session

Copilot AI reviewed Jun 1, 2026

View reviewed changes

more descriptive comment

e1f06db

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

aledudek reviewed Jun 2, 2026

View reviewed changes

Comment thread projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp Outdated

aosewski and others added 8 commits June 3, 2026 09:07

review fixes

e5a44a5

Merge branch 'users/bidlekm/cktile_compile_fix' of github.com:ROCm/ro…

36d169d

…cm-libraries into users/bidlekm/cktile_compile_fix

aosewski merged commit 2ebe353 into users/shumway/ck/exp-kpack Jun 3, 2026
5 of 9 checks passed

aosewski deleted the users/bidlekm/cktile_compile_fix branch June 3, 2026 15:18

aosewski mentioned this pull request Jun 3, 2026

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

Merged

chris-tsiaousis-hpc added a commit that referenced this pull request Jun 4, 2026

Import the changes from PRs #7538, #7274 and #7928

7e21ee1

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Navigation Menu

Uh oh!

Synchronizing Arg structure based on the CKTile changes - #7928

Synchronizing Arg structure based on the CKTile changes#7928
aosewski merged 14 commits into
users/shumway/ck/exp-kpackfrom
users/bidlekm/cktile_compile_fix

bidlekm commented Jun 1, 2026 •

edited by aosewski

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aledudek left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

bidlekm commented Jun 1, 2026 • edited by aosewski Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Resolves

Additional fixes folded in (beyond the CK Tile interface-drift sync)

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!

aledudek left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bidlekm commented Jun 1, 2026 •

edited by aosewski

Loading