Adding DqDkDv d64 variants fp16 + bf16 batch - #7804
Merged
aosewski merged 7 commits intoJun 2, 2026
Merged
Conversation
bidlekm
added
organization: ROCm
organization: streamhpc
Contributor
There was a problem hiding this comment.
Pull request overview
This PR exposes new d64 batch-mode dqdkdv FMHA BWD kernel variants (FP16 and BF16) through the rocm_ck wrapper layer, wiring existing CK Tile kernels so they can be discovered, built, packed, and selected from host code.
Changes:
- Added new
.hipentrypoints forfmha_bwd_dqdkdv_{fp16,bf16}_d64_batch. - Extended the
dqdkdvdevice bridge to select the correctd64tile configuration (while retainingd128support). - Updated registry, pack manifest, CMake build lists, and compatibility/registry tests to include the new variants.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/experimental/rocm_ck/tests/test_fmha_bwd_compat.cpp | Adds frozen-baseline + registry lookup coverage for the new d64 batch variants; updates unregistered-case and expected variant count. |
| projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp | Wires d64 tile geometry via conditional BlockTile selection and updates guards/comments accordingly. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/rocm_fmha_bwd_registry.hpp | Registers the new FP16/BF16 d64 batch variants for host-side selection. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/pack.py | Adds the new variants to the kpack manifest so they’re packed into the archive. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/fmha_bwd_dqdkdv_fp16_d64_batch.hip | New device entrypoint that launches the FP16 d64 batch variant via runFmhaBwdDQDKDV. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/fmha_bwd_dqdkdv_bf16_d64_batch.hip | New device entrypoint that launches the BF16 d64 batch variant via runFmhaBwdDQDKDV. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/CMakeLists.txt | Includes the new dqdkdv d64 batch variants in the per-variant build list. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
|
please resolve merge conflict. |
The host grid is sized as ceil(seqlen_k / K.block_n0) while the kernel maps blockIdx.x to the K-tile offset via kN0 (BlockTile[1]). If a future hdim row wires a different bn0 (e.g. the gfx9 d256 row uses bn0=64), the two diverge silently: too few blocks leaves dK/dV tail rows unwritten, too many wastes seqlen-guarded work. Assert block_n0 == BlockTile kN0 so the mismatch is a compile error rather than a wrong-grid launch. Also reflow the over-length conditional_t BlockTile alias to satisfy clang-format-18. Addresses review findings #4 and #6 on PR #7804. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The tile-geometry and block_per_cu comments still claimed "d128 only" after d64 was wired. Clarify that the wired d64 and d128 rows both use bn0=128 (so the hardcoded block_n0 is correct for both, guarded by the device static_assert) and that production would derive these per hdim. Comment-only; no behavior change. Addresses review finding #7 on PR #7804. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Registry_DqDkDv_ReturnsNullForUnregistered probes d96 to assert findVariant() returns null. If a d96 variant is ever registered, that probe silently passes for the wrong reason. Add a loop over ALL_DQDKDV_VARIANTS asserting no d96 dqdkdv entry exists, so the premise fails loudly and forces a new probe instead. Also fix the file-header NOTE: drop the dead dqdkdv_spec.hpp:200-204 line reference and explain that block_n0 is identical (128) for d64 and d128, so hdim_q/hdim_v gate tile selection (verified by the device static_assert). Collapse the four d64 test initializers per clang-format-18. Addresses review findings #5, #6, #9, #10 on PR #7804. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The base branch (users/shumway/ck/exp-kpack) refactored the dqdkdv bridge to a table-driven design: dqdkdv_dev.hpp now builds BlockTile from a consteval kTile lookup, and dqdkdv_spec.hpp derives block_n0, block_size, and block_per_cu from GFX9_FP16_DQDKDV_BASE_TILES (which already includes the d64 row) instead of hardcoding them. It also added d32/d96/d256 batch variants. Conflict resolution: - dqdkdv_dev.hpp, dqdkdv_spec.hpp: take base. Its table-driven approach supersedes this PR's std::conditional_t<kIsD64,...> and the earlier review fixes built on the old structure (block_n0 is now derived from the same table as BlockTile, so the hand-added block_n0==kN0 assert and the stale-comment edits are no longer needed). - CMakeLists.txt, pack.py, rocm_fmha_bwd_registry.hpp: union -- keep base's d32/d96/d128/d256 entries and re-add this PR's d64 batch variants (the unique P3-5 deliverable; base wired the d64 tile but not the user-facing variant). - test_fmha_bwd_compat.cpp: take base (already has d64 spec tests via the tile table), add Registry_DqDkDv_FindsFP16/BF16D64Batch, and bump the dqdkdv variant count 28 -> 30. The d64 .hip kernels remain from this PR. Host compat layer verified: registry resolves both d64 variants, count is 30, and d64 derives block_n0=128 / block_size=256 / block_per_cu=1 from the base table. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Exposes new
d64batch-modedqdkdvkernel variants in therocm_ckFMHA BWD wrapper layer. The underlying CK Tile kernels already exist; this PR adds the wrapper wiring so they are reachable from host code: new.hipentrypoints, registry entries, CMake build inclusion,pack.pymanifest entries, host-side compatibility test coverage, and the requireddqdkdvdevice bridge update for thed64tile configuration.Workplan task: ROCm_CK FMHA BWD P3-5
What this PR contains
fmha_bwd_dqdkdv_bf16_d64_batch.hip.hipentrypointfmha_bwd_dqdkdv_fp16_d64_batch.hip.hipentrypointCMakeLists.txtdqdkdv d64 batchvariants to the buildpack.pyrocm_fmha_bwd_registry.hppd64batch variantsdqdkdv_dev.hppdqdkdvdevice bridge to support thed64tile configurationtest_fmha_bwd_compat.cppd64batch variantsTest plan
ninja rocm_ck_testsctest -L rocm_ck --output-on-failure: 64/64 passctest -L compile_fail --output-on-failure: 45/45 pass