diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py
index 32a54a7884..3b44c31a2f 100644
--- a/tests/pytorch/attention/test_attention_with_cp.py
+++ b/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",
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
index 78e599199f..e9948cb9d8 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
+++ b/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,