feat: [CK Tile] mxfp8 support for qr async pipeline by zsotakal · Pull Request #6526 · ROCm/rocm-libraries

zsotakal · 2026-04-17T13:53:07Z

Closes: #6526

Motivation

Adds MXFP8 support to the QR async FMHA forward pipeline by enabling MX quant scale paths, loosening padding constraints, and updating tests/codegen to cover the new combinations.

Technical Details

Changes:

Extend BlockFmhaPipelineQRKSVSAsync to support MX quant scales (Q/K/V/P scale tiles) and adjust padding/alignment behavior accordingly.
Update FMHA fwd tests to include head-dim adjustment helpers and add a new “General” parameterized test suite.
Update codegen pipeline constraints/selection for qr_async, and adjust tile-window vectorization logic for packed types.

Test Plan

Running CI with extended set of tests

Test Result

CI passing

Submission Checklist

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

Copilot

Pull request overview

Note

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

Adds MXFP8 support to the QR async FMHA forward pipeline by enabling MX quant scale paths, loosening padding constraints, and updating tests/codegen to cover the new combinations.

Changes:

Extend BlockFmhaPipelineQRKSVSAsync to support MX quant scales (Q/K/V/P scale tiles) and adjust padding/alignment behavior accordingly.
Update FMHA fwd tests to include head-dim adjustment helpers and add a new “General” parameterized test suite.
Update codegen pipeline constraints/selection for qr_async, and adjust tile-window vectorization logic for packed types.

Reviewed changes

Copilot reviewed 4 out of 5 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp	Adds `adjust_hdim()` per datatype and introduces a new generalized parameterized test suite; uses `adjust_hdim` in padding tests.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp	Implements MX scale loading/casting paths in the QR async pipeline; relaxes hardcoded padding assumptions; updates alignment and loop structure.
projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp	Changes vectorized thread-buffer sizing/indexing logic to accommodate packed types.
projects/composablekernel/example/ck_tile/01_fmha/script/benchmark_fwd.sh	Fixes indentation for the last benchmark invocation line.
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py	Adjusts `qr_async` head-dim constraints and pipeline enumeration; adds QR async pipelines for MXFP8 in column layout.

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

JP-Fernando

This work adds MXFP8 support to the qr_ks_vs_async pipeline. And it looks good to me! Just two inconsequential comments.

github-actions · 2026-06-16T01:25:24Z

This pull request has been inactive for 25 days and will be marked as stale.

If you would like to keep this PR open, please:

Add new commits
Add a comment explaining why it should remain open

This PR will be automatically closed in 5 days if no further activity occurs.

github-actions · 2026-06-22T01:23:18Z

This pull request has been automatically closed due to inactivity (30 days with no updates).

If you'd like to continue working on this, feel free to reopen the PR or create a new one.

therock-pr-bot · 2026-07-07T07:28:54Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

therock-pr-bot · 2026-07-07T07:33:03Z

🎉 All checks passed! This PR is ready for review.

poyenc

LGTM. mxfp4 Fix A is in, and I confirmed the V-window advance and dcheck→true codegen changes are correct. Approving.

feat: [CK Tile] mxfp8 support for qr async pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes: #6526 ## Motivation Adds MXFP8 support to the QR async FMHA forward pipeline by enabling MX quant scale paths, loosening padding constraints, and updating tests/codegen to cover the new combinations. ## Technical Details Changes: - Extend BlockFmhaPipelineQRKSVSAsync to support MX quant scales (Q/K/V/P scale tiles) and adjust padding/alignment behavior accordingly. - Update FMHA fwd tests to include head-dim adjustment helpers and add a new “General” parameterized test suite. - Update codegen pipeline constraints/selection for qr_async, and adjust tile-window vectorization logic for packed types. ## Test Plan Running CI with extended set of tests ## Test Result CI passing ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

This reverts commit 3e01710.

This reverts commit 723ff1b.

This reverts commit 3e01710.

This reverts commit 723ff1b.

…6526)" (#9461) This reverts commit 3e01710. ## Motivation The PR is identified breaking the CI RUN_CK_TILE_FMHA_TESTS. http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2fComposable%20Kernel/detail/develop/1762/pipeline/305 and Aiter test_mha.py http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2FComposable%20Kernel/detail/users%2Fjialuo%2Fck%2Ffix_ci_aiter_error/33/pipeline/246 ## Technical Details  ## Test Plan  ## Test Result  ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

This reverts commit 723ff1b.

…6526)" (#9461) This reverts commit 3e01710. ## Motivation The PR is identified breaking the CI RUN_CK_TILE_FMHA_TESTS. http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2fComposable%20Kernel/detail/develop/1762/pipeline/305 and Aiter test_mha.py http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2FComposable%20Kernel/detail/users%2Fjialuo%2Fck%2Ffix_ci_aiter_error/33/pipeline/246 ## Technical Details  ## Test Plan  ## Test Result  ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

feat: [CK Tile] mxfp8 support for qr async pipeline Closes: #6526 ## Motivation Adds MXFP8 support to the QR async FMHA forward pipeline by enabling MX quant scale paths, loosening padding constraints, and updating tests/codegen to cover the new combinations. ## Technical Details Changes: - Extend BlockFmhaPipelineQRKSVSAsync to support MX quant scales (Q/K/V/P scale tiles) and adjust padding/alignment behavior accordingly. - Update FMHA fwd tests to include head-dim adjustment helpers and add a new “General” parameterized test suite. - Update codegen pipeline constraints/selection for qr_async, and adjust tile-window vectorization logic for packed types. ## Test Plan Running CI with extended set of tests ## Test Result CI passing ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Zoltán Lakatos <zoltan.lakatos@streamhpc.com>

zsotakal self-assigned this Apr 17, 2026

zsotakal added project: composablekernel organization: streamhpc contributors from streamhpc labels Apr 17, 2026

zsotakal force-pushed the users/zlakatos/ck/fmha-qr-async-mxfp8 branch 2 times, most recently from 8c5c014 to 5d65659 Compare April 20, 2026 14:28

zsotakal force-pushed the users/zlakatos/ck/fmha-qr-async-mxfp8 branch from 296d5c9 to e76717a Compare May 13, 2026 08:05

Zoltán Lakatos added 7 commits May 20, 2026 12:02

async pipeline implementation

343bcab

npad for qr async

b312578

fix adync trload v3 compilation

14313ff

turn off mxfp4 for qr async

9c91592

codegen cleanup

315bfa3

removed failing tests

3207f1c

revert benchmark script

6685a90

zsotakal force-pushed the users/zlakatos/ck/fmha-qr-async-mxfp8 branch from e76717a to 6685a90 Compare May 20, 2026 12:04

zsotakal requested a review from Copilot May 20, 2026 12:05

Copilot started reviewing on behalf of zsotakal May 20, 2026 12:11 View session

Copilot AI reviewed May 20, 2026

View reviewed changes

static assert to check division by PackedSize

0d6d8b6

zsotakal changed the title ~~mxfp8 support for qr async pipeline~~ [CK Tile] mxfp8 support for qr async pipeline May 21, 2026

Merge branch 'develop' into users/zlakatos/ck/fmha-qr-async-mxfp8

095fb80

krithalith requested review from JP-Fernando and ex-rzr May 21, 2026 12:14

JP-Fernando approved these changes May 21, 2026

View reviewed changes

Comment thread ...ts/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp

Comment thread ...ts/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp Outdated

github-actions Bot added the Stale PR has no activity for 25+ days label Jun 16, 2026

github-actions Bot closed this Jun 22, 2026

zsotakal reopened this Jun 25, 2026

zsotakal and others added 2 commits June 25, 2026 09:08

Merge branch 'develop' into users/zlakatos/ck/fmha-qr-async-mxfp8

c67d966

remove unintentional comment

64868ef

github-actions Bot removed the Stale PR has no activity for 25+ days label Jun 26, 2026

zsotakal marked this pull request as ready for review June 26, 2026 11:09

zsotakal requested a review from a team as a code owner June 26, 2026 11:09

poyenc reviewed Jul 3, 2026

View reviewed changes

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

fix unintented narrowing of mxfp4 pipelines

2cbbe3d

zsotakal changed the title ~~[CK Tile] mxfp8 support for qr async pipeline~~ feat: [CK Tile] mxfp8 support for qr async pipeline Jul 7, 2026

Merge branch 'develop' into users/zlakatos/ck/fmha-qr-async-mxfp8

db582f5

poyenc approved these changes Jul 9, 2026

View reviewed changes

Merge branch 'develop' into users/zlakatos/ck/fmha-qr-async-mxfp8

af67cde

zsotakal merged commit 3e01710 into develop Jul 13, 2026
40 checks passed

zsotakal deleted the users/zlakatos/ck/fmha-qr-async-mxfp8 branch July 13, 2026 10:43

JiaLuo-CAN added a commit that referenced this pull request Jul 14, 2026

Revert "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

89efe7b

This reverts commit 3e01710.

JiaLuo-CAN added a commit that referenced this pull request Jul 15, 2026

Revert "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

723ff1b

This reverts commit 3e01710.

JiaLuo-CAN added a commit that referenced this pull request Jul 15, 2026

Reapply "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

78394ea

This reverts commit 723ff1b.

JiaLuo-CAN added a commit that referenced this pull request Jul 15, 2026

Revert "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

1114d9f

This reverts commit 3e01710.

JiaLuo-CAN added a commit that referenced this pull request Jul 15, 2026

Revert "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

a97febc

This reverts commit 3e01710.

JiaLuo-CAN added a commit that referenced this pull request Jul 15, 2026

Reapply "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

2a84463

This reverts commit 723ff1b.

JiaLuo-CAN added a commit that referenced this pull request Jul 16, 2026

Reapply "feat: [CK Tile] mxfp8 support for qr async pipeline (#6526)"

d845bfe

This reverts commit 723ff1b.

Uh oh!

Conversation

zsotakal commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

JP-Fernando left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jun 16, 2026

Uh oh!

github-actions Bot commented Jun 22, 2026

Uh oh!

Uh oh!

therock-pr-bot Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Jul 7, 2026

Uh oh!

poyenc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zsotakal commented Apr 17, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 7, 2026 •

edited

Loading