modify the cu_seqlens_q_padded&cu_seqlens_k_padded acquisition way, add the parameter in python api - #1615
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refactors the handling of padded cumulative sequence length pointers in the FMHA (Fused Multi-Head Attention) v3 API. The key change is moving the cu_seqlens_q_padded and cu_seqlens_k_padded parameters from being passed as separate function arguments to being part of the mha_fwd_args structure, accessed via cu_seqlen_q_ptr and cu_seqlen_k_ptr fields.
Key Changes:
- Removed
seqstart_q_padding_ptrandseqstart_k_padding_ptrparameters from internal v3 dispatcher functions - Added
cu_seqlens_q_paddedandcu_seqlens_k_paddedas optional parameters to Python API functions - Modified C++ implementation to retrieve padded pointers from the
mha_fwd_argsstruct instead of separate function parameters
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| hsa/gfx950/fmha_v3_fwd/codegen.py | Updated dispatcher function signatures to remove padding pointer parameters; modified argument assignment to use struct fields |
| hsa/gfx942/fmha_v3_fwd/codegen.py | Updated dispatcher function signatures to remove padding pointer parameters; modified argument assignment to use struct fields |
| csrc/py_itfs_cu/asm_mha_varlen_fwd.cu | Added cu_seqlens_q_padded and cu_seqlens_k_padded parameters to function signature and passed them to args construction |
| csrc/include/torch/mha_v3_varlen_fwd.h | Added cu_seqlens_q_padded and cu_seqlens_k_padded parameters to function declaration |
| csrc/include/rocm_ops.hpp | Added Python bindings for cu_seqlens_q_padded and cu_seqlens_k_padded parameters |
| csrc/include/mha_fwd.h | Removed default parameters from fmha_fwd_v3 function declarations in gfx942 and gfx950 namespaces |
| csrc/cpp_itfs/mha_fwd_generate.py | Updated v3 function call to remove padding pointer arguments |
| aiter/ops/mha.py | Added cu_seqlens_q_padded and cu_seqlens_k_padded parameters to Python API; removed check that prevented using these parameters |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
wangye805
requested changes
Dec 11, 2025
wangye805
reviewed
Dec 12, 2025
wangye805
reviewed
Dec 12, 2025
minmengdie
force-pushed
the
mmd/fix/fwd_v3
branch
from
December 13, 2025 10:35
574c56d to
520ad94
Compare
wangye805
approved these changes
Dec 13, 2025
valarLip
approved these changes
Dec 16, 2025
valarLip
approved these changes
Dec 16, 2025
zhuyuhua-v
pushed a commit
that referenced
this pull request
Dec 17, 2025
…dd the parameter in python api (#1615) * modify the cu_seqlens_q_padded cu_seqlens_k_padded * fix * fix the ck parameter * delete the log * exchange the ptr_qseq_padding and ptr_qseq * batch mode not support padded
ZhangLirong-amd
pushed a commit
that referenced
this pull request
Dec 29, 2025
…dd the parameter in python api (#1615) * modify the cu_seqlens_q_padded cu_seqlens_k_padded * fix * fix the ck parameter * delete the log * exchange the ptr_qseq_padding and ptr_qseq * batch mode not support padded
valarLip
pushed a commit
that referenced
this pull request
Mar 18, 2026
…dd the parameter in python api (#1615) * modify the cu_seqlens_q_padded cu_seqlens_k_padded * fix * fix the ck parameter * delete the log * exchange the ptr_qseq_padding and ptr_qseq * batch mode not support padded
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
Technical Details
Test Plan
python3 reproduce_loadpt.py
bash smoke_test_fwd_v3.sh
Test Result
Submission Checklist