diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8ebce563ce..7422b0290a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 785e438cda..091c6fd07d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/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" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index a62ca73187..995dd2e90a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7be94a6fa1..b60d9a11a2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/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, diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9c22c56bd1..046019ee58 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -5,6 +5,7 @@ """Python interface for fused attention extensions""" import math +from enum import IntEnum from typing import Tuple, List, Union, Optional import torch import transformer_engine_torch as tex @@ -97,11 +98,65 @@ "learnable": NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, } -FusedAttnBackend = { - "F16_arbitrary_seqlen": NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - "FP8": NVTE_Fused_Attn_Backend.NVTE_FP8, - "No_Backend": NVTE_Fused_Attn_Backend.NVTE_No_Backend, -} + +class FusedAttnBackend(IntEnum): + """Fused attention sub-backends. + + This is the canonical fused-attention backend enum for + ``transformer_engine.pytorch``. It mirrors the backend + ``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. + """ + + No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) + F16_arbitrary_seqlen = int(NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) + FP8 = int(NVTE_Fused_Attn_Backend.NVTE_FP8) + + @classmethod + def cast( + cls, backend: "Union[FusedAttnBackend, NVTE_Fused_Attn_Backend]" + ) -> "FusedAttnBackend": + """Normalize a backend value to the canonical ``FusedAttnBackend`` member. + + The pybind ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` enum is + accepted as input for backward compatibility and mapped to the matching + ``FusedAttnBackend`` member. + """ + if isinstance(backend, cls): + return backend + return cls(int(backend)) + + def __eq__(self, other: object) -> bool: + # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` + # is a pybind11 enum. Compare by integer value so the two enums stay + # equivalent regardless of the pybind11 version (the pybind ``__eq__`` + # handles the reverse order). + if isinstance(other, NVTE_Fused_Attn_Backend): + return int(self) == int(other) + return int.__eq__(self, other) + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + def __hash__(self) -> int: + return int.__hash__(self) + + +# Fail fast at import time if a new enumerator is added on the C++ side +# without being mirrored above. +assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( + "FusedAttnBackend in python is out of sync with" + " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." + " Please make sure TE C++ and python are in sync." +) BACKEND_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 @@ -124,7 +179,7 @@ def fused_attn_fwd( k: torch.Tensor, v: torch.Tensor, fake_dtype: torch.dtype, - fused_attention_backend: tex.NVTE_Fused_Attn_Backend, + fused_attention_backend: FusedAttnBackend, attn_bias: torch.Tensor = None, cu_seqlens_q_padded: torch.Tensor = None, cu_seqlens_kv_padded: torch.Tensor = None, @@ -176,7 +231,7 @@ def fused_attn_fwd( fake_dtype : DType data type of Q, K and V - in case of high precision, fake dtype in case of FP8; in torch.dtype - fused_attention_backend : tex.NVTE_Fused_Attn_Backend + fused_attention_backend : FusedAttnBackend please see FusedAttention module for details on supported backends. attn_bias : torch.Tensor, default = None input tensor Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; @@ -285,6 +340,8 @@ def fused_attn_fwd( f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." ) + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) if fused_attention_backend == FusedAttnBackend["No_Backend"]: raise ValueError( "Fused attention does not support this input combination:" @@ -397,7 +454,7 @@ def fused_attn_bwd( d_o: torch.Tensor, fake_dtype: torch.dtype, aux_ctx_tensors: List[torch.Tensor], - fused_attention_backend: tex.NVTE_Fused_Attn_Backend, + fused_attention_backend: FusedAttnBackend, cu_seqlens_q_padded: torch.Tensor = None, cu_seqlens_kv_padded: torch.Tensor = None, s_quantizer: Quantizer = None, @@ -453,7 +510,7 @@ def fused_attn_bwd( aux_ctx_tensors : List[torch.Tensor] auxiliary output tensors of the forward pass when its is_training is True, e.g. aux_ctx_tensors = [S, Max, rng_state] - fused_attention_backend : tex.NVTE_Fused_Attn_Backend + fused_attention_backend : FusedAttnBackend please see FusedAttention module for details on supported backends. cu_seqlens_q_padded : torch.Tensor, default = None cumulative sequence offsets for Q; shape [batch_size + 1] @@ -538,6 +595,8 @@ def fused_attn_bwd( d = q.size(-1) attn_scale = 1.0 / math.sqrt(d) + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) if fused_attention_backend == FusedAttnBackend["No_Backend"]: raise ValueError( "Fused attention backward does not support this input combination:" diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index ffbfbc1fdd..39162f8311 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -83,6 +83,7 @@ def _get_device_compute_capability(device: torch.device) -> Tuple[int, int]: return (props.major, props.minor) +@torch.compiler.assume_constant_result def get_device_compute_capability() -> Tuple[int, int]: """CUDA compute capability of current GPU""" return _get_device_compute_capability(torch.cuda.current_device()) @@ -663,7 +664,7 @@ def is_non_tn_fp8_gemm_supported() -> bool: @functools.lru_cache(maxsize=None) -def get_cudnn_version() -> Tuple[int, int, int]: +def _get_cudnn_version() -> Tuple[int, int, int]: """Runtime cuDNN version (major, minor, patch)""" import transformer_engine.pytorch.cpp_extensions as ext @@ -674,6 +675,12 @@ def get_cudnn_version() -> Tuple[int, int, int]: return (major, minor, patch) +@torch.compiler.assume_constant_result +def get_cudnn_version() -> Tuple[int, int, int]: + """Runtime cuDNN version (major, minor, patch)""" + return _get_cudnn_version() + + def canonicalize_device(device: Optional[torch.device | str]) -> torch.device: """Canonicalize PyTorch device