Add attn sink by LJ-underdog · Pull Request #2892 · ROCm/composable_kernel

LJ-underdog · 2025-09-22T09:07:32Z

Proposed changes

Please describe the motivation behind the pull request, whether it enables a new feature or fixes a bug. If there are associated pull requests or issues, please link them to the pull request.

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

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

example/ck_tile/01_fmha/script/correct_test_fwd_sink.sh

example/ck_tile/01_fmha/fmha_fwd.hpp

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp

asleepzzz

aiter side is not ready

This reverts commit 9fa4e8d.

* enable attn sink Signed-off-by: JL-underdog <Jun.Lin@amd.com> * update attn_sink script Signed-off-by: JL-underdog <Jun.Lin@amd.com> * fix some error Signed-off-by: JL-underdog <Jun.Lin@amd.com> * clang-format Signed-off-by: JL-underdog <Jun.Lin@amd.com> * update fmha_bwd mask Signed-off-by: JL-underdog <Jun.Lin@amd.com> * update fmha_bwd_kernel'mask Signed-off-by: JL-underdog <Jun.Lin@amd.com> * update block_fmha_pipeline_qr_ks_vs.hpp Signed-off-by: JL-underdog <Jun.Lin@amd.com> * fix ci error Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * fix format error Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * Update block_fmha_bwd_pipeline_default_policy.hpp * Update fmha_fwd_runner.hpp * Update block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp * Update fmha_fwd_runner.hpp * Update fmha_fwd_runner.hpp * Update fmha_fwd_runner.hpp * update splitkv_pipline Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * update splitkv&pagedkv pipeline Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * add sink test Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * update attn_sink result log Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * update smoke_test_fwd_sink.sh Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * update test file Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * update test script Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * Update block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp * use constexpr kHasSink for sink in fmha pipeline Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * update by pre-commit Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update include/ck_tile/ops/fmha/kernel/fmha_fwd_pagedkv_kernel.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update fmha_fwd.py * Update example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Remove causal mask setting logic from mask.hpp Removed the mask setting logic for causal masks. * fix ci error that some usage of lamada not support in c++17 Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * Update remod.py * add smoke sink test Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * Update fmha_pagedkv_prefill.py * Update FmhaFwdPipeline parameters in fmha_fwd.py * update block_fmha_pipeline_qr_ks_vs_async_trload.hpp Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * fix c++17 unsupprot error Signed-off-by: LJ-underdog <Jun.Lin@amd.com> * Update block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp * Fix formatting of sink_seq_end assignment * Fix indentation for sink_seq_end assignment * Update block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp --------- Signed-off-by: JL-underdog <Jun.Lin@amd.com> Signed-off-by: LJ-underdog <Jun.Lin@amd.com> Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

This reverts commit 9fa4e8d.

This reverts commit 5adaa20.

* Revert "Revert "Add attn sink (#2892)" (#3250)" This reverts commit 5adaa20. * fix conflict Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Add F_sink parameter to FmhaFwdPipeline * Update tile_fmha_traits.hpp * Refactor pipeline creation in fmha_fwd.py Updated the pipeline creation logic to include 'sink' parameter in product combinations and adjusted the FmhaFwdPipeline calls accordingly. * Update fmha_fwd.py * Update fmha_fwd.py * Update example/ck_tile/01_fmha/script/correct_test_fwd_sink.sh Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * update CHANGELOG.md Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update CHANGELOG with new features and support * Update fmha_fwd.hpp * Update CHANGELOG.md * Update smoke_test_fwd_sink.sh * Update correct_test_fwd_sink.sh * Update smoke_test_fwd_sink.sh --------- Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

LJ-underdog added 18 commits September 22, 2025 02:40

enable attn sink

a092cf2

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

update attn_sink script

4e6f2de

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

fix some error

e89fd09

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

14bef45

clang-format

628082d

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

80be924

update fmha_bwd mask

77a1107

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

33db67c

update fmha_bwd_kernel'mask

27b79f1

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

update block_fmha_pipeline_qr_ks_vs.hpp

8c659bc

Signed-off-by: JL-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

ba22123

Merge branch 'develop' into lj/attn_sink

f83007d

fix ci error

7b51f1b

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

9438f5f

fix format error

a946d48

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

Update block_fmha_bwd_pipeline_default_policy.hpp

4f06d30

Update fmha_fwd_runner.hpp

42af7c2

Update block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp

00e8ae8

DDEle reviewed Sep 29, 2025

View reviewed changes

example/ck_tile/01_fmha/script/correct_test_fwd_sink.sh Show resolved Hide resolved

example/ck_tile/01_fmha/fmha_fwd.hpp Show resolved Hide resolved

LJ-underdog added 5 commits September 29, 2025 10:30

Update fmha_fwd_runner.hpp

aa9b1b1

Update fmha_fwd_runner.hpp

171d16c

Update fmha_fwd_runner.hpp

3290f92

update splitkv_pipline

edaef37

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

update splitkv&pagedkv pipeline

3ea15ba

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

LJ-underdog changed the title ~~Add attn sink to qr_ks_vs_pipeline~~ Add attn sink Oct 23, 2025

LJ-underdog added 5 commits October 23, 2025 15:39

add sink test

d4d9698

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

update attn_sink result log

e116fa0

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

fd7f676

update smoke_test_fwd_sink.sh

dd6f264

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

update test file

648cbf3

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

LJ-underdog added 12 commits November 14, 2025 14:57

Merge branch 'develop' into lj/attn_sink

afc92b7

Update remod.py

0c7cb3d

add smoke sink test

5ce1524

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

Merge branch 'develop' into lj/attn_sink

14c05cb

Update fmha_pagedkv_prefill.py

2fe5243

Update FmhaFwdPipeline parameters in fmha_fwd.py

4f5116f

update block_fmha_pipeline_qr_ks_vs_async_trload.hpp

3259192

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

fix c++17 unsupprot error

79d0e45

Signed-off-by: LJ-underdog <Jun.Lin@amd.com>

Update block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp

bd06e29

Fix formatting of sink_seq_end assignment

0debfd9

Fix indentation for sink_seq_end assignment

6aee384

Update block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp

e2b2254

poyenc reviewed Nov 18, 2025

View reviewed changes

include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp Show resolved Hide resolved

poyenc approved these changes Nov 18, 2025

View reviewed changes

LJ-underdog added 5 commits November 18, 2025 15:29

Merge branch 'develop' into lj/attn_sink

38fc5d6

Merge branch 'develop' into lj/attn_sink

7671121

Merge branch 'develop' into lj/attn_sink

e414747

Merge branch 'develop' into lj/attn_sink

c24ccab

Merge branch 'develop' into lj/attn_sink

660efeb

LJ-underdog merged commit 9fa4e8d into develop Nov 20, 2025
21 checks passed

LJ-underdog deleted the lj/attn_sink branch November 20, 2025 11:24

asleepzzz reviewed Nov 20, 2025

View reviewed changes

asleepzzz added a commit that referenced this pull request Nov 20, 2025

Revert "Add attn sink (#2892)"

557d273

This reverts commit 9fa4e8d.

asleepzzz mentioned this pull request Nov 20, 2025

Revert "Add attn sink" #3250

Merged

illsilin pushed a commit that referenced this pull request Nov 20, 2025

Revert "Add attn sink (#2892)" (#3250)

5adaa20

This reverts commit 9fa4e8d.

AviralGoelAMD pushed a commit that referenced this pull request Nov 28, 2025

Revert "Add attn sink (#2892)" (#3250)

edb9097

This reverts commit 9fa4e8d.

LJ-underdog added a commit that referenced this pull request Dec 8, 2025

Revert "Revert "Add attn sink (#2892)" (#3250)"

fc03494

This reverts commit 5adaa20.

Copilot AI mentioned this pull request Dec 9, 2025

Add attention sink support for FMHA FWD #3368

Merged

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add attn sink#2892

Add attn sink#2892
LJ-underdog merged 62 commits intodevelopfrom
lj/attn_sink

LJ-underdog commented Sep 22, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

asleepzzz left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

LJ-underdog commented Sep 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

Checklist

Discussion

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

asleepzzz left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

LJ-underdog commented Sep 22, 2025 •

edited

Loading