Fix THD P2P pad detection and tail-zero gating

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -1825,29 +1825,45 @@ def backward(ctx, d_out, *_args):
             dq = dq[..., : d_out.shape[-1]]
             dk = dk[..., : d_out.shape[-1]]
             dv = dv[..., : d_out.shape[-1]]
-            # Zero-fill positions beyond cu_seqlens_*_padded[-1] in dQ/dK/dV for THD.
-            # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV.
-            # Sync-free `arange + mask` so capture and eager paths run the same code.
             _qkv_format = ctx.qkv_layout.split("_")[0].replace("3", "").replace("2", "")
             if _qkv_format == "thd":
-                if (
-                    cu_seqlens_q_padded is not None
-                    and isinstance(dq, torch.Tensor)
-                    and dq.shape[0] > 0
-                ):
-                    q_pad_mask = (
-                        torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
-                    )
-                    dq[q_pad_mask] = 0
-                if cu_seqlens_kv_padded is not None:
-                    kv_actual_t = cu_seqlens_kv_padded[-1]
-                    for d_tensor in (dk, dv):
-                        if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0:
-                            kv_pad_mask = (
-                                torch.arange(d_tensor.shape[0], device=d_tensor.device)
-                                >= kv_actual_t
-                            )
-                            d_tensor[kv_pad_mask] = 0
+                if is_graph_capturing():
+                    if (
+                        cu_seqlens_q_padded is not None
+                        and isinstance(dq, torch.Tensor)
+                        and dq.shape[0] > 0
+                    ):
+                        q_pad_mask = (
+                            torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
+                        )
+                        dq[q_pad_mask] = 0
+                    if cu_seqlens_kv_padded is not None:
+                        kv_actual_t = cu_seqlens_kv_padded[-1]
+                        for d_tensor in (dk, dv):
+                            if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0:
+                                kv_pad_mask = (
+                                    torch.arange(d_tensor.shape[0], device=d_tensor.device)
+                                    >= kv_actual_t
+                                )
+                                d_tensor[kv_pad_mask] = 0
+                else:
+                    if (
+                        cu_seqlens_q_padded is not None
+                        and isinstance(dq, torch.Tensor)
+                        and dq.shape[0] > 0
+                    ):
+                        q_end = int(cu_seqlens_q_padded[-1])
+                        if q_end < dq.shape[0]:
+                            dq[q_end:].fill_(0)
+                    if cu_seqlens_kv_padded is not None:
+                        kv_end = int(cu_seqlens_kv_padded[-1])
+                        for d_tensor in (dk, dv):
+                            if (
+                                isinstance(d_tensor, torch.Tensor)
+                                and d_tensor.shape[0] > 0
+                                and kv_end < d_tensor.shape[0]
+                            ):
+                                d_tensor[kv_end:].fill_(0)
         else:
             with get_nvtx_range_context("FusedAttnFunc.backward"):
                 # get nominal data type of dq, dk, dv

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -3107,8 +3107,6 @@ def backward(ctx, dout, *_args):
             and cu_seqlens_kv_padded is not None
         ):
             if is_graph_capturing():
-                # arange+mask under capture: `tensor[scalar_tensor:]` slicing would
-                # force a GPU->CPU sync that is forbidden during CUDA graph capture.
                 q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1]
                 kv_pad_mask = (
                     torch.arange(dk.shape[0], device=dk.device) >= cu_seqlens_kv_padded[-1]
@@ -3117,10 +3115,14 @@ def backward(ctx, dout, *_args):
                 dk[kv_pad_mask] = 0
                 dv[kv_pad_mask] = 0
             else:
-                # Pre-existing TE eager-mode behaviour.
-                dq[cu_seqlens_q_padded[-1] :].fill_(0)
-                dk[cu_seqlens_kv_padded[-1] :].fill_(0)
-                dv[cu_seqlens_kv_padded[-1] :].fill_(0)
+                q_end = int(cu_seqlens_q_padded[-1])
+                kv_end = int(cu_seqlens_kv_padded[-1])
+                if q_end < dq.shape[0]:
+                    dq[q_end:].fill_(0)
+                if kv_end < dk.shape[0]:
+                    dk[kv_end:].fill_(0)
+                if kv_end < dv.shape[0]:
+                    dv[kv_end:].fill_(0)
 
         if ctx.fp8 and ctx.is_input_fp8:
             dq, dk, dv, _, _ = combine_and_quantize(ctx.qkv_layout, dq, dk, dv, ctx.dQKV_quantizer)
@@ -3138,9 +3140,7 @@ def backward(ctx, dout, *_args):
                 ctx.dP_quantizer,
             )
 
-        # Partial-gradient reduction can write THD inter-sequence padding.
-        # Clean it while gradients and per-step sequence metadata share sequence order.
-        if ctx.qkv_format == "thd":
+        if ctx.qkv_format == "thd" and ctx.pad_between_seqs:
             _zero_thd_padding((dq,), cu_seqlens_q_per_step[0], cu_seqlens_q_padded)
             _zero_thd_padding((dk, dv), cu_seqlens_kv_per_step[0], cu_seqlens_kv_padded)
 

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -2682,23 +2682,25 @@ def forward(
                 else None
             )
 
-            # Default pad_between_seqs auto-detect. For THD, infer presence of
-            # inter-sequence padding from whether padded cu_seqlens were supplied --
-            # sync-free, and stable across eager and CUDA graph capture (the auto-detect
-            # must return the same value in both modes for backend selection to match).
-            # If padded cu_seqlens are the *same object* as the unpadded ones, no real
-            # inter-sequence padding exists (only THD tail padding) -- treat as False so
-            # FlashAttention v2/v4 remain eligible.
             if pad_between_seqs is None:
                 if qkv_format == "thd":
-                    if (
-                        cu_seqlens_q_padded is cu_seqlens_q
-                        and cu_seqlens_kv_padded is cu_seqlens_kv
-                    ):
-                        pad_between_seqs = False
+                    if is_graph_capturing():
+                        if (
+                            cu_seqlens_q_padded is cu_seqlens_q
+                            and cu_seqlens_kv_padded is cu_seqlens_kv
+                        ):
+                            pad_between_seqs = False
+                        else:
+                            pad_between_seqs = (
+                                cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None
+                            )
                     else:
                         pad_between_seqs = (
-                            cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None
+                            cu_seqlens_q_padded is not None
+                            and not torch.equal(cu_seqlens_q_padded[:-1], cu_seqlens_q[:-1])
+                        ) or (
+                            cu_seqlens_kv_padded is not None
+                            and not torch.equal(cu_seqlens_kv_padded[:-1], cu_seqlens_kv[:-1])
                         )
                 else:
                     pad_between_seqs = False