Skip to content

[FMHA] Enable page size 16 for batch prefill kernel #3568

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
Jeff-Huang merged 2 commits into from
Jan 15, 2026
Merged

[FMHA] Enable page size 16 for batch prefill kernel #3568

Jeff-Huang merged 2 commits into from
Jan 15, 2026
+62 −28

Conversation

@Jeff-Huang
Copy link
Contributor

@Jeff-Huang Jeff-Huang commented Jan 14, 2026

Proposed changes

  • Enable support for page_size=16 in the FMHA batch prefill kernel.
  • Remove the restriction that page size must be a multiple of the tile size (kN0).

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

@Jeff-Huang Jeff-Huang force-pushed the batch_prefill_page_size_16_vectorized branch from 9e4ed9c to c9b1a8e Compare January 15, 2026 00:11
@Jeff-Huang Jeff-Huang self-assigned this Jan 15, 2026
@poyenc poyenc requested a review from Copilot January 15, 2026 05:02
Copilot AI reviewed Jan 15, 2026
View reviewed changes
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

@Jeff-Huang Jeff-Huang force-pushed the batch_prefill_page_size_16_vectorized branch from c9b1a8e to ef5e6f1 Compare January 15, 2026 06:41
poyenc
poyenc reviewed Jan 15, 2026
View reviewed changes
include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp Outdated
OffsetVecType& kv_offset_vec,
index_t global_seq_offset = 0)
{
static constexpr index_t kLog2PageSize = []() constexpr {
Copy link
Contributor

@poyenc poyenc Jan 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can omit the parentheses () if no parameters are needed.

poyenc
poyenc previously approved these changes Jan 15, 2026
View reviewed changes
@Jeff-Huang Jeff-Huang force-pushed the batch_prefill_page_size_16_vectorized branch from ef5e6f1 to 1e7458c Compare January 15, 2026 07:03
@Jeff-Huang
Refactor batch prefill KV offset logic to simplify template arguments
af868e4 
- Remove redundant `kLog2PageSize` and `kIsVTileFitsInPage` from template args.
- Add static assert to forbid `page_size=1` with vectorized layout.
@Jeff-Huang Jeff-Huang force-pushed the batch_prefill_page_size_16_vectorized branch from 1e7458c to af868e4 Compare January 15, 2026 07:27
@Jeff-Huang Jeff-Huang merged commit 993d3e2 into develop Jan 15, 2026
21 checks passed
@Jeff-Huang Jeff-Huang deleted the batch_prefill_page_size_16_vectorized branch January 15, 2026 14:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@poyenc poyenc poyenc approved these changes

@illsilin illsilin Awaiting requested review from illsilin illsilin is a code owner

@carlushuang carlushuang Awaiting requested review from carlushuang carlushuang is a code owner

@qianfengz qianfengz Awaiting requested review from qianfengz qianfengz is a code owner

@aosewski aosewski Awaiting requested review from aosewski aosewski is a code owner

@geyyer geyyer Awaiting requested review from geyyer geyyer is a code owner

@bartekxk bartekxk Awaiting requested review from bartekxk bartekxk is a code owner

@andriy-ca andriy-ca Awaiting requested review from andriy-ca andriy-ca is a code owner

@afagaj afagaj Awaiting requested review from afagaj afagaj is a code owner

@asleepzzz asleepzzz Awaiting requested review from asleepzzz asleepzzz is a code owner

@tenpercent tenpercent Awaiting requested review from tenpercent tenpercent is a code owner

@ThomasNing ThomasNing Awaiting requested review from ThomasNing ThomasNing is a code owner

@coderfeli coderfeli Awaiting requested review from coderfeli coderfeli is a code owner

@cgmillette cgmillette Awaiting requested review from cgmillette cgmillette is a code owner

@shumway shumway Awaiting requested review from shumway shumway is a code owner

@vidyasagar-amd vidyasagar-amd Awaiting requested review from vidyasagar-amd vidyasagar-amd is a code owner

@vpietila-amd vpietila-amd Awaiting requested review from vpietila-amd vpietila-amd is a code owner

@Snektron Snektron Awaiting requested review from Snektron Snektron is a code owner

Assignees

@Jeff-Huang Jeff-Huang

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Jeff-Huang @poyenc