Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
poyenc wants to merge 17 commits into
base: develop
Choose a base branch
from
+36 −1
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
9a47f21
Fix wrong sink_size for no_mask cases
poyenc Jan 2, 2026
57d5599
Avoid invalid instances
poyenc Jan 2, 2026
9de34c9
Filter out redundant instances in fmha_fwd.py
poyenc Jan 2, 2026
84c5f7b
Filter out redundant instances in more codegen scripts
poyenc Jan 2, 2026
b7da954
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 4, 2026
6542982
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 5, 2026
f20b16e
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 6, 2026
e459f43
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 8, 2026
692ffd9
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 8, 2026
d739507
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 10, 2026
2cc4fea
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 13, 2026
7975c94
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 13, 2026
4c37878
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 21, 2026
7425250
Update comment
poyenc Jan 21, 2026
26f3a0b
Update comment in fmha_fwd_splitkv.py
poyenc Jan 21, 2026
ac46bdc
Update comment in fmha_pagedkv_prefill.py
poyenc Jan 21, 2026
5844a1f
Merge branch 'develop' into poyenc/fix-mask-decode
poyenc Jan 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -809,6 +809,11 @@ def check_hdim(problem_ctx: ProblemContext, kernel_ctx: KernelContext) -> bool:
False
return True

all_mask_keys = list(get_mask_map("simplified").keys()) + list(
get_mask_map("generic").keys()
)
no_mask_keys = [mask_key for mask_key in all_mask_keys if "no" in mask_key]

def check_feature(
problem_ctx: ProblemContext, kernel_ctx: KernelContext
) -> bool:
Expand All @@ -821,6 +826,12 @@ def check_feature(
or kernel_ctx.pipeline.F_logits == "f"
):
return False
# 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
Comment on lines +829 to +834
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
return True

return [check_mode, check_hdim, check_feature]
Expand Down
9 changes: 9 additions & 0 deletions example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -880,6 +880,11 @@ def get_fwd_splitkv_blobs(

factories = get_factories_for_targets(targets, get_factory)

all_mask_keys = list(get_mask_map("simplified").keys()) + list(
get_mask_map("generic").keys()
)
no_mask_keys = [mask_key for mask_key in all_mask_keys if "no" in mask_key]

for factory, dtype in itertools.product(factories, FWD_DTYPE_MAP.keys()):
d = factory.get_hdim_tile_size_dict(dtype)
if d is None:
Expand All @@ -899,6 +904,10 @@ def get_fwd_splitkv_blobs(
or pipeline.F_logits == "f"
):
continue
# sink_size is only meaningful when mask is applied
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
if pipeline.F_mask in no_mask_keys and pipeline.F_sink == "t":
continue

k = Kernel(
F_arch=factory.arch,
F_idx=0,
Expand Down
9 changes: 9 additions & 0 deletions example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,11 @@ def get_fwd_blobs(

factories = get_factories_for_targets(targets, get_factory)

all_mask_keys = list(get_mask_map("simplified").keys()) + list(
get_mask_map("generic").keys()
)
no_mask_keys = [mask_key for mask_key in all_mask_keys if "no" in mask_key]

for factory, dtype in itertools.product(factories, FWD_DTYPE_MAP.keys()):
d = factory.get_hdim_tile_size_dict(dtype)
if d is None:
Expand All @@ -666,6 +671,10 @@ def get_fwd_blobs(
or pipeline.F_logits == "f"
):
continue
# sink_size is only meaningful when mask is applied
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot uses AI. Check for mistakes.
if pipeline.F_mask in no_mask_keys and pipeline.F_sink == "t":
continue

k = FmhaFwdKernel(
F_arch=factory.arch,
F_idx=0,
Expand Down
5 changes: 4 additions & 1 deletion example/ck_tile/01_fmha/mask.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,10 @@ struct mask_info
}
else if(str == "0")
{
tmp.type = mask_enum::no_mask;
tmp.type = mask_enum::no_mask;
tmp.left = -1;
tmp.right = -1;
tmp.sink = 0;
}
else if(str == "1" || str == "t")
{
Expand Down
3 changes: 3 additions & 0 deletions include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ struct BlockFmhaPipelineProblem
static constexpr auto QScaleEnum = Traits::QScaleEnum;
static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
static constexpr bool kHasSink = Traits::kHasSink;
static_assert(FmhaMask::IsMasking || !kHasSink);
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
};

template <typename QDataType_,
Expand Down Expand Up @@ -174,6 +175,7 @@ struct BlockFmhaFwdPagedKVPipelineProblem
static constexpr bool kIsPagedKV = Traits::kIsPagedKV;
static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
static constexpr bool kHasSink = Traits::kHasSink;
static_assert(FmhaMask::IsMasking || !kHasSink);
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
};

template <typename QDataType_,
Expand Down Expand Up @@ -228,6 +230,7 @@ struct BlockFmhaFwdSplitKVPipelineProblem
static constexpr bool kMergeNumHeadGroupsSeqLenQ = Traits::kMergeNumHeadGroupsSeqLenQ;
static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
static constexpr bool kHasSink = Traits::kHasSink;
static_assert(FmhaMask::IsMasking || !kHasSink);
Copy link

Copilot AI Jan 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
};

// extract tile size attributes to remove dependency on traits
Expand Down