[CK_TILE][FMHA] Enable gpt-oss sink

[LJ-underdog]

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

[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.

[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]

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

[LJ-underdog]

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

no see any performance regression in my local test.