Fix THD P2P pad detection and tail-zero gating

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -2067,22 +2067,45 @@ def _fused_attn_backward_impl(
         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 = args.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

@@ -3109,8 +3109,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]
@@ -3119,10 +3117,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)
@@ -3140,9 +3142,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

@@ -4,6 +4,7 @@
 
 """Attention."""
 from contextlib import nullcontext
+import functools
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -113,6 +114,45 @@
 _THD_POLICY_VALIDATION_CACHE_LIMIT = 64
 _thd_policy_validation_cache = []
 
+_PAD_BETWEEN_SEQS_EAGER_CACHE_LIMIT = 64
+_pad_between_seqs_eager_store = {}
+
+
+def _pad_between_seqs_cache_key(
+    cu_seqlens_q, cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded
+):
+    def _key(t):
+        if t is None:
+            return (0, 0)
+        return (id(t), _get_thd_policy_tensor_version(t))
+
+    return (
+        _key(cu_seqlens_q),
+        _key(cu_seqlens_kv),
+        _key(cu_seqlens_q_padded),
+        _key(cu_seqlens_kv_padded),
+    )
+
+
+@functools.lru_cache(maxsize=_PAD_BETWEEN_SEQS_EAGER_CACHE_LIMIT)
+def _lookup_pad_between_seqs_cache(key):
+    return _pad_between_seqs_eager_store[key]
+
+
+def _store_pad_between_seqs_cache(key, tensors, decision):
+    overwrite = key in _pad_between_seqs_eager_store
+    _pad_between_seqs_eager_store[key] = (tensors, decision)
+    evicted = False
+    while len(_pad_between_seqs_eager_store) > _PAD_BETWEEN_SEQS_EAGER_CACHE_LIMIT:
+        _pad_between_seqs_eager_store.pop(next(iter(_pad_between_seqs_eager_store)))
+        evicted = True
+    if overwrite or evicted:
+        _lookup_pad_between_seqs_cache.cache_clear()
+    try:
+        _lookup_pad_between_seqs_cache(key)
+    except KeyError:
+        pass
+
 
 def _get_thd_policy_attention_backend(
     policy: Dict[str, Any],
@@ -2706,23 +2746,47 @@ 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
+                    tensors = (
+                        cu_seqlens_q,
+                        cu_seqlens_kv,
+                        cu_seqlens_q_padded,
+                        cu_seqlens_kv_padded,
+                    )
+                    if is_graph_capturing():
+                        cache_key = _pad_between_seqs_cache_key(*tensors)
+                        try:
+                            cached_tensors, pad_between_seqs = _lookup_pad_between_seqs_cache(
+                                cache_key
+                            )
+                        except KeyError:
+                            cached_tensors = None
+                        if cached_tensors is None or any(
+                            t is not ct for t, ct in zip(tensors, cached_tensors)
+                        ):
+                            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])
+                        )
+                        _store_pad_between_seqs_cache(
+                            _pad_between_seqs_cache_key(*tensors),
+                            tensors,
+                            pad_between_seqs,
                         )
                 else:
                     pad_between_seqs = False