Conversation
chris-tsiaousis-hpc
added
the
organization: streamhpc
aosewski
approved these changes
May 11, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
This PR upgrades the FMHA BWD dQ/dK/dV configuration schema by replacing the lossy has_mask boolean with a FmhaMaskType enum that preserves CK Tile’s distinct mask families (top-left causal, bottom-right causal, and generic sliding-window), enabling upcoming variant work without silent aliasing.
Changes:
- Introduces
FmhaMaskType(value-pinned tock_tile::GenericAttentionMaskEnum) and replaceshas_maskacross DqDkDv spec/config/registry/runner/docs/tests. - Adds a shared
hasMask()predicate and compile-time validation to reject out-of-range enum values (with a new compile-fail test). - Updates variant matching to compare
mask_typeexactly (preventing bottom-right/generic requests from incorrectly matching top-left variants).
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/common.hpp | Adds FmhaMaskType enum (pinned integer values) and documentation. |
| projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_spec.hpp | Replaces has_mask with mask_type, adds hasMask() helper, validates enum range, updates scalar-slot predicate. |
| projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp | Switches device bridge gating from K.has_mask to hasMask(K) and keeps runtime MASK_TYPE wiring. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/rocm_fmha_bwd_registry.hpp | Updates registered variants and findVariant() to match mask_type exactly. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/main.cpp | Populates MASK_TYPE scalar from variant.spec.mask_type and uses hasMask() for gating. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/pack.py | Updates kpack variant metadata to use mask_type instead of has_mask. |
| projects/composablekernel/experimental/rocm_ck/examples/06_rocm_ck_fmha_bwd/INSIGHTS.md | Updates documentation to reflect hasMask(K) gating and mask_type terminology. |
| projects/composablekernel/experimental/rocm_ck/tests/test_fmha_bwd_dqdkdv.cpp | Updates tests for mask_type, adds coverage for all enum values + hasMask() agreement. |
| projects/composablekernel/experimental/rocm_ck/tests/test_fmha_bwd_compat.cpp | Updates frozen baselines and adds negative registry match test for unregistered mask types. |
| projects/composablekernel/experimental/rocm_ck/tests/compile_fail/dqdkdv_invalid_mask_type.cpp | Adds compile-fail coverage for out-of-range mask_type values. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
chris-tsiaousis-hpc
force-pushed
the
users/chris-tsiaousis-hpc/ck/rocm-ck-ae-1-mask-type-enum
branch
from
May 12, 2026 07:55
2a81cf7 to
8924eb7
Compare
chris-tsiaousis-hpc
marked this pull request as ready for review
May 15, 2026 14:32
bidlekm
reviewed
May 19, 2026
Comment on lines
+15
to
+19
| using namespace rocm_ck; | ||
|
|
||
| constexpr auto bad = makeSpec(FmhaBwdDQDKDVConfig{ | ||
| .signature = {.dtype = DataType::FP16, .hdim_q = 128, .hdim_v = 128, .mode = FmhaMode::BATCH}, | ||
| .algorithm = {.mask_type = static_cast<FmhaMaskType>(99), .pad_hdim_q = 8, .pad_hdim_v = 8}}); |
Contributor
There was a problem hiding this comment.
Nice one! Since the enum has to maintained parallel to the CK Tile one, it is nice to have a test to see what happens if one's enum value is different from the "approved" ones.
bidlekm
reviewed
May 19, 2026
| // This is the same class of predicate bug fixed for `has_bias_grad` in 7abd4. | ||
| TEST(FmhaBwdCompat, Registry_DqDkDv_DisambiguatesMaskType) | ||
| { | ||
| const auto* v_tl = findVariant(FmhaBwdDQDKDVConfig{ |
Contributor
There was a problem hiding this comment.
makeSpec tests for valid arguments, but this findVariant function does not. It seems like it doesnt check for mask type anywhere, just selects an implementation based on other signature and algorithm fields. Does it make sense to do these tests then? Should the algorithm and signature be validated before calling this function?
bidlekm
approved these changes
May 19, 2026
3 tasks
chris-tsiaousis-hpc
force-pushed
the
users/chris-tsiaousis-hpc/ck/rocm-ck-ae-1-mask-type-enum
branch
from
May 28, 2026 07:53
8924eb7 to
5ad8d70
Compare
Add a four-valued enum that names every attention-mask family CK Tile's GenericAttentionMaskEnum exposes (NO_MASK, TOP_LEFT_CAUSAL, BOTTOM_RIGHT_CAUSAL, GENERIC). Integer values are pinned to match ck_tile::GenericAttentionMaskEnum so the device bridge can forward the spec-time enum to the MASK_TYPE kernel scalar via a static_cast without a translation table. This is the vocabulary commit; no spec, registry, or runner change yet -- those are in the follow-up that replaces FmhaBwdDQDKDVAlgorithm's boolean has_mask field with this enum. Workplan task: AE-1 (https://github.com/orgs/ROCm/projects/178). Test plan - Header-only addition; existing 65 host + compile-fail tests continue to pass (build is unchanged). Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>
Migrate FmhaBwdDQDKDVAlgorithm and FmhaBwdDQDKDVSpec from a single boolean has_mask flag to the FmhaMaskType enum added in the previous commit. The boolean was a 1-of-2 collapse that lost the top-left-vs-bottom-right distinction CK Tile already supports and could not express GENERIC sliding-window masks at all -- both blocked the upcoming _swa / _cmask_br variants (P1-1b / P1-7). Workplan task: AE-1 (https://github.com/orgs/ROCm/projects/178). Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>
chris-tsiaousis-hpc
force-pushed
the
users/chris-tsiaousis-hpc/ck/rocm-ck-ae-1-mask-type-enum
branch
from
June 2, 2026 13:03
57dcc70 to
220a6bb
Compare
… sites The AE-1 has_mask->mask_type rename was rebased onto a base that later gained new has_mask call sites (group-mode tests from #7534, the consteval suite from #7504, a group_cmask registry entry, and the group_cmask kpack TOC entry). Those sites still referenced the removed has_mask field, so the host test TUs no longer compiled. Convert them with the PR's established idiom: * .has_mask = true -> .mask_type = FmhaMaskType::TOP_LEFT_CAUSAL * EXPECT_FALSE(k.has_mask) -> EXPECT_EQ(k.mask_type, NO_MASK) * EXPECT_TRUE(k.has_mask) -> family-specific EXPECT_EQ (TL / BR / GENERIC) * cross-variant has_mask compare-> hasMask() equality + EXPECT_NE(mask_type) * pack.py "has_mask": True -> "mask_type": "top_left_causal" Also refresh two consteval comments that still described the pre-AE-1 findVariant() aliasing behaviour (mask_type now disambiguates the family). Verified host compile-time gate with standalone clang++ -std=c++20: test_fmha_bwd_{compat,consteval,dqdkdv,common,validate_args,convert_dq, ograd_dot_o} all compile clean; compile_fail/dqdkdv_invalid_mask_type still correctly rejects the out-of-range enum. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aosewski
deleted the
users/chris-tsiaousis-hpc/ck/rocm-ck-ae-1-mask-type-enum
branch
June 3, 2026 10:30
chris-tsiaousis-hpc
added a commit
that referenced
this pull request
Jun 4, 2026
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
Two stacked commits implementing AE-1 from the FMHA BWD project board: replace the boolean
has_maskflag onFmhaBwdDQDKDVAlgorithm/FmhaBwdDQDKDVSpecwith aFmhaMaskTypeenum that mirrorsck_tile::GenericAttentionMaskEnum.The boolean was a 1-of-2 collapse that lost the top-left-vs-bottom-right distinction CK Tile already supports and could not express
GENERICsliding-window masks at all — both block the upcoming_swa/_cmask_brvariants tracked by P1-1b / P1-7.Workplan task: https://github.com/orgs/ROCm/projects/178/views/1?pane=issue&itemId=183664589
Commits
af3d56f—rocm-ck: [AE-1] Add FmhaMaskType enum to common.hppVocabulary-only addition. Integer values pinned to
ck_tile::GenericAttentionMaskEnumso the device bridge canstatic_caststraight into theMASK_TYPEkernel scalar without a translation table.2a81cf7—rocm-ck: [AE-1] Replace has_mask bool with mask_type enum on DqDkDvAtomic API rename across spec / dev bridge / registry / runner / kpack TOC (
pack.py) / docs / tests, plus a newcompile_fail/dqdkdv_invalid_mask_type.cppfor out-of-range enum rejection.findVariant()compares the enum directly so futureGENERIC/BOTTOM_RIGHT_CAUSALrequests are not silently aliased onto the top-left variant — same class of bug as thehas_bias_gradaliasing fixed in7abd4.Test plan
ninja rocm_ck_testsbuilds 19 host targets clean.ctest -L rocm_ck: 19/19 pass — including the newMakeSpecMaskTypes(covers all 4 enum values +hasMask()agreement) andRegistry_DqDkDv_RejectsUnregisteredMaskType(verifiesBOTTOM_RIGHT_CAUSAL/GENERICrequests returnnullptrinstead of aliasing ontocmask).ctest -L compile_fail: 46/46 pass (45 pre-existing + newdqdkdv_invalid_mask_type).clang-format-18, copyright headers, exec-bit,ck_tile remod.Out of scope
_cmask_br/_swaare P1-1b / P1-7 follow-ups that consume this enum.OGradDotO/ConvertDQspecs unchanged (they don't carry a mask field).