@@ -13,7 +13,10 @@
     get_cu_seqlens_on_cp_rank,
     get_thd_partitioned_indices,
 )
-from transformer_engine.pytorch.attention.dot_product_attention.utils import combine_and_quantize
+from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    combine_and_quantize,
+    get_thd_padding_mask,
+)
 from transformer_engine.pytorch import DType
 from test_attention_with_cp import (
     model_configs_flash_attn,
@@ -699,11 +702,18 @@ def run_dpa_with_cp(
             num_pads_kv = (cu_seqlens_kv_padded - cu_seqlens_kv)[1:] - (
                 cu_seqlens_kv_padded - cu_seqlens_kv
             )[:-1]
-            # FA3 leaves garbage at padding positions despite seqused_q/k (tile spillover).
-            # Forward out_ can't be pre-zeroed because FA3's custom op returns out_ as an
-            # output rather than mutating it in-place, triggering PyTorch's aliasing constraint.
-            # Backward dq/dk/dv CAN be pre-zeroed because FA3 marks them as mutated inputs.
+            # FA3/FA4 kernels may leave physically padded rows unwritten. The CP
+            # implementation must clean them after partial-output/gradient merges.
             if fa_pad_between_seqs == "True":
+                out_padding_mask = get_thd_padding_mask(
+                    out_.shape[0], cu_seqlens_q, cu_seqlens_q_padded
+                )
+                nnz = torch.count_nonzero(out_[out_padding_mask]).item()
+                assert nnz == 0, (
+                    f"out_ has {nnz} nonzero values in THD padding — "
+                    "context_parallel.py should zero padding positions"
+                )
+
                 # out_ is a view inside the CP custom autograd Function, so in-place
                 # zeroing is blocked by PyTorch. Clone to break the view relationship.
                 out_ = out_.clone()
 
@@ -33,11 +33,14 @@
     _attention_backends,
 )
 from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+    AttentionParams,
     FlashAttentionUtils,
     FusedAttentionParams,
     _get_supported_versions,
     check_set_window_size,
+    get_attention_backend,
     get_fused_attn_spec,
+    get_thd_padding_mask,
 )
 from transformer_engine.pytorch.attention import RotaryPositionEmbedding
 import transformer_engine.pytorch.cpp_extensions as ext
@@ -212,6 +215,147 @@ def test_fused_attn_backend_message():
     assert message != ""
 
 
+@pytest.mark.parametrize(
+    "cu_seqlens,cu_seqlens_padded,expected",
+    [
+        ([0, 3, 8], [0, 4, 12], [False] * 3 + [True] + [False] * 5 + [True] * 3),
+        ([0, 3, 8], [0, 4, 12], [False] * 3 + [True] + [False] * 5 + [True] * 7),
+        ([0, 3, 8], [0, 3, 8], [False] * 8 + [True] * 4),
+        ([0, 0, 3, 3, 5], [0, 0, 4, 4, 8], [False] * 3 + [True] + [False] * 2 + [True] * 4),
+    ],
+)
+def test_thd_padding_mask_capacity(cu_seqlens, cu_seqlens_padded, expected):
+    """Cover sequence gaps, reserved buffer tails, and zero-length sequences."""
+    actual = get_thd_padding_mask(
+        len(expected),
+        torch.tensor(cu_seqlens, dtype=torch.int32, device="cuda"),
+        torch.tensor(cu_seqlens_padded, dtype=torch.int32, device="cuda"),
+    )
+    torch.testing.assert_close(actual, torch.tensor(expected, dtype=torch.bool, device="cuda"))
+
+
+@pytest.mark.parametrize("fa_version", [3, 4])
+def test_dpa_flash_thd_padding_capacity_cuda_graph(fa_version, monkeypatch):
+    """Check padded outputs and gradients, including buffer tails, without CP.
+
+    This regression exercises eager execution and raw CUDA Graph capture. CP graph
+    capture uses TE's make_graphed_callables and is outside this test's scope.
+    """
+    for key, value in {
+        "NVTE_FLASH_ATTN": "1",
+        "NVTE_FUSED_ATTN": "0",
+        "NVTE_UNFUSED_ATTN": "0",
+        "NVTE_FLASH_ATTN_V2": "0",
+        "NVTE_FLASH_ATTN_V3": str(int(fa_version == 3)),
+        "NVTE_FLASH_ATTN_V4": str(int(fa_version == 4)),
+    }.items():
+        monkeypatch.setenv(key, value)
+    use_flash_attention, flash_attention_backend, *_ = get_attention_backend(
+        AttentionParams(
+            qkv_dtype=torch.bfloat16,
+            qkv_layout="thd_thd_thd",
+            batch_size=2,
+            num_heads=4,
+            num_gqa_groups=4,
+            max_seqlen_q=128,
+            max_seqlen_kv=128,
+            num_tokens_q=128,
+            num_tokens_kv=192,
+            head_dim_qk=64,
+            head_dim_v=64,
+            attn_mask_type="padding",
+            window_size=(-1, -1),
+            core_attention_bias_requires_grad=False,
+            pad_between_seqs=True,
+            attention_dropout=0.0,
+            context_parallel=False,
+            deterministic=_deterministic,
+            is_training=True,
+            cuda_graph=True,
+            softmax_scale=0.125,
+        )
+    )
+    if not use_flash_attention:
+        pytest.skip(f"FlashAttention {fa_version} does not support this CUDA Graph configuration.")
+    assert flash_attention_backend.major == fa_version
+    # Restore the complete backend cache along with the environment at test teardown.
+    for key, value in _attention_backends.items():
+        monkeypatch.setitem(_attention_backends, key, value)
+    _attention_backends["backend_selection_requires_update"] = True
+    reset_rng_states()
+    cu_q, cu_kv, padded_q, padded_kv = [
+        torch.tensor(offsets, dtype=torch.int32, device="cuda")
+        for offsets in ([0, 61, 96], [0, 45, 118], [0, 64, 112], [0, 64, 160])
+    ]
+    # Expected valid token positions are independent of get_thd_padding_mask.
+    q_rows = torch.tensor(list(range(61)) + list(range(64, 99)), device="cuda")
+    kv_rows = torch.tensor(list(range(45)) + list(range(64, 137)), device="cuda")
+    qkv = [
+        torch.randn(tokens, 4, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+        for tokens in (128, 192, 192)
+    ]
+    grad_output = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+    block = DotProductAttention(
+        4,
+        64,
+        qkv_format="thd",
+        attn_mask_type="padding",
+        attention_type="cross",
+        attention_dropout=0.0,
+    ).cuda()
+
+    def run(inputs, grad, padded):
+        """Run the selected FlashAttention backend and differentiate Q/K/V."""
+        output = block(
+            *inputs,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_kv=cu_kv,
+            cu_seqlens_q_padded=padded_q if padded else None,
+            cu_seqlens_kv_padded=padded_kv if padded else None,
+            max_seqlen_q=128,
+            max_seqlen_kv=128,
+            pad_between_seqs=padded,
+        )
+        assert _attention_backends["use_flash_attention"]
+        assert _attention_backends["flash_attention_backend"].major == fa_version
+        assert not _attention_backends["use_fused_attention"]
+        assert not _attention_backends["use_unfused_attention"]
+        return (output, *torch.autograd.grad(output, inputs, grad))
+
+    def check(actual):
+        """Compare valid token positions with compact attention and require exact-zero padding."""
+        compact_qkv = [
+            tensor.detach()[rows].requires_grad_()
+            for tensor, rows in zip(qkv, (q_rows, kv_rows, kv_rows))
+        ]
+        expected = run(compact_qkv, grad_output[q_rows], False)
+        for tensor, reference, rows in zip(actual, expected, (q_rows, q_rows, kv_rows, kv_rows)):
+            assert tensor.shape[0] == (128 if rows is q_rows else 192)
+            torch.testing.assert_close(tensor[rows], reference, atol=1.5e-2, rtol=1.5e-2)
+            padding = torch.ones(tensor.shape[0], dtype=torch.bool, device="cuda")
+            padding[rows] = False
+            assert torch.count_nonzero(tensor[padding]).item() == 0
+
+    check(run(qkv, grad_output, True))
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            run(qkv, grad_output, True)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        captured = run(qkv, grad_output, True)
+    for _ in range(2):
+        # Replays must consume fresh inputs and upstream gradients at the same addresses.
+        with torch.no_grad():
+            for tensor in (*qkv, grad_output):
+                tensor.normal_()
+        graph.replay()
+        check(captured)
+
+
 # Define F16 data types to test
 param_types = [torch.float16]
 if is_bf16_available():
@@ -2166,6 +2310,22 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
 
     if backend in ["UnfusedDotProductAttention"]:
         return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset)
+    if backend == "FlashAttention" and qkv_format == "thd" and pad_between_seqs:
+        tensors_and_boundaries = [("out", out, cu_seqlens_q, cu_seqlens_q_after_pad)]
+        if is_training:
+            tensors_and_boundaries.extend(
+                [
+                    ("dq", q_grad, cu_seqlens_q, cu_seqlens_q_after_pad),
+                    ("dk", k_grad, cu_seqlens_kv, cu_seqlens_kv_after_pad),
+                    ("dv", v_grad, cu_seqlens_kv, cu_seqlens_kv_after_pad),
+                ]
+            )
+        for tensor_name, tensor, cu_seqlens, cu_seqlens_padded in tensors_and_boundaries:
+            padding_mask = get_thd_padding_mask(tensor.shape[0], cu_seqlens, cu_seqlens_padded)
+            assert (
+                torch.count_nonzero(tensor[padding_mask]).item() == 0
+            ), f"{backend} left nonzero values in {tensor_name} padding"
+
     if backend in ["FusedAttention", "FlashAttention"]:
         if qkv_format == "thd" and pad_between_seqs:
             out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype)
 
@@ -909,6 +909,33 @@ def _unalias_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
     return cu_q, cu_kv
 
 
+def _mask_thd_padding(tensor: torch.Tensor, padding_mask: torch.Tensor) -> torch.Tensor:
+    """Return ``tensor`` with invalid THD token positions zeroed."""
+    mask = padding_mask.view(padding_mask.shape[0], *([1] * (tensor.ndim - 1)))
+    return tensor.masked_fill(mask, 0)
+
+
+class _MaskTHDPaddingGrad(torch.autograd.Function):
+    """Identity in forward and zero THD padding positions in backward.
+
+    FA3 and FA4 may leave padding positions unwritten in their gradient buffers.
+    Keeping the forward as an identity avoids copying Q/K/V, while the custom
+    backward guarantees zero padding positions before gradients reach the caller.
+    """
+
+    @staticmethod
+    def forward(ctx, tensor: torch.Tensor, padding_mask: torch.Tensor) -> torch.Tensor:
+        """Save the padding mask and return the input without copying."""
+        ctx.save_for_backward(padding_mask)
+        return tensor
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        """Zero padding positions before propagating the input gradient."""
+        (padding_mask,) = ctx.saved_tensors
+        return _mask_thd_padding(grad_output, padding_mask), None
+
+
 class FlashAttention(torch.nn.Module):
     """Dot product attention, using HazyResearch flash-attn package:
     https://github.com/Dao-AILab/flash-attention
@@ -1158,6 +1185,21 @@ def forward(
         use_flash_attn_3 = (
             flash_attention_backend is not None and flash_attention_backend.major == 3
         )
+        q_padding_mask = None
+        if pad_between_seqs and qkv_format == "thd" and not context_parallel:
+            assert cu_seqlens_q_padded is not None and cu_seqlens_kv_padded is not None, (
+                "cu_seqlens_q_padded and cu_seqlens_kv_padded are required when "
+                "pad_between_seqs=True"
+            )
+            q_padding_mask = dpa_utils.get_thd_padding_mask(
+                query_layer.shape[0], cu_seqlens_q, cu_seqlens_q_padded
+            )
+            kv_padding_mask = dpa_utils.get_thd_padding_mask(
+                key_layer.shape[0], cu_seqlens_kv, cu_seqlens_kv_padded
+            )
+            query_layer = _MaskTHDPaddingGrad.apply(query_layer, q_padding_mask)
+            key_layer = _MaskTHDPaddingGrad.apply(key_layer, kv_padding_mask)
+            value_layer = _MaskTHDPaddingGrad.apply(value_layer, kv_padding_mask)
         if (
             use_flash_attn_4
             and (10, 0) <= get_device_compute_capability() < (12, 0)
@@ -1301,6 +1343,8 @@ def forward(
                     )
                     if isinstance(output, (List, Tuple)):
                         output = output[0]
+                    if q_padding_mask is not None:
+                        output = _mask_thd_padding(output, q_padding_mask)
                 elif not use_flash_attn_3:
                     fa_optional_forward_kwargs = {}
                     if fa_utils.v2_3_plus:
@@ -1422,6 +1466,8 @@ def convert_to_torch_float8(tensor, dtype):
 
                     if fp8:
                         output = output.to(dtype=torch_orig_dtype)
+                    if q_padding_mask is not None:
+                        output = _mask_thd_padding(output, q_padding_mask)
                     if fp8 and fp8_output:
                         O_quantizer = quantizers["scaling_fwd"][META_O]
                         output = O_quantizer(output)
 
@@ -2252,6 +2252,9 @@ def forward(
 
         # Save the rank-local output for backward before restoring the A2A layout.
         out = out.view(post_a2a_o_shape)
+        if qkv_format == "thd" and pad_between_seqs and not use_fused_attention:
+            # Partial-output correction can write FA3/FA4 padding rows.
+            _zero_thd_padding((out,), cu_seqlens_q_per_step[0], cu_seqlens_q_padded)
         out_part = out.to(fwd_nominal_dtype)
 
         if cp_size_a2a > 1:
@@ -4906,6 +4909,10 @@ def forward(
             aux_ctx_tensors = [softmax_lse, rng_state]
             out_part = out_
 
+            # Clean FA3/FA4 padding before inverse A2A changes sequence order.
+            if qkv_format == "thd" and pad_between_seqs:
+                _zero_thd_padding((out_,), cu_seqlens_q, cu_seqlens_q_padded)
+
         # a2a: split s and gather h
         # [b, s, h//cp, d] -> [b*s//cp, h, d]
         # [s, b, h//cp, d] -> [s//cp*b, h, d]
@@ -5304,6 +5311,11 @@ def backward(ctx, dout, *_args):
                     **fa_backward_kwargs,
                 )
 
+            # Clean FA3/FA4 padding before inverse A2A changes sequence order.
+            if ctx.dqkv_format == "thd" and ctx.pad_between_seqs:
+                _zero_thd_padding((dq,), cu_seqlens_q, cu_seqlens_q_padded)
+                _zero_thd_padding((dk, dv), cu_seqlens_kv, cu_seqlens_kv_padded)
+
         # dq, dk, dv:
         # FP8DS: torch.uint8
         # FP8CS/MXFP8/F16: torch.float16 or torch.bfloat16
 
@@ -1939,11 +1939,26 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
 
 @torch.no_grad()
 def get_thd_padding_mask(num_tokens, cu_seqlens, cu_seqlens_padded):
-    """Identify inter-sequence padding in a flattened packed THD buffer."""
+    """Return a boolean mask identifying padding token positions in a THD buffer.
+
+    ``num_tokens`` is the full physical buffer capacity, which may exceed
+    ``cu_seqlens_padded[-1]``. ``cu_seqlens`` describes cumulative valid token
+    counts, while ``cu_seqlens_padded`` describes physical sequence boundaries.
+    Inter-sequence gaps and all positions at or beyond the final padded boundary
+    are marked True. For example, a 16-token buffer with ``cu_seqlens=[0, 3, 8]``
+    and ``cu_seqlens_padded=[0, 4, 12]`` includes tail padding at positions [12, 16).
+
+    This helper is shared by non-CP and CP attention. Its operations stay on
+    device so that padding cleanup remains CUDA-graph-capturable in both paths.
+    """
     rows = torch.arange(num_tokens, device=cu_seqlens_padded.device)
     sequence = torch.searchsorted(cu_seqlens_padded[1:], rows, right=True)
-    valid_end = cu_seqlens_padded[sequence] + cu_seqlens[sequence + 1] - cu_seqlens[sequence]
-    return rows >= valid_end
+    valid_ends = cu_seqlens_padded[:-1] + cu_seqlens[1:] - cu_seqlens[:-1]
+    # Tail token positions map to sequence == batch_size. Append a zero-length sequence
+    # at the final physical boundary so that index is valid, even for an empty
+    # batch. All operations stay on device for CUDA Graph capture and replay.
+    valid_ends = torch.cat((valid_ends, cu_seqlens_padded[-1:]))
+    return rows >= valid_ends[sequence]
 
 
 @torch.no_grad()