Skip to content

Commit 6f9598d

Browse files
committed
fmha: clarify IsOutOfSinkBound predicate comment
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
1 parent 8bac412 commit 6f9598d

1 file changed

Lines changed: 11 additions & 6 deletions

File tree

include/ck_tile/ops/fmha/block/block_masking.hpp

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -242,12 +242,17 @@ struct GenericAttentionMask
242242
index_t x_start = -y + i_y + 1;
243243
index_t x_end = min(i_y + x, x_total);
244244

245-
// Sink un-mask must respect causal/right-window: a query at row i_y
246-
// sees keys in [..., i_y + x), so a sink column i_x is only attended
247-
// when i_x < i_y + x. The previous (i_y + x) > 1 guard erroneously
248-
// gated on the row index instead of the column index, which let
249-
// queries 1..sink-1 attend to their own future sink positions and
250-
// forced query 0 to fall back to the plain causal mask.
245+
// Sink un-mask predicate, clause by clause:
246+
// i_x < sink : the column lives inside the StreamLLM sink prefix.
247+
// i_x < i_y + x : the column is not in the masked-out future of the
248+
// window (= < x_end modulo the min with x_total);
249+
// without this, queries <= sink-1 would be allowed
250+
// to look at later sink rows than they should under
251+
// causality / right-window.
252+
// y < y_total : the local window doesn't already span everything
253+
// (otherwise sink un-mask is meaningless).
254+
// i_y < x_total : the query row is in-range vs. the key sequence
255+
// (handles seqlen_q > seqlen_k padding).
251256
if constexpr(IsLocal)
252257
{
253258
if((i_x < sink) && (i_x < i_y + x) && (y < y_total) && i_y < x_total)

0 commit comments

Comments
 (0)