From f6ed27c446829e0874281ec7744359c082188d07 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 27 Jul 2026 11:07:09 +0200
Subject: [PATCH 01/39] [PyTorch][torch.compile] Support for
 DotProductAttention

Make DotProductAttention itself traceable, so that torch.compile captures
the whole module -- input unpacking, qkv layout, backend selection and the
backend -- on the FlashAttention and UnfusedDotProductAttention backends.
FusedAttention stays an eager island, as before.

DotProductAttention.forward no longer carries @no_torch_dynamo. What that
uncovered, and the fixes:

- init_fp8_metadata unconditionally called get_fp8_recipe(), which builds a
  default Recipe when no autocast is active -- not traceable. It now returns
  early when FP8 is off everywhere, which is also what the rest of the
  function did in that case.
- get_qkv_layout detects packed q/k/v from data pointers and storage
  offsets, neither of which dynamo can trace (and neither is meaningful on
  fake tensors). While tracing, the layout is built from the format instead;
  packed inputs are declared explicitly via qkv_layer/kv_layer, and
  non-contiguous q/k/v are made contiguous so the layout stays truthful.
- get_full_cu_seqlens keyed its cache on torch.is_inference_mode_enabled(),
  a fundamental graph break; while tracing it builds the tensor directly.
- get_padding_mask read sequence lengths on the host (a device sync, and
  data-dependent under torch.compile). It is now vectorized.
- The fused sub-backend enum member does not survive a graph break: dynamo
  reconstructs the assume_constant_result value by re-emitting the call that
  produced it, so the resumed frame received the function itself instead of
  the enum -- which then reached fused_attn_fwd. get_attention_backend now
  keeps a plain int while tracing and casts back to FusedAttnBackend in
  eager mode only.
- Backend selection logging is skipped while tracing (logging.Logger methods
  graph-break, and int() of the sub-backend enum is untraceable too).

Tests: test_dpa_torch_compile covers 7 configurations (bshd/sbhd/thd,
causal/no_mask/sliding window, GQA, cross attention, packed qkv_layer)
against both backends with fullgraph=True, comparing forward and backward
against eager; plus a CUDA-graphs (reduce-overhead) test and a test that the
FusedAttention path stays correct around its graph break.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 243 ++++++++++++++++++
 .../dot_product_attention.py                  |  47 +++-
 .../attention/dot_product_attention/utils.py  |  58 ++++-
 .../pytorch/cpp_extensions/fused_attn.py      |  11 +-
 4 files changed, 330 insertions(+), 29 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d720ca7a4a..ed2ad90c76 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -550,6 +550,249 @@ def fn(q, k, v, extra):
         assert v.grad is not None
 
 
+# ---------------------------------------------------------------------------
+# DotProductAttention under torch.compile
+# ---------------------------------------------------------------------------
+
+
+# Each config is a model configuration that both the FlashAttention and the
+# UnfusedDotProductAttention backend support, so that every one of them runs
+# against both backends.
+_DPA_COMPILE_CONFIGS = {
+    "self_bshd_causal": dict(qkv_format="bshd", attn_mask_type="causal"),
+    "self_sbhd_no_mask": dict(qkv_format="sbhd", attn_mask_type="no_mask"),
+    "self_bshd_swa": dict(qkv_format="bshd", attn_mask_type="causal", window_size=(16, 0)),
+    "gqa_bshd_causal": dict(
+        qkv_format="bshd", attn_mask_type="causal", num_heads=8, num_gqa_groups=2
+    ),
+    "cross_bshd_no_mask": dict(
+        qkv_format="bshd", attn_mask_type="no_mask", attention_type="cross", max_seqlen_kv=256
+    ),
+    "self_thd_padding_causal": dict(qkv_format="thd", attn_mask_type="padding_causal"),
+    "packed_qkv_bs3hd": dict(qkv_format="bshd", attn_mask_type="causal", packed="bs3hd"),
+}
+
+
+def _dpa_config(**overrides):
+    cfg = dict(
+        batch_size=2,
+        num_heads=4,
+        num_gqa_groups=None,
+        head_dim=64,
+        max_seqlen_q=128,
+        max_seqlen_kv=128,
+        qkv_format="bshd",
+        attn_mask_type="causal",
+        attention_type="self",
+        window_size=None,
+        packed=None,
+    )
+    cfg.update(overrides)
+    if cfg["num_gqa_groups"] is None:
+        cfg["num_gqa_groups"] = cfg["num_heads"]
+    return cfg
+
+
+def _make_dpa(cfg, dtype: torch.dtype) -> te.DotProductAttention:
+    return te.DotProductAttention(
+        num_attention_heads=cfg["num_heads"],
+        kv_channels=cfg["head_dim"],
+        num_gqa_groups=cfg["num_gqa_groups"],
+        attention_dropout=0.0,
+        qkv_format=cfg["qkv_format"],
+        attn_mask_type=cfg["attn_mask_type"],
+        window_size=cfg["window_size"],
+        attention_type=cfg["attention_type"],
+    ).to(dtype=dtype, device="cuda")
+
+
+def _make_dpa_inputs(cfg, dtype: torch.dtype):
+    """Build the (args, kwargs) that `DotProductAttention.forward` is called
+    with, plus the list of tensors whose gradients the test compares."""
+    b = cfg["batch_size"]
+    s_q, s_kv = cfg["max_seqlen_q"], cfg["max_seqlen_kv"]
+    h, g, d = cfg["num_heads"], cfg["num_gqa_groups"], cfg["head_dim"]
+    fmt = cfg["qkv_format"]
+
+    def _shape(s, heads):
+        return {
+            "bshd": (b, s, heads, d),
+            "sbhd": (s, b, heads, d),
+            "thd": (b * s, heads, d),
+        }[fmt]
+
+    def _randn(shape):
+        return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+
+    kwargs = {}
+    if cfg["packed"] is not None:
+        # Declarative packed QKV: q/k/v are derived from one buffer by DPA
+        # itself, and the layout comes from the declaration -- the only packed
+        # layout that torch.compile supports (pointer-based detection of packed
+        # q/k/v is untraceable).
+        assert cfg["packed"] == "bs3hd" and fmt == "bshd" and h == g
+        qkv = _randn((b, s_q, 3, h, d))
+        kwargs["qkv_layer"] = qkv
+        kwargs["qkv_interleave_dim"] = -3
+        grad_tensors = [qkv]
+        args = ()
+    else:
+        q, k, v = _randn(_shape(s_q, h)), _randn(_shape(s_kv, g)), _randn(_shape(s_kv, g))
+        grad_tensors = [q, k, v]
+        args = (q, k, v)
+
+    if fmt == "thd":
+        # All sequences have the maximum length, i.e. no padding. max_seqlen_*
+        # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
+        # (a device sync, and an unbacked SymInt under torch.compile).
+        kwargs["cu_seqlens_q"] = torch.arange(
+            0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda"
+        )
+        kwargs["cu_seqlens_kv"] = torch.arange(
+            0, (b + 1) * s_kv, step=s_kv, dtype=torch.int32, device="cuda"
+        )
+        kwargs["max_seqlen_q"] = s_q
+        kwargs["max_seqlen_kv"] = s_kv
+
+    return args, kwargs, grad_tensors
+
+
+def _force_dpa_backend(monkeypatch, backend: str) -> None:
+    """Restrict DotProductAttention to a single backend."""
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+    from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+        FlashAttentionUtils,
+    )
+
+    if backend == "flash" and not FlashAttentionUtils.is_installed:
+        pytest.skip("flash-attn is not installed")
+
+    for name, var in (
+        ("flash", "NVTE_FLASH_ATTN"),
+        ("fused", "NVTE_FUSED_ATTN"),
+        ("unfused", "NVTE_UNFUSED_ATTN"),
+    ):
+        monkeypatch.setenv(var, "1" if name == backend else "0")
+    # Backend selection is cached on the attention params only, so the env vars
+    # above are not enough to invalidate it.
+    _attention_backends["backend_selection_requires_update"] = True
+
+
+def _assert_dpa_backend(backend: str) -> None:
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+
+    assert _attention_backends[f"use_{backend}_attention"], (
+        f"expected the {backend} backend to run, selected:"
+        f" flash={_attention_backends['use_flash_attention']},"
+        f" fused={_attention_backends['use_fused_attention']},"
+        f" unfused={_attention_backends['use_unfused_attention']}"
+    )
+
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
+def test_dpa_torch_compile(monkeypatch, backend, config):
+    """`DotProductAttention` must trace under `torch.compile(fullgraph=True)`
+    on the FlashAttention and UnfusedDotProductAttention backends, and match
+    eager in forward and backward.
+
+    `fullgraph=True` makes the test fail on any graph break, i.e. it covers the
+    whole module: input unpacking, qkv layout, backend selection and the
+    backend itself. The FusedAttention backend is not covered -- it is an eager
+    island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS[config])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+
+    ref = module(*args, **kwargs)
+    ref.sum().backward()
+    _assert_dpa_backend(backend)
+    ref_grads = [t.grad.clone() for t in grads]
+    for t in grads:
+        t.grad = None
+
+    torch._dynamo.reset()
+    # Force backend selection to be re-run (and traced) inside the compiled
+    # region instead of being served from the cache the eager call populated.
+    _force_dpa_backend(monkeypatch, backend)
+    compiled = torch.compile(module, fullgraph=True)
+
+    out = compiled(*args, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend(backend)
+    torch.testing.assert_close(out, ref)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        assert tensor.grad is not None
+        torch.testing.assert_close(tensor.grad, ref_grad)
+
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
+    """`mode="reduce-overhead"`: forward and backward of DotProductAttention
+    are captured into CUDA graphs and replayed on subsequent iterations."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
+
+    for _ in range(3):
+        args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+        out = compiled(*args, **kwargs)
+        out.sum().backward()
+        torch.cuda.synchronize()
+        assert torch.isfinite(out).all()
+        for tensor in grads:
+            assert tensor.grad is not None
+    _assert_dpa_backend(backend)
+
+
+def test_dpa_torch_compile_fused_backend(monkeypatch):
+    """The FusedAttention backend is an eager island, so it graph-breaks rather
+    than compiles; DotProductAttention must still be correct around it."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, "fused")
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+
+    ref = module(*args, **kwargs)
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+
+    if not _attention_backends["use_fused_attention"]:
+        pytest.skip("FusedAttention does not support this config on this device")
+    ref.sum().backward()
+    ref_grads = [t.grad.clone() for t in grads]
+    for t in grads:
+        t.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, "fused")
+    out = torch.compile(module)(*args, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend("fused")
+    torch.testing.assert_close(out, ref)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        torch.testing.assert_close(tensor.grad, ref_grad)
+
+
 # ---------------------------------------------------------------------------
 # get_attention_backend under torch.compile
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4cc4cab1b8..0b21adb15c 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,6 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -715,6 +714,22 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
+        fp8 = FP8GlobalStateManager.is_fp8_enabled()
+        fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
+        fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
+        if not (fp8 or fp8_calibration or fp8_parameters):
+            # Nothing below has an effect when FP8 is off everywhere -- the recipe
+            # juggling ends in the "turn off and return" branch. Bailing out here
+            # also keeps DotProductAttention traceable by torch.compile: with no
+            # autocast active, get_fp8_recipe() below builds a default Recipe,
+            # whose construction is not traceable.
+            self.fast_setattr("fp8_parameters", fp8_parameters)
+            self.fast_setattr("fp8", fp8)
+            self.fast_setattr("fp8_calibration", fp8_calibration)
+            self.fp8_meta["fp8_checkpoint"] = False
+            self.fast_setattr("fp8_initialized", False)
+            return
+
         # global recipe set in autocast()
         fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
         if fp8_recipe.custom():
@@ -1090,7 +1105,6 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @no_torch_dynamo(recursive=False)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
@@ -1820,18 +1834,23 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    if use_flash_attention:
-                        self.logger.info(
-                            "Running with FlashAttention backend (version %s)",
-                            flash_attention_backend,
-                        )
-                    elif use_fused_attention:
-                        self.logger.info(
-                            "Running with FusedAttention backend (sub-backend %s)",
-                            int(fused_attention_backend),
-                        )
-                    elif use_unfused_attention:
-                        self.logger.info("Running with UnfusedDotProductAttention backend")
+                    # Backend selection is only logged in eager mode: dynamo
+                    # graph-breaks on logging.Logger methods, and even the
+                    # arguments below are untraceable (int() of the sub-backend
+                    # enum). Same restriction as in get_attention_backend.
+                    if not torch.compiler.is_compiling():
+                        if use_flash_attention:
+                            self.logger.info(
+                                "Running with FlashAttention backend (version %s)",
+                                flash_attention_backend,
+                            )
+                        elif use_fused_attention:
+                            self.logger.info(
+                                "Running with FusedAttention backend (sub-backend %s)",
+                                fused_attention_backend,
+                            )
+                        elif use_unfused_attention:
+                            self.logger.info("Running with UnfusedDotProductAttention backend")
                 else:
                     use_flash_attention = _attention_backends["use_flash_attention"]
                     flash_attention_backend = _attention_backends["flash_attention_backend"]
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 8a4c6ce286..4a8b7beb14 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -361,11 +361,17 @@ def _get_fused_attn_backend(
     *args,
 ):
     """Constant-foldable tex.get_fused_attn_backend: the result depends only on
-    the attention config, and the python-side enum keeps it traceable by
-    torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax
-    are taken as their string keys and resolved to the pybind enums here, so
-    that every argument is a python literal or a python enum."""
-    return FusedAttnBackend.cast(
+    the attention config. Layout/bias/mask/softmax are taken as their string
+    keys and resolved to the pybind enums here, so that every argument is a
+    python literal or a python enum.
+
+    Returns a plain int rather than a FusedAttnBackend member: dynamo
+    reconstructs the result of an assume_constant_result call by re-emitting the
+    call, which is only valid inside the frame that made it. An int survives a
+    graph break because it is baked into the graph as a literal, while an enum
+    member comes out of the reconstruction corrupted (see the cast at the call
+    site, which restores the enum)."""
+    return int(
         tex.get_fused_attn_backend(
             is_training,
             q_type,
@@ -397,6 +403,9 @@ def get_attention_backend(
         Whether the `FusedAttention` backend has been selected.
     fused_attention_backend : FusedAttnBackend
         If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
+        Under `torch.compile` this is the sub-backend's plain integer value instead of the enum
+        member (see the cast below); `FusedAttnBackend.cast` accepts both, as does every consumer
+        in `transformer_engine.pytorch`.
     use_unfused_attention : bool
         Whether the `UnfusedDotProductAttention` backend has been selected.
     available_backends : List[bool]
@@ -1455,11 +1464,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             cuda_graph,
             deterministic,
         )
-        if fused_attention_backend == FusedAttnBackend["No_Backend"]:
+        if fused_attention_backend == FusedAttnBackend.No_Backend.value:
             logger.debug("Disabling FusedAttention as no backend supports the provided input")
             use_fused_attention = False
             fused_attention_backend = None
-        elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]:
+        elif (
+            has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value
+        ):
             logger.debug(
                 "Disabling FusedAttention for score_mod because sub-backend %s is not "
                 "F16/BF16 arbitrary-seqlen",
@@ -1509,7 +1520,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             use_fused_attention = False
             fused_attention_backend = None
         if (
-            fused_attention_backend == FusedAttnBackend["FP8"]
+            fused_attention_backend == FusedAttnBackend.FP8.value
             and is_training
             and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0))
         ):
@@ -1520,7 +1531,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             use_fused_attention = False
             fused_attention_backend = None
         if (
-            fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
+            fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value
             and is_training
             and (
                 device_compute_capability < (9, 0)
@@ -1589,6 +1600,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         use_flash_attention_4 = False
     use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
     available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention]
+    if fused_attention_backend is not None and not torch.compiler.is_compiling():
+        # The fused sub-backend is kept as a plain int while tracing: it has to
+        # survive a graph break (DotProductAttention hands it to FusedAttention,
+        # which is an eager island) and an enum member does not -- dynamo
+        # reconstructs the constant-folded member into the resumed frame as the
+        # function that produced it. Eager callers get the enum, as before.
+        fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend)
     if use_flash_attention_2:
         flash_attention_backend = FlashAttentionUtils.version
     if use_flash_attention_3:
@@ -2025,7 +2043,10 @@ def _get_cu_seqlens(batch_size, max_seqlen, device):
             device=device,
         )
 
-    if is_in_onnx_export_mode():
+    if is_in_onnx_export_mode() or torch.compiler.is_compiling():
+        # torch.is_inference_mode_enabled() (the cache key below) is a fundamental
+        # graph break under torch.compile; building the tensor is a single arange
+        # that the compiler captures into the graph anyway.
         return _get_cu_seqlens(batch_size, max_seqlen, device)
 
     is_inference = torch.is_inference_mode_enabled()
@@ -2538,7 +2559,22 @@ def run_iteratively(q, k, v):
 
         return qkv_layout
 
-    if not is_in_onnx_export_mode():
+    if torch.compiler.is_compiling():
+        # The detection in run_iteratively is untraceable by dynamo -- it reads
+        # UntypedStorage.data_ptr and Tensor.storage_offset -- and pointers are
+        # meaningless on the fake tensors used while tracing anyway. Under
+        # torch.compile, packed q/k/v must therefore be declared explicitly, via
+        # DotProductAttention's qkv_layer/kv_layer arguments (which bypass this
+        # function entirely); q/k/v passed separately are treated as three
+        # independent tensors. Views into a packed buffer are non-contiguous, so
+        # making the inputs contiguous keeps the returned layout truthful for
+        # every backend.
+        q, k, v = [x if x.is_contiguous() else x.contiguous() for x in [q, k, v]]
+        if is_same_q_kv_format:
+            qkv_layout = "_".join([qkv_format] * 3)
+        else:
+            qkv_layout = q_format + "_" + kv_format + "_" + kv_format
+    elif not is_in_onnx_export_mode():
         qkv_layout = run_iteratively(q, k, v)
     else:
         qkv_layout = "not_supported"
diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
index 046019ee58..3a4e5aac66 100644
--- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py
+++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
@@ -107,10 +107,13 @@ class FusedAttnBackend(IntEnum):
     ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum
     value-for-value, and instances of the two enums compare equal when they
     share the same integer value. Unlike the pybind enum, a plain-python
-    ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold
-    cleanly and instances safely cross the ``assume_constant_result`` boundary
-    in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``)
-    works the same way as with the dict this used to be.
+    ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member
+    constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works
+    the same way as with the dict this used to be.
+
+    Members do not survive a graph break, though, so ``get_attention_backend``
+    hands out plain ints while tracing -- ``cast`` normalizes both back to a
+    member.
     """
 
     No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend)

From 3053d76f97dba1e4b1abba8ed009c9ff176a4310 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 12:51:11 +0200
Subject: [PATCH 02/39] Delegate the no-FP8 case of DPA's init_fp8_metadata to
 the base class

The override hand-rolled the state that TransformerEngineBaseModule already
sets when FP8 is off everywhere (the three flags, fp8_checkpoint and
fp8_initialized). Call super() instead, and read the flags straight off
FP8GlobalStateManager.quantization_state, as the other hot paths do.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 20 ++++++-------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 0b21adb15c..41aef33280 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -714,20 +714,12 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
-        fp8 = FP8GlobalStateManager.is_fp8_enabled()
-        fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
-        fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
-        if not (fp8 or fp8_calibration or fp8_parameters):
-            # Nothing below has an effect when FP8 is off everywhere -- the recipe
-            # juggling ends in the "turn off and return" branch. Bailing out here
-            # also keeps DotProductAttention traceable by torch.compile: with no
-            # autocast active, get_fp8_recipe() below builds a default Recipe,
-            # whose construction is not traceable.
-            self.fast_setattr("fp8_parameters", fp8_parameters)
-            self.fast_setattr("fp8", fp8)
-            self.fast_setattr("fp8_calibration", fp8_calibration)
-            self.fp8_meta["fp8_checkpoint"] = False
-            self.fast_setattr("fp8_initialized", False)
+        qstate = FP8GlobalStateManager.quantization_state
+        if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
+            # Without FP8 the recipe juggling below is a no-op, and its
+            # get_fp8_recipe() call would build a default Recipe -- which
+            # torch.compile cannot trace.
+            super().init_fp8_metadata(num_gemms=num_gemms)
             return
 
         # global recipe set in autocast()

From d6b285ed1cf5ff6d036acadd0cc7b517d0af2002 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 15:42:46 +0200
Subject: [PATCH 03/39] Run DPA eagerly for undeclared packed q/k/v instead of
 guessing the layout

get_qkv_layout recognizes packed q/k/v by inspecting data pointers and storage
offsets, which dynamo cannot trace. Rather than assume such inputs are three
separate tensors -- a layout that would not describe memory -- detect the case
from what is traceable (a strided tensor that does not own its whole storage)
and run the whole module as an eager island, with a warning. Declaring the
packing via qkv_layer/kv_layer keeps the call on the compiled path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 41 ++++++++++++++++
 .../dot_product_attention.py                  | 49 +++++++++++++++++++
 .../attention/dot_product_attention/utils.py  | 33 ++++++++-----
 3 files changed, 112 insertions(+), 11 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ed2ad90c76..af67ea8605 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -759,6 +759,47 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("interleave_dim", [-3, -2])
+def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
+    """Packed q/k/v passed as plain views cannot be recognized while tracing (the
+    detection reads data pointers), so DotProductAttention runs that detection
+    eagerly instead -- correct results, at the cost of a graph break. Declaring
+    the packing via qkv_layer/kv_layer keeps it on the compiled path, which
+    test_dpa_torch_compile covers."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
+    h, d = cfg["num_heads"], cfg["head_dim"]
+    shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
+    qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+    q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
+
+    ref = module(q, k, v)
+    ref.sum().backward()
+    ref_grad = qkv.grad.clone()
+    qkv.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    with pytest.warns(UserWarning, match="Falling back to eager execution"):
+        out = torch.compile(module)(q, k, v)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(out, ref)
+    torch.testing.assert_close(qkv.grad, ref_grad)
+
+    # The same graph break is an error when the user asked for a full graph.
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    with pytest.raises(Exception, match="torch.compiler.disable"):
+        torch.compile(module, fullgraph=True)(q, k, v)
+
+
 def test_dpa_torch_compile_fused_backend(monkeypatch):
     """The FusedAttention backend is an eager island, so it graph-breaks rather
     than compiles; DotProductAttention must still be correct around it."""
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 41aef33280..d00d2c3849 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -4,6 +4,7 @@
 
 """Attention."""
 from contextlib import nullcontext
+from functools import wraps
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -42,6 +43,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
+from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -195,6 +197,50 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
+def _needs_eager_dpa(
+    _self,
+    query_layer: Optional[torch.Tensor] = None,
+    key_layer: Optional[torch.Tensor] = None,
+    value_layer: Optional[torch.Tensor] = None,
+    *_args,
+    qkv_layer: Optional[torch.Tensor] = None,
+    kv_layer: Optional[torch.Tensor] = None,
+    **_kwargs,
+) -> bool:
+    """Whether this DotProductAttention call has to run outside the graph, i.e.
+    whether it passes packed q/k/v without declaring them."""
+    if qkv_layer is not None or kv_layer is not None:
+        return False
+    return dpa_utils.qkv_layout_needs_detection(query_layer, key_layer, value_layer)
+
+
+def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
+    """Decorator running the wrapped method eagerly when `needs_eager` says the
+    call is unsupported on the compiled path."""
+
+    def decorator(fn):
+        # The warning belongs inside the dynamo-disabled function: warnings.warn
+        # graph-breaks on its own, masking the break that matters.
+        @no_torch_dynamo()
+        def eager_fn(*args, **kwargs):
+            warnings.warn(
+                f"Falling back to eager execution under torch.compile: {reason} is"
+                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                stacklevel=3,
+            )
+            return fn(*args, **kwargs)
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if torch.compiler.is_compiling() and needs_eager(*args, **kwargs):
+                return eager_fn(*args, **kwargs)
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def _unpack_packed_qkv(
     qkv_layer: Optional[torch.Tensor],
     kv_layer: Optional[torch.Tensor],
@@ -1097,6 +1143,9 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
+    @_eager_under_compile_if(
+        _needs_eager_dpa, "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+    )
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 4a8b7beb14..6ea15942d1 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -59,7 +59,7 @@
 )
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
-from transformer_engine.pytorch.jit import jit_fuser
+from transformer_engine.pytorch.jit import jit_fuser, no_torch_dynamo
 
 # NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0
 _NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0"))
@@ -2386,6 +2386,20 @@ def get_qkv_format(
     return qkv_format, q_format, kv_format
 
 
+def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool:
+    """Whether the layout of these q/k/v can only be told by inspecting memory.
+
+    True for tensors that may be slices of a packed buffer, i.e. strided ones
+    that do not own their whole storage.
+    """
+    return any(
+        x is not None
+        and not x.is_contiguous()
+        and x.untyped_storage().size() != x.numel() * x.element_size()
+        for x in qkv
+    )
+
+
 def get_qkv_layout(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -2560,16 +2574,13 @@ def run_iteratively(q, k, v):
         return qkv_layout
 
     if torch.compiler.is_compiling():
-        # The detection in run_iteratively is untraceable by dynamo -- it reads
-        # UntypedStorage.data_ptr and Tensor.storage_offset -- and pointers are
-        # meaningless on the fake tensors used while tracing anyway. Under
-        # torch.compile, packed q/k/v must therefore be declared explicitly, via
-        # DotProductAttention's qkv_layer/kv_layer arguments (which bypass this
-        # function entirely); q/k/v passed separately are treated as three
-        # independent tensors. Views into a packed buffer are non-contiguous, so
-        # making the inputs contiguous keeps the returned layout truthful for
-        # every backend.
-        q, k, v = [x if x.is_contiguous() else x.contiguous() for x in [q, k, v]]
+        # run_iteratively reads data pointers and storage offsets, which dynamo
+        # cannot trace; unpacked q/k/v need no detection anyway.
+        assert not qkv_layout_needs_detection(q, k, v), (
+            "q/k/v may be views into a packed buffer, whose layout cannot be detected under"
+            " torch.compile. Pass the packed buffer explicitly via DotProductAttention's"
+            " qkv_layer/kv_layer (with qkv_interleave_dim)."
+        )
         if is_same_q_kv_format:
             qkv_layout = "_".join([qkv_format] * 3)
         else:

From 1cd6704912eff4d3fc2fe4841f49aa85bcfa431a Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 16:23:11 +0200
Subject: [PATCH 04/39] Log backend selection in eager only; shorten comments

Log arguments are evaluated regardless of the logger, so the no-op logger
from #3189 is not enough here -- skip the whole block while tracing.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/dot_product_attention.py       | 9 ++++-----
 .../pytorch/attention/dot_product_attention/utils.py     | 5 ++---
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d00d2c3849..fa261f0c61 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1875,10 +1875,9 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    # Backend selection is only logged in eager mode: dynamo
-                    # graph-breaks on logging.Logger methods, and even the
-                    # arguments below are untraceable (int() of the sub-backend
-                    # enum). Same restriction as in get_attention_backend.
+                    # Logged in eager only: dynamo graph-breaks on logging.Logger
+                    # methods, and log arguments are evaluated regardless of the
+                    # logger, so a no-op logger would not be enough.
                     if not torch.compiler.is_compiling():
                         if use_flash_attention:
                             self.logger.info(
@@ -1888,7 +1887,7 @@ def forward(
                         elif use_fused_attention:
                             self.logger.info(
                                 "Running with FusedAttention backend (sub-backend %s)",
-                                fused_attention_backend,
+                                int(fused_attention_backend),
                             )
                         elif use_unfused_attention:
                             self.logger.info("Running with UnfusedDotProductAttention backend")
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 6ea15942d1..84868f7740 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2044,9 +2044,8 @@ def _get_cu_seqlens(batch_size, max_seqlen, device):
         )
 
     if is_in_onnx_export_mode() or torch.compiler.is_compiling():
-        # torch.is_inference_mode_enabled() (the cache key below) is a fundamental
-        # graph break under torch.compile; building the tensor is a single arange
-        # that the compiler captures into the graph anyway.
+        # The cache only saves an allocation, and its key -- is_inference_mode_enabled()
+        # -- is a fundamental graph break.
         return _get_cu_seqlens(batch_size, max_seqlen, device)
 
     is_inference = torch.is_inference_mode_enabled()

From 42a9e6cf209a3a10134a92bf2aec28795513e0af Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 09:04:37 +0200
Subject: [PATCH 05/39] Fix lint: unused import and keyword-arg-before-vararg

- utils.py no longer needs no_torch_dynamo: the eager island moved to
  dot_product_attention.py.
- _needs_eager_dpa reads q/k/v out of *args/**kwargs instead of naming them
  before *args, which pylint flags (W1113).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 25 +++++++++----------
 .../attention/dot_product_attention/utils.py  |  2 +-
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index fa261f0c61..d842a8e55b 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,21 +197,20 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-def _needs_eager_dpa(
-    _self,
-    query_layer: Optional[torch.Tensor] = None,
-    key_layer: Optional[torch.Tensor] = None,
-    value_layer: Optional[torch.Tensor] = None,
-    *_args,
-    qkv_layer: Optional[torch.Tensor] = None,
-    kv_layer: Optional[torch.Tensor] = None,
-    **_kwargs,
-) -> bool:
+def _needs_eager_dpa(*args, **kwargs) -> bool:
     """Whether this DotProductAttention call has to run outside the graph, i.e.
-    whether it passes packed q/k/v without declaring them."""
-    if qkv_layer is not None or kv_layer is not None:
+    whether it passes packed q/k/v without declaring them.
+
+    Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
+    are read from either side: `args[0]` is `self`, followed by the positional
+    query/key/value.
+    """
+    if kwargs.get("qkv_layer") is not None or kwargs.get("kv_layer") is not None:
         return False
-    return dpa_utils.qkv_layout_needs_detection(query_layer, key_layer, value_layer)
+    qkv = list(args[1:4]) + [
+        kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+    ]
+    return dpa_utils.qkv_layout_needs_detection(*qkv)
 
 
 def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 84868f7740..95b5a7f15f 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -59,7 +59,7 @@
 )
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
-from transformer_engine.pytorch.jit import jit_fuser, no_torch_dynamo
+from transformer_engine.pytorch.jit import jit_fuser
 
 # NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0
 _NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0"))

From af11f8f13b0d5e3e99aef2f6067634faf94da8cb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 09:59:30 +0200
Subject: [PATCH 06/39] Compare against eager under CUDA graphs, and tighten
 the tolerances

The CUDA-graphs test only checked that the output was finite and that
gradients existed. A replay that reuses stale buffers -- the failure mode
that test is there to catch -- satisfies both, so it now compares against
eager on fresh inputs every iteration.

The remaining comparisons used the dtype defaults (rtol=1.6e-2 for bfloat16),
which is far looser than what these paths deliver: compiled and eager agree
bit-for-bit on every config and backend. FlashAttention now has to match
exactly, as it runs the same kernel either way; the unfused backend is
allowed a single bfloat16 rounding, since inductor may reassociate the
softmax sums.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 45 +++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 9 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index af67ea8605..050f0f25a5 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -693,6 +693,19 @@ def _assert_dpa_backend(backend: str) -> None:
     )
 
 
+def _dpa_tolerances(backend: str) -> dict:
+    """How closely the compiled backend has to match eager.
+
+    FlashAttention calls the same kernel either way, so it has to match exactly.
+    The unfused backend is compiled by inductor, which may fuse the softmax
+    chain differently and reassociate its sums; allow a single bfloat16 rounding
+    (eps = 2**-8) for that, which is still 4x tighter than the dtype default.
+    """
+    if backend == "flash":
+        return {"rtol": 0.0, "atol": 0.0}
+    return {"rtol": 2**-8, "atol": 1e-5}
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
 def test_dpa_torch_compile(monkeypatch, backend, config):
@@ -729,10 +742,11 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    torch.testing.assert_close(out, ref)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
     for tensor, ref_grad in zip(grads, ref_grads):
         assert tensor.grad is not None
-        torch.testing.assert_close(tensor.grad, ref_grad)
+        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -748,14 +762,26 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
+    tols = _dpa_tolerances(backend)
     for _ in range(3):
+        # Fresh inputs every iteration: a replay that reuses stale buffers would
+        # still produce finite values and non-None gradients, so only comparing
+        # against eager catches it.
         args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
-        out = compiled(*args, **kwargs)
+        # CUDA graphs hand back tensors owned by their memory pool, which the
+        # next replay overwrites.
+        out = compiled(*args, **kwargs).clone()
         out.sum().backward()
         torch.cuda.synchronize()
-        assert torch.isfinite(out).all()
+        compiled_grads = [t.grad.clone() for t in grads]
         for tensor in grads:
-            assert tensor.grad is not None
+            tensor.grad = None
+
+        ref = module(*args, **kwargs)
+        ref.sum().backward()
+        torch.testing.assert_close(out, ref, **tols)
+        for compiled_grad, tensor in zip(compiled_grads, grads):
+            torch.testing.assert_close(compiled_grad, tensor.grad, **tols)
     _assert_dpa_backend(backend)
 
 
@@ -790,8 +816,9 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     out.sum().backward()
     torch.cuda.synchronize()
 
-    torch.testing.assert_close(out, ref)
-    torch.testing.assert_close(qkv.grad, ref_grad)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
+    torch.testing.assert_close(qkv.grad, ref_grad, **tols)
 
     # The same graph break is an error when the user asked for a full graph.
     torch._dynamo.reset()
@@ -829,9 +856,9 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     torch.cuda.synchronize()
 
     _assert_dpa_backend("fused")
-    torch.testing.assert_close(out, ref)
+    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
     for tensor, ref_grad in zip(grads, ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad)
+        torch.testing.assert_close(tensor.grad, ref_grad, rtol=0.0, atol=0.0)
 
 
 # ---------------------------------------------------------------------------

From 0b64c9695abc27aab434326c258e0edc5c4dac02 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:16:55 +0200
Subject: [PATCH 07/39] Support one cu_seqlens tensor for both q and kv under
 torch.compile

Self attention on thd naturally passes the same tensor as cu_seqlens_q and
cu_seqlens_kv, and flash-attn's varlen entry point forwards both to the same
autograd.Function -- which dynamo refuses to trace ("duplicate tensor input"),
failing the compilation under fullgraph=True.

Hand the varlen call a copy of one of them while tracing. The clone is
b + 1 int32 elements, only on that path, and only under torch.compile.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 43 +++++++++++++++++++
 .../dot_product_attention/backends.py         | 15 ++++---
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 050f0f25a5..71eb58a10d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -785,6 +785,49 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
+    """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
+    and kv, which flash-attn hands to two inputs of the same autograd.Function --
+    something dynamo cannot trace. The other tests build two tensors, so this
+    covers the shape a caller is most likely to write."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_thd_padding_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
+    h, d = cfg["num_heads"], cfg["head_dim"]
+    q, k, v = [
+        torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
+    ]
+    cu_seqlens = torch.arange(0, (b + 1) * s, step=s, dtype=torch.int32, device="cuda")
+    kwargs = dict(
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_kv=cu_seqlens,  # the same tensor, not a copy
+        max_seqlen_q=s,
+        max_seqlen_kv=s,
+    )
+
+    ref = module(q, k, v, **kwargs)
+    ref.sum().backward()
+    ref_grads = [t.grad.clone() for t in (q, k, v)]
+    for t in (q, k, v):
+        t.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    out = torch.compile(module, fullgraph=True)(q, k, v, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend(backend)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
+    for tensor, ref_grad in zip((q, k, v), ref_grads):
+        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("interleave_dim", [-3, -2])
 def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 6a60160f77..f0382be30c 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -1122,12 +1122,15 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        fa_optional_forward_args_thd.append(
-                            cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
-                        )
-                        fa_optional_forward_args_thd.append(
-                            cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
-                        )
+                        cu_q = cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
+                        cu_kv = cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
+                        if cu_q is cu_kv and torch.compiler.is_compiling():
+                            # Self attention passes one tensor for both, and dynamo
+                            # cannot trace autograd.Function.apply with the same
+                            # tensor on two of its inputs.
+                            cu_kv = cu_kv.clone()
+                        fa_optional_forward_args_thd.append(cu_q)
+                        fa_optional_forward_args_thd.append(cu_kv)
                         fa_optional_forward_args_thd.append(max_seqlen_q)
                         fa_optional_forward_args_thd.append(max_seqlen_kv)
                 if use_flash_attn_4:

From 7f480c6d07a9d04bca718f581e6a5e6b3e9bb113 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:21:49 +0200
Subject: [PATCH 08/39] Apply the shared-cu_seqlens workaround to the
 FlashAttention v4 path too

Pull the deduplication into a helper and use it wherever both cumulative
sequence lengths are handed to flash-attn: the varlen path shared by v2 and
v3 training, and v4's varlen kwargs. The v3 KV-cache path derives the kv
lengths by subtraction, so its two arguments are distinct by construction.

v4 is not installed in the environment this was tested in, so that call site
is covered by construction rather than by a test; the helper is a no-op
unless the two arguments are the same object.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/backends.py         | 28 +++++++++++++------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index f0382be30c..2c47ff8a4b 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -793,6 +793,18 @@ def backward(
         return dq, dk, dv
 
 
+def _distinct_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
+    """Make q's and kv's cumulative sequence lengths distinct objects.
+
+    Self attention passes one tensor for both, and flash-attn forwards them to
+    two inputs of the same autograd.Function -- which dynamo cannot trace. In
+    eager the duplicate is harmless, so nothing is copied there.
+    """
+    if cu_q is cu_kv and torch.compiler.is_compiling():
+        return cu_q, cu_kv.clone()
+    return cu_q, cu_kv
+
+
 class FlashAttention(torch.nn.Module):
     """Dot product attention, using HazyResearch flash-attn package:
     https://github.com/Dao-AILab/flash-attention
@@ -1122,13 +1134,10 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        cu_q = cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
-                        cu_kv = cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
-                        if cu_q is cu_kv and torch.compiler.is_compiling():
-                            # Self attention passes one tensor for both, and dynamo
-                            # cannot trace autograd.Function.apply with the same
-                            # tensor on two of its inputs.
-                            cu_kv = cu_kv.clone()
+                        cu_q, cu_kv = _distinct_cu_seqlens(
+                            cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q,
+                            cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv,
+                        )
                         fa_optional_forward_args_thd.append(cu_q)
                         fa_optional_forward_args_thd.append(cu_kv)
                         fa_optional_forward_args_thd.append(max_seqlen_q)
@@ -1141,8 +1150,9 @@ def forward(
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
                     if func is flash_attn_varlen_func_v4:
-                        fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q
-                        fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv
+                        cu_q, cu_kv = _distinct_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
+                        fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q
+                        fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv
                         fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
                         fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv
                     output = func(

From adc96c48a9b455678acbda96e367e8a0c63d16b2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:58:31 +0200
Subject: [PATCH 09/39] Run FP8 attention eagerly under torch.compile

FP8 attention brings quantizers, fp8_meta mutation and Float8Tensor across
the graph boundary, and none of it is exercised by this PR's tests. Rather
than trace it untested, route it to the eager island that packed q/k/v
already use.

The predicate is deliberately narrow: it asks the recipe for fp8_dpa/fp8_mha
rather than whether an autocast is active, so FP8 GEMMs with attention in
high precision -- the common training setup -- stay on the compiled path.

The decorator now takes a predicate that returns the reason, so the warning
names which of the two paths was taken.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 34 ++++++++++++++
 .../dot_product_attention.py                  | 44 +++++++++++--------
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 71eb58a10d..61e78e58b0 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 import abc
+import warnings
 
 import pytest
 import torch
@@ -785,6 +786,39 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"])
+def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
+    """FP8 attention is not supported on the compiled path and falls back to
+    eager. FP8 elsewhere in the model with attention in high precision -- the
+    common training setup -- must stay compiled."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, "unfused")
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, _ = _make_dpa_inputs(cfg, dtype)
+    fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
+
+    def fn(*args, **kwargs):
+        with te.autocast(enabled=True, recipe=fp8_recipe):
+            return module(*args, **kwargs)
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, "unfused")
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        try:
+            torch.compile(fn)(*args, **kwargs)
+        except Exception:  # pylint: disable=broad-except
+            # FP8 attention is not available on every device; what matters here
+            # is which path was taken, which the warning below reports.
+            pass
+        fell_back = any("Falling back to eager" in str(w.message) for w in caught)
+
+    assert fell_back == fp8_attention
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d842a8e55b..f82aa91c13 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,31 +197,39 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-def _needs_eager_dpa(*args, **kwargs) -> bool:
-    """Whether this DotProductAttention call has to run outside the graph, i.e.
-    whether it passes packed q/k/v without declaring them.
+def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
+    """Why this DotProductAttention call has to run outside the graph, or None.
 
     Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
     are read from either side: `args[0]` is `self`, followed by the positional
     query/key/value.
     """
-    if kwargs.get("qkv_layer") is not None or kwargs.get("kv_layer") is not None:
-        return False
-    qkv = list(args[1:4]) + [
-        kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
-    ]
-    return dpa_utils.qkv_layout_needs_detection(*qkv)
+    # FP8 GEMMs with the attention itself in high precision -- the common
+    # training setup -- stay on the compiled path; only FP8 attention bails out.
+    qstate = FP8GlobalStateManager.quantization_state
+    fp8_recipe = qstate.fp8_recipe
+    if qstate.fp8_enabled and fp8_recipe is not None:
+        if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
+            return "FP8 attention"
+
+    if kwargs.get("qkv_layer") is None and kwargs.get("kv_layer") is None:
+        qkv = list(args[1:4]) + [
+            kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+        ]
+        if dpa_utils.qkv_layout_needs_detection(*qkv):
+            return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+    return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
-    """Decorator running the wrapped method eagerly when `needs_eager` says the
-    call is unsupported on the compiled path."""
+def _eager_under_compile_if(needs_eager: Callable[..., Optional[str]]):
+    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
+    a reason why the call is unsupported on the compiled path."""
 
     def decorator(fn):
         # The warning belongs inside the dynamo-disabled function: warnings.warn
         # graph-breaks on its own, masking the break that matters.
         @no_torch_dynamo()
-        def eager_fn(*args, **kwargs):
+        def eager_fn(reason, *args, **kwargs):
             warnings.warn(
                 f"Falling back to eager execution under torch.compile: {reason} is"
                 " unsupported on the compiled path (graph-breaks under fullgraph=True).",
@@ -231,8 +239,10 @@ def eager_fn(*args, **kwargs):
 
         @wraps(fn)
         def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling() and needs_eager(*args, **kwargs):
-                return eager_fn(*args, **kwargs)
+            if torch.compiler.is_compiling():
+                reason = needs_eager(*args, **kwargs)
+                if reason is not None:
+                    return eager_fn(reason, *args, **kwargs)
             return fn(*args, **kwargs)
 
         return wrapper
@@ -1142,9 +1152,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @_eager_under_compile_if(
-        _needs_eager_dpa, "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
-    )
+    @_eager_under_compile_if(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,

From 737904f09cc1957db238e73a68c7e1edeb231158 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:23:04 +0200
Subject: [PATCH 10/39] Drive the compile tests off ModelConfig instead of a
 hand-picked intersection

The configurations were chosen so that both backends could run all of them,
which meant nothing backend-specific was ever exercised: no biases, no
arbitrary masks, no MLA head dims, no sink softmax. Describe them with the
ModelConfig the eager attention tests already use, and let
get_available_attention_backends() decide which backend runs which, so a
configuration only one backend supports is covered rather than avoided.

Five such configurations are added. One of them immediately showed that the
tolerance for the unfused backend was too tight: an off-by-one softmax
differs from eager by two bfloat16 roundings on a single element out of
65536, because inductor reassociates the softmax sums. Only FlashAttention,
which runs the same kernel either way, keeps the exact comparison.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 208 +++++++++++++++++-----------
 1 file changed, 125 insertions(+), 83 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 61e78e58b0..77c6c31a3b 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -4,6 +4,7 @@
 
 import abc
 import warnings
+from typing import Optional
 
 import pytest
 import torch
@@ -37,7 +38,7 @@
     MXFP8Quantizer,
     NVFP4Quantizer,
 )
-from utils import recipe_id
+from utils import ModelConfig, get_available_attention_backends, recipe_id
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
     UnfusedDotProductAttention,
 )
@@ -556,93 +557,117 @@ def fn(q, k, v, extra):
 # ---------------------------------------------------------------------------
 
 
-# Each config is a model configuration that both the FlashAttention and the
-# UnfusedDotProductAttention backend support, so that every one of them runs
-# against both backends.
+# Model configurations, described with the same ModelConfig the eager attention
+# tests use. Which backends can run each of them is not hardcoded here --
+# get_available_attention_backends() answers that, so configurations only one
+# backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
+# rather than avoided.
 _DPA_COMPILE_CONFIGS = {
-    "self_bshd_causal": dict(qkv_format="bshd", attn_mask_type="causal"),
-    "self_sbhd_no_mask": dict(qkv_format="sbhd", attn_mask_type="no_mask"),
-    "self_bshd_swa": dict(qkv_format="bshd", attn_mask_type="causal", window_size=(16, 0)),
-    "gqa_bshd_causal": dict(
-        qkv_format="bshd", attn_mask_type="causal", num_heads=8, num_gqa_groups=2
+    # name: (ModelConfig, qkv_format, packed layout or None)
+    "self_bshd_causal": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", None),
+    "self_sbhd_no_mask": (ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd", None),
+    "self_bshd_swa": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0)),
+        "bshd",
+        None,
     ),
-    "cross_bshd_no_mask": dict(
-        qkv_format="bshd", attn_mask_type="no_mask", attention_type="cross", max_seqlen_kv=256
+    "gqa_bshd_causal": (
+        ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal"),
+        "bshd",
+        None,
+    ),
+    "cross_bshd_no_mask": (
+        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask"),
+        "bshd",
+        None,
+    ),
+    "self_thd_padding_causal": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
+        "thd",
+        None,
+    ),
+    "packed_qkv_bs3hd": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", "bs3hd"),
+    # Configurations below are supported by one backend only, or take a code
+    # path of their own inside DotProductAttention.
+    "alibi_bshd_causal": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi"),
+        "bshd",
+        None,
+    ),
+    "post_scale_bias_bshd": (
+        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss"),
+        "bshd",
+        None,
+    ),
+    "arbitrary_mask_bshd": (ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary"), "bshd", None),
+    "mla_bshd_causal": (
+        ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal"),
+        "bshd",
+        None,
+    ),
+    "sink_softmax_bshd": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one"),
+        "bshd",
+        None,
     ),
-    "self_thd_padding_causal": dict(qkv_format="thd", attn_mask_type="padding_causal"),
-    "packed_qkv_bs3hd": dict(qkv_format="bshd", attn_mask_type="causal", packed="bs3hd"),
 }
 
 
-def _dpa_config(**overrides):
-    cfg = dict(
-        batch_size=2,
-        num_heads=4,
-        num_gqa_groups=None,
-        head_dim=64,
-        max_seqlen_q=128,
-        max_seqlen_kv=128,
-        qkv_format="bshd",
-        attn_mask_type="causal",
-        attention_type="self",
-        window_size=None,
-        packed=None,
-    )
-    cfg.update(overrides)
-    if cfg["num_gqa_groups"] is None:
-        cfg["num_gqa_groups"] = cfg["num_heads"]
-    return cfg
+def _qkv_layout(qkv_format: str, packed: Optional[str]) -> str:
+    return packed if packed is not None else "_".join([qkv_format] * 3)
 
 
-def _make_dpa(cfg, dtype: torch.dtype) -> te.DotProductAttention:
+def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.DotProductAttention:
     return te.DotProductAttention(
-        num_attention_heads=cfg["num_heads"],
-        kv_channels=cfg["head_dim"],
-        num_gqa_groups=cfg["num_gqa_groups"],
-        attention_dropout=0.0,
-        qkv_format=cfg["qkv_format"],
-        attn_mask_type=cfg["attn_mask_type"],
-        window_size=cfg["window_size"],
-        attention_type=cfg["attention_type"],
+        num_attention_heads=config.num_heads,
+        kv_channels=config.kv_channels,
+        num_gqa_groups=config.num_gqa_groups,
+        attention_dropout=config.dropout_p,
+        qkv_format=qkv_format,
+        attn_mask_type=config.attn_mask_type,
+        window_size=config.window_size,
+        attention_type=config.attn_type,
+        softmax_type=config.softmax_type,
     ).to(dtype=dtype, device="cuda")
 
 
-def _make_dpa_inputs(cfg, dtype: torch.dtype):
+def _make_dpa_inputs(config: ModelConfig, qkv_format: str, packed: Optional[str], dtype):
     """Build the (args, kwargs) that `DotProductAttention.forward` is called
     with, plus the list of tensors whose gradients the test compares."""
-    b = cfg["batch_size"]
-    s_q, s_kv = cfg["max_seqlen_q"], cfg["max_seqlen_kv"]
-    h, g, d = cfg["num_heads"], cfg["num_gqa_groups"], cfg["head_dim"]
-    fmt = cfg["qkv_format"]
+    b = config.batch_size
+    s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv
+    h, g = config.num_heads, config.num_gqa_groups
+    d_qk, d_v = config.head_dim_qk, config.head_dim_v
 
-    def _shape(s, heads):
+    def _shape(s, heads, head_dim):
         return {
-            "bshd": (b, s, heads, d),
-            "sbhd": (s, b, heads, d),
-            "thd": (b * s, heads, d),
-        }[fmt]
+            "bshd": (b, s, heads, head_dim),
+            "sbhd": (s, b, heads, head_dim),
+            "thd": (b * s, heads, head_dim),
+        }[qkv_format]
 
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
     kwargs = {}
-    if cfg["packed"] is not None:
+    if packed is not None:
         # Declarative packed QKV: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
-        # layout that torch.compile supports (pointer-based detection of packed
-        # q/k/v is untraceable).
-        assert cfg["packed"] == "bs3hd" and fmt == "bshd" and h == g
-        qkv = _randn((b, s_q, 3, h, d))
+        # layout that torch.compile supports.
+        assert packed == "bs3hd" and qkv_format == "bshd" and h == g and d_qk == d_v
+        qkv = _randn((b, s_q, 3, h, d_qk))
         kwargs["qkv_layer"] = qkv
         kwargs["qkv_interleave_dim"] = -3
         grad_tensors = [qkv]
         args = ()
     else:
-        q, k, v = _randn(_shape(s_q, h)), _randn(_shape(s_kv, g)), _randn(_shape(s_kv, g))
+        q = _randn(_shape(s_q, h, d_qk))
+        k = _randn(_shape(s_kv, g, d_qk))
+        v = _randn(_shape(s_kv, g, d_v))
         grad_tensors = [q, k, v]
         args = (q, k, v)
 
-    if fmt == "thd":
+    if qkv_format == "thd":
         # All sequences have the maximum length, i.e. no padding. max_seqlen_*
         # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
         # (a device sync, and an unbacked SymInt under torch.compile).
@@ -655,9 +680,25 @@ def _randn(shape):
         kwargs["max_seqlen_q"] = s_q
         kwargs["max_seqlen_kv"] = s_kv
 
+    if config.attn_mask_type == "arbitrary":
+        kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda")
+    if config.attn_bias_type != "no_bias":
+        kwargs["core_attention_bias_type"] = config.attn_bias_type
+    if config.attn_bias_type == "post_scale_bias":
+        kwargs["core_attention_bias"] = torch.randn(1, h, s_q, s_kv, dtype=dtype, device="cuda")
+
     return args, kwargs, grad_tensors
 
 
+def _skip_unsupported(config: ModelConfig, qkv_layout: str, backend: str, dtype) -> None:
+    """Skip configurations the backend under test cannot run at all."""
+    available, _, _ = get_available_attention_backends(config, dtype, qkv_layout)
+    flash_supported, _, unfused_supported = available
+    supported = {"flash": flash_supported, "unfused": unfused_supported}[backend]
+    if not supported:
+        pytest.skip(f"the {backend} backend does not support this configuration")
+
+
 def _force_dpa_backend(monkeypatch, backend: str) -> None:
     """Restrict DotProductAttention to a single backend."""
     from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
@@ -698,13 +739,13 @@ def _dpa_tolerances(backend: str) -> dict:
     """How closely the compiled backend has to match eager.
 
     FlashAttention calls the same kernel either way, so it has to match exactly.
-    The unfused backend is compiled by inductor, which may fuse the softmax
-    chain differently and reassociate its sums; allow a single bfloat16 rounding
-    (eps = 2**-8) for that, which is still 4x tighter than the dtype default.
+    The unfused backend is compiled by inductor, which fuses the softmax chain
+    differently and reassociates its sums: with an off-by-one softmax that costs
+    up to two bfloat16 roundings on individual elements, so it keeps the dtype
+    defaults. A systematic difference -- a wrong mask or scale -- lands far
+    outside them anyway.
     """
-    if backend == "flash":
-        return {"rtol": 0.0, "atol": 0.0}
-    return {"rtol": 2**-8, "atol": 1e-5}
+    return {"rtol": 0.0, "atol": 0.0} if backend == "flash" else {}
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -719,11 +760,12 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     backend itself. The FusedAttention backend is not covered -- it is an eager
     island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS[config])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS[config]
+    _skip_unsupported(model_config, _qkv_layout(qkv_format, packed), backend, dtype)
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
 
     ref = module(*args, **kwargs)
     ref.sum().backward()
@@ -755,10 +797,10 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
 
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
@@ -768,7 +810,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
-        args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+        args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
         # CUDA graphs hand back tensors owned by their memory pool, which the
         # next replay overwrites.
         out = compiled(*args, **kwargs).clone()
@@ -793,11 +835,11 @@ def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
     eager. FP8 elsewhere in the model with attention in high precision -- the
     common training setup -- must stay compiled."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "unfused")
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, _ = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, _ = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
     fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
 
     def fn(*args, **kwargs):
@@ -826,12 +868,12 @@ def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     something dynamo cannot trace. The other tests build two tensors, so this
     covers the shape a caller is most likely to write."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_thd_padding_causal"])
+    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_thd_padding_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
-    h, d = cfg["num_heads"], cfg["head_dim"]
+    module = _make_dpa(model_config, qkv_format, dtype)
+    b, s = model_config.batch_size, model_config.max_seqlen_q
+    h, d = model_config.num_heads, model_config.head_dim_qk
     q, k, v = [
         torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
     ]
@@ -871,12 +913,12 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     the packing via qkv_layer/kv_layer keeps it on the compiled path, which
     test_dpa_torch_compile covers."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
-    h, d = cfg["num_heads"], cfg["head_dim"]
+    module = _make_dpa(model_config, qkv_format, dtype)
+    b, s = model_config.batch_size, model_config.max_seqlen_q
+    h, d = model_config.num_heads, model_config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
@@ -908,11 +950,11 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     """The FusedAttention backend is an eager island, so it graph-breaks rather
     than compiles; DotProductAttention must still be correct around it."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "fused")
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
 
     ref = module(*args, **kwargs)
     from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (

From 9026a41c3303a1e602b6be6e99d2cefe5e4da64d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:31:36 +0200
Subject: [PATCH 11/39] Scale the unfused comparison to the tensor, and take
 tolerances from utils

An off-by-one softmax exposed that comparing the unfused backend elementwise
does not hold up: inductor's reassociation of the softmax sums produces an
error proportional to those sums, which for gradient elements near zero is
far outside any relative tolerance -- while being one bfloat16 rounding at
the tensor's own scale.

Derive the absolute tolerance from the reference tensor's magnitude, and take
the dtype tolerances from tests/pytorch/utils.py rather than restating them.
FlashAttention and FusedAttention keep the exact comparison.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 53 ++++++++++++++++-------------
 1 file changed, 29 insertions(+), 24 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 77c6c31a3b..b88222be60 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -38,7 +38,7 @@
     MXFP8Quantizer,
     NVFP4Quantizer,
 )
-from utils import ModelConfig, get_available_attention_backends, recipe_id
+from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
     UnfusedDotProductAttention,
 )
@@ -735,17 +735,26 @@ def _assert_dpa_backend(backend: str) -> None:
     )
 
 
-def _dpa_tolerances(backend: str) -> dict:
-    """How closely the compiled backend has to match eager.
+def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None:
+    """Assert a compiled result matches the eager one.
 
-    FlashAttention calls the same kernel either way, so it has to match exactly.
-    The unfused backend is compiled by inductor, which fuses the softmax chain
-    differently and reassociates its sums: with an off-by-one softmax that costs
-    up to two bfloat16 roundings on individual elements, so it keeps the dtype
-    defaults. A systematic difference -- a wrong mask or scale -- lands far
-    outside them anyway.
+    FlashAttention and FusedAttention run the same kernel either way, so they
+    have to match exactly. Inductor reassociates the unfused backend's softmax
+    sums, and that error scales with the magnitude of those sums rather than
+    with each output element -- so the absolute tolerance comes from the
+    tensor's own scale, an elementwise relative tolerance being meaningless for
+    elements near zero.
     """
-    return {"rtol": 0.0, "atol": 0.0} if backend == "flash" else {}
+    if backend != "unfused":
+        torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+        return
+    tols = dtype_tols(dtype)
+    torch.testing.assert_close(
+        actual,
+        expected,
+        rtol=tols["rtol"],
+        atol=max(tols["atol"], 1e-3 * expected.abs().max().item()),
+    )
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -785,11 +794,10 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
     for tensor, ref_grad in zip(grads, ref_grads):
         assert tensor.grad is not None
-        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -805,7 +813,6 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
-    tols = _dpa_tolerances(backend)
     for _ in range(3):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
@@ -822,9 +829,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
 
         ref = module(*args, **kwargs)
         ref.sum().backward()
-        torch.testing.assert_close(out, ref, **tols)
+        _assert_matches_eager(out, ref, backend, dtype)
         for compiled_grad, tensor in zip(compiled_grads, grads):
-            torch.testing.assert_close(compiled_grad, tensor.grad, **tols)
+            _assert_matches_eager(compiled_grad, tensor.grad, backend, dtype)
     _assert_dpa_backend(backend)
 
 
@@ -898,10 +905,9 @@ def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
     for tensor, ref_grad in zip((q, k, v), ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -935,9 +941,8 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     out.sum().backward()
     torch.cuda.synchronize()
 
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
-    torch.testing.assert_close(qkv.grad, ref_grad, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
+    _assert_matches_eager(qkv.grad, ref_grad, backend, dtype)
 
     # The same graph break is an error when the user asked for a full graph.
     torch._dynamo.reset()
@@ -975,9 +980,9 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     torch.cuda.synchronize()
 
     _assert_dpa_backend("fused")
-    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
+    _assert_matches_eager(out, ref, "fused", dtype)
     for tensor, ref_grad in zip(grads, ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad, rtol=0.0, atol=0.0)
+        _assert_matches_eager(tensor.grad, ref_grad, "fused", dtype)
 
 
 # ---------------------------------------------------------------------------

From bd2c5066fd871f4b3cad37166c6a441250c0c489 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:53:47 +0200
Subject: [PATCH 12/39] Merge the per-backend compile tests, and give padding
 masks real padding

The configurations built cu_seqlens with every sequence at the maximum length,
so a padding mask was all-False: the padding branches ran, but on data that
could not tell a difference. They now use varying sequence lengths, which
immediately showed that FlashAttention on bshd/sbhd with a padding mask does
not trace: it packs the tensors first, and get_indices built the index list
with a Python comprehension over sequence lengths held in a GPU tensor -- a
host synchronization per sequence in eager, untraceable while compiling.
It is now built on the device, verified exhaustively against the previous
implementation (10304 cases, 0 mismatches) and free of synchronization.

The separate tests for one shared cu_seqlens tensor and for the fused backend
folded into the main one: the first is a configuration, and the second is the
backend axis, which now runs flash, fused and unfused with the graph break
that FusedAttention's eager island implies.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 306 +++++++-----------
 .../attention/dot_product_attention/utils.py  |  24 +-
 2 files changed, 136 insertions(+), 194 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index b88222be60..397db816f4 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -562,68 +562,61 @@ def fn(q, k, v, extra):
 # get_available_attention_backends() answers that, so configurations only one
 # backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
 # rather than avoided.
+def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
+    return dict(
+        model_config=model_config,
+        qkv_format=qkv_format,
+        packed=packed,
+        share_cu_seqlens=share_cu_seqlens,
+    )
+
+
 _DPA_COMPILE_CONFIGS = {
-    # name: (ModelConfig, qkv_format, packed layout or None)
-    "self_bshd_causal": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", None),
-    "self_sbhd_no_mask": (ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd", None),
-    "self_bshd_swa": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0)),
-        "bshd",
-        None,
-    ),
-    "gqa_bshd_causal": (
-        ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal"),
-        "bshd",
-        None,
+    "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")),
+    "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"),
+    "self_bshd_swa": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0))),
+    "gqa_bshd_causal": _cfg(ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal")),
+    "cross_bshd_no_mask": _cfg(
+        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask")
     ),
-    "cross_bshd_no_mask": (
-        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask"),
-        "bshd",
-        None,
+    "self_bshd_padding_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal")),
+    "self_thd_padding_causal": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd"
     ),
-    "self_thd_padding_causal": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
-        "thd",
-        None,
+    # Self attention naturally passes one cu_seqlens tensor for both q and kv,
+    # which flash-attn hands to two inputs of the same autograd.Function.
+    "self_thd_shared_cu_seqlens": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True
     ),
-    "packed_qkv_bs3hd": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", "bs3hd"),
+    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="bs3hd"),
     # Configurations below are supported by one backend only, or take a code
     # path of their own inside DotProductAttention.
-    "alibi_bshd_causal": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi"),
-        "bshd",
-        None,
+    "alibi_bshd_causal": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi")
     ),
-    "post_scale_bias_bshd": (
-        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss"),
-        "bshd",
-        None,
+    "post_scale_bias_bshd": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss")
     ),
-    "arbitrary_mask_bshd": (ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary"), "bshd", None),
-    "mla_bshd_causal": (
-        ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal"),
-        "bshd",
-        None,
-    ),
-    "sink_softmax_bshd": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one"),
-        "bshd",
-        None,
+    "arbitrary_mask_bshd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary")),
+    "mla_bshd_causal": _cfg(ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal")),
+    "sink_softmax_bshd": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
 }
 
 
-def _qkv_layout(qkv_format: str, packed: Optional[str]) -> str:
-    return packed if packed is not None else "_".join([qkv_format] * 3)
+def _qkv_layout(spec: dict) -> str:
+    return spec["packed"] or "_".join([spec["qkv_format"]] * 3)
 
 
-def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.DotProductAttention:
+def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
+    config = spec["model_config"]
     return te.DotProductAttention(
         num_attention_heads=config.num_heads,
         kv_channels=config.kv_channels,
         num_gqa_groups=config.num_gqa_groups,
         attention_dropout=config.dropout_p,
-        qkv_format=qkv_format,
+        qkv_format=spec["qkv_format"],
         attn_mask_type=config.attn_mask_type,
         window_size=config.window_size,
         attention_type=config.attn_type,
@@ -631,25 +624,51 @@ def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.Do
     ).to(dtype=dtype, device="cuda")
 
 
-def _make_dpa_inputs(config: ModelConfig, qkv_format: str, packed: Optional[str], dtype):
+def _cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
+    cu = torch.zeros(seqlens.numel() + 1, dtype=torch.int32, device="cuda")
+    cu[1:] = torch.cumsum(seqlens, dim=0)
+    return cu
+
+
+def _make_dpa_inputs(spec: dict, dtype: torch.dtype):
     """Build the (args, kwargs) that `DotProductAttention.forward` is called
     with, plus the list of tensors whose gradients the test compares."""
+    config = spec["model_config"]
+    qkv_format, packed = spec["qkv_format"], spec["packed"]
     b = config.batch_size
     s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv
     h, g = config.num_heads, config.num_gqa_groups
     d_qk, d_v = config.head_dim_qk, config.head_dim_v
+    padded = "padding" in config.attn_mask_type
 
-    def _shape(s, heads, head_dim):
+    kwargs = {}
+    if padded:
+        # Sequences shorter than the maximum, so the padding mask is not
+        # degenerate: with all sequences full it would be all-False, and any
+        # difference in how masking is compiled would be invisible.
+        seqlens_q = torch.randint(1, s_q, [b], dtype=torch.int32, device="cuda")
+        seqlens_kv = (
+            seqlens_q
+            if config.attn_type == "self"
+            else (torch.randint(1, s_kv, [b], dtype=torch.int32, device="cuda"))
+        )
+        cu_q = _cu_seqlens(seqlens_q)
+        cu_kv = cu_q if spec["share_cu_seqlens"] else _cu_seqlens(seqlens_kv)
+        kwargs.update(cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, max_seqlen_q=s_q, max_seqlen_kv=s_kv)
+        t_q, t_kv = int(cu_q[-1]), int(cu_kv[-1])
+    else:
+        t_q, t_kv = b * s_q, b * s_kv
+
+    def _shape(s, t, heads, head_dim):
         return {
             "bshd": (b, s, heads, head_dim),
             "sbhd": (s, b, heads, head_dim),
-            "thd": (b * s, heads, head_dim),
+            "thd": (t, heads, head_dim),
         }[qkv_format]
 
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
-    kwargs = {}
     if packed is not None:
         # Declarative packed QKV: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
@@ -658,27 +677,12 @@ def _randn(shape):
         qkv = _randn((b, s_q, 3, h, d_qk))
         kwargs["qkv_layer"] = qkv
         kwargs["qkv_interleave_dim"] = -3
-        grad_tensors = [qkv]
-        args = ()
+        grad_tensors, args = [qkv], ()
     else:
-        q = _randn(_shape(s_q, h, d_qk))
-        k = _randn(_shape(s_kv, g, d_qk))
-        v = _randn(_shape(s_kv, g, d_v))
-        grad_tensors = [q, k, v]
-        args = (q, k, v)
-
-    if qkv_format == "thd":
-        # All sequences have the maximum length, i.e. no padding. max_seqlen_*
-        # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
-        # (a device sync, and an unbacked SymInt under torch.compile).
-        kwargs["cu_seqlens_q"] = torch.arange(
-            0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda"
-        )
-        kwargs["cu_seqlens_kv"] = torch.arange(
-            0, (b + 1) * s_kv, step=s_kv, dtype=torch.int32, device="cuda"
-        )
-        kwargs["max_seqlen_q"] = s_q
-        kwargs["max_seqlen_kv"] = s_kv
+        q = _randn(_shape(s_q, t_q, h, d_qk))
+        k = _randn(_shape(s_kv, t_kv, g, d_qk))
+        v = _randn(_shape(s_kv, t_kv, g, d_v))
+        grad_tensors, args = [q, k, v], (q, k, v)
 
     if config.attn_mask_type == "arbitrary":
         kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda")
@@ -690,11 +694,17 @@ def _randn(shape):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(config: ModelConfig, qkv_layout: str, backend: str, dtype) -> None:
+def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
     """Skip configurations the backend under test cannot run at all."""
-    available, _, _ = get_available_attention_backends(config, dtype, qkv_layout)
-    flash_supported, _, unfused_supported = available
-    supported = {"flash": flash_supported, "unfused": unfused_supported}[backend]
+    available, _, _ = get_available_attention_backends(
+        spec["model_config"], dtype, _qkv_layout(spec)
+    )
+    flash_supported, fused_supported, unfused_supported = available
+    supported = {
+        "flash": flash_supported,
+        "fused": fused_supported,
+        "unfused": unfused_supported,
+    }[backend]
     if not supported:
         pytest.skip(f"the {backend} backend does not support this configuration")
 
@@ -757,39 +767,22 @@ def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) ->
     )
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
-def test_dpa_torch_compile(monkeypatch, backend, config):
-    """`DotProductAttention` must trace under `torch.compile(fullgraph=True)`
-    on the FlashAttention and UnfusedDotProductAttention backends, and match
-    eager in forward and backward.
-
-    `fullgraph=True` makes the test fail on any graph break, i.e. it covers the
-    whole module: input unpacking, qkv layout, backend selection and the
-    backend itself. The FusedAttention backend is not covered -- it is an eager
-    island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS[config]
-    _skip_unsupported(model_config, _qkv_layout(qkv_format, packed), backend, dtype)
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
-
+def _compare_compiled_to_eager(
+    module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs
+) -> None:
+    """Run the module eagerly and compiled on the same inputs, and compare both
+    the output and every input gradient."""
     ref = module(*args, **kwargs)
     ref.sum().backward()
-    _assert_dpa_backend(backend)
     ref_grads = [t.grad.clone() for t in grads]
-    for t in grads:
-        t.grad = None
+    for tensor in grads:
+        tensor.grad = None
 
     torch._dynamo.reset()
     # Force backend selection to be re-run (and traced) inside the compiled
     # region instead of being served from the cache the eager call populated.
     _force_dpa_backend(monkeypatch, backend)
-    compiled = torch.compile(module, fullgraph=True)
-
-    out = compiled(*args, **kwargs)
+    out = torch.compile(module, **compile_kwargs)(*args, **kwargs)
     out.sum().backward()
     torch.cuda.synchronize()
 
@@ -800,15 +793,39 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
         _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
+@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
+@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
+def test_dpa_torch_compile(monkeypatch, backend, config):
+    """`DotProductAttention` under `torch.compile` must match eager in forward
+    and backward, for every backend that supports the configuration.
+
+    FlashAttention and UnfusedDotProductAttention are compiled with
+    `fullgraph=True`, so the test fails on any graph break -- covering the whole
+    module: input unpacking, qkv layout, backend selection and the backend
+    itself. FusedAttention is an eager island (`@no_torch_dynamo` on its
+    forward), so it is expected to graph-break and only has to stay correct.
+    """
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS[config]
+    _skip_unsupported(spec, backend, dtype)
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    _compare_compiled_to_eager(
+        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=backend != "fused"
+    )
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(model_config, qkv_format, dtype)
+    module = _make_dpa(spec, dtype)
 
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
@@ -817,7 +834,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
-        args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
+        args, kwargs, grads = _make_dpa_inputs(spec, dtype)
         # CUDA graphs hand back tensors owned by their memory pool, which the
         # next replay overwrites.
         out = compiled(*args, **kwargs).clone()
@@ -842,11 +859,11 @@ def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
     eager. FP8 elsewhere in the model with attention in high precision -- the
     common training setup -- must stay compiled."""
     dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "unfused")
 
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, _ = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
+    module = _make_dpa(spec, dtype)
+    args, kwargs, _ = _make_dpa_inputs(spec, dtype)
     fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
 
     def fn(*args, **kwargs):
@@ -868,48 +885,6 @@ def fn(*args, **kwargs):
     assert fell_back == fp8_attention
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
-    """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
-    and kv, which flash-attn hands to two inputs of the same autograd.Function --
-    something dynamo cannot trace. The other tests build two tensors, so this
-    covers the shape a caller is most likely to write."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_thd_padding_causal"]
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    b, s = model_config.batch_size, model_config.max_seqlen_q
-    h, d = model_config.num_heads, model_config.head_dim_qk
-    q, k, v = [
-        torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
-    ]
-    cu_seqlens = torch.arange(0, (b + 1) * s, step=s, dtype=torch.int32, device="cuda")
-    kwargs = dict(
-        cu_seqlens_q=cu_seqlens,
-        cu_seqlens_kv=cu_seqlens,  # the same tensor, not a copy
-        max_seqlen_q=s,
-        max_seqlen_kv=s,
-    )
-
-    ref = module(q, k, v, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in (q, k, v)]
-    for t in (q, k, v):
-        t.grad = None
-
-    torch._dynamo.reset()
-    _force_dpa_backend(monkeypatch, backend)
-    out = torch.compile(module, fullgraph=True)(q, k, v, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_dpa_backend(backend)
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip((q, k, v), ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
-
-
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("interleave_dim", [-3, -2])
 def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
@@ -919,12 +894,13 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     the packing via qkv_layer/kv_layer keeps it on the compiled path, which
     test_dpa_torch_compile covers."""
     dtype = torch.bfloat16
-    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    config = spec["model_config"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(model_config, qkv_format, dtype)
-    b, s = model_config.batch_size, model_config.max_seqlen_q
-    h, d = model_config.num_heads, model_config.head_dim_qk
+    module = _make_dpa(spec, dtype)
+    b, s = config.batch_size, config.max_seqlen_q
+    h, d = config.num_heads, config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
@@ -951,40 +927,6 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
         torch.compile(module, fullgraph=True)(q, k, v)
 
 
-def test_dpa_torch_compile_fused_backend(monkeypatch):
-    """The FusedAttention backend is an eager island, so it graph-breaks rather
-    than compiles; DotProductAttention must still be correct around it."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
-    _force_dpa_backend(monkeypatch, "fused")
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
-
-    ref = module(*args, **kwargs)
-    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
-        _attention_backends,
-    )
-
-    if not _attention_backends["use_fused_attention"]:
-        pytest.skip("FusedAttention does not support this config on this device")
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for t in grads:
-        t.grad = None
-
-    torch._dynamo.reset()
-    _force_dpa_backend(monkeypatch, "fused")
-    out = torch.compile(module)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_dpa_backend("fused")
-    _assert_matches_eager(out, ref, "fused", dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, "fused", dtype)
-
-
 # ---------------------------------------------------------------------------
 # get_attention_backend under torch.compile
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 95b5a7f15f..383afebd69 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2005,21 +2005,21 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
     tensor of shape [batch_size * max_seqlen, 1, 1] containing the indices for
     the valid tokens in a batch.
     """
+    # Built with device-side ops only: reading the sequence lengths on the host
+    # would synchronize the device once per sequence.
     bs = len(cu_seqlens) - 1
     seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
-    indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)]
-    indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda")
-
-    num_nonzeros = indices.shape[0]
-    pad_amount = bs * max_seqlen - num_nonzeros
-    indices = F.pad(
-        input=indices,
-        pad=(0, 0, 0, 0, 0, pad_amount),
-        mode="constant",
-        value=float(bs * max_seqlen),
+    positions = torch.arange(max_seqlen, device=cu_seqlens.device)
+    valid = (positions.unsqueeze(0) < seqlens.unsqueeze(1)).flatten()
+    # Sorting the mask is stable, so the valid positions come first in order,
+    # and the invalid tail is replaced by the out-of-range padding index.
+    ordered = torch.argsort(~valid, stable=True)
+    indices = torch.where(
+        torch.arange(bs * max_seqlen, device=cu_seqlens.device) < valid.sum(),
+        ordered,
+        bs * max_seqlen,
     )
-
-    return indices
+    return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1)
 
 
 def get_full_cu_seqlens(

From 725931ec24db7425e0f008444b65b89571692a07 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 12:11:58 +0200
Subject: [PATCH 13/39] Run eagerly when max_seqlen has to be derived from
 cu_seqlens

On thd inputs without max_seqlen, DotProductAttention reads the sequence
lengths off cu_seqlens to derive it. Unlike the padding mask and the packing
indices, this one cannot be moved to the device: max_seqlen sizes tensors
downstream, so it has to be a Python int, and reading it is a device
synchronization -- a data-dependent value dynamo cannot trace. Route those
calls to the eager island, so passing max_seqlen is what keeps a call
compiled.

Arguments are now looked up by name or position, so the predicate sees them
whichever way the caller passed them.

The two eager-fallback tests merge into one: they set up different inputs but
assert the same thing -- a warning, results matching eager, and an error under
fullgraph.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 70 +++++++++++++------
 .../dot_product_attention.py                  | 44 ++++++++++--
 2 files changed, 88 insertions(+), 26 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 397db816f4..ea21cdd2ed 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -885,46 +885,76 @@ def fn(*args, **kwargs):
     assert fell_back == fp8_attention
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-@pytest.mark.parametrize("interleave_dim", [-3, -2])
-def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
-    """Packed q/k/v passed as plain views cannot be recognized while tracing (the
-    detection reads data pointers), so DotProductAttention runs that detection
-    eagerly instead -- correct results, at the cost of a graph break. Declaring
-    the packing via qkv_layer/kv_layer keeps it on the compiled path, which
-    test_dpa_torch_compile covers."""
-    dtype = torch.bfloat16
-    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+def _packed_views_inputs(spec, dtype, interleave_dim=-3):
+    """Packed q/k/v handed over as plain views, without declaring the packing."""
     config = spec["model_config"]
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(spec, dtype)
     b, s = config.batch_size, config.max_seqlen_q
     h, d = config.num_heads, config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
+    return (q, k, v), {}, [qkv]
+
+
+def _thd_without_max_seqlen_inputs(spec, dtype):
+    """thd inputs that leave max_seqlen to be derived from cu_seqlens."""
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    del kwargs["max_seqlen_q"], kwargs["max_seqlen_kv"]
+    return args, kwargs, grads
+
+
+_EAGER_FALLBACK_CASES = {
+    # name: (config name, input builder)
+    "packed_views_bs3hd": (
+        "self_bshd_causal",
+        lambda spec, dtype: _packed_views_inputs(spec, dtype, -3),
+    ),
+    "packed_views_bsh3d": (
+        "self_bshd_causal",
+        lambda spec, dtype: _packed_views_inputs(spec, dtype, -2),
+    ),
+    "thd_without_max_seqlen": ("self_thd_padding_causal", _thd_without_max_seqlen_inputs),
+}
 
-    ref = module(q, k, v)
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("case", _EAGER_FALLBACK_CASES.keys())
+def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
+    """Calls that cannot be traced run as an eager island instead, with a
+    warning: recognizing packed q/k/v takes the data pointers dynamo cannot
+    read, and deriving max_seqlen off cu_seqlens is a device synchronization.
+    Both keep working, and both are avoidable -- by declaring the packing via
+    qkv_layer/kv_layer, or by passing max_seqlen."""
+    dtype = torch.bfloat16
+    config_name, make_inputs = _EAGER_FALLBACK_CASES[case]
+    spec = _DPA_COMPILE_CONFIGS[config_name]
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = make_inputs(spec, dtype)
+
+    ref = module(*args, **kwargs)
     ref.sum().backward()
-    ref_grad = qkv.grad.clone()
-    qkv.grad = None
+    ref_grads = [t.grad.clone() for t in grads]
+    for tensor in grads:
+        tensor.grad = None
 
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.warns(UserWarning, match="Falling back to eager execution"):
-        out = torch.compile(module)(q, k, v)
+        out = torch.compile(module)(*args, **kwargs)
     out.sum().backward()
     torch.cuda.synchronize()
 
     _assert_matches_eager(out, ref, backend, dtype)
-    _assert_matches_eager(qkv.grad, ref_grad, backend, dtype)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
-    # The same graph break is an error when the user asked for a full graph.
+    # The same eager island is an error when the user asked for a full graph.
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.raises(Exception, match="torch.compiler.disable"):
-        torch.compile(module, fullgraph=True)(q, k, v)
+        torch.compile(module, fullgraph=True)(*args, **kwargs)
 
 
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index f82aa91c13..4e0697fcd7 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,12 +197,32 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
+# Positions of the `DotProductAttention.forward` arguments the predicate below
+# reads, so that they are found whether they were passed by name or by position.
+_FORWARD_ARG_POSITIONS = {
+    "query_layer": 1,
+    "key_layer": 2,
+    "value_layer": 3,
+    "qkv_format": 5,
+    "max_seqlen_q": 10,
+    "max_seqlen_kv": 11,
+    "qkv_layer": 28,
+    "kv_layer": 29,
+}
+
+
+def _forward_arg(name: str, args: tuple, kwargs: dict) -> Any:
+    if name in kwargs:
+        return kwargs[name]
+    position = _FORWARD_ARG_POSITIONS[name]
+    return args[position] if len(args) > position else None
+
+
 def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
     """Why this DotProductAttention call has to run outside the graph, or None.
 
-    Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
-    are read from either side: `args[0]` is `self`, followed by the positional
-    query/key/value.
+    Takes `DotProductAttention.forward`'s arguments as they were passed, with
+    `args[0]` being `self`.
     """
     # FP8 GEMMs with the attention itself in high precision -- the common
     # training setup -- stay on the compiled path; only FP8 attention bails out.
@@ -212,9 +232,21 @@ def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
-    if kwargs.get("qkv_layer") is None and kwargs.get("kv_layer") is None:
-        qkv = list(args[1:4]) + [
-            kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+    qkv_format = _forward_arg("qkv_format", args, kwargs) or args[0].qkv_format
+    if qkv_format == "thd" and (
+        _forward_arg("max_seqlen_q", args, kwargs) is None
+        or _forward_arg("max_seqlen_kv", args, kwargs) is None
+    ):
+        # Deriving it reads the sequence lengths off cu_seqlens, which is a
+        # device synchronization and a data-dependent value while tracing.
+        return "deriving max_seqlen from cu_seqlens"
+
+    if (
+        _forward_arg("qkv_layer", args, kwargs) is None
+        and _forward_arg("kv_layer", args, kwargs) is None
+    ):
+        qkv = [
+            _forward_arg(name, args, kwargs) for name in ("query_layer", "key_layer", "value_layer")
         ]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"

From b2ab7ebad5034690f51ee8d6064ff0f3fbfa20fb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 13:54:50 +0200
Subject: [PATCH 14/39] Share the run-and-compare code between the compile
 tests

Both tests ran a callable, went backward, copied the gradients out and cleared
them. That is now _run_and_capture(), with _assert_run_matches() comparing two
of its results, which halves the CUDA-graphs test.

The two stay separate tests: what CUDA graphs risk -- a replay reading stale
buffers -- does not depend on the attention configuration, so folding the mode
into the configuration matrix would multiply the cases without covering
anything new, and would report a replay bug as a failure of some unrelated
configuration.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 59 ++++++++++++++++-------------
 1 file changed, 33 insertions(+), 26 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ea21cdd2ed..4be4ec9fea 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -767,30 +767,48 @@ def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) ->
     )
 
 
+def _run_and_capture(fn, args, kwargs, grads):
+    """Run `fn` and backward through its output, returning the output and a copy
+    of every input gradient.
+
+    The output is cloned because CUDA graphs hand back tensors owned by their
+    memory pool, which the next replay overwrites, and the gradients are cleared
+    so that the same inputs can be reused for the other run.
+    """
+    out = fn(*args, **kwargs).clone()
+    out.sum().backward()
+    torch.cuda.synchronize()
+    captured = []
+    for tensor in grads:
+        assert tensor.grad is not None
+        captured.append(tensor.grad.clone())
+        tensor.grad = None
+    return out, captured
+
+
+def _assert_run_matches(actual, expected, backend: str, dtype: torch.dtype) -> None:
+    """Compare an (output, gradients) pair produced by `_run_and_capture`."""
+    (out, out_grads), (ref, ref_grads) = actual, expected
+    _assert_matches_eager(out, ref, backend, dtype)
+    for out_grad, ref_grad in zip(out_grads, ref_grads):
+        _assert_matches_eager(out_grad, ref_grad, backend, dtype)
+
+
 def _compare_compiled_to_eager(
     module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs
 ) -> None:
     """Run the module eagerly and compiled on the same inputs, and compare both
     the output and every input gradient."""
-    ref = module(*args, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for tensor in grads:
-        tensor.grad = None
+    eager = _run_and_capture(module, args, kwargs, grads)
 
     torch._dynamo.reset()
     # Force backend selection to be re-run (and traced) inside the compiled
     # region instead of being served from the cache the eager call populated.
     _force_dpa_backend(monkeypatch, backend)
-    out = torch.compile(module, **compile_kwargs)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
+    compiled = _run_and_capture(torch.compile(module, **compile_kwargs), args, kwargs, grads)
 
     _assert_dpa_backend(backend)
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        assert tensor.grad is not None
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
+    _assert_run_matches(compiled, eager, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
@@ -835,20 +853,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
         args, kwargs, grads = _make_dpa_inputs(spec, dtype)
-        # CUDA graphs hand back tensors owned by their memory pool, which the
-        # next replay overwrites.
-        out = compiled(*args, **kwargs).clone()
-        out.sum().backward()
-        torch.cuda.synchronize()
-        compiled_grads = [t.grad.clone() for t in grads]
-        for tensor in grads:
-            tensor.grad = None
-
-        ref = module(*args, **kwargs)
-        ref.sum().backward()
-        _assert_matches_eager(out, ref, backend, dtype)
-        for compiled_grad, tensor in zip(compiled_grads, grads):
-            _assert_matches_eager(compiled_grad, tensor.grad, backend, dtype)
+        replayed = _run_and_capture(compiled, args, kwargs, grads)
+        eager = _run_and_capture(module, args, kwargs, grads)
+        _assert_run_matches(replayed, eager, backend, dtype)
     _assert_dpa_backend(backend)
 
 

From 50d19848ba97743cfe06ed619984792fbe9b3ecb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 13:57:46 +0200
Subject: [PATCH 15/39] Use the shared run-and-compare code in the remaining
 compile tests

The eager-fallback test hand-rolled the same run, backward and gradient copy.

The backend-level unfused test only checked that the output was finite and
that gradients existed, which a CUDA-graph replay reading stale buffers also
satisfies -- the weakness the DotProductAttention CUDA-graphs test was given
an eager comparison for. It now compares against eager as well.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 28 +++++++++-------------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 4be4ec9fea..8e3617af92 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -542,14 +542,13 @@ def fn(q, k, v, extra):
     compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead")
 
     for _ in range(3):
+        # Compared against eager rather than only checked for finiteness: a
+        # replay reusing stale buffers would pass the latter.
         q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True)
-        out = compiled(q, k, v, extra)
-        out.sum().backward()
-        torch.cuda.synchronize()
-        assert torch.isfinite(out).all()
-        assert q.grad is not None
-        assert k.grad is not None
-        assert v.grad is not None
+        inputs = (q, k, v, extra)
+        replayed = _run_and_capture(compiled, inputs, {}, [q, k, v])
+        eager = _run_and_capture(fn, inputs, {}, [q, k, v])
+        _assert_run_matches(replayed, eager, "unfused", dtype)
 
 
 # ---------------------------------------------------------------------------
@@ -940,22 +939,13 @@ def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = make_inputs(spec, dtype)
 
-    ref = module(*args, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for tensor in grads:
-        tensor.grad = None
+    eager = _run_and_capture(module, args, kwargs, grads)
 
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.warns(UserWarning, match="Falling back to eager execution"):
-        out = torch.compile(module)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
+        fell_back = _run_and_capture(torch.compile(module), args, kwargs, grads)
+    _assert_run_matches(fell_back, eager, backend, dtype)
 
     # The same eager island is an error when the user asked for a full graph.
     torch._dynamo.reset()

From f19c52d6e58f0c4eeca4a557df39c91a718f64fc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:05:10 +0200
Subject: [PATCH 16/39] Read the predicate's arguments by name, not by
 hardcoded position

The positions of forward's parameters were written out as a table of indices,
which anyone inserting a parameter into a thirty-argument signature would
silently invalidate: the predicate would go on reading whatever now sits at
index 5 or 10 and decide the fallback on it, with nothing to notice.

The decorator already has the function it wraps, so it takes the parameter
names from its signature and hands the predicate the call's arguments keyed by
name.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 57 +++++++------------
 1 file changed, 20 insertions(+), 37 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4e0697fcd7..4d7c66d22d 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -5,6 +5,7 @@
 """Attention."""
 from contextlib import nullcontext
 from functools import wraps
+import inspect
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -197,32 +198,11 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-# Positions of the `DotProductAttention.forward` arguments the predicate below
-# reads, so that they are found whether they were passed by name or by position.
-_FORWARD_ARG_POSITIONS = {
-    "query_layer": 1,
-    "key_layer": 2,
-    "value_layer": 3,
-    "qkv_format": 5,
-    "max_seqlen_q": 10,
-    "max_seqlen_kv": 11,
-    "qkv_layer": 28,
-    "kv_layer": 29,
-}
-
-
-def _forward_arg(name: str, args: tuple, kwargs: dict) -> Any:
-    if name in kwargs:
-        return kwargs[name]
-    position = _FORWARD_ARG_POSITIONS[name]
-    return args[position] if len(args) > position else None
-
-
-def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
+def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     """Why this DotProductAttention call has to run outside the graph, or None.
 
-    Takes `DotProductAttention.forward`'s arguments as they were passed, with
-    `args[0]` being `self`.
+    `call` maps `DotProductAttention.forward`'s parameter names to the arguments
+    this call passed, including `self`.
     """
     # FP8 GEMMs with the attention itself in high precision -- the common
     # training setup -- stay on the compiled path; only FP8 attention bails out.
@@ -232,32 +212,33 @@ def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
-    qkv_format = _forward_arg("qkv_format", args, kwargs) or args[0].qkv_format
+    qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
-        _forward_arg("max_seqlen_q", args, kwargs) is None
-        or _forward_arg("max_seqlen_kv", args, kwargs) is None
+        call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None
     ):
         # Deriving it reads the sequence lengths off cu_seqlens, which is a
         # device synchronization and a data-dependent value while tracing.
         return "deriving max_seqlen from cu_seqlens"
 
-    if (
-        _forward_arg("qkv_layer", args, kwargs) is None
-        and _forward_arg("kv_layer", args, kwargs) is None
-    ):
-        qkv = [
-            _forward_arg(name, args, kwargs) for name in ("query_layer", "key_layer", "value_layer")
-        ]
+    if call.get("qkv_layer") is None and call.get("kv_layer") is None:
+        qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
     return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[..., Optional[str]]):
+def _eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
     """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path."""
+    a reason why the call is unsupported on the compiled path.
+
+    `needs_eager` is handed the call's arguments keyed by parameter name, taken
+    from the wrapped signature, so it does not have to care whether the caller
+    passed them by name or by position.
+    """
 
     def decorator(fn):
+        parameter_names = list(inspect.signature(fn).parameters)
+
         # The warning belongs inside the dynamo-disabled function: warnings.warn
         # graph-breaks on its own, masking the break that matters.
         @no_torch_dynamo()
@@ -272,7 +253,9 @@ def eager_fn(reason, *args, **kwargs):
         @wraps(fn)
         def wrapper(*args, **kwargs):
             if torch.compiler.is_compiling():
-                reason = needs_eager(*args, **kwargs)
+                call = dict(zip(parameter_names, args))
+                call.update(kwargs)
+                reason = needs_eager(call)
                 if reason is not None:
                     return eager_fn(reason, *args, **kwargs)
             return fn(*args, **kwargs)

From fbc7146e78d465c118ed7312d58983460cf7559f Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:12:38 +0200
Subject: [PATCH 17/39] Move the eager-fallback decorator to jit.py

Nothing about it is specific to attention: it takes a predicate, and runs the
wrapped method as an eager island whenever that predicate gives a reason. It
belongs next to no_torch_dynamo, which it builds on, rather than buried in the
attention module where the modules that will need it next -- MultiheadAttention
and TransformerLayer, both of which still have blockers of their own -- would
not find it.

What stays in the attention module is the predicate, which is where the
attention-specific knowledge lives.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 44 +------------------
 transformer_engine/pytorch/jit.py             | 42 +++++++++++++++++-
 2 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4d7c66d22d..5f6628c951 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -4,8 +4,6 @@
 
 """Attention."""
 from contextlib import nullcontext
-from functools import wraps
-import inspect
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -44,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import no_torch_dynamo
+from transformer_engine.pytorch.jit import eager_under_compile_if
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -227,44 +225,6 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
-    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path.
-
-    `needs_eager` is handed the call's arguments keyed by parameter name, taken
-    from the wrapped signature, so it does not have to care whether the caller
-    passed them by name or by position.
-    """
-
-    def decorator(fn):
-        parameter_names = list(inspect.signature(fn).parameters)
-
-        # The warning belongs inside the dynamo-disabled function: warnings.warn
-        # graph-breaks on its own, masking the break that matters.
-        @no_torch_dynamo()
-        def eager_fn(reason, *args, **kwargs):
-            warnings.warn(
-                f"Falling back to eager execution under torch.compile: {reason} is"
-                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
-                stacklevel=3,
-            )
-            return fn(*args, **kwargs)
-
-        @wraps(fn)
-        def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling():
-                call = dict(zip(parameter_names, args))
-                call.update(kwargs)
-                reason = needs_eager(call)
-                if reason is not None:
-                    return eager_fn(reason, *args, **kwargs)
-            return fn(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
 def _unpack_packed_qkv(
     qkv_layer: Optional[torch.Tensor],
     kv_layer: Optional[torch.Tensor],
@@ -1167,7 +1127,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @_eager_under_compile_if(_needs_eager_dpa)
+    @eager_under_compile_if(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 1b93b8254c..2d0ec02d8d 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -3,9 +3,11 @@
 # See LICENSE for license information.
 
 """NVFuser functions and JIT utilities"""
+import inspect
 import os
+import warnings
 from functools import wraps
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 
 from .torch_version import torch_version
@@ -77,6 +79,44 @@ def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
         return lambda func: func
 
 
+def eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
+    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
+    a reason why the call is unsupported on the compiled path.
+
+    `needs_eager` is handed the call's arguments keyed by parameter name, taken
+    from the wrapped signature, so it does not have to care whether the caller
+    passed them by name or by position. Returning None keeps the call compiled.
+    """
+
+    def decorator(fn):
+        parameter_names = list(inspect.signature(fn).parameters)
+
+        # The warning belongs inside the dynamo-disabled function: warnings.warn
+        # graph-breaks on its own, masking the break that matters.
+        @no_torch_dynamo()
+        def eager_fn(reason, *args, **kwargs):
+            warnings.warn(
+                f"Falling back to eager execution under torch.compile: {reason} is"
+                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                stacklevel=3,
+            )
+            return fn(*args, **kwargs)
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if torch.compiler.is_compiling():
+                call = dict(zip(parameter_names, args))
+                call.update(kwargs)
+                reason = needs_eager(call)
+                if reason is not None:
+                    return eager_fn(reason, *args, **kwargs)
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def set_jit_fusion_options() -> None:
     """Set PyTorch JIT layer fusion options."""
     # flags required to enable jit fusion kernels

From 6b246303b847db73c1a71312915bccec1c3514e0 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:40:19 +0200
Subject: [PATCH 18/39] Skip FusedAttention rather than special-casing it in
 the compile tests

The test compiled every backend with fullgraph=True except fused, which was
passed fullgraph=False because its forward is an eager island. That put an
assumption about one backend into the comparison logic.

FusedAttention is skipped instead, next to the configurations a backend cannot
run, with the reason spelled out. Everything else is compiled the same way, and
supporting fused later is deleting the skip.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 8e3617af92..e32304af1b 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -694,7 +694,12 @@ def _randn(shape):
 
 
 def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
-    """Skip configurations the backend under test cannot run at all."""
+    """Skip what the backend under test cannot run, or cannot compile."""
+    if backend == "fused":
+        # FusedAttention's forward carries @no_torch_dynamo, so there is nothing
+        # to compile: it runs as an eager island. Drop this skip once it traces,
+        # and the tests below cover it as they do the others.
+        pytest.skip("FusedAttention is an eager island and does not compile")
     available, _, _ = get_available_attention_backends(
         spec["model_config"], dtype, _qkv_layout(spec)
     )
@@ -813,14 +818,13 @@ def _compare_compiled_to_eager(
 @pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
 @pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
 def test_dpa_torch_compile(monkeypatch, backend, config):
-    """`DotProductAttention` under `torch.compile` must match eager in forward
-    and backward, for every backend that supports the configuration.
-
-    FlashAttention and UnfusedDotProductAttention are compiled with
-    `fullgraph=True`, so the test fails on any graph break -- covering the whole
-    module: input unpacking, qkv layout, backend selection and the backend
-    itself. FusedAttention is an eager island (`@no_torch_dynamo` on its
-    forward), so it is expected to graph-break and only has to stay correct.
+    """`DotProductAttention` under `torch.compile(fullgraph=True)` must match
+    eager in forward and backward, for every backend that supports the
+    configuration.
+
+    `fullgraph=True` makes the test fail on any graph break, so it covers the
+    whole module: input unpacking, qkv layout, backend selection and the backend
+    itself.
     """
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS[config]
@@ -830,7 +834,7 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = _make_dpa_inputs(spec, dtype)
     _compare_compiled_to_eager(
-        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=backend != "fused"
+        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True
     )
 
 

From 9a204b8e0a97237a85ee7a99ea85c7091b0fc6e4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:39:18 +0200
Subject: [PATCH 19/39] Cover the compiled path around FusedAttention

FusedAttention is skipped from the backend matrix because it does not compile,
but it is the default on Hopper and Blackwell, and this change does affect it:
what used to be one eager island is now a traced prologue, a graph break at the
backend, and a traced remainder. Whatever crosses that break has to survive it,
which the sub-backend enum did not.

One test covers it, without a fullgraph and outside the matrix. Verified to
fail with the enum restored -- TypeError: int() argument must be ... not
'function', from cuDNN's entry point.

Also rename _distinct_cu_seqlens: it returns its arguments untouched in eager
and whenever they already differ, so it promised more than it does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 24 ++++++++++++++++---
 .../dot_product_attention/backends.py         |  8 +++----
 .../dot_product_attention.py                  |  3 ---
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index e32304af1b..d20a0e70dc 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -693,9 +693,10 @@ def _randn(shape):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
-    """Skip what the backend under test cannot run, or cannot compile."""
-    if backend == "fused":
+def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) -> None:
+    """Skip what the backend under test cannot run, or -- for a test that
+    compiles it -- cannot be compiled."""
+    if compiled and backend == "fused":
         # FusedAttention's forward carries @no_torch_dynamo, so there is nothing
         # to compile: it runs as an eager island. Drop this skip once it traces,
         # and the tests below cover it as they do the others.
@@ -838,6 +839,23 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     )
 
 
+def test_dpa_torch_compile_around_fused(monkeypatch):
+    """FusedAttention itself is an eager island, but everything around it is
+    compiled: DotProductAttention traces up to the backend call, breaks the
+    graph there and resumes afterwards. What crosses that break has to survive
+    it -- the sub-backend enum did not, and reached cuDNN as the function that
+    produced it."""
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    _skip_unsupported(spec, "fused", dtype, compiled=False)
+    _force_dpa_backend(monkeypatch, "fused")
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    # No fullgraph: the graph break at the eager island is the point here.
+    _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype)
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 2c47ff8a4b..04bcabe7c4 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -793,8 +793,8 @@ def backward(
         return dq, dk, dv
 
 
-def _distinct_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
-    """Make q's and kv's cumulative sequence lengths distinct objects.
+def _maybe_unshare_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
+    """Copy kv's cumulative sequence lengths when they are q's as well.
 
     Self attention passes one tensor for both, and flash-attn forwards them to
     two inputs of the same autograd.Function -- which dynamo cannot trace. In
@@ -1134,7 +1134,7 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        cu_q, cu_kv = _distinct_cu_seqlens(
+                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(
                             cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q,
                             cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv,
                         )
@@ -1150,7 +1150,7 @@ def forward(
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
                     if func is flash_attn_varlen_func_v4:
-                        cu_q, cu_kv = _distinct_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
+                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
                         fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q
                         fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv
                         fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 5f6628c951..34aa6cf403 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -746,9 +746,6 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
 
         qstate = FP8GlobalStateManager.quantization_state
         if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
-            # Without FP8 the recipe juggling below is a no-op, and its
-            # get_fp8_recipe() call would build a default Recipe -- which
-            # torch.compile cannot trace.
             super().init_fp8_metadata(num_gemms=num_gemms)
             return
 

From 2188f183a41bd0cdd6dc662d03169c10a6685c4c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:47:26 +0200
Subject: [PATCH 20/39] Log backend selection through the no-op logger, as
 get_attention_backend does

Skipping the block with is_compiling() introduced a second mechanism for a
problem this module already solves one way: get_attention_backend swaps in a
no-op logger while tracing. The skip was needed while the sub-backend argument
was an enum that did not survive tracing, and is not any more.

The logger, which was private to the utils module, is now named without the
underscore, since it is used from outside it.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 35 ++++++++++---------
 .../attention/dot_product_attention/utils.py  |  4 +--
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 34aa6cf403..d72eb6e98d 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1854,22 +1854,25 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    # Logged in eager only: dynamo graph-breaks on logging.Logger
-                    # methods, and log arguments are evaluated regardless of the
-                    # logger, so a no-op logger would not be enough.
-                    if not torch.compiler.is_compiling():
-                        if use_flash_attention:
-                            self.logger.info(
-                                "Running with FlashAttention backend (version %s)",
-                                flash_attention_backend,
-                            )
-                        elif use_fused_attention:
-                            self.logger.info(
-                                "Running with FusedAttention backend (sub-backend %s)",
-                                int(fused_attention_backend),
-                            )
-                        elif use_unfused_attention:
-                            self.logger.info("Running with UnfusedDotProductAttention backend")
+                    # logging.Logger methods graph-break under torch.compile, so
+                    # selection is only logged in eager -- as in
+                    # get_attention_backend. Note the arguments below are
+                    # evaluated either way, so they have to stay traceable.
+                    logger = (
+                        dpa_utils.no_op_logger if torch.compiler.is_compiling() else self.logger
+                    )
+                    if use_flash_attention:
+                        logger.info(
+                            "Running with FlashAttention backend (version %s)",
+                            flash_attention_backend,
+                        )
+                    elif use_fused_attention:
+                        logger.info(
+                            "Running with FusedAttention backend (sub-backend %s)",
+                            int(fused_attention_backend),
+                        )
+                    elif use_unfused_attention:
+                        logger.info("Running with UnfusedDotProductAttention backend")
                 else:
                     use_flash_attention = _attention_backends["use_flash_attention"]
                     flash_attention_backend = _attention_backends["flash_attention_backend"]
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 383afebd69..ccf46fda22 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -346,7 +346,7 @@ def error(self, *args, **kwargs):
         """No-op."""
 
 
-_no_op_logger = _NoOpLogger()
+no_op_logger = _NoOpLogger()
 
 
 @torch.compiler.assume_constant_result
@@ -461,7 +461,7 @@ def get_attention_backend(
     if torch.compiler.is_compiling():
         # logging.Logger methods graph-break under torch.compile; backend
         # selection logs are only emitted in eager mode.
-        logger = _no_op_logger
+        logger = no_op_logger
     else:
         logger = logging.getLogger("DotProductAttention")
         logger.setLevel(AttentionLogging._log_level)

From 10d141b3fc70fe5602a8a15bd84a71d25732fcef Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:51:46 +0200
Subject: [PATCH 21/39] Return the fused sub-backend as an int, in eager as
 well

get_attention_backend cast it back to a FusedAttnBackend member unless it was
tracing, which made its return type depend on the execution mode -- something
every reader of that function then has to keep in mind, for no benefit inside
transformer_engine: fused_attn_fwd/bwd call cast() on whatever they are given,
and every comparison against a member works with an int on either side.

It now returns an int always, and the docstring says so.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/utils.py    | 17 +++++------------
 .../pytorch/cpp_extensions/fused_attn.py        |  2 +-
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index ccf46fda22..9e055ef6e6 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -401,11 +401,11 @@ def get_attention_backend(
         Whether the `FlashAttention` backend has been selected.
     use_fused_attention : bool
         Whether the `FusedAttention` backend has been selected.
-    fused_attention_backend : FusedAttnBackend
-        If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
-        Under `torch.compile` this is the sub-backend's plain integer value instead of the enum
-        member (see the cast below); `FusedAttnBackend.cast` accepts both, as does every consumer
-        in `transformer_engine.pytorch`.
+    fused_attention_backend : int
+        If `use_fused_attention = True`, the integer value of one of `FusedAttention`'s three
+        sub-backends, else `None`. It is not a `FusedAttnBackend` member because that does not
+        survive a graph break under `torch.compile`; `FusedAttnBackend.cast` turns it into one,
+        and comparing it against a member works either way.
     use_unfused_attention : bool
         Whether the `UnfusedDotProductAttention` backend has been selected.
     available_backends : List[bool]
@@ -1600,13 +1600,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         use_flash_attention_4 = False
     use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
     available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention]
-    if fused_attention_backend is not None and not torch.compiler.is_compiling():
-        # The fused sub-backend is kept as a plain int while tracing: it has to
-        # survive a graph break (DotProductAttention hands it to FusedAttention,
-        # which is an eager island) and an enum member does not -- dynamo
-        # reconstructs the constant-folded member into the resumed frame as the
-        # function that produced it. Eager callers get the enum, as before.
-        fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend)
     if use_flash_attention_2:
         flash_attention_backend = FlashAttentionUtils.version
     if use_flash_attention_3:
diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
index 3a4e5aac66..9a33df7634 100644
--- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py
+++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
@@ -112,7 +112,7 @@ class FusedAttnBackend(IntEnum):
     the same way as with the dict this used to be.
 
     Members do not survive a graph break, though, so ``get_attention_backend``
-    hands out plain ints while tracing -- ``cast`` normalizes both back to a
+    returns the sub-backend as a plain int and ``cast`` turns it back into a
     member.
     """
 

From 93a201d42e02f01bba21a67ad9c3e6f44f78ce05 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:57:12 +0200
Subject: [PATCH 22/39] Rename eager_under_compile_if to fallback_to_eager_when

The old name read as a run-on, and its '_if' suffix suggested a boolean
condition where the argument actually returns the reason. The new one reads as
a sentence where it is used and matches the warning the decorator emits.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py  | 4 ++--
 transformer_engine/pytorch/jit.py                             | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d72eb6e98d..ba6c5a4769 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import eager_under_compile_if
+from transformer_engine.pytorch.jit import fallback_to_eager_when
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -1124,7 +1124,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @eager_under_compile_if(_needs_eager_dpa)
+    @fallback_to_eager_when(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 2d0ec02d8d..f6bbf55bde 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -79,7 +79,7 @@ def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
         return lambda func: func
 
 
-def eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
+def fallback_to_eager_when(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
     """Decorator running the wrapped method eagerly whenever `needs_eager` gives
     a reason why the call is unsupported on the compiled path.
 

From 0866319fe05e2d17d1f768ca817919c53fa9f08d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:06:44 +0200
Subject: [PATCH 23/39] Fold the eager fallback into no_torch_dynamo as a
 `when` predicate

It was a decorator of its own, which meant a second name for what is the same
thing this module already does: disable dynamo for a frame. It is now the
conditional form of no_torch_dynamo -- @no_torch_dynamo(when=predicate) -- so
there is one decorator to find and one docstring to read.

The predicate receives the call's arguments keyed by parameter name and returns
the reason this call cannot be traced, or None to have it traced.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  |  4 +-
 transformer_engine/pytorch/jit.py             | 94 +++++++++----------
 2 files changed, 48 insertions(+), 50 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index ba6c5a4769..774f6a44d1 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import fallback_to_eager_when
+from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -1124,7 +1124,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @fallback_to_eager_when(_needs_eager_dpa)
+    @no_torch_dynamo(when=_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index f6bbf55bde..164610f5a5 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -7,7 +7,7 @@
 import os
 import warnings
 from functools import wraps
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Callable, Optional, Tuple
 import torch
 
 from .torch_version import torch_version
@@ -51,22 +51,58 @@ def wrapper(*args, **kwargs):
 if torch.__version__ >= "2":
     import torch._dynamo
 
-    def no_torch_dynamo(recursive=True):
-        """Decorator to disable Torch Dynamo, except during ONNX export."""
+    def no_torch_dynamo(recursive=True, when=None):
+        """Decorator to disable Torch Dynamo, except during ONNX export.
 
-        def decorator(f):
+        `when` makes it conditional: it is called with the arguments of the call
+        keyed by parameter name, and returns the reason this particular call
+        cannot be traced, or None to have it traced as usual. The reason is
+        reported once per distinct message.
+        """
+
+        def _disable(f):
             # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True
-            disabled_f = (
+            return (
                 torch._dynamo.disable(f, recursive=recursive)
                 if torch.__version__ >= "2.1"
                 else torch._dynamo.disable(f)
             )
 
-            @wraps(f)
-            def wrapper(*args, **kwargs):
-                if is_in_onnx_export_mode():
+        def decorator(f):
+            disabled_f = _disable(f)
+            if when is None:
+
+                @wraps(f)
+                def wrapper(*args, **kwargs):
+                    if is_in_onnx_export_mode():
+                        return f(*args, **kwargs)
+                    return disabled_f(*args, **kwargs)
+
+            else:
+                parameter_names = list(inspect.signature(f).parameters)
+
+                # The warning belongs inside the disabled function: warnings.warn
+                # graph-breaks on its own, which would mask the break that matters.
+                def report_and_run(reason, *args, **kwargs):
+                    warnings.warn(
+                        f"Falling back to eager execution under torch.compile: {reason} is"
+                        " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                        stacklevel=3,
+                    )
                     return f(*args, **kwargs)
-                return disabled_f(*args, **kwargs)
+
+                disabled_report_and_run = _disable(report_and_run)
+
+                @wraps(f)
+                def wrapper(*args, **kwargs):  # pylint: disable=function-redefined
+                    if is_in_onnx_export_mode() or not torch.compiler.is_compiling():
+                        return f(*args, **kwargs)
+                    call = dict(zip(parameter_names, args))
+                    call.update(kwargs)
+                    reason = when(call)
+                    if reason is None:
+                        return f(*args, **kwargs)
+                    return disabled_report_and_run(reason, *args, **kwargs)
 
             return wrapper
 
@@ -74,49 +110,11 @@ def wrapper(*args, **kwargs):
 
 else:
     # Fallback for PyTorch < 2.0: no-op decorator
-    def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
+    def no_torch_dynamo(recursive=True, when=None):  # pylint: disable=unused-argument
         """No-op decorator for PyTorch < 2.0."""
         return lambda func: func
 
 
-def fallback_to_eager_when(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
-    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path.
-
-    `needs_eager` is handed the call's arguments keyed by parameter name, taken
-    from the wrapped signature, so it does not have to care whether the caller
-    passed them by name or by position. Returning None keeps the call compiled.
-    """
-
-    def decorator(fn):
-        parameter_names = list(inspect.signature(fn).parameters)
-
-        # The warning belongs inside the dynamo-disabled function: warnings.warn
-        # graph-breaks on its own, masking the break that matters.
-        @no_torch_dynamo()
-        def eager_fn(reason, *args, **kwargs):
-            warnings.warn(
-                f"Falling back to eager execution under torch.compile: {reason} is"
-                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
-                stacklevel=3,
-            )
-            return fn(*args, **kwargs)
-
-        @wraps(fn)
-        def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling():
-                call = dict(zip(parameter_names, args))
-                call.update(kwargs)
-                reason = needs_eager(call)
-                if reason is not None:
-                    return eager_fn(reason, *args, **kwargs)
-            return fn(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
 def set_jit_fusion_options() -> None:
     """Set PyTorch JIT layer fusion options."""
     # flags required to enable jit fusion kernels

From a80a41aa6d89d8922d27a4e725c6172e8b61c3ce Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:13:25 +0200
Subject: [PATCH 24/39] Cover the declared packed layouts beyond bs3hd

Packing was pinned to one layout by an assert in the input builder, so of the
twelve declarations DotProductAttention accepts -- qkv_layer or kv_layer,
interleaved at -3 or -2, in each of the three formats -- exactly one was
tested. The builder now takes what is packed and where it is interleaved, and
derives the shapes and the layout string the way DotProductAttention does.

Three configurations are added, so both packed inputs and both interleave
dimensions are covered.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 69 +++++++++++++++++++++++++----
 1 file changed, 60 insertions(+), 9 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d20a0e70dc..f65bb8cb95 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -561,15 +561,34 @@ def fn(q, k, v, extra):
 # get_available_attention_backends() answers that, so configurations only one
 # backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
 # rather than avoided.
-def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
+def _cfg(
+    model_config,
+    qkv_format="bshd",
+    packed=None,
+    interleave_dim=-3,
+    share_cu_seqlens=False,
+):
+    """A model configuration plus what DotProductAttention is handed it as.
+
+    `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at
+    `interleave_dim`, or None for separate q/k/v.
+    """
     return dict(
         model_config=model_config,
         qkv_format=qkv_format,
         packed=packed,
+        interleave_dim=interleave_dim,
         share_cu_seqlens=share_cu_seqlens,
     )
 
 
+def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str:
+    """bshd + 3 @ -3 -> bs3hd; bshd + 2 @ -2 -> bsh2d; as DotProductAttention
+    derives it from the declaration."""
+    position = len(qkv_format) + interleave_dim + 1
+    return qkv_format[:position] + str(packed_dim) + qkv_format[position:]
+
+
 _DPA_COMPILE_CONFIGS = {
     "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")),
     "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"),
@@ -587,7 +606,17 @@ def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
     "self_thd_shared_cu_seqlens": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True
     ),
-    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="bs3hd"),
+    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv"),
+    "packed_qkv_bsh3d": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv", interleave_dim=-2
+    ),
+    "packed_kv_bshd_bs2hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="kv"),
+    "packed_kv_thd_th2d": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
+        "thd",
+        packed="kv",
+        interleave_dim=-2,
+    ),
     # Configurations below are supported by one backend only, or take a code
     # path of their own inside DotProductAttention.
     "alibi_bshd_causal": _cfg(
@@ -605,7 +634,12 @@ def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
 
 
 def _qkv_layout(spec: dict) -> str:
-    return spec["packed"] or "_".join([spec["qkv_format"]] * 3)
+    qkv_format, packed = spec["qkv_format"], spec["packed"]
+    if packed is None:
+        return "_".join([qkv_format] * 3)
+    if packed == "qkv":
+        return _packed_layout(qkv_format, 3, spec["interleave_dim"])
+    return f"{qkv_format}_{_packed_layout(qkv_format, 2, spec['interleave_dim'])}"
 
 
 def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
@@ -669,14 +703,31 @@ def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
     if packed is not None:
-        # Declarative packed QKV: q/k/v are derived from one buffer by DPA
+        # Declarative packed inputs: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
         # layout that torch.compile supports.
-        assert packed == "bs3hd" and qkv_format == "bshd" and h == g and d_qk == d_v
-        qkv = _randn((b, s_q, 3, h, d_qk))
-        kwargs["qkv_layer"] = qkv
-        kwargs["qkv_interleave_dim"] = -3
-        grad_tensors, args = [qkv], ()
+        assert h == g and d_qk == d_v, "packed inputs require uniform heads and head dims"
+        interleave_dim = spec["interleave_dim"]
+
+        def _packed(seqlen, tokens, packed_dim):
+            leading = {
+                "bshd": (b, seqlen),
+                "sbhd": (seqlen, b),
+                "thd": (tokens,),
+            }[qkv_format]
+            trailing = (packed_dim, h, d_qk) if interleave_dim == -3 else (h, packed_dim, d_qk)
+            return _randn(leading + trailing)
+
+        kwargs["qkv_interleave_dim"] = interleave_dim
+        if packed == "qkv":
+            qkv = _packed(s_q, t_q, 3)
+            kwargs["qkv_layer"] = qkv
+            grad_tensors, args = [qkv], ()
+        else:
+            q = _randn(_shape(s_q, t_q, h, d_qk))
+            kv = _packed(s_kv, t_kv, 2)
+            kwargs["kv_layer"] = kv
+            grad_tensors, args = [q, kv], (q,)
     else:
         q = _randn(_shape(s_q, t_q, h, d_qk))
         k = _randn(_shape(s_kv, t_kv, g, d_qk))

From 9eb1449daf1d40e349d27ad50835169f02c9e9c1 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:21:13 +0200
Subject: [PATCH 25/39] Let ONNX export keep its own path in get_qkv_layout

The ONNX exporter runs through dynamo, so torch.compiler.is_compiling() is
true during an export as well -- measured: inside torch.onnx.export(dynamo=True)
both it and is_in_onnx_export_mode() are set. The compile branch was written
first and took the export over: the layout came from the format instead of the
detection ONNX relies on, tensors were no longer made contiguous, and packed
views hit an assert where they used to be copied.

Export mode is the narrower context, so it is checked first. This restores the
previous behaviour for export while leaving the compiled path as it was.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/utils.py   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 9e055ef6e6..a2d260c6b2 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2565,7 +2565,11 @@ def run_iteratively(q, k, v):
 
         return qkv_layout
 
-    if torch.compiler.is_compiling():
+    if is_in_onnx_export_mode():
+        # Checked first: the ONNX exporter runs through dynamo, so it also sets
+        # is_compiling(), and it has its own handling below.
+        qkv_layout = "not_supported"
+    elif torch.compiler.is_compiling():
         # run_iteratively reads data pointers and storage offsets, which dynamo
         # cannot trace; unpacked q/k/v need no detection anyway.
         assert not qkv_layout_needs_detection(q, k, v), (
@@ -2577,10 +2581,8 @@ def run_iteratively(q, k, v):
             qkv_layout = "_".join([qkv_format] * 3)
         else:
             qkv_layout = q_format + "_" + kv_format + "_" + kv_format
-    elif not is_in_onnx_export_mode():
-        qkv_layout = run_iteratively(q, k, v)
     else:
-        qkv_layout = "not_supported"
+        qkv_layout = run_iteratively(q, k, v)
     if qkv_layout == "not_supported":
         # force q,k,v to be contiguous and run get_layout again
         q, k, v = [x.contiguous() for x in [q, k, v]]

From 5c5b451264637a5726868043bee551be71cb88d4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:31:08 +0200
Subject: [PATCH 26/39] Guard the assumption the argument binding rests on

Arguments are matched to parameter names by position, which is only right
while the signature has no *args in between -- and would misalign silently if
one were added. It is checked when the decorator is applied.

The docstring now also says that an argument the call left out is simply
absent, so a predicate reading it with .get() sees its default only as long as
that default is None.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/jit.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 164610f5a5..62e6950154 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -55,9 +55,11 @@ def no_torch_dynamo(recursive=True, when=None):
         """Decorator to disable Torch Dynamo, except during ONNX export.
 
         `when` makes it conditional: it is called with the arguments of the call
-        keyed by parameter name, and returns the reason this particular call
-        cannot be traced, or None to have it traced as usual. The reason is
-        reported once per distinct message.
+        keyed by parameter name -- an argument the call left out is absent, so
+        reading it with .get() yields its default as long as that default is
+        None -- and returns the reason this particular call cannot be traced, or
+        None to have it traced as usual. The reason is reported once per
+        distinct message.
         """
 
         def _disable(f):
@@ -79,7 +81,13 @@ def wrapper(*args, **kwargs):
                     return disabled_f(*args, **kwargs)
 
             else:
-                parameter_names = list(inspect.signature(f).parameters)
+                parameters = inspect.signature(f).parameters
+                # Arguments are matched to names by position, which only holds
+                # without a *args in between.
+                assert not any(
+                    p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values()
+                ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes"
+                parameter_names = list(parameters)
 
                 # The warning belongs inside the disabled function: warnings.warn
                 # graph-breaks on its own, which would mask the break that matters.

From 95030f9a3140e7b62067bbf903610562a55462f2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 13:49:53 +0200
Subject: [PATCH 27/39] Check that CUDA graphs were actually captured, and
 generalize a docstring

fullgraph=True says dynamo did not break the graph, but inductor can still
decline to capture -- a mutated input, a CPU scalar -- and run the thing
normally, which would leave the CUDA-graphs test passing while measuring
nothing. Its skip counter is asserted to be zero. Measured on this branch:
nothing is skipped, the tree has one root, and interleaving the eager
reference call between replays does not disturb it.

The tolerance docstring described which backends are compared exactly by
naming them; it now states the rule -- a backend whose kernel is the same
either way -- so it still holds once FusedAttention traces.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f65bb8cb95..6ae070b70d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -21,6 +21,8 @@
 except ImportError:
     _opaque_available = False
 
+from torch._dynamo.utils import counters
+
 import transformer_engine.pytorch as te
 import transformer_engine_torch as tex
 from transformer_engine.common import recipe
@@ -804,12 +806,11 @@ def _assert_dpa_backend(backend: str) -> None:
 def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None:
     """Assert a compiled result matches the eager one.
 
-    FlashAttention and FusedAttention run the same kernel either way, so they
-    have to match exactly. Inductor reassociates the unfused backend's softmax
-    sums, and that error scales with the magnitude of those sums rather than
-    with each output element -- so the absolute tolerance comes from the
-    tensor's own scale, an elementwise relative tolerance being meaningless for
-    elements near zero.
+    A backend that calls the same kernel either way has to match exactly:
+    compiling changes what surrounds it, not its arithmetic. The unfused
+    backend is instead built from PyTorch ops, which inductor fuses and
+    reassociates; that error scales with the softmax sums rather than with each
+    output element, hence an absolute tolerance taken from the tensor's scale.
     """
     if backend != "unfused":
         torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
@@ -918,6 +919,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     module = _make_dpa(spec, dtype)
 
     torch._dynamo.reset()
+    counters.clear()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
     for _ in range(3):
@@ -929,6 +931,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         eager = _run_and_capture(module, args, kwargs, grads)
         _assert_run_matches(replayed, eager, backend, dtype)
     _assert_dpa_backend(backend)
+    # Without this, inductor declining to capture -- a mutated input, a CPU
+    # scalar -- would leave the test passing while measuring nothing.
+    assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)

From b0e56d450bac0300da046cc7e56ea6b098ec20b3 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 14:51:20 +0200
Subject: [PATCH 28/39] Skip lazy compilation while already tracing

Filling lazy_compile's closure cell is a side effect Dynamo rejects inside
a higher-order op, so compiling a module whose autograd.Function backward
makes the first call to a jit_fuser'd function fails with

    Unsupported: HOP: Unsafe side effect
      Attempted to mutate CellVariable()

An earlier eager run hides this by filling the cell beforehand. The nested
torch.compile is inlined into the outer graph either way, so calling the
function directly while tracing is equivalent.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/jit.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 62e6950154..22e9c9ae3e 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -28,6 +28,8 @@ def lazy_compile(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         nonlocal compiled_func
+        if torch.compiler.is_compiling():
+            return func(*args, **kwargs)
         if compiled_func is None:
             compiled_func = torch.compile(func)
         return compiled_func(*args, **kwargs)

From 2cf695bb59a551efb8e6ad49dbe3ef85af2d54f4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 16:49:43 +0200
Subject: [PATCH 29/39] Run context parallel attention eagerly

The context parallel implementation uses P2P communication and its own CUDA
stream, neither of which dynamo traces, and FlashAttention reaches it from a
forward that is no longer excluded from the graph.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 774f6a44d1..0ed01a9c9e 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -210,6 +210,9 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
+    if call["self"].cp_group is not None:
+        return "context parallelism"
+
     qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
         call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None

From b89446f4d04341a91708ffe5c3e8e532bfa05bff Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 13:04:46 +0200
Subject: [PATCH 30/39] Keep DotProductAttention compilable under a CUDA RNG
 states tracker

Forking the tracker swaps the global CUDA generator state, which dynamo
refuses to trace. With attention_dropout == 0 the dropout draws nothing
from the generator, so the fork is a no-op -- use nullcontext instead and
keep the call in the graph. Calls that do use dropout under a tracker run
as an eager island.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/dot_product_attention.py     | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 0ed01a9c9e..6bb10aec56 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -225,6 +225,11 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
         qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+
+    if call["self"].rng_states_tracker is not None and call["self"].attention_dropout > 0:
+        # Forking the tracker swaps the global CUDA generator state, which Dynamo
+        # refuses to trace. With no dropout there is nothing to fork for.
+        return "attention dropout under a CUDA RNG states tracker"
     return None
 
 
@@ -577,7 +582,10 @@ def __init__(
         else:
             self.rng_states_tracker = get_rng_state_tracker()
             set_all_rng_states(self.rng_states_tracker.get_states())
-            attention_dropout_ctx = self.rng_states_tracker.fork
+            # Forking only matters if the dropout actually draws from the generator.
+            attention_dropout_ctx = (
+                self.rng_states_tracker.fork if attention_dropout > 0 else nullcontext
+            )
 
         if softmax_scale is None:
             softmax_scale = 1.0 / math.sqrt(

From fb50154cc461711e5b9daae88ee2dcee0f3e71ca Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:14:50 +0200
Subject: [PATCH 31/39] Run checkpointed attention eagerly

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 6bb10aec56..d4cba79ea9 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -213,6 +213,9 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     if call["self"].cp_group is not None:
         return "context parallelism"
 
+    if call.get("checkpoint_core_attention", False):
+        return "activation checkpointing of the attention"
+
     qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
         call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None

From 4b1ebb2806c9e8d3e5907e9d2442e2ec61934a40 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:48:26 +0200
Subject: [PATCH 32/39] Support KV caching under torch.compile

Wrap the KV cache kernels in custom ops: dynamo cannot trace pybind calls
into transformer_engine_torch, so a decoding step used to break the graph.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 69 +++++++++++++--
 .../dot_product_attention/backends.py         | 14 +--
 .../pytorch/attention/inference.py            | 13 +--
 .../pytorch/attention/kv_cache_ops.py         | 88 +++++++++++++++++++
 4 files changed, 167 insertions(+), 17 deletions(-)
 create mode 100644 transformer_engine/pytorch/attention/kv_cache_ops.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 6ae070b70d..504204a6bf 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -569,11 +569,13 @@ def _cfg(
     packed=None,
     interleave_dim=-3,
     share_cu_seqlens=False,
+    kv_cache=False,
 ):
     """A model configuration plus what DotProductAttention is handed it as.
 
     `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at
-    `interleave_dim`, or None for separate q/k/v.
+    `interleave_dim`, or None for separate q/k/v. `kv_cache` runs the call as a
+    decoding step against an InferenceParams KV cache.
     """
     return dict(
         model_config=model_config,
@@ -581,6 +583,7 @@ def _cfg(
         packed=packed,
         interleave_dim=interleave_dim,
         share_cu_seqlens=share_cu_seqlens,
+        kv_cache=kv_cache,
     )
 
 
@@ -632,6 +635,14 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     "sink_softmax_bshd": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
+    # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
+    # cache length divisible by 256.
+    "kv_cache_bshd": _cfg(
+        ModelConfig(
+            2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"
+        ),
+        kv_cache=True,
+    ),
 }
 
 
@@ -646,6 +657,21 @@ def _qkv_layout(spec: dict) -> str:
 
 def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
     config = spec["model_config"]
+    if spec["kv_cache"]:
+        # KV caching addresses the cache by layer number, and a decoding step
+        # runs in inference mode.
+        return (
+            te.DotProductAttention(
+                num_attention_heads=config.num_heads,
+                kv_channels=config.kv_channels,
+                num_gqa_groups=config.num_gqa_groups,
+                qkv_format=spec["qkv_format"],
+                attn_mask_type=config.attn_mask_type,
+                layer_number=1,
+            )
+            .to(dtype=dtype, device="cuda")
+            .eval()
+        )
     return te.DotProductAttention(
         num_attention_heads=config.num_heads,
         kv_channels=config.kv_channels,
@@ -704,6 +730,28 @@ def _shape(s, t, heads, head_dim):
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
+    if spec["kv_cache"]:
+        from collections import OrderedDict
+        from transformer_engine.pytorch.attention import InferenceParams
+
+        inference_params = InferenceParams(
+            max_batch_size=b,
+            max_sequence_length=s_kv,
+            num_heads_kv=g,
+            head_dim_k=d_qk,
+            dtype=dtype,
+            qkv_format=qkv_format,
+        )
+        inference_params.allocate_memory(1)
+        inference_params.pre_step(OrderedDict((i, s_q) for i in range(b)))
+        qkv = [
+            torch.randn(_shape(s_q, t_q, heads, d_qk), dtype=dtype, device="cuda")
+            for heads in (h, g, g)
+        ]
+        # The sequence lengths come from the cache, not from cu_seqlens, and a
+        # decoding step has no gradients to compare.
+        return tuple(qkv), {"inference_params": inference_params}, []
+
     if packed is not None:
         # Declarative packed inputs: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
@@ -746,7 +794,9 @@ def _packed(seqlen, tokens, packed_dim):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) -> None:
+def _skip_unsupported(
+    spec: dict, backend: str, dtype, compiled: bool = True, inference_params=None
+) -> None:
     """Skip what the backend under test cannot run, or -- for a test that
     compiles it -- cannot be compiled."""
     if compiled and backend == "fused":
@@ -755,7 +805,11 @@ def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) ->
         # and the tests below cover it as they do the others.
         pytest.skip("FusedAttention is an eager island and does not compile")
     available, _, _ = get_available_attention_backends(
-        spec["model_config"], dtype, _qkv_layout(spec)
+        spec["model_config"],
+        dtype,
+        _qkv_layout(spec),
+        inference_params=inference_params,
+        is_training=inference_params is None,
     )
     flash_supported, fused_supported, unfused_supported = available
     supported = {
@@ -833,7 +887,8 @@ def _run_and_capture(fn, args, kwargs, grads):
     so that the same inputs can be reused for the other run.
     """
     out = fn(*args, **kwargs).clone()
-    out.sum().backward()
+    if grads:
+        out.sum().backward()
     torch.cuda.synchronize()
     captured = []
     for tensor in grads:
@@ -881,11 +936,11 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     """
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS[config]
-    _skip_unsupported(spec, backend, dtype)
-    _force_dpa_backend(monkeypatch, backend)
-
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    _skip_unsupported(spec, backend, dtype, inference_params=kwargs.get("inference_params"))
+    _force_dpa_backend(monkeypatch, backend)
+
     _compare_compiled_to_eager(
         module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True
     )
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 04bcabe7c4..825e4cd412 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -50,6 +50,10 @@
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
 from transformer_engine.pytorch.distributed import get_distributed_world_size
+from transformer_engine.pytorch.attention.kv_cache_ops import (
+    convert_bshd_to_thd,
+    convert_thd_to_bshd,
+)
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     attn_forward_func_with_cp,
@@ -437,7 +441,7 @@ def _forward(
 
         if qkv_format == "thd_2bshd":
             batch_size = key_layer.shape[0]
-            query_layer = tex.convert_thd_to_bshd(
+            query_layer = convert_thd_to_bshd(
                 query_layer,
                 cu_seqlens_q,
                 batch_size,
@@ -1029,7 +1033,7 @@ def forward(
                 # convert from bshd to thd_2bshd for flash_attn_varlen_func/_with_kvcache;
                 # kernel assumes tensor is contiguous
                 if isinstance(query_layer, Float8Tensor):
-                    query_layer._data = tex.convert_bshd_to_thd(
+                    query_layer._data = convert_bshd_to_thd(
                         query_layer._data,
                         cu_seqlens_q,
                         batch_size * context_len,
@@ -1038,7 +1042,7 @@ def forward(
                         query_layer, data=query_layer._data, shape=query_layer._data.shape
                     )
                 else:
-                    query_layer = tex.convert_bshd_to_thd(
+                    query_layer = convert_bshd_to_thd(
                         query_layer,
                         cu_seqlens_q,
                         batch_size * context_len,
@@ -1287,7 +1291,7 @@ def convert_to_torch_float8(tensor, dtype):
             # all KV caching cases use thd_2bshd for calculation
             # convert results back to bshd from thd_2bshd
             if isinstance(query_layer, Float8Tensor):
-                output._data = tex.convert_thd_to_bshd(
+                output._data = convert_thd_to_bshd(
                     output._data,
                     cu_seqlens_q,
                     batch_size,
@@ -1295,7 +1299,7 @@ def convert_to_torch_float8(tensor, dtype):
                 )
                 output = Float8Tensor.make_like(output, data=output._data, shape=output._data.shape)
             else:
-                output = tex.convert_thd_to_bshd(
+                output = convert_thd_to_bshd(
                     output,
                     cu_seqlens_q,
                     batch_size,
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index 08e50aad8b..b42be244af 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -11,7 +11,10 @@
 import torch
 
 import transformer_engine_torch as tex
-from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
+from transformer_engine.pytorch.attention.kv_cache_ops import (
+    copy_to_kv_cache,
+    QKV_FORMAT_VALUE,
+)
 
 __all__ = ["InferenceParams", "KVCacheManager", "NonPagedKVCacheManager", "PagedKVCacheManager"]
 
@@ -549,7 +552,7 @@ def step(
             batch_size = new_k.shape[1]
             ctx_len = new_k.shape[0]
 
-        tex.copy_to_kv_cache(
+        copy_to_kv_cache(
             new_k,
             new_v,
             k_cache,
@@ -557,7 +560,7 @@ def step(
             self.batch_indices,
             cu_new_seqlens,
             cu_cached_seqlens,
-            QKVFormat[qkv_format],
+            QKV_FORMAT_VALUE[qkv_format],
             batch_size,
             ctx_len,
             self.max_seqlen,
@@ -807,7 +810,7 @@ def step(
             batch_size = new_k.shape[1]
             ctx_len = new_k.shape[0]
 
-        tex.copy_to_kv_cache(
+        copy_to_kv_cache(
             new_k,
             new_v,
             k_cache,
@@ -815,7 +818,7 @@ def step(
             self.page_table,
             cu_new_seqlens,
             cu_cached_seqlens,
-            QKVFormat[qkv_format],
+            QKV_FORMAT_VALUE[qkv_format],
             batch_size,
             ctx_len,
             self.max_seqlen,
diff --git a/transformer_engine/pytorch/attention/kv_cache_ops.py b/transformer_engine/pytorch/attention/kv_cache_ops.py
new file mode 100644
index 0000000000..105e8cac3f
--- /dev/null
+++ b/transformer_engine/pytorch/attention/kv_cache_ops.py
@@ -0,0 +1,88 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""KV cache kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
+
+import torch
+import transformer_engine_torch as tex
+
+from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
+
+# The ops take the format's value rather than the pybind enum: converting the
+# enum inside a traced region makes dynamo recurse until it gives up.
+QKV_FORMAT_VALUE = {name: int(fmt) for name, fmt in QKVFormat.items()}
+_QKV_FORMAT_BY_VALUE = {int(fmt): fmt for fmt in QKVFormat.values()}
+
+
+@torch.library.custom_op(
+    "te_kv_cache::copy_to_kv_cache",
+    mutates_args=("k_cache", "v_cache"),
+    device_types="cuda",
+)
+def copy_to_kv_cache(
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    cu_new_lens: torch.Tensor,
+    cu_cached_lens: torch.Tensor,
+    qkv_format: int,
+    b: int,
+    max_ctx_len: int,
+    max_seq_len: int,
+    max_pages_per_seq: int,
+    is_non_paged: bool,
+) -> None:
+    """Copy new key/value tokens into the KV cache."""
+    tex.copy_to_kv_cache(
+        new_k,
+        new_v,
+        k_cache,
+        v_cache,
+        page_table,
+        cu_new_lens,
+        cu_cached_lens,
+        _QKV_FORMAT_BY_VALUE[qkv_format],
+        b,
+        max_ctx_len,
+        max_seq_len,
+        max_pages_per_seq,
+        is_non_paged,
+    )
+
+
+@copy_to_kv_cache.register_fake
+def _copy_to_kv_cache_fake(*_args, **_kwargs) -> None:
+    return None
+
+
+@torch.library.custom_op("te_kv_cache::convert_bshd_to_thd", mutates_args=(), device_types="cuda")
+def convert_bshd_to_thd(tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int) -> torch.Tensor:
+    """Convert a tensor from bshd to thd."""
+    return tex.convert_bshd_to_thd(tensor, cu_seqlens, t)
+
+
+@convert_bshd_to_thd.register_fake
+def _convert_bshd_to_thd_fake(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int
+) -> torch.Tensor:
+    del cu_seqlens
+    return tensor.new_empty((t, *tensor.shape[2:]))
+
+
+@torch.library.custom_op("te_kv_cache::convert_thd_to_bshd", mutates_args=(), device_types="cuda")
+def convert_thd_to_bshd(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int
+) -> torch.Tensor:
+    """Convert a tensor from thd to bshd."""
+    return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len)
+
+
+@convert_thd_to_bshd.register_fake
+def _convert_thd_to_bshd_fake(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int
+) -> torch.Tensor:
+    del cu_seqlens
+    return tensor.new_empty((b, max_seq_len, *tensor.shape[1:]))

From eb5ad558d37a8fe822832eaae8257d529f195bbf Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:59:46 +0200
Subject: [PATCH 33/39] Compile the fused sbh3d QKV split

Wrap fa_prepare_fwd/bwd like the KV cache kernels, and cover the layout
that reaches them: no configuration exercised it, compiled or eager.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  8 ++++++
 .../{kv_cache_ops.py => custom_ops.py}        | 28 ++++++++++++++++++-
 .../dot_product_attention/backends.py         |  8 ++++--
 .../pytorch/attention/inference.py            |  2 +-
 4 files changed, 41 insertions(+), 5 deletions(-)
 rename transformer_engine/pytorch/attention/{kv_cache_ops.py => custom_ops.py} (70%)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 504204a6bf..d9a5130f40 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -635,6 +635,14 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     "sink_softmax_bshd": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
+    # FlashAttention takes a fused sbhd->bshd split for sbh3d with a head
+    # dimension of 128 and at least 512 tokens, and a plain transpose otherwise.
+    "packed_qkv_sbh3d_fused_split": _cfg(
+        ModelConfig(4, 128, 4, 128, attn_mask_type="no_mask"),
+        "sbhd",
+        packed="qkv",
+        interleave_dim=-2,
+    ),
     # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
     # cache length divisible by 256.
     "kv_cache_bshd": _cfg(
diff --git a/transformer_engine/pytorch/attention/kv_cache_ops.py b/transformer_engine/pytorch/attention/custom_ops.py
similarity index 70%
rename from transformer_engine/pytorch/attention/kv_cache_ops.py
rename to transformer_engine/pytorch/attention/custom_ops.py
index 105e8cac3f..f364633607 100644
--- a/transformer_engine/pytorch/attention/kv_cache_ops.py
+++ b/transformer_engine/pytorch/attention/custom_ops.py
@@ -2,7 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""KV cache kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
+"""Attention kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
 
 import torch
 import transformer_engine_torch as tex
@@ -86,3 +86,29 @@ def _convert_thd_to_bshd_fake(
 ) -> torch.Tensor:
     del cu_seqlens
     return tensor.new_empty((b, max_seq_len, *tensor.shape[1:]))
+
+
+@torch.library.custom_op("te_attention::fa_prepare_fwd", mutates_args=(), device_types="cuda")
+def fa_prepare_fwd(qkvi: torch.Tensor) -> torch.Tensor:
+    """Split interleaved sbh3d QKV into bshd q/k/v."""
+    return tex.fa_prepare_fwd(qkvi)
+
+
+@fa_prepare_fwd.register_fake
+def _fa_prepare_fwd_fake(qkvi: torch.Tensor) -> torch.Tensor:
+    # qkvi is the q view into the packed buffer, and its strides cover all of it.
+    s, b, n, h = qkvi.shape
+    return qkvi.new_empty((3, b, s, n, h))
+
+
+@torch.library.custom_op("te_attention::fa_prepare_bwd", mutates_args=(), device_types="cuda")
+def fa_prepare_bwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Pack bshd gradients back into an interleaved sbh3d buffer."""
+    return tex.fa_prepare_bwd(q, k, v)
+
+
+@fa_prepare_bwd.register_fake
+def _fa_prepare_bwd_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    del k, v
+    b, s, n, h = q.shape
+    return q.new_empty((s, b, n, 3 * h))
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 825e4cd412..4a6ff10b15 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -50,9 +50,11 @@
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
 from transformer_engine.pytorch.distributed import get_distributed_world_size
-from transformer_engine.pytorch.attention.kv_cache_ops import (
+from transformer_engine.pytorch.attention.custom_ops import (
     convert_bshd_to_thd,
     convert_thd_to_bshd,
+    fa_prepare_bwd,
+    fa_prepare_fwd,
 )
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
@@ -777,7 +779,7 @@ def forward(
         # All inputs received are non-contiguous tensors.
         # The `query_layer` tensor is used to access the
         # full memory region of the QKV tensor.
-        qkv = tex.fa_prepare_fwd(query_layer)
+        qkv = fa_prepare_fwd(query_layer)
         q, k, v = split_tensor_along_dim(qkv, 0, 3)
         query_layer = torch.squeeze(q, 0)
         key_layer = torch.squeeze(k, 0)
@@ -792,7 +794,7 @@ def backward(
         dv: torch.Tensor,
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         # pylint: disable=missing-function-docstring
-        dqkv = tex.fa_prepare_bwd(dq, dk, dv)
+        dqkv = fa_prepare_bwd(dq, dk, dv)
         dq, dk, dv = split_tensor_along_dim(dqkv, -1, 3)
         return dq, dk, dv
 
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index b42be244af..0db6fdca2e 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -11,7 +11,7 @@
 import torch
 
 import transformer_engine_torch as tex
-from transformer_engine.pytorch.attention.kv_cache_ops import (
+from transformer_engine.pytorch.attention.custom_ops import (
     copy_to_kv_cache,
     QKV_FORMAT_VALUE,
 )

From 7da1a66d053c3058313aaf05893ef8cfd818a374 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 15:23:54 +0200
Subject: [PATCH 34/39] Run FlashAttention 4 eagerly

Version 4 registers no custom ops, unlike 2 and 3: it builds its kernels
through the CUTLASS DSL as it runs. Disable dynamo on its entry points, so
only the calls that reach it become eager islands.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/backends.py  | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 4a6ff10b15..a5d2b05141 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -174,11 +174,16 @@
     flash_attn_varlen_func_v4 = None
 else:
     from flash_attn.cute.interface import (  # pylint: disable=ungrouped-imports,no-name-in-module
-        flash_attn_func as flash_attn_func_v4,
-        flash_attn_varlen_func as flash_attn_varlen_func_v4,
+        flash_attn_func as _flash_attn_func_v4,
+        flash_attn_varlen_func as _flash_attn_varlen_func_v4,
         _validate_head_dims as _fa4_validate_head_dims,
     )
 
+    # Unlike versions 2 and 3, FlashAttention 4 registers no custom ops: it builds
+    # its kernels through the CUTLASS DSL as it runs. Keep it an eager island.
+    flash_attn_func_v4 = no_torch_dynamo()(_flash_attn_func_v4)
+    flash_attn_varlen_func_v4 = no_torch_dynamo()(_flash_attn_varlen_func_v4)
+
     fa_utils.v4_validate_head_dims = _fa4_validate_head_dims
     fa_utils.set_flash_attention_4_params()
 

From d1beb9b91cc16c8d80804da4c92d4a20ea2617b7 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 15:49:15 +0200
Subject: [PATCH 35/39] Pin the sequence lengths backend selection bakes in

A second sequence length makes dynamo's automatic dynamic turn max_seqlen
symbolic, and _get_fused_attn_backend, being assume_constant_result, cannot
take a symbolic argument -- so the call died once the length changed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py  | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d4cba79ea9..383eb3a457 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -12,6 +12,7 @@
 
 import torch
 import torch.nn.functional as F
+from torch.fx.experimental.symbolic_shapes import guard_scalar
 from torch.nn.parameter import Parameter
 
 from transformer_engine.common.recipe import (
@@ -1535,6 +1536,9 @@ def forward(
                     batch_size = query_layer.shape[0]
                     max_seqlen_q = query_layer.shape[1] if max_seqlen_q is None else max_seqlen_q
                     max_seqlen_kv = key_layer.shape[1] if max_seqlen_kv is None else max_seqlen_kv
+                # Backend selection bakes these in, which it cannot do symbolically.
+                max_seqlen_q = guard_scalar(max_seqlen_q)
+                max_seqlen_kv = guard_scalar(max_seqlen_kv)
             if qkv_format == "thd":
                 assert all(
                     len(x.shape) == 3 for x in (query_layer, key_layer, value_layer)

From 4af8c30341d0bae89467b889ea1809d478c6d082 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 16:03:17 +0200
Subject: [PATCH 36/39] Capture CUDA graphs with a KV cache

The cache is mutated in place, so inductor declined to capture until its
buffers were marked static.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py               | 5 +++--
 transformer_engine/pytorch/attention/inference.py | 6 ++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d9a5130f40..d5138ec5ac 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -972,11 +972,12 @@ def test_dpa_torch_compile_around_fused(monkeypatch):
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
-def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
+@pytest.mark.parametrize("config", ["self_bshd_causal", "kv_cache_bshd"])
+def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS[config]
     _force_dpa_backend(monkeypatch, backend)
 
     module = _make_dpa(spec, dtype)
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index 0db6fdca2e..ae9e02dd7d 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -462,6 +462,9 @@ def allocate_memory(self, layer_number):
             dtype=self.dtype,
             device=torch.cuda.current_device(),
         )
+        # Mutated in place, so CUDA graphs are only captured if the addresses are fixed.
+        torch._dynamo.mark_static_address(k_cache)
+        torch._dynamo.mark_static_address(v_cache)
         self.cache[layer_number] = (k_cache, v_cache)
 
     def pre_step(
@@ -659,6 +662,9 @@ def allocate_memory(self, layer_number):
             dtype=self.dtype,
             device=torch.cuda.current_device(),
         )
+        # Mutated in place, so CUDA graphs are only captured if the addresses are fixed.
+        torch._dynamo.mark_static_address(k_cache)
+        torch._dynamo.mark_static_address(v_cache)
         self.cache[layer_number] = (k_cache, v_cache)
 
     def print_cache(self):

From bd94fff1c155700d5d2397a6ee3fd4f020290377 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 16:16:59 +0200
Subject: [PATCH 37/39] Test generation against a KV cache

A prefill and three single-token steps, each compared to eager, with and
without CUDA graphs: the cache carries state, so a step that updates it
wrongly only shows in the step after.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 74 +++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d5138ec5ac..ab7cddaf54 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -1000,6 +1000,80 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config):
     assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("paged", [False, True], ids=["non_paged", "paged"])
+@pytest.mark.parametrize("cuda_graphs", [False, True], ids=["default", "cudagraphs"])
+def test_dpa_torch_compile_kv_cache_decoding(monkeypatch, backend, paged, cuda_graphs):
+    """Generation against a KV cache, one prefill and three single-token steps.
+
+    The cache carries state from one step to the next, so a step that updates it
+    wrongly is only visible in the step after -- hence comparing every step, and
+    a cache of its own for each of the two runs.
+    """
+    from collections import OrderedDict
+    from transformer_engine.pytorch.attention import InferenceParams
+
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS["kv_cache_bshd"]
+    config = spec["model_config"]
+    b, ctx_len = config.batch_size, config.max_seqlen_q
+    h, g, d = config.num_heads, config.num_gqa_groups, config.head_dim_qk
+
+    def make_cache():
+        kwargs = dict(
+            max_batch_size=b,
+            max_sequence_length=config.max_seqlen_kv,
+            num_heads_kv=g,
+            head_dim_k=d,
+            dtype=dtype,
+            qkv_format=spec["qkv_format"],
+        )
+        if paged:
+            # One page per sequence, and FlashAttention 2 wants it divisible by 256.
+            kwargs.update(is_paged=True, page_size=config.max_seqlen_kv, total_num_pages=b)
+        inference_params = InferenceParams(**kwargs)
+        inference_params.allocate_memory(1)
+        return inference_params
+
+    _skip_unsupported(spec, backend, dtype, inference_params=make_cache())
+
+    gen = torch.Generator(device="cuda").manual_seed(1234)
+    steps = [
+        [
+            torch.randn(b, seqlen, heads, d, device="cuda", dtype=dtype, generator=gen)
+            for heads in (h, g, g)
+        ]
+        for seqlen in (ctx_len, 1, 1, 1)
+    ]
+
+    _force_dpa_backend(monkeypatch, backend)
+    module = _make_dpa(spec, dtype)
+
+    def generate(fn):
+        inference_params = make_cache()
+        outputs = []
+        with torch.no_grad():
+            for args in steps:
+                inference_params.pre_step(OrderedDict((i, args[0].shape[1]) for i in range(b)))
+                _force_dpa_backend(monkeypatch, backend)
+                # Cloned because a CUDA graph replay overwrites what it returned.
+                outputs.append(fn(*args, inference_params=inference_params).clone())
+        return outputs
+
+    eager = generate(module)
+
+    torch._dynamo.reset()
+    counters.clear()
+    compile_kwargs = {"mode": "reduce-overhead"} if cuda_graphs else {}
+    compiled = generate(torch.compile(module, fullgraph=True, **compile_kwargs))
+
+    _assert_dpa_backend(backend)
+    for out, ref in zip(compiled, eager):
+        _assert_matches_eager(out, ref, backend, dtype)
+    if cuda_graphs:
+        assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 @pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"])
 def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):

From 688e10c0ac0ecb1299feebf1e3905de9fe9a684c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 4 Aug 2026 14:20:22 +0000
Subject: [PATCH 38/39] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/pytorch/test_torch_compile.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ab7cddaf54..5876fe857f 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -646,9 +646,7 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
     # cache length divisible by 256.
     "kv_cache_bshd": _cfg(
-        ModelConfig(
-            2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"
-        ),
+        ModelConfig(2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"),
         kv_cache=True,
     ),
 }

From 43eba11be959c3405df2d3618985233b0f170ac0 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 21:55:38 +0200
Subject: [PATCH 39/39] Drop the imports the custom ops made unused

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/backends.py          | 1 -
 transformer_engine/pytorch/attention/inference.py                | 1 -
 2 files changed, 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index a5d2b05141..71afed25c5 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -15,7 +15,6 @@
 
 import torch
 import torch.nn.functional as F
-import transformer_engine_torch as tex
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
     split_tensor_along_dim,
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index ae9e02dd7d..f98d1693a6 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -10,7 +10,6 @@
 
 import torch
 
-import transformer_engine_torch as tex
 from transformer_engine.pytorch.attention.custom_ops import (
     copy_to_kv_cache,
     QKV_FORMAT_VALUE,