[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends

transformer_engine/pytorch/attention/custom_ops.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Attention 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:]))
+
+
+@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))

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,
@@ -50,6 +49,12 @@
 )
 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.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 (
     attn_forward_func_with_cp,
@@ -168,11 +173,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()
 
@@ -437,7 +447,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,
@@ -773,7 +783,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)
@@ -788,11 +798,23 @@ 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
 
 
+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
+    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
@@ -1017,7 +1039,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,
@@ -1026,7 +1048,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,
@@ -1122,12 +1144,12 @@ 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_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,
                         )
+                        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:
@@ -1138,8 +1160,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 = _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
                         fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv
                     output = func(
@@ -1274,15 +1297,15 @@ 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,
                     context_len,
                 )
                 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,

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 (
@@ -196,6 +197,46 @@ 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(call: Dict[str, Any]) -> Optional[str]:
+    """Why this DotProductAttention call has to run outside the graph, or None.
+
+    `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.
+    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 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
+    ):
+        # 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 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"
+
+    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
+
+
 def _unpack_packed_qkv(
     qkv_layer: Optional[torch.Tensor],
     kv_layer: Optional[torch.Tensor],
@@ -545,7 +586,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(
@@ -715,6 +759,11 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
+        qstate = FP8GlobalStateManager.quantization_state
+        if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
+            super().init_fp8_metadata(num_gemms=num_gemms)
+            return
+
         # global recipe set in autocast()
         fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
         if fp8_recipe.custom():
@@ -1090,7 +1139,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @no_torch_dynamo(recursive=False)
+    @no_torch_dynamo(when=_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
@@ -1487,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)
@@ -1820,18 +1872,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
+                    # 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:
-                        self.logger.info(
+                        logger.info(
                             "Running with FlashAttention backend (version %s)",
                             flash_attention_backend,
                         )
                     elif use_fused_attention:
-                        self.logger.info(
+                        logger.info(
                             "Running with FusedAttention backend (sub-backend %s)",
                             int(fused_attention_backend),
                         )
                     elif use_unfused_attention:
-                        self.logger.info("Running with UnfusedDotProductAttention backend")
+                        logger.info("Running with UnfusedDotProductAttention backend")
                 else:
                     use_flash_attention = _attention_backends["use_flash_attention"]
                     flash_attention_backend = _attention_backends["flash_attention_backend"]