[Pytorch][Attention][CP][CI] Make 3221 changes robust (#3438)

* Support no-load-balance THD with FlashAttention

FlashAttention 2 backward needs zeroed gradient outputs for deterministic grouped-query attention in this path. Permit unpadded FlashAttention backends because inter-sequence padding requires no backend-specific handling, and align the autograd return arity with the forward inputs.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Clarify no-load-balance backend constraints

Document why FA2 initializes backward gradient buffers and why inter-sequence padding remains limited to FusedAttention. This makes the supported backend matrix explicit without changing behavior.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Add no-load-balance FlashAttention regression coverage

Reuse the focused CP test across supported FlashAttention generations and force FA2 from L1 in a fresh process. Backend controls are read at import time, so a separate invocation covers deterministic FA2 GQA backward without adding runner plumbing.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Use backend selection for FlashAttention CP test

Query the existing backend selector for the exact THD context-parallel configuration instead of duplicating architecture and version gates in pytest. This keeps capability decisions in one place while allowing version-forced test invocations to exercise production selection logic.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Clarify forced FA2 regression coverage

Document why the focused deterministic case runs in a fresh process and how the version flags force FA2, while the full runs retain normal backend selection.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Rely on existing L1 attention runs

The general context-parallel test is already collected by both L1 attention invocations. Since this suite selects FA2, a separate forced deterministic invocation duplicates coverage without exercising a distinct path.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: sudhakarsingh27

tests/pytorch/attention/test_attention_with_cp.py

@@ -739,12 +739,22 @@ def test_cp_with_fused_attention_no_load_balance(cp_pool):
     )
 
 
-@pytest.mark.skipif(
-    get_device_compute_capability() != (9, 0) or not FlashAttentionUtils.v3_is_installed,
-    reason="FlashAttention 3 requires sm90 and an installed FA3 package.",
-)
-def test_cp_with_flash_attention_3_no_load_balance(cp_pool):
-    """Check the supported unpadded FlashAttention 3 path."""
+def test_cp_with_flash_attention_no_load_balance(cp_pool):
+    """Check the supported unpadded FlashAttention path."""
+    config = copy.deepcopy(model_configs_flash_attn["cp_2_0"])
+    config.context_parallel = True
+    config.cp_comm_type = "all_gather"
+    config.attn_mask_type = "padding_causal"
+    available_backends, _, _ = get_available_attention_backends(
+        config,
+        qkv_dtype=torch.bfloat16,
+        qkv_layout="thd_thd_thd",
+        pad_between_seqs=False,
+        is_training=True,
+        deterministic=_deterministic,
+    )
+    if not available_backends[0]:
+        pytest.skip("FlashAttention is unavailable.")
     _submit(
         cp_pool(2),
         dtype="bf16",

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -4184,6 +4184,13 @@ def backward(ctx, dout, *_args):
                     fa_backward_kwargs["deterministic"] = ctx.deterministic
                 if fa_utils.v2_6_0_plus:
                     fa_backward_kwargs["softcap"] = 0.0
+                if (
+                    ctx.qkv_format == "thd"
+                    and ctx.load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
+                ):
+                    # FA2 GQA backward accumulates into expanded K/V gradient buffers.
+                    # Initialize them before accumulation in this partitioning path.
+                    fa_backward_kwargs["zero_tensors"] = True
 
         local_seq_chunk_ids = (
             [rank]
@@ -4563,6 +4570,7 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -5452,7 +5460,7 @@ def attn_forward_func_with_cp(
     assigns one contiguous physical-buffer chunk to each rank and uses one attention
     step per rank. Logical sequences remain isolated by ``cu_seqlens``. This strategy
     requires THD, all-gather, causal self-attention without a sliding window, and
-    FusedAttention, or FlashAttention 3 with ``pad_between_seqs=False``. Input producers
+    FusedAttention, or FlashAttention with ``pad_between_seqs=False``. Input producers
     must use the same strategy when partitioning inputs with
     :func:`get_batch_on_this_cp_rank` or :func:`get_thd_partitioned_indices`.
 
@@ -5540,12 +5548,11 @@ def attn_forward_func_with_cp(
         assert (
             cp_comm_type == "all_gather"
         ), "No-load-balance THD partitioning requires cp_comm_type='all_gather'."
+        # Backend selection is handled by the caller. FlashAttention does not yet
+        # support inter-sequence padding with this partitioning strategy.
         assert (
-            use_fused_attention or use_flash_attn_3
-        ), "No-load-balance THD partitioning requires FusedAttention or FlashAttention 3."
-        assert not (
-            use_flash_attn_3 and pad_between_seqs
-        ), "No-load-balance THD partitioning with FlashAttention 3 does not support padding yet."
+            not pad_between_seqs or use_fused_attention
+        ), "No-load-balance THD partitioning only supports padding with FusedAttention."
         assert "causal" in attn_mask_type and window_size == (
             -1,
             0,