Skip to content

[Bugfix][PyTorch][FlashAttention][CP] Honor requested FA padding in CP tests #3269

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

Merged
sudhakarsingh27 merged 2 commits into from
Jul 29, 2026
Merged

[Bugfix][PyTorch][FlashAttention][CP] Honor requested FA padding in CP tests #3269

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
775d986
[PyTorch] Honor requested FA padding in CP tests
sudhakarsingh27 Jul 28, 2026
1c6d1ac
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 7 additions & 13 deletions tests/pytorch/attention/run_attention_with_cp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,11 @@ def run_dpa_with_cp(
logging.root.setLevel(log_level)
# When is_training is False, gradient outputs are None.
is_training = is_training == "True"
pad_between_seqs = None
if qkv_format == "thd":
# Keep this in sync with generate_input_shapes so DPA gets the explicit
# padding state without a GPU-to-CPU sync.
pad_between_seqs = kernel_backend == "FusedAttention" or fa_pad_between_seqs == "True"

# set up environment variables and config
if deterministic == "True":
Expand Down Expand Up @@ -411,13 +416,7 @@ def run_dpa_with_cp(
cu_seqlens_kv=cu_seqlens_kv,
cu_seqlens_q_padded=cu_seqlens_q_padded,
cu_seqlens_kv_padded=cu_seqlens_kv_padded,
# Test runner sets cu_seqlens_q == cu_seqlens_q_padded for the
# FlashAttention path, i.e. no inter-sequence padding. Declare this
# explicitly so the sync-free auto-detect (which conservatively
# picks True when padded cu_seqlens are present) does not disable FA.
pad_between_seqs=(
(kernel_backend != "FlashAttention") if qkv_format == "thd" else None
),
pad_between_seqs=pad_between_seqs,
fp8_output=fp8_mha,
)
if config.return_max_logit:
Expand Down Expand Up @@ -535,12 +534,7 @@ def run_dpa_with_cp(
cu_seqlens_kv=cu_seqlens_kv,
cu_seqlens_q_padded=cu_seqlens_q_padded,
cu_seqlens_kv_padded=cu_seqlens_kv_padded,
# See note above (non-CP branch): same explicit declaration so
# FlashAttention isn't disabled by the conservative sync-free
# auto-detect when this test path constructs no inter-seq padding.
pad_between_seqs=(
(kernel_backend != "FlashAttention") if qkv_format == "thd" else None
),
pad_between_seqs=pad_between_seqs,
fp8_output=fp8_mha,
)
if config.return_max_logit:
Expand Down
Loading