[ck_tile/fmha] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels - #3732
Closed
jundaf2 wants to merge 3 commits into
Closed
[ck_tile/fmha] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels#3732jundaf2 wants to merge 3 commits into
jundaf2 wants to merge 3 commits into
Conversation
Add F_sink ∈ {t,f} to the fp8bf16 codegen loop so that fp8bf16 +
kv_blockscale variants also produce kernel instances supporting
StreamLLM sink token. The underlying qr_async pipeline already gates
sink logic via constexpr kHasSink, so this only widens the dispatch
table; non-sink callers keep using the existing _nsink_* blobs.
Made-with: Cursor
The previous sink un-mask predicate ((i_y + x) > 1) gated on the query
row index instead of the key column, with two consequences:
- For causal masking (x = 1) and sink_size = S, queries 1..S-1
attended to their own future sink positions (i_x in (i_y, sink)),
breaking autoregressive correctness.
- Query 0 was *excluded* from sink un-mask entirely and fell back to
the plain causal mask, so it could not attend to any sink position
other than column 0.
Replace the row-side guard with the column-side bound (i_x < i_y + x),
which lets sink un-mask precisely those sink columns that lie inside
the query's causal/right-window past, restoring StreamLLM semantics.
Made-with: Cursor
Replace the historical fix-rationale note above the sink un-mask check with a clause-by-clause explanation of i_x < sink, i_x < i_y + x, y < y_total, and i_y < x_total so future readers can map the predicate to the StreamLLM sink + sliding-window semantics directly. Co-authored-by: AI Coding Agent Made-with: Cursor
jundaf2
requested review from
Snektron,
ThomasNing,
afagaj,
andriy-ca,
aosewski,
asleepzzz,
bartekxk,
carlushuang,
cgmillette,
coderfeli,
geyyer,
illsilin,
poyenc,
qianfengz,
shumway,
tenpercent,
vidyasagar-amd and
vpietila-amd
as code owners
April 28, 2026 14:56
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes ck_tile FMHA StreamLLM sink behavior for sliding-window batch-prefill with fp8 KV / bf16 compute by aligning codegen + masking logic so sink-aware kernels are emitted and respect the right-window boundary.
Changes:
- Correct
IsOutOfSinkBoundsink un-mask gating to respect the per-row right-window bound (i_x < i_y + x) instead of a row-index-based predicate. - Update
SimplifiedGenericAttentionMaskto apply the same corrected sink un-mask condition. - Extend fp8bf16 batch-prefill codegen to emit both
has_sinkvariants (F_sink=t/f) so runtime dispatch can select sink-aware kernels.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| include/ck_tile/ops/fmha/block/block_masking.hpp | Fixes sink un-mask predicate to respect right-window/causality; updates and expands explanatory comments; mirrors fix in simplified mask. |
| example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py | Emits fp8bf16 batch-prefill pipelines with F_sink variants to enable dispatch to sink-aware kernels. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+250
to
+251
| // to look at later sink rows than they should under | ||
| // causality / right-window. |
There was a problem hiding this comment.
In the sink un-mask predicate explanation, the text says queries could look at later "sink rows". Since i_x is the key column index here, this should refer to later sink columns/positions (keys), not rows, to avoid confusion when reasoning about causal/right-window behavior.
Suggested change
| // to look at later sink rows than they should under | |
| // causality / right-window. | |
| // to look at later sink columns/positions (keys) | |
| // than they should under causality / right-window. |
Contributor
|
Reopen as ROCm/rocm-libraries#6914 |
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 related fixes to
ck_tileFMHA so that StreamLLM-sink + sliding-windowbatch-prefill works correctly for fp8 KV / bf16 compute.
Review the commits in this order:
fmha: emit sink kernels for fp8bf16 batch_prefillExtends
example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.pysothe fp8(KV) / bf16(QO) batch-prefill codegen also emits the
mask=mask_enum::generic_with_sinkvariant. Without this the runtimecould not dispatch to a sink-aware kernel for the fp8bf16 path.
fmha: respect right-window in IsOutOfSinkBoundThe sink un-mask in
GenericAttentionMask::IsOutOfSinkBound(local-maskbranch) used
(i_y + x) > 1as the gate, which conditioned on the rowindex instead of the column index. As a result, queries
1..sink-1could attend to future sink positions (violating causal / right-window),
while query
0fell back to the plain causal mask. The fix replaces theguard with
i_x < i_y + xso every query only sees sink columns up toits own right-window boundary.
fmha: clarify IsOutOfSinkBound predicate commentDoc-only follow-up that rewrites the comment above the predicate as a
clause-by-clause explanation (
i_x < sink,i_x < i_y + x,y < y_total,i_y < x_total).Test plan
op_tests/test_batch_prefill.py(fp8 + bf16_dequantmodes with
sink=4,win_left=1023,softcap=0.0,sal=True)now passes for all parametrized shapes.
diff only adds the
generic_with_sinkvariant for fp8bf16; existingkernel object lists unaffected.