[CK_TILE] fmha: Add backward pass support for padded inputs #1212
+484
−226
Conversation
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
Jeff-Huang
force-pushed
the
SWDEV-558893
branch
2 times, most recently
from
7486448 to
67ca00e
Compare
Introduces support for padded sequence lengths in the backward pass of the variable-length flash attention (fmha_v3_varlen_bwd).
- Updated Python and C++ function signatures to accept optional `cu_seqlens_q_padded` and `cu_seqlens_k_padded` arguments.
- Modified the underlying CUDA kernels and code generation scripts to pass padding information via the new `seqlen_q_ptr` and `seqlen_k_ptr` fields in
the CK `fmha_bwd_args` struct.
- Modified the underlying kernels and code generation scripts to correctly handle pointers for both padded and unpadded sequence data.
- Added comprehensive gradient verification to the test suite (`test_mha_varlen.py`) to ensure the correctness of the backward pass with various
padding scenarios.
Jeff-Huang
force-pushed
the
SWDEV-558893
branch
from
67ca00e to
6b0d7ae
Compare
Refactor the FMHA forward and backward pass to align with the updated padding API in `composable_kernel`. - Argument Simplification: Removed the manual calculation of `seqlen_q` and `seqlen_k` from `cu_seqlens` in the `mha.cu` interface. The underlying kernels now handle this logic. - API Alignment: Updated the arguments passed to `aiter::mha_fwd` and `aiter::mha_bwd` to match the new `composable_kernel` API. This involves passing `cu_seqlen` pointers directly. - Kernel Interface Update: Modified the `codegen.py` scripts for `gfx942` and `gfx950` to reflect the changes in the kernel's function signatures and argument handling for padded and unpadded sequence lengths.
Jeff-Huang
force-pushed
the
SWDEV-558893
branch
from
6b0d7ae to
cad028e
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Introduces support for padded sequence lengths in the backward pass of the variable-length flash attention (fmha_v3_varlen_bwd).
cu_seqlens_q_paddedandcu_seqlens_k_paddedarguments.seqlen_q_ptrandseqlen_k_ptrfields in the CKfmha_bwd_argsstruct.test_mha_varlen.py) to ensure the correctness of the backward pass with various padding scenarios.Motivation
Technical Details
Test Plan
Test Result
Submission Checklist