[PyTorch][torch.compile] Make get_attention_backend traceable without graph breaks

tests/pytorch/test_torch_compile.py

@@ -389,6 +389,111 @@ def fn(inp):
     out.sum().backward()
 
 
+# ---------------------------------------------------------------------------
+# get_attention_backend under torch.compile
+# ---------------------------------------------------------------------------
+
+
+# Scalars in AttentionParams must stay concrete: assume_constant_result cannot
+# convert symbolic scalars (dynamo's automatic dynamic would make changed ints
+# symbolic on recompilation), so pin them static explicitly.
+@torch._dynamo.config.patch(specialize_int=True, specialize_float=True, recompile_limit=32)
+def test_get_attention_backend_traceable(monkeypatch):
+    """get_attention_backend must trace under torch.compile(fullgraph=True)
+    without graph breaks. The compiled selection must stay consistent with
+    eager when NVTE_* env vars flip (dynamo guards on os.environ) and when
+    attention params change, and the baked tex.get_fused_attn_backend result
+    must drive the selection."""
+    from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils
+
+    def fn(x, params):
+        (
+            use_flash_attention,
+            _,
+            use_fused_attention,
+            fused_attention_backend,
+            use_unfused_attention,
+            _,
+        ) = dpa_utils.get_attention_backend(params)
+        # Encode the full selection (enabled backends + fused sub-backend) in
+        # the tensor value: without a tensor op dynamo skips the frame entirely
+        # (nothing gets compiled or guarded), and the output makes compiled vs
+        # eager selection directly comparable.
+        return (
+            x
+            + (1 if use_flash_attention else 0)
+            + (2 if use_fused_attention else 0)
+            + (4 if use_unfused_attention else 0)
+            + (8 * int(fused_attention_backend) if fused_attention_backend is not None else 0)
+        )
+
+    # Dynamo only guards os.environ entries that exist at trace time (reads of
+    # absent keys are not guarded yet), so set the vars explicitly.
+    for env_var, value in (
+        ("NVTE_FLASH_ATTN", "1"),
+        ("NVTE_FUSED_ATTN", "1"),
+        ("NVTE_UNFUSED_ATTN", "1"),
+        ("NVTE_FP8_DPA_BWD", "1"),
+        ("NVTE_DPA_FP8CS_O_in_F16", "1"),
+        ("NVTE_DPA_FP8_RECIPE", ""),
+        ("NVTE_UnfusedDPA_Emulate_FP8", "0"),
+    ):
+        monkeypatch.setenv(env_var, value)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+    x = torch.zeros(8, device="cuda")
+    params = dpa_utils.AttentionParams()
+
+    torch.testing.assert_close(compiled(x, params), fn(x, params))
+
+    # Flip env vars one by one: the compiled function must recompile (guards
+    # on os.environ) and keep matching eager.
+    for env_var in ("NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN", "NVTE_FLASH_ATTN"):
+        monkeypatch.setenv(env_var, "0")
+        torch.testing.assert_close(compiled(x, params), fn(x, params))
+        monkeypatch.setenv(env_var, "1")
+
+    # FP8 attention (fp8_dpa recipe): covers the FP8-only branch (run_config
+    # env reads, recipe filters, get_fp8_te_dtype). Flipping an FP8-only env
+    # var (emulation enables UnfusedDotProductAttention) must recompile too.
+    fp8_params = dpa_utils.AttentionParams(
+        fp8=True, fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}
+    )
+    torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params))
+    monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1")
+    torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params))
+    monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "0")
+
+    # Changing attention params (ints, layout string, dtype) must recompile
+    # and keep matching eager, still with no graph break.
+    for changed_params in (
+        dpa_utils.AttentionParams(head_dim_qk=128, head_dim_v=128),
+        dpa_utils.AttentionParams(max_seqlen_q=512, max_seqlen_kv=512),
+        dpa_utils.AttentionParams(qkv_layout="bshd_bshd_bshd"),
+        dpa_utils.AttentionParams(qkv_dtype=torch.float16),
+    ):
+        torch.testing.assert_close(compiled(x, changed_params), fn(x, changed_params))
+
+    # The baked probe result must drive the selection: report no fused
+    # sub-backend and expect UnfusedDotProductAttention (flash disabled, so the
+    # outcome is deterministic). Use a fresh frame: already-compiled frames
+    # keep the previously baked constant (assume_constant_result installs no
+    # guard on the wrapped function).
+    monkeypatch.setenv("NVTE_FLASH_ATTN", "0")
+    monkeypatch.setattr(
+        dpa_utils.tex,
+        "get_fused_attn_backend",
+        lambda *args: dpa_utils.FusedAttnBackend["No_Backend"],
+    )
+
+    def fn_no_backend(x, params):
+        return fn(x, params)
+
+    compiled_no_backend = torch.compile(fn_no_backend, fullgraph=True)
+    torch.testing.assert_close(compiled_no_backend(x, params), x + 4.0)
+
+
 # ---------------------------------------------------------------------------
 # Value-opaque quantizers
 # ---------------------------------------------------------------------------

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -1995,7 +1995,7 @@ def forward(
         attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
         window_size: Optional[Tuple[int, int]] = None,
         bottom_right_diagonal: Optional[bool] = None,
-        fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend,
+        fused_attention_backend: FusedAttnBackend = FusedAttnBackend["No_Backend"],
         core_attention_bias_type: str = "no_bias",
         core_attention_bias: Optional[torch.Tensor] = None,
         fast_zero_fill: bool = True,
@@ -2017,7 +2017,7 @@ def forward(
     ) -> torch.Tensor:
         """fused attention fprop"""
         assert (
-            fused_attention_backend != tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend
+            fused_attention_backend != FusedAttnBackend["No_Backend"]
         ), "No fused attention backend supports this input combination!"
         assert all(
             x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, QuantizedTensorStorage)
@@ -2110,15 +2110,15 @@ def forward(
         use_FAv2_bwd = (
             self.use_FAv2_bwd
             and (core_attention_bias_type == "no_bias")
-            and (fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen)
+            and (fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"])
         )
 
         if fp8:
             fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
             if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None:
                 fp8_recipe = fp8_meta["local_recipes"][0]
-            assert fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8, (
-                f"cuDNN attention sub-backend {int(tex.NVTE_Fused_Attn_Backend.NVTE_FP8)}"
+            assert fused_attention_backend == FusedAttnBackend["FP8"], (
+                f"cuDNN attention sub-backend {int(FusedAttnBackend['FP8'])}"
                 " is required for FP8 attention!"
             )
             assert fp8_meta is not None, "FP8 metadata fp8_meta is required for FP8 attention!"
@@ -2138,8 +2138,7 @@ def forward(
 
         if context_parallel:
             assert (
-                fp8
-                or fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
+                fp8 or fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
             ), f"{fused_attention_backend} does not work with context parallelism!"
             assert core_attention_bias_type not in [
                 "alibi"

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -3180,7 +3180,7 @@ def forward(
         fp8_meta_kwargs = {}
         if fp8:
             assert use_fused_attention, "FP8 is only supported with FusedAttention backend!"
-            fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8
+            fused_attn_backend = FusedAttnBackend["FP8"]
             if not is_input_fp8 and not fp8_recipe.mxfp8():
                 q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize(
                     qkv_layout, q, k, v, QKV_quantizer
@@ -3190,7 +3190,7 @@ def forward(
             fp8_meta_kwargs["s_quantizer"] = S_quantizer
             fp8_meta_kwargs["o_quantizer"] = O_quantizer
         elif use_fused_attention:
-            fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
+            fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"]
         orig_q_shape, _, orig_v_shape = q.shape, k.shape, v.shape
         orig_o_shape = orig_q_shape[:-1] + orig_v_shape[-1:]
 
@@ -3901,14 +3901,14 @@ def backward(ctx, dout, *_args):
                             softmax_lse_per_step[i],
                             rng_states[i],
                         ]
-                        fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
+                        fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"]
                         fp8_meta_kwargs = {}
                         new_qkv_layout = ctx.qkv_layout
                         do_format = ctx.o_format
                         qkv_scale_inv_format = None
                         do_scale_inv_format = None
                         if ctx.fp8:
-                            fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8
+                            fused_attn_backend = FusedAttnBackend["FP8"]
                             fp8_meta_kwargs["s_quantizer"] = ctx.S_quantizer
                             fp8_meta_kwargs["dp_quantizer"] = ctx.dP_quantizer
                             fp8_meta_kwargs["dqkv_quantizer"] = ctx.dQKV_quantizer

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -5,6 +5,7 @@
 """
 Utils/Helper classes and methods for attention
 """
+
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -324,6 +325,53 @@ def __eq__(self, other):
         return True
 
 
+class _NoOpLogger:
+    """
+    Stand-in for the "DotProductAttention" logger used when get_attention_backend
+    is traced by torch.compile. logging.Logger methods are not traceable by dynamo
+    (they cause graph breaks), while this class's no-op methods are inlined away.
+    """
+
+    def debug(self, *args, **kwargs):
+        """No-op."""
+
+    def warning(self, *args, **kwargs):
+        """No-op."""
+
+
+_no_op_logger = _NoOpLogger()
+
+
+@torch.compiler.assume_constant_result
+def _get_fused_attn_backend(
+    is_training,
+    q_type,
+    kv_type,
+    qkv_layout,
+    bias_type,
+    attn_mask_type,
+    softmax_type,
+    *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(
+        tex.get_fused_attn_backend(
+            is_training,
+            q_type,
+            kv_type,
+            QKVLayout[qkv_layout],
+            AttnBiasType[bias_type],
+            AttnMaskType[attn_mask_type],
+            SoftmaxType[softmax_type],
+            *args,
+        )
+    )
+
+
 def get_attention_backend(
     attention_params: AttentionParams = None,
 ):
@@ -340,7 +388,7 @@ 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 : tex.NVTE_Fused_Attn_Backend
+    fused_attention_backend : FusedAttnBackend
         If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
     use_unfused_attention : bool
         Whether the `UnfusedDotProductAttention` backend has been selected.
@@ -387,17 +435,29 @@ def get_attention_backend(
     has_score_mod = attention_params.has_score_mod
     has_score_mod_bprop = attention_params.has_score_mod_bprop
 
+    # NOTE: environment variables in this function are read with
+    # os.environ.get, NOT os.getenv, on purpose: dynamo installs guards on
+    # os.environ reads (so changing an NVTE_* variable triggers recompilation
+    # under torch.compile), while os.getenv reads are unguarded and would bake
+    # stale values into compiled graphs. New code must follow suit.
+
     # Run config
-    logger = logging.getLogger("DotProductAttention")
-    logger.setLevel(AttentionLogging._log_level)
-    if not logger.hasHandlers():
-        logger.addHandler(AttentionLogging._stream_handler)
+    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
+    else:
+        logger = logging.getLogger("DotProductAttention")
+        logger.setLevel(AttentionLogging._log_level)
+        if not logger.hasHandlers():
+            logger.addHandler(AttentionLogging._stream_handler)
     device_compute_capability = get_device_compute_capability()
     cudnn_version = get_cudnn_version()
     run_config = {
         "transformer_engine_version": te.__version__,
-        "compute_capability": "sm"
-        + str(10 * device_compute_capability[0] + device_compute_capability[1]),
+        "compute_capability": (
+            "sm" + str(10 * device_compute_capability[0] + device_compute_capability[1])
+        ),
         "cuda_version": torch.version.cuda,
         "flash_attn_version": (
             str(FlashAttentionUtils.version)
@@ -423,27 +483,27 @@ def get_attention_backend(
     # Add FP8 environment variables to config
     if fp8:
         # all FP8 recipes: 1: (FP8 fwd, FP8 bwd), 0: (FP8 fwd, F16 bwd)
-        run_config["NVTE_FP8_DPA_BWD"] = int(os.getenv("NVTE_FP8_DPA_BWD", "1"))
+        run_config["NVTE_FP8_DPA_BWD"] = int(os.environ.get("NVTE_FP8_DPA_BWD", "1"))
         # Float8CurrentScaling: 1: use F16 O in bwd, 0: use FP8 O in bwd
-        run_config["NVTE_DPA_FP8CS_O_in_F16"] = int(os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1"))
+        run_config["NVTE_DPA_FP8CS_O_in_F16"] = int(os.environ.get("NVTE_DPA_FP8CS_O_in_F16", "1"))
         # switch recipe to "F16", "DelayedScaling", or "Float8CurrentScaling"
-        _dpa_fp8_recipe = os.getenv("NVTE_DPA_FP8_RECIPE", "")
+        _dpa_fp8_recipe = os.environ.get("NVTE_DPA_FP8_RECIPE", "")
         run_config["NVTE_DPA_FP8_RECIPE"] = _dpa_fp8_recipe
         if _dpa_fp8_recipe != "":
             # config new recipe if switched
-            run_config["NVTE_DPA_FP8_FORMAT"] = os.getenv("NVTE_DPA_FP8_FORMAT", "HYBRID")
-            run_config["NVTE_DPA_FP8DS_AMAX_ALGO"] = os.getenv(
+            run_config["NVTE_DPA_FP8_FORMAT"] = os.environ.get("NVTE_DPA_FP8_FORMAT", "HYBRID")
+            run_config["NVTE_DPA_FP8DS_AMAX_ALGO"] = os.environ.get(
                 "NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent"
             )
             run_config["NVTE_DPA_FP8DS_AMAX_HISTLEN"] = int(
-                os.getenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1")
+                os.environ.get("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1")
             )
             run_config["NVTE_DPA_FP8DS_REDUCE_AMAX"] = int(
-                os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1")
+                os.environ.get("NVTE_DPA_FP8DS_REDUCE_AMAX", "1")
             )
         # UnfusedDotProductAttention: 1: allow FP8 emulation, 0: do not allow
         run_config["NVTE_UnfusedDPA_Emulate_FP8"] = int(
-            os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0")
+            os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0")
         )
     logger.debug("Running with config=%s", run_config)
 
@@ -454,13 +514,13 @@ def get_attention_backend(
     qkv_format, q_format, kv_format = get_qkv_format(qkv_layout, inference_params)
 
     # Filter: Environment variables
-    use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1"))
+    use_flash_attention = int(os.environ.get("NVTE_FLASH_ATTN", "1"))
     use_flash_attention_2 = use_flash_attention
     use_flash_attention_3 = use_flash_attention
     use_flash_attention_4 = use_flash_attention
     flash_attention_backend = None
-    use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1"))
-    use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1"))
+    use_fused_attention = int(os.environ.get("NVTE_FUSED_ATTN", "1"))
+    use_unfused_attention = int(os.environ.get("NVTE_UNFUSED_ATTN", "1"))
     if not use_flash_attention_2 and FlashAttentionUtils.is_installed:
         logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0")
     if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
@@ -585,7 +645,8 @@ def _disable_all_flash_attention() -> None:
             use_flash_attention_3 = False
         if use_unfused_attention:
             allow_emulation = (
-                os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode()
+                os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1"
+                or is_in_onnx_export_mode()
             )
             if not allow_emulation:
                 logger.debug("Disabling UnfusedDotProductAttention for FP8 attention")
@@ -1363,14 +1424,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         if fp8 and fp8_meta["recipe"].fp8_dpa:
             q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True)
             kv_type = q_type
-        fused_attention_backend = tex.get_fused_attn_backend(
+        # NOTE: under torch.compile the numeric args below must not be symbolic
+        # (assume_constant_result requires concrete values); ints/floats made
+        # dynamic by automatic dynamic currently graph break here.
+        fused_attention_backend = _get_fused_attn_backend(
             is_training,
             q_type,
             kv_type,
-            QKVLayout[qkv_layout],
-            AttnBiasType[fu_core_attention_bias_type],
-            AttnMaskType[attn_mask_type],
-            SoftmaxType[softmax_type],
+            qkv_layout,
+            fu_core_attention_bias_type,
+            attn_mask_type,
+            softmax_type,
             attention_dropout,
             num_heads,
             num_gqa_groups,