[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels by LJ-underdog · Pull Request #6914 · ROCm/rocm-libraries

LJ-underdog · 2026-04-29T03:46:15Z

Summary

Two related fixes to ck_tile FMHA so that StreamLLM-sink + sliding-window
batch-prefill works correctly for fp8 KV / bf16 compute.

Review the commits in this order:

fmha: emit sink kernels for fp8bf16 batch_prefill
Extends example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py so
the fp8(KV) / bf16(QO) batch-prefill codegen also emits the
mask=mask_enum::generic_with_sink variant. Without this the runtime
could not dispatch to a sink-aware kernel for the fp8bf16 path.
fmha: respect right-window in IsOutOfSinkBound
The sink un-mask in GenericAttentionMask::IsOutOfSinkBound (local-mask
branch) used (i_y + x) > 1 as the gate, which conditioned on the row
index instead of the column index. As a result, queries 1..sink-1
could attend to future sink positions (violating causal / right-window),
while query 0 fell back to the plain causal mask. The fix replaces the
guard with i_x < i_y + x so every query only sees sink columns up to
its own right-window boundary.
fmha: clarify IsOutOfSinkBound predicate comment
Doc-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

Repro on aiter op_tests/test_batch_prefill.py (fp8 + bf16_dequant
modes with sink=4, win_left=1023, softcap=0.0, sal=True)
now passes for all parametrized shapes.
Existing fp16/bf16 batch-prefill paths (no sink) unchanged — codegen
diff only adds the generic_with_sink variant for fp8bf16; existing
kernel object lists unaffected.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

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

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR fixes StreamLLM sink + sliding-window behavior in ck_tile FMHA masking and updates fp8(KV)/bf16(QO) batch-prefill codegen to emit sink-enabled kernels for correct runtime dispatch.

Changes:

Fix sink “unmask” gating under right-window/local masking by using a column-based predicate (i_x < i_y + x) instead of a row-based gate.
Emit fp8bf16 batch-prefill pipeline variants with F_sink enabled (t/f) so sink-aware kernels can be generated and dispatched.
Improve inline documentation explaining the sink unmask predicate.

Reviewed changes

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

File	Description
projects/composablekernel/include/ck_tile/ops/fmha/block/block_masking.hpp	Adjusts sink unmask predicate to respect right-window and clarifies predicate documentation.
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py	Extends fp8bf16 batch-prefill pipeline generation to include sink-enabled variants.

Comments suppressed due to low confidence (1)

projects/composablekernel/include/ck_tile/ops/fmha/block/block_masking.hpp:269

i_x < i_y + x is described as “< x_end modulo the min with x_total”, but it does not actually enforce the min(..., x_total) part. If i_y + x > x_total and sink > x_total, this predicate can unmask sink columns i_x that are already out of range w.r.t. x_total, while the rest of the masking logic uses x_end as the true upper bound. Consider using i_x < x_end in the sink-unmask condition (it matches the effective right boundary in both the padded and non-padded cases).

        index_t x_start = -y + i_y + 1;
        index_t x_end   = min(i_y + x, x_total);

        // Sink un-mask predicate, clause by clause:
        //   i_x < sink       : the column lives inside the StreamLLM sink prefix.
        //   i_x < i_y + x    : the column is not in the masked-out future of the
        //                      window (= < x_end modulo the min with x_total);
        //                      without this, queries <= sink-1 would be allowed
        //                      to look at later sink rows than they should under
        //                      causality / right-window.
        //   y < y_total      : the local window doesn't already span everything
        //                      (otherwise sink un-mask is meaningless).
        //   i_y < x_total    : the query row is in-range vs. the key sequence
        //                      (handles seqlen_q > seqlen_k padding).
        if constexpr(IsLocal)
        {
            if((i_x < sink) && (i_x < i_y + x) && (y < y_total) && i_y < x_total)
                return false;
            else
                return i_x < x_start || i_x >= x_end;
        }
        else
        {
            if((i_x < sink) && (i_x < i_y + x) && (y < y_total) && i_y < x_total)
                return false;
            else
                return i_x >= x_end || i_y >= y_total;
        }

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

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

In get_pipelines() for both fp16/bf16 and fp8bf16 dispatch tables, the sink axis is iterated together with the mask axis via itertools.product, which produces (sink="t", mask="no") and (sink="t", mask="s_no") rows that the underlying qr_async pipeline cannot meaningfully run: kHasSink only takes effect when a real causal/right-window mask is in play, so those instances would compile to dead kernels that no caller ever dispatches. Add a `continue` guard in both branches that skips the sink="t" rows whenever mask is one of the no-op variants ("no", "s_no"), shrinking the emitted kernel matrix without losing any dispatchable shape. Addresses review feedback from poyenc on PR #6914. Co-authored-by: AI Coding Agent Made-with: Cursor

Jeff-Huang · 2026-04-29T17:34:21Z

LGTM

[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels (#6914) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Two related fixes to `ck_tile` FMHA so that StreamLLM-sink + sliding-window batch-prefill works correctly for fp8 KV / bf16 compute. Review the commits in this order: 1. `fmha: emit sink kernels for fp8bf16 batch_prefill` Extends `example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py` so the fp8(KV) / bf16(QO) batch-prefill codegen also emits the `mask=mask_enum::generic_with_sink` variant. Without this the runtime could not dispatch to a sink-aware kernel for the fp8bf16 path. 2. `fmha: respect right-window in IsOutOfSinkBound` The sink un-mask in `GenericAttentionMask::IsOutOfSinkBound` (local-mask branch) used `(i_y + x) > 1` as the gate, which conditioned on the row index instead of the column index. As a result, queries `1..sink-1` could attend to *future* sink positions (violating causal / right-window), while query `0` fell back to the plain causal mask. The fix replaces the guard with `i_x < i_y + x` so every query only sees sink columns up to its own right-window boundary. 3. `fmha: clarify IsOutOfSinkBound predicate comment` Doc-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 - [x] Repro on aiter `op_tests/test_batch_prefill.py` (fp8 + bf16_dequant modes with `sink=4`, `win_left=1023`, `softcap=0.0`, `sal=True`) now passes for all parametrized shapes. - [x] Existing fp16/bf16 batch-prefill paths (no sink) unchanged — codegen diff only adds the `generic_with_sink` variant for fp8bf16; existing kernel object lists unaffected. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…7272) ## Summary In `fmha_bwd_runner.hpp`, the `sink_host` `HostTensor` is allocated with first dimension `shape_batch` (= 1 in group mode), but the reference forward loop accesses `sink_host(wb, i_h)` with `wb ∈ [0, batch-1]`. For any `wb >= 1` this is an out-of-bounds heap read, silently corrupting the reference forward math chain (`lse_host`, `o_host`) and turning the bwd-side `d_sink_head_acc` reference into non-deterministic garbage. `HostTensor::operator()` does not bounds check, so the OOB is not caught at runtime. This manifests as intermittent `tile_example_fmha_bwd` failures (25–67% fail rate) when `-sink_grad=1` is combined with `-mode=1` (group mode), with bit-exact but spurious `max_err` values like 4.27 / 14.6. ## Fix One-line: allocate `sink_host` with `batch` (the real per-batch dim) instead of `shape_batch`, mirroring how `sink_host` is accessed by the loop. ```diff - sink_grad ? std::array<ck_tile::index_t, 2>{shape_batch, nhead} + sink_grad ? std::array<ck_tile::index_t, 2>{batch, nhead} Repro tile_example_fmha_bwd -b=2 -h=2 -s=516 -s_k=253 -prec=bf16 -d=72 \ -bias=n -dbias=0 -p_drop=0 -iperm=1 -operm=1 -deterministic=0 \ -v=3 -mode=1 -kname=1 -sink_grad=1 Verification - 0/30 fail on the repro config after fix - Baselines (before fix): - sink=1, mask=n: 25% fail rate (p ≈ 1.8e-4) - sink=1, mask=t: 67% fail rate (p ≈ 6e-15) Attribution Shape bug introduced together with sink_grad in #5504. Unrelated to #6914 (which is a fwd-only fix on a different code path) ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Signed-off-by: junlin12 <junlin12@amd.com> Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com>

Bumps 3rdparty/composable_kernel from fdf4bb7fc → 33b62ed08, the mirror merge of ROCm/rocm-libraries#6914. Also brings in the following rocm-libraries PRs that landed on the mirror between these two pins: #6152, #7046, #6932, #6912, #6972, #6574, #6741, #6209, #6701. Co-authored-by: Linjun-AMD <Jun.Lin@amd.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: LJ-underdog <105184542+LJ-underdog@users.noreply.github.com>

…batch_prefill sink kernels (#6914) ## Summary Two related fixes to `ck_tile` FMHA so that StreamLLM-sink + sliding-window batch-prefill works correctly for fp8 KV / bf16 compute. Review the commits in this order: 1. `fmha: emit sink kernels for fp8bf16 batch_prefill` Extends `example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py` so the fp8(KV) / bf16(QO) batch-prefill codegen also emits the `mask=mask_enum::generic_with_sink` variant. Without this the runtime could not dispatch to a sink-aware kernel for the fp8bf16 path. 2. `fmha: respect right-window in IsOutOfSinkBound` The sink un-mask in `GenericAttentionMask::IsOutOfSinkBound` (local-mask branch) used `(i_y + x) > 1` as the gate, which conditioned on the row index instead of the column index. As a result, queries `1..sink-1` could attend to *future* sink positions (violating causal / right-window), while query `0` fell back to the plain causal mask. The fix replaces the guard with `i_x < i_y + x` so every query only sees sink columns up to its own right-window boundary. 3. `fmha: clarify IsOutOfSinkBound predicate comment` Doc-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 - [x] Repro on aiter `op_tests/test_batch_prefill.py` (fp8 + bf16_dequant modes with `sink=4`, `win_left=1023`, `softcap=0.0`, `sal=True`) now passes for all parametrized shapes. - [x] Existing fp16/bf16 batch-prefill paths (no sink) unchanged — codegen diff only adds the `generic_with_sink` variant for fp8bf16; existing kernel object lists unaffected. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: fengjunda.aml <fengjunda.aml@bytedance.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: root <root@smci350-rck-g03-f12-31.rck.dcgpu>

…7272) ## Summary In `fmha_bwd_runner.hpp`, the `sink_host` `HostTensor` is allocated with first dimension `shape_batch` (= 1 in group mode), but the reference forward loop accesses `sink_host(wb, i_h)` with `wb ∈ [0, batch-1]`. For any `wb >= 1` this is an out-of-bounds heap read, silently corrupting the reference forward math chain (`lse_host`, `o_host`) and turning the bwd-side `d_sink_head_acc` reference into non-deterministic garbage. `HostTensor::operator()` does not bounds check, so the OOB is not caught at runtime. This manifests as intermittent `tile_example_fmha_bwd` failures (25–67% fail rate) when `-sink_grad=1` is combined with `-mode=1` (group mode), with bit-exact but spurious `max_err` values like 4.27 / 14.6. ## Fix One-line: allocate `sink_host` with `batch` (the real per-batch dim) instead of `shape_batch`, mirroring how `sink_host` is accessed by the loop. ```diff - sink_grad ? std::array<ck_tile::index_t, 2>{shape_batch, nhead} + sink_grad ? std::array<ck_tile::index_t, 2>{batch, nhead} Repro tile_example_fmha_bwd -b=2 -h=2 -s=516 -s_k=253 -prec=bf16 -d=72 \ -bias=n -dbias=0 -p_drop=0 -iperm=1 -operm=1 -deterministic=0 \ -v=3 -mode=1 -kname=1 -sink_grad=1 Verification - 0/30 fail on the repro config after fix - Baselines (before fix): - sink=1, mask=n: 25% fail rate (p ≈ 1.8e-4) - sink=1, mask=t: 67% fail rate (p ≈ 6e-15) Attribution Shape bug introduced together with sink_grad in #5504. Unrelated to #6914 (which is a fwd-only fix on a different code path) ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Signed-off-by: junlin12 <junlin12@amd.com> Co-authored-by: Max Podkorytov <4273004+tenpercent@users.noreply.github.com>

[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels (#6914) ## Summary Two related fixes to `ck_tile` FMHA so that StreamLLM-sink + sliding-window batch-prefill works correctly for fp8 KV / bf16 compute. Review the commits in this order: 1. `fmha: emit sink kernels for fp8bf16 batch_prefill` Extends `example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py` so the fp8(KV) / bf16(QO) batch-prefill codegen also emits the `mask=mask_enum::generic_with_sink` variant. Without this the runtime could not dispatch to a sink-aware kernel for the fp8bf16 path. 2. `fmha: respect right-window in IsOutOfSinkBound` The sink un-mask in `GenericAttentionMask::IsOutOfSinkBound` (local-mask branch) used `(i_y + x) > 1` as the gate, which conditioned on the row index instead of the column index. As a result, queries `1..sink-1` could attend to *future* sink positions (violating causal / right-window), while query `0` fell back to the plain causal mask. The fix replaces the guard with `i_x < i_y + x` so every query only sees sink columns up to its own right-window boundary. 3. `fmha: clarify IsOutOfSinkBound predicate comment` Doc-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 - [x] Repro on aiter `op_tests/test_batch_prefill.py` (fp8 + bf16_dequant modes with `sink=4`, `win_left=1023`, `softcap=0.0`, `sal=True`) now passes for all parametrized shapes. - [x] Existing fp16/bf16 batch-prefill paths (no sink) unchanged — codegen diff only adds the `generic_with_sink` variant for fp8bf16; existing kernel object lists unaffected. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: fengjunda.aml <fengjunda.aml@bytedance.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: root <root@smci350-rck-g03-f12-31.rck.dcgpu>

jundaf2 added 3 commits April 29, 2026 03:38

github-actions Bot added the project: composablekernel label Apr 29, 2026

LJ-underdog marked this pull request as ready for review April 29, 2026 04:39

LJ-underdog requested a review from a team as a code owner April 29, 2026 04:39

LJ-underdog requested a review from poyenc April 29, 2026 04:39

assistant-librarian Bot added the organization: ROCm label Apr 29, 2026

LJ-underdog changed the title ~~Ck/lj/add batch prefill swa ck try~~ [CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels Apr 29, 2026

LJ-underdog requested review from Jeff-Huang and Copilot April 29, 2026 05:54

Copilot started reviewing on behalf of LJ-underdog April 29, 2026 06:00 View session

Copilot AI reviewed Apr 29, 2026

View reviewed changes

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/block/block_masking.hpp

Comment thread projects/composablekernel/include/ck_tile/ops/fmha/block/block_masking.hpp Outdated

Apply suggestion from @Copilot

6bc0cd2

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

poyenc reviewed Apr 29, 2026

View reviewed changes

Comment thread projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py

DDEle mentioned this pull request Apr 29, 2026

[ck_tile/fmha] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels ROCm/composable_kernel#3732

Closed

2 tasks

poyenc approved these changes Apr 30, 2026

View reviewed changes

poyenc enabled auto-merge (squash) April 30, 2026 03:33

Jeff-Huang approved these changes Apr 30, 2026

View reviewed changes

LJ-underdog requested a review from asleepzzz April 30, 2026 07:12

Merge branch 'develop' into ck/lj/add_batch_prefill_swa_ck_try

f5848ce

asleepzzz approved these changes May 7, 2026

View reviewed changes

poyenc merged commit b791478 into develop May 7, 2026
35 checks passed

poyenc deleted the ck/lj/add_batch_prefill_swa_ck_try branch May 7, 2026 02:39

arthurliu1998 mentioned this pull request May 7, 2026

bump composable_kernel to include rocm-libraries PR #6914 ROCm/aiter#3073

Merged

LJ-underdog mentioned this pull request May 11, 2026

[ck_tile][fmha_bwd] Fix sink_host OOB in group mode reference runner #7272

Merged

1 task

Navigation Menu

Uh oh!

[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels - #6914

[CK_TILE][FMHA] Fix sink un-mask under right-window and emit fp8bf16 batch_prefill sink kernels#6914
poyenc merged 6 commits into
developfrom
ck/lj/add_batch_prefill_swa_ck_try

LJ-underdog commented Apr 29, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Jeff-Huang commented Apr 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

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

Uh oh!

Summary

Test plan

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Jeff-Huang commented Apr 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

LJ-underdog commented Apr 29, 2026 •

edited

Loading