[CK_TILE][FMHA] Enable gpt-oss sink by LJ-underdog · Pull Request #3490 · ROCm/composable_kernel

LJ-underdog · 2025-12-25T07:10:58Z

Proposed changes

This PR enables support for "gptoss sink" tokens in the FMHA (Fused Multi-Head Attention) implementation by adding a sink_ptr parameter throughout the codebase. The sink feature allows attention mechanisms to maintain a virtual sink token that affects softmax normalization.

Key Changes:

Added sink_ptr parameter to kernel argument structures and pipeline operators
Modified initialization logic to conditionally set sink values based on infinity checks
Updated validation logic to account for sink tokens in softmax computation

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

None

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

Copilot

Pull request overview

This PR enables support for "gptoss sink" tokens in the FMHA (Fused Multi-Head Attention) implementation by adding a sink_ptr parameter throughout the codebase. The sink feature allows attention mechanisms to maintain a virtual sink token that affects softmax normalization.

Key Changes:

Added sink_ptr parameter to kernel argument structures and pipeline operators
Modified initialization logic to conditionally set sink values based on infinity checks
Updated validation logic to account for sink tokens in softmax computation

Reviewed changes

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

Show a summary per file

File	Description
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp`	Added `sink_v` parameter and conditional initialization logic for m/l tiles based on sink value
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp`	Added `sink_v` parameter to both operator() overloads with conditional tile initialization
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp`	Added `sink_v` parameter to operator() methods with conditional tile initialization
`include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp`	Added `sink_v` parameter with split-aware conditional initialization; includes commented debug prints
`include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp`	Added `sink_v` parameter with split-aware conditional initialization
`include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp`	Added `sink_v` parameter to both operator() overloads with conditional tile initialization
`include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp`	Added `sink_ptr` to Kargs, MakeKargs methods; compute per-head sink value from pointer
`include/ck_tile/ops/fmha/kernel/fmha_fwd_pagedkv_kernel.hpp`	Added `sink_ptr` to Kargs, MakeKargs methods; compute per-head sink value from pointer
`include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp`	Added `sink_ptr` to Kargs and all MakeKargs overloads; compute per-head sink value; removed extraneous whitespace
`include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp`	Added `sink_ptr` to Kargs and MakeKargsImpl methods; compute per-head sink value from pointer
`example/ck_tile/01_fmha/fmha_fwd_runner.hpp`	Added `init_sink_value` parameter, sink tensor allocation/initialization, and validation logic for sink tokens
`example/ck_tile/01_fmha/fmha_fwd.hpp`	Added `sink_ptr` field to all args structures and threaded through kargs creation
`example/ck_tile/01_fmha/example_fmha_fwd.cpp`	Added `init_sink` command-line argument

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

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

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

Copilot

Pull request overview

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

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

poyenc · 2026-01-13T08:15:42Z

Do we see any performance regression after adding the extra sink_ptr kernel argument?

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

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

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

LJ-underdog · 2026-01-14T02:47:31Z

Do we see any performance regression after adding the extra sink_ptr kernel argument?

no see any performance regression in my local test.

* Enable gptoss sink Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * add gptoss sink test Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * update CHANGELOG.md Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * fix test args error Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update test_fmha_fwd.cpp * update sink test Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Revert "update sink test" This reverts commit 970b4f1. * update sink test Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * update valid sink_v in splitkv pipeline Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp * Update example_fmha_fwd.cpp * fix lse error Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * fix clangformat error Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * fix aiter scale error Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update block_fmha_pipeline_qr_ks_vs.hpp * div scale_s for sink_value Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update fmha_fwd_runner.hpp * update sink_value with bias Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp * Fix typo in dropout parameter in fmha_batch_prefill_kernel * Update block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp * Update example_fmha_fwd.cpp * Update include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp 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> * optimized some code Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * fix splitkv error Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * update sink reference Signed-off-by: Linjun-AMD <Jun.Lin@amd.com> * Update fmha_fwd_runner.hpp * 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> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>

Enable gptoss sink

4bc109f

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

LJ-underdog requested review from ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, geyyer, illsilin, poyenc, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners December 25, 2025 07:11

LJ-underdog requested a review from Copilot December 25, 2025 07:11

Merge branch 'develop' into gptoss_sink

b86a860

Copilot AI reviewed Dec 25, 2025

View reviewed changes

LJ-underdog and others added 4 commits December 25, 2025 15:15

Update include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipel…

6711cec

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

Update include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipel…

1e3c54d

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

add gptoss sink test

7c9cb83

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

update CHANGELOG.md

5163868

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

LJ-underdog requested review from a team and ddembeckAMD as code owners December 25, 2025 07:51

LJ-underdog added 4 commits December 29, 2025 09:57

Merge branch 'develop' into gptoss_sink

4d73213

fix test args error

5ab683b

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

Update test_fmha_fwd.cpp

5c0e07a

update sink test

970b4f1

Signed-off-by: Linjun-AMD <Jun.Lin@amd.com>

poyenc requested a review from Copilot January 13, 2026 07:45

Merge branch 'develop' into gptoss_sink

26a3e9f

Copilot started reviewing on behalf of poyenc January 13, 2026 07:45 View session

poyenc changed the title ~~Enable gptoss sink~~ [CK_TILE][FMHA] Enable gptoss sink Jan 13, 2026

poyenc assigned LJ-underdog Jan 13, 2026

Copilot AI reviewed Jan 13, 2026

View reviewed changes

Comment thread include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp Outdated

Comment thread include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp Outdated

poyenc changed the title ~~[CK_TILE][FMHA] Enable gptoss sink~~ [CK_TILE][FMHA] Enable gpt-oss sink Jan 13, 2026