From 775d9868acbe72b101b19bf9791cfa908ff56b7c Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 28 Jul 2026 10:44:03 -0700 Subject: [PATCH 1/2] [PyTorch] Honor requested FA padding in CP tests The CP test runner creates inter-sequence padding when the FlashAttention padding case is requested, but it unconditionally told DPA that FlashAttention THD inputs had no padding. That mismatch left CP backward padding uninitialized.\n\nDerive the explicit padding state from the same condition used to generate the inputs and reuse it for both reference and CP calls, preserving the sync-free CUDA-graph path for non-padding cases. Signed-off-by: Sudhakar Singh --- .../attention/run_attention_with_cp.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 82b9df262f..0bd8b7b074 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -209,6 +209,13 @@ 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": @@ -411,13 +418,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: @@ -535,12 +536,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: From 1c6d1acce1c2a41125d8dbc254b722e3f3452ad0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:28:11 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/run_attention_with_cp.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 0bd8b7b074..332c4e3ae2 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -213,9 +213,7 @@ def run_dpa_with_cp( 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" - ) + pad_between_seqs = kernel_backend == "FusedAttention" or fa_pad_between_seqs == "True" # set up environment variables and config if deterministic == "True":