Added custom FMHA codegen receipt for TransformerEngine - #6867
Merged
Conversation
wangye805
reviewed
Apr 27, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a dedicated “700” receipt to constrain FMHA codegen output for TransformerEngine, aiming to reduce kernel count, compile time, and memory footprint compared to the broader “600” receipt.
Changes:
- Documented a new
700receipt in the generator CLI help text. - Added
receipt == 700filtering rules for FMHA forward kernel selection. - Added
receipt == 700filtering rules for FMHA backward blob generation.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| projects/composablekernel/example/ck_tile/01_fmha/generate.py | Adds CLI help text describing the new TE-specific receipt and its invariants. |
| projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py | Introduces TE-specific forward-kernel compatibility rule under receipt 700. |
| projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py | Introduces TE-specific backward blob filtering under receipt 700. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
poyenc
reviewed
Apr 29, 2026
poyenc
left a comment
poyenc
left a comment
Contributor
There was a problem hiding this comment.
Review of receipt 700 addition for TransformerEngine integration.
poyenc
approved these changes
May 4, 2026
poyenc
left a comment
poyenc
left a comment
Contributor
There was a problem hiding this comment.
LGTM. All review comments addressed — zero-output guards added for unsupported APIs, and remaining questions answered satisfactorily.
poyenc
approved these changes
May 11, 2026
poyenc
approved these changes
May 12, 2026
poyenc
approved these changes
May 13, 2026
poyenc
enabled auto-merge (squash)
May 14, 2026 02:02
poyenc
approved these changes
May 14, 2026
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
May 14, 2026
Added custom FMHA codegen receipt for TransformerEngine (#6867) ## Motivation TE uses AITER to build static MHA libraries, which ultimately rely on CK kernels. We use the `600` receipt which generates more kernels than TE truly needs. This bespoke receipt allows us to minimize the kernel count, compile time, and memory footprint of our MHA library. ## Technical Details Extended the receipt mechanism to include a custom `700` receipt for TE's needs ## Test Plan Test by building TE using the same receipt profile ## Test Result Build validated in TE using a custom feature branches of AITER/CK to temporarily apply the patch ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
aledudek
pushed a commit
that referenced
this pull request
May 20, 2026
## Motivation TE uses AITER to build static MHA libraries, which ultimately rely on CK kernels. We use the `600` receipt which generates more kernels than TE truly needs. This bespoke receipt allows us to minimize the kernel count, compile time, and memory footprint of our MHA library. ## Technical Details Extended the receipt mechanism to include a custom `700` receipt for TE's needs ## Test Plan Test by building TE using the same receipt profile ## Test Result Build validated in TE using a custom feature branches of AITER/CK to temporarily apply the patch ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
May 27, 2026
Added custom FMHA codegen receipt for TransformerEngine (#6867) ## Motivation TE uses AITER to build static MHA libraries, which ultimately rely on CK kernels. We use the `600` receipt which generates more kernels than TE truly needs. This bespoke receipt allows us to minimize the kernel count, compile time, and memory footprint of our MHA library. ## Technical Details Extended the receipt mechanism to include a custom `700` receipt for TE's needs ## Test Plan Test by building TE using the same receipt profile ## Test Result Build validated in TE using a custom feature branches of AITER/CK to temporarily apply the patch ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> Co-authored-by: Po Yen Chen <PoYen.Chen@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
TE uses AITER to build static MHA libraries, which ultimately rely on CK kernels. We use the
600receipt which generates more kernels than TE truly needs. This bespoke receipt allows us to minimize the kernel count, compile time, and memory footprint of our MHA library.Technical Details
Extended the receipt mechanism to include a custom
700receipt for TE's needsTest Plan
Test by building TE using the same receipt profile
Test Result
Build validated in TE using a custom feature branches of AITER/CK to temporarily apply the patch
Submission Checklist