[CK_TILE][FMHA] Fix uninitialized sink_size in mask_info::decode() and filter redundant no-mask+sink instances by poyenc · Pull Request #3504 · ROCm/composable_kernel

poyenc · 2026-01-02T09:33:04Z

Problem

When mask_info::decode() parses "0" (no_mask), it only set the type field but left left, right, and sink uninitialized. This caused:

Uninitialized sink could be arbitrary garbage value
traits.has_sink = (mask.sink > 0) in fmha_fwd_runner.hpp:882 might evaluate to true
Even without masking, code could incorrectly dispatch to kHasSink=true instantiations

Solution

Runtime fix: Initialize left=-1, right=-1, sink=0 when decoding no_mask in mask.hpp
Compile-time check: Add static_assert(FmhaMask::IsMasking || !kHasSink) to pipeline problems
Reduce instantiations: Filter out F_mask=no_mask + F_sink=true combinations in codegen scripts:
- fmha_fwd.py
- fmha_fwd_splitkv.py
- fmha_pagedkv_prefill.py

Impact

✅ Fixes potential runtime dispatch bug
✅ Prevents invalid template combinations at compile-time
✅ Reduces binary size by eliminating redundant kernel instantiations
✅ Improves build time

Testing

Existing tests should pass
Invalid combinations now caught at compile-time

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

I have added tests relevant to the introduced functionality, and the unit tests are passing locally
I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
I have added inline documentation which enables the maintainers with understanding the motivation
I have removed the stale documentation which is no longer relevant after this pull request
(If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
I have run clang-format on all changed files
Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

Copilot

Pull request overview

Fixes an FMHA runtime dispatch hazard caused by uninitialized sink_size for no_mask, adds a compile-time guard against invalid sink+no-mask template combinations, and reduces redundant kernel instantiations in codegen.

Changes:

Initialize left/right/sink when decoding no_mask in mask_info::decode().
Add static_assert(FmhaMask::IsMasking || !kHasSink) to prevent invalid pipeline instantiations.
Filter out no_mask + sink=true combinations in FMHA fwd-related codegen scripts.

Reviewed changes

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

Show a summary per file

File	Description
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp	Adds compile-time validation to prevent `kHasSink=true` when masking is disabled.
example/ck_tile/01_fmha/mask.hpp	Fixes uninitialized fields for `no_mask` decoding (prevents bogus runtime `has_sink`).
example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py	Skips generating redundant/invalid `no_mask + sink` kernel variants.
example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py	Skips generating redundant/invalid `no_mask + sink` kernel variants.
example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py	Adds compatibility filtering to avoid `no_mask + sink` kernels in fwd generation.

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

Copilot · 2026-01-21T05:58:20Z

    static constexpr auto QScaleEnum        = Traits::QScaleEnum;
    static constexpr index_t kBlockPerCu    = Traits::kBlockPerCu;
    static constexpr bool kHasSink          = Traits::kHasSink;
+    static_assert(FmhaMask::IsMasking || !kHasSink);


The new static_assert has no diagnostic message, while other static_asserts in this file provide one (e.g., lines 108–123). Adding a short message (e.g., that sink requires masking) would make template instantiation failures much easier to understand.

Copilot · 2026-01-21T05:58:20Z

    static constexpr bool kIsPagedKV        = Traits::kIsPagedKV;
    static constexpr index_t kBlockPerCu    = Traits::kBlockPerCu;
    static constexpr bool kHasSink          = Traits::kHasSink;
+    static_assert(FmhaMask::IsMasking || !kHasSink);


The new static_assert has no diagnostic message, while other static_asserts in this file provide one (e.g., lines 108–123). Adding a short message (e.g., that sink requires masking) would make template instantiation failures much easier to understand.

Copilot · 2026-01-21T05:58:21Z

    static constexpr bool kMergeNumHeadGroupsSeqLenQ = Traits::kMergeNumHeadGroupsSeqLenQ;
    static constexpr index_t kBlockPerCu             = Traits::kBlockPerCu;
    static constexpr bool kHasSink                   = Traits::kHasSink;
+    static_assert(FmhaMask::IsMasking || !kHasSink);


The new static_assert has no diagnostic message, while other static_asserts in this file provide one (e.g., lines 108–123). Adding a short message (e.g., that sink requires masking) would make template instantiation failures much easier to understand.

Copilot · 2026-01-21T05:58:21Z

+            # sink_size is only meaningful when mask is applied
+            if (
+                kernel_ctx.pipeline.F_mask in no_mask_keys
+                and kernel_ctx.pipeline.F_sink == "t"
+            ):
+                return False


This comment says sink_size is only meaningful when no masking is applied, but the condition directly below filters out the no-mask + sink=true combination. Please update the comment to match the logic (i.e., sink is only meaningful when masking is enabled).

Copilot · 2026-01-21T05:58:21Z

                    or pipeline.F_logits == "f"
                ):
                    continue
+                # sink_size is only meaningful when mask is applied


This comment says sink_size is only meaningful when no masking is applied, but the condition directly below filters out the no-mask + sink=true combination. Please update the comment to match the logic (i.e., sink is only meaningful when masking is enabled).

Suggested change

# sink_size is only meaningful when mask is applied

# sink_size is only meaningful when masking is enabled, so disallow sink when no mask is applied

Copilot · 2026-01-21T05:58:22Z

                    or pipeline.F_logits == "f"
                ):
                    continue
+                # sink_size is only meaningful when mask is applied


This comment says sink_size is only meaningful when no masking is applied, but the condition directly below filters out the no-mask + sink=true combination. Please update the comment to match the logic (i.e., sink is only meaningful when masking is enabled).

Suggested change

# sink_size is only meaningful when mask is applied

# sink_size is only meaningful when masking is enabled; disallow sink when no mask is used

LJ-underdog · 2026-01-21T06:12:47Z

LGTM @asleepzzz Please approve it.

ammallya · 2026-02-03T22:22:25Z

Imported to ROCm/rocm-libraries

Fix wrong sink_size for no_mask cases

9a47f21

poyenc self-assigned this Jan 2, 2026

poyenc requested review from ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, geyyer, illsilin, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners January 2, 2026 09:33

poyenc marked this pull request as draft January 2, 2026 09:57

poyenc added 3 commits January 2, 2026 04:33

Avoid invalid instances

57d5599

Filter out redundant instances in fmha_fwd.py

9de34c9

Filter out redundant instances in more codegen scripts

84c5f7b

poyenc marked this pull request as ready for review January 2, 2026 14:21

poyenc changed the title ~~Fix mask_info initialization for no_mask cases~~ [CK_TILE][FMHA] Fix uninitialized sink_size in mask_info::decode() and filter redundant no-mask+sink instances Jan 2, 2026

poyenc added 7 commits January 4, 2026 10:40

Merge branch 'develop' into poyenc/fix-mask-decode

b7da954

Merge branch 'develop' into poyenc/fix-mask-decode

6542982

Merge branch 'develop' into poyenc/fix-mask-decode

f20b16e

Merge branch 'develop' into poyenc/fix-mask-decode

e459f43

Merge branch 'develop' into poyenc/fix-mask-decode

692ffd9

Merge branch 'develop' into poyenc/fix-mask-decode

d739507

Merge branch 'develop' into poyenc/fix-mask-decode

2cc4fea

poyenc added 2 commits January 13, 2026 20:32

Merge branch 'develop' into poyenc/fix-mask-decode

7975c94

Merge branch 'develop' into poyenc/fix-mask-decode

4c37878

poyenc requested review from Snektron and vpietila-amd as code owners January 21, 2026 00:17

poyenc requested a review from Copilot January 21, 2026 05:52

Update comment

7425250

Copilot started reviewing on behalf of poyenc January 21, 2026 05:54 View session

poyenc added 3 commits January 21, 2026 13:55

Update comment in fmha_fwd_splitkv.py

26f3a0b

Update comment in fmha_pagedkv_prefill.py

ac46bdc

Merge branch 'develop' into poyenc/fix-mask-decode

5844a1f

Copilot AI reviewed Jan 21, 2026

View reviewed changes

LJ-underdog and others added 3 commits January 22, 2026 12:41

Merge branch 'develop' into poyenc/fix-mask-decode

98e5432

Merge branch 'develop' into poyenc/fix-mask-decode

f153799

Merge branch 'develop' into poyenc/fix-mask-decode

1dd3225

assistant-librarian Bot mentioned this pull request Feb 3, 2026

[CK_TILE][FMHA] Fix uninitialized sink_size in mask_info::decode() and filter redundant no-mask+sink instances ROCm/rocm-libraries#4306

Closed

7 tasks

ammallya closed this Feb 3, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[CK_TILE][FMHA] Fix uninitialized sink_size in mask_info::decode() and filter redundant no-mask+sink instances#3504

[CK_TILE][FMHA] Fix uninitialized sink_size in mask_info::decode() and filter redundant no-mask+sink instances#3504
poyenc wants to merge 20 commits into
ROCm:developfrom
poyenc:poyenc/fix-mask-decode

poyenc commented Jan 2, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

Copilot AI Jan 21, 2026

Uh oh!

LJ-underdog commented Jan 21, 2026 •

edited

Loading

Uh oh!

ammallya commented Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	# sink_size is only meaningful when mask is applied
	# sink_size is only meaningful when masking is enabled, so disallow sink when no mask is applied

Conversation

poyenc commented Jan 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Solution

Impact

Testing

Checklist

Discussion

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

ammallya commented Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

poyenc commented Jan 2, 2026 •

edited

Loading

LJ-underdog commented Jan 21, 2026 •

edited

Loading