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

[chris-tsiaousis-hpc]

No description provided.

[Copilot]

Pull request overview

This PR ports the FMHA backward (BWD) infrastructure from an experimental branch (users/shumway/ck/exp-kpack) into develop. It adds three kernel families — OGradDotO, DqDkDv, and ConvertDQ — exposed through a layered spec/api/dev header convention (consteval factories, host-only helpers, device-only CK Tile bridges), a host-side variant registry, comprehensive GoogleTest unit/death tests, and CMake-driven compile-fail tests for the consteval validators.

Changes:

• New rocm_ck/ops/fmha_bwd/* header trio per kernel family plus shared common.hpp/grid_dim.hpp/ck_type_map.hpp, with consteval makeSpec() validators and named tensor/scalar slot constants.
• New variant registry header and findVariant() / name-lookup helpers for programmatic kernel selection.
• New unit, consteval, death, compatibility, and compile-fail tests wired into the existing CMake test targets.

Reviewed changes

Copilot reviewed 43 out of 43 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
include/rocm_ck/grid_dim.hpp	HIP-free GridDim type.
include/rocm_ck/ck_type_map.hpp	Device-only DataType → CK Tile type map.
include/rocm_ck/ops/fmha_bwd/common.hpp	Shared FmhaMode / FmhaBiasType enums.
include/rocm_ck/ops/fmha_bwd/ograd_dot_o_{spec,api,dev}.hpp	OGradDotO kernel spec, host helpers, and CK Tile bridge.
include/rocm_ck/ops/fmha_bwd/dqdkdv_{spec,api,dev}.hpp	DqDkDv kernel spec, host helpers/validator, and CK Tile bridge.
include/rocm_ck/ops/fmha_bwd/convert_dq_{spec,api,dev}.hpp	ConvertDQ kernel spec, host helpers, and CK Tile bridge.
tests/unit/fmha_bwd_registry.hpp	Variant registry + consteval name lookup.
tests/unit/unit_fmha_bwd_common.cpp	Tests for shared enums and GridDim.
tests/unit/unit_fmha_bwd_compat.cpp	Frozen-baseline + registry tests across all three families.
tests/unit/unit_fmha_bwd_consteval.cpp	SWA/CMaskBR consteval spec/slot tests.
tests/unit/unit_fmha_bwd_convert_dq.cpp	ConvertDQ spec/slot/grid tests.
tests/unit/unit_fmha_bwd_dqdkdv.cpp	DqDkDv spec/slot/grid tests.
tests/unit/unit_fmha_bwd_ograd_dot_o.cpp	OGradDotO spec/slot/grid tests.
tests/unit/unit_fmha_bwd_validate_args.cpp	Death tests for the debug-mode DqDkDv validateArgs.
tests/compile_fail/*.cpp + CMakeLists.txt	Compile-fail coverage for consteval validators.
tests/CMakeLists.txt	Adds new unit sources to the unit-test target.

---

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

[bidlekm]

LGTM! The FMHA BWD has a lot of pr-s merged, do you think it would be better to push this to production like it is now, or to push it with the merges?

[chris-tsiaousis-hpc]

No additional code changes needed on the productionize branch. The only new exp-kpack content since our port (96bd1a114bd) that touches headers/non-compat-tests was #7534 , and it's already present.
Since we don't want to migrate the examples to develop I think the other PRs aren't needed.

Specifically the new PRs to port (post 96bd1a1, touching FMHA headers/tests):

• d8e300c — P1-8: multi-tile grid-size coverage tests
• 0750803 — P1-1b: _swa and _cmask_br variants
• 3041557 — P1-3: bf16 bias variants + frozen baselines
• 40924b3 + 5ef7463 — P4-1 arch dispatch + revert (net zero)
• cc118ce — bf16 ograddoto, dqdkdv
• 6832e29 — Removing test
• a136983 — dqdkdv group mode dropout, cmask, bias
• 09d06dc — convertdq variants
• 0d7ec62 — P1-7: sliding-window configs
• 12cbe67 — OGradDotO fp16 multidim variants
• 1a217a6 — P3-1: Tile-config table design
• 9484d2e — DqDkDv d32/d96/d256 variants
• 5bcb8f8 — DqDkDv d64 variants
• cd41900 — AE-1: FmhaMaskType enum

[aosewski]

Thanks for getting this onto develop! I went through it carefully and diffed it against exp-kpack — overall this is a clean, well-structured port, and the ops/fmha_bwd/ layout + test wiring fit develop's conventions nicely. No blocking correctness issues; a few things worth a look before merge, plus inline notes on specific lines.

On "does this carry everything from exp-kpack?" — I verified the headers are a faithful snapshot of exp-kpack @ #7534 (2026-05-28): the _dev/_api headers are byte-identical, and _spec/common.hpp differ only by the datatype_utils.hpp→datatype.hpp include renames. So nothing was dropped — it's just that three FMHA commits landed upstream after the snapshot and aren't here yet:

• #7538 — consteval tile-config table (06-01)
• #7274 — FmhaMaskType enum (06-03)
• #7928 — CK-Tile Arg-struct sync (06-03)

Could we make the snapshot scope explicit in the title/description and track those three as catch-up? One call worth making now: #7274 replaces the public has_mask bool with a FmhaMaskType enum and it merged today — re-cutting to include it would avoid a breaking field change on develop later.

Docs — the PR adds a whole new ops/ subtree, but README.md is unchanged: its Directory-layout tree and Status blurb are still flat/GEMM-only, and it documents a single unified Signature model that FMHA intentionally breaks. Worth updating; and a short ops/fmha_bwd/README.md carrying the exp INSIGHTS.md rationale (per-kernel Signature, PhysicalTensor-rejected, pad_hdim=8, TrLoad/IsWG32 future) would stop a future maintainer from "fixing" intentional choices.

Device-bridge catch-up (heads-up, not blocking) — the _dev.hpp headers are dormant on develop (no .hip instantiations ship), but they're one CK-Tile-sync generation behind tip (pre-#7928/#6152/#5504), so their Kargs aggregate-inits would mis-map against current CK Tile. They'll need a rebase onto exp-tip before any kernel instantiation lands — the __AMDGCN_WAVEFRONT_SIZE assert (inline note) is the first thing that'll trip on clang ≥ 23.

Test note — the mask int → ck_tile::GenericAttentionMaskEnum contract is live at runtime (read from the MASK_TYPE scalar slot) but nothing pins the canonical integer values the host packer must write; a tiny mapping test would lock it down.
A handful of smaller Suggestions (unused <cstdint>, grid_dim.hpp Role: header, include sorting, slot-count canaries, asymmetric-hdim handling, etc.) I've kept out of the thread to reduce noise — happy to share the full list. Nice work overall; the structure is solid. 🙏