From 17f80bd2bfa7f4308131c3e6bf15c538316ee859 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 12:24:57 +0200 Subject: [PATCH 1/9] Make get_attention_backend traceable by torch.compile without graph breaks - Read NVTE_* env vars via os.environ.get instead of os.getenv so dynamo installs guards on the values (os.getenv reads are not guarded and would bake stale backend selections into compiled graphs). - Wrap tex.get_fused_attn_backend in a torch.compiler.assume_constant_result helper so the pybind call does not graph-break. - Mark get_device_compute_capability/get_cudnn_version with assume_constant_result for the same reason. - Use a no-op logger when compiling (logging.Logger methods graph-break) and skip debug-log blocks that call int()/str() on pybind enums. - Add test in tests/pytorch/test_torch_compile.py checking fullgraph=True tracing and recompilation on NVTE_* env var changes. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 49 +++++++ .../attention/dot_product_attention/utils.py | 126 ++++++++++++------ transformer_engine/pytorch/utils.py | 9 +- 3 files changed, 140 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..ed2e536ecc 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -384,3 +384,52 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# get_attention_backend under torch.compile +# --------------------------------------------------------------------------- + + +def test_get_attention_backend_traceable(monkeypatch): + """get_attention_backend must trace under torch.compile(fullgraph=True), + i.e. without any graph break, and NVTE_* env var reads must be guarded so + that changing an env var triggers recompilation instead of silently + reusing a stale backend selection.""" + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + + attention_params = dpa_utils.AttentionParams() + + def fn(x): + ( + use_flash_attention, + _, + use_fused_attention, + _, + use_unfused_attention, + _, + ) = dpa_utils.get_attention_backend(attention_params) + return ( + x + + (1 if use_flash_attention else 0) + + (2 if use_fused_attention else 0) + + (4 if use_unfused_attention 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 in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): + monkeypatch.setenv(env_var, "1") + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + x = torch.zeros(8, device="cuda") + torch.testing.assert_close(compiled(x), fn(x)) + + # 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), fn(x)) + monkeypatch.setenv(env_var, "1") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9913b78dfc..797eef3017 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -324,6 +324,34 @@ 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(*args): + """ + Wrapper for tex.get_fused_attn_backend. The result only depends on the + attention configuration (not on tensor values), so it is marked with + assume_constant_result to keep get_attention_backend traceable by + torch.compile without a graph break on the C extension call. + """ + return tex.get_fused_attn_backend(*args) + + def get_attention_backend( attention_params: AttentionParams = None, ): @@ -388,10 +416,15 @@ def get_attention_backend( has_score_mod_bprop = attention_params.has_score_mod_bprop # 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 = { @@ -423,27 +456,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 +487,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 +618,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") @@ -1359,7 +1393,7 @@ 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( + fused_attention_backend = _get_fused_attn_backend( is_training, q_type, kv_type, @@ -1385,11 +1419,13 @@ 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 elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: - logger.debug( - "Disabling FusedAttention for score_mod because sub-backend %s is not " - "F16/BF16 arbitrary-seqlen", - int(fused_attention_backend), - ) + if not torch.compiler.is_compiling(): + # int() on a pybind enum is not traceable by dynamo + logger.debug( + "Disabling FusedAttention for score_mod because sub-backend %s is not " + "F16/BF16 arbitrary-seqlen", + int(fused_attention_backend), + ) use_fused_attention = False fused_attention_backend = None # Filter: Determinism @@ -1521,19 +1557,21 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_flash_attention_4: flash_attention_backend = FlashAttentionUtils.fa4_version - logger.debug( - "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s," - " UnfusedDotProductAttention=%s}", - bool(available_backends[0]), - (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), - bool(available_backends[1]), - ( - f" (sub-backend {int(fused_attention_backend)})" - if fused_attention_backend is not None - else "" - ), - bool(available_backends[2]), - ) + if not torch.compiler.is_compiling(): + # int()/str() on the backend objects are not traceable by dynamo + logger.debug( + "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s," + " UnfusedDotProductAttention=%s}", + bool(available_backends[0]), + (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), + bool(available_backends[1]), + ( + f" (sub-backend {int(fused_attention_backend)})" + if fused_attention_backend is not None + else "" + ), + bool(available_backends[2]), + ) # Select FusedAttention for performance if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0): @@ -1549,14 +1587,16 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_unfused_attention = False elif use_fused_attention: use_unfused_attention = False - selected_backend = "NoBackend" - if use_flash_attention: - selected_backend = f"FlashAttention ({str(flash_attention_backend)})" - elif use_fused_attention: - selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" - elif use_unfused_attention: - selected_backend = "UnfusedDotProductAttention" - logger.debug("Selected backend = %s.", selected_backend) + if not torch.compiler.is_compiling(): + # int()/str() on the backend objects are not traceable by dynamo + selected_backend = "NoBackend" + if use_flash_attention: + selected_backend = f"FlashAttention ({str(flash_attention_backend)})" + elif use_fused_attention: + selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" + elif use_unfused_attention: + selected_backend = "UnfusedDotProductAttention" + logger.debug("Selected backend = %s.", selected_backend) return ( use_flash_attention, 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 From 36bb8e86881bf73bd088ce618e6acce88a760d75 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 13:18:09 +0200 Subject: [PATCH 2/9] Return int from the fused-attn backend probe and test symbolic-int path Comparing the pybind enum returned through assume_constant_result against module-level enum values generates guards dynamo cannot evaluate (crash when the comparison is true, i.e. when cuDNN rejects the config). The wrapper now returns a plain int, comparisons use precomputed int values, and the enum for callers is reconstructed by a second assume_constant_result helper that is never compared during tracing. Also document that os.environ.get (vs os.getenv) is intentional, and drop the guard_scalar specialization of numeric args: symbolic scalars (automatic dynamic) now graph break at the probe instead of forcing a full recompile per seqlen value; the test covers that path without fullgraph and checks the selection stays correct. A second test monkeypatches tex.get_fused_attn_backend to verify the baked result is trace-time-only and actually drives selection. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 92 +++++++++++++++++++ .../attention/dot_product_attention/utils.py | 76 +++++++++++---- 2 files changed, 148 insertions(+), 20 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index ed2e536ecc..d07d03bb4e 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -433,3 +433,95 @@ def fn(x): monkeypatch.setenv(env_var, "0") torch.testing.assert_close(compiled(x), fn(x)) monkeypatch.setenv(env_var, "1") + + +def test_get_attention_backend_fused_backend_constant(monkeypatch): + """The tex.get_fused_attn_backend call inside get_attention_backend is + wrapped with assume_constant_result. Check that under torch.compile it is + (1) invoked at trace time only (not on every run), (2) still consulted + (with correct results) when attention params change and dynamo turns the + changed ints into symbolic scalars via automatic dynamic, and (3) its + return value actually drives the backend selection (verified by + monkeypatching it to report no available sub-backend).""" + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + + # Restrict candidates to FusedAttention vs UnfusedDotProductAttention so + # the selection outcome encodes the tex.get_fused_attn_backend result. + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "1") + + def fn(x, params): + ( + use_flash_attention, + _, + use_fused_attention, + _, + use_unfused_attention, + _, + ) = dpa_utils.get_attention_backend(params) + return ( + x + + (1 if use_flash_attention else 0) + + (2 if use_fused_attention else 0) + + (4 if use_unfused_attention else 0) + ) + + x = torch.zeros(8, device="cuda") + params = dpa_utils.AttentionParams() + if fn(x, params)[0].item() != 2.0: + pytest.skip("FusedAttention not available for the default attention params") + + calls = [] + real_get_backend = tex.get_fused_attn_backend + + def counting_get_backend(*args): + calls.append(args) + return real_get_backend(*args) + + monkeypatch.setattr(dpa_utils.tex, "get_fused_attn_backend", counting_get_backend) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + out = compiled(x, params) + torch.testing.assert_close(out, x + 2.0) + calls_after_trace = len(calls) + assert calls_after_trace >= 1, "tex.get_fused_attn_backend not invoked during tracing" + + # Baked as a constant: running the compiled function again must not call it. + compiled(x, params) + assert len(calls) == calls_after_trace + + # Changing attention params: the changed ints (head_dim, seqlens) become + # symbolic scalars on recompilation (dynamo's automatic dynamic). + # assume_constant_result cannot convert symbolic scalars to constants and + # graph breaks on them, so this phase compiles without fullgraph and only + # requires that the selection stays correct (tex consulted again, eagerly). + compiled_dyn = torch.compile(fn) + 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), + ): + num_calls = len(calls) + out = compiled_dyn(x, changed_params) + assert len(calls) > num_calls, f"tex not consulted for {changed_params}" + torch.testing.assert_close(out, fn(x, changed_params)) + + # The baked result must drive the selection: report no fused sub-backend + # and expect UnfusedDotProductAttention instead of FusedAttention. Use a + # fresh frame: already-compiled frames keep the previously baked constant + # (assume_constant_result installs no guard on the wrapped function). + 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) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 797eef3017..e5df286934 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 @@ -342,14 +343,30 @@ def warning(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend(*args): +def _get_fused_attn_backend_int(*args): """ - Wrapper for tex.get_fused_attn_backend. The result only depends on the - attention configuration (not on tensor values), so it is marked with - assume_constant_result to keep get_attention_backend traceable by - torch.compile without a graph break on the C extension call. + Wrapper for tex.get_fused_attn_backend returning the backend as a plain int. + The result only depends on the attention configuration (not on tensor + values), so it is marked with assume_constant_result to keep + get_attention_backend traceable by torch.compile without a graph break on + the C extension call. It returns an int rather than the pybind enum: an int + is a python literal that safely crosses the compile boundary, while + comparing a baked pybind enum object against module-level enum values + generates guards dynamo cannot evaluate. """ - return tex.get_fused_attn_backend(*args) + return int(tex.get_fused_attn_backend(*args)) + + +@torch.compiler.assume_constant_result +def _fused_attn_backend_from_int(backend_int): + """Reconstruct the NVTE_Fused_Attn_Backend enum from its int value.""" + return tex.NVTE_Fused_Attn_Backend(backend_int) + + +# Plain-int values of the FusedAttnBackend enums, precomputed eagerly so that +# backend comparisons inside get_attention_backend are traceable (see +# _get_fused_attn_backend_int). +_FUSED_ATTN_BACKEND_INT = {name: int(backend) for name, backend in FusedAttnBackend.items()} def get_attention_backend( @@ -415,6 +432,12 @@ 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 if torch.compiler.is_compiling(): # logging.Logger methods graph-break under torch.compile; backend @@ -429,8 +452,9 @@ def get_attention_backend( 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) @@ -1385,6 +1409,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False # Filter: cuDNN support fused_attention_backend = None + fused_attention_backend_int = None if use_fused_attention: # ``DType`` is implicitly convertible to ``transformer_engine::DType`` # on the C++ side, so pass it straight to the pybind function. @@ -1393,7 +1418,11 @@ 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 = _get_fused_attn_backend( + # NOTE: under torch.compile the numeric args below may be symbolic + # (dynamo's automatic dynamic on recompilation, or dynamic shapes); + # assume_constant_result requires concrete values and currently graph + # breaks on symbolic scalars. + fused_attention_backend_int = _get_fused_attn_backend_int( is_training, q_type, kv_type, @@ -1414,20 +1443,24 @@ 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"]: + fused_attention_backend = _fused_attn_backend_from_int(fused_attention_backend_int) + if fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["No_Backend"]: 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"]: - if not torch.compiler.is_compiling(): - # int() on a pybind enum is not traceable by dynamo - logger.debug( - "Disabling FusedAttention for score_mod because sub-backend %s is not " - "F16/BF16 arbitrary-seqlen", - int(fused_attention_backend), - ) + fused_attention_backend_int = None + elif ( + has_score_mod + and fused_attention_backend_int != _FUSED_ATTN_BACKEND_INT["F16_arbitrary_seqlen"] + ): + logger.debug( + "Disabling FusedAttention for score_mod because sub-backend %s is not " + "F16/BF16 arbitrary-seqlen", + fused_attention_backend_int, + ) use_fused_attention = False fused_attention_backend = None + fused_attention_backend_int = None # Filter: Determinism # backend | deterministic # --------------------------------------------- @@ -1469,8 +1502,9 @@ 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 + fused_attention_backend_int = None if ( - fused_attention_backend == FusedAttnBackend["FP8"] + fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["FP8"] and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1480,8 +1514,9 @@ 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 + fused_attention_backend_int = None if ( - fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["F16_arbitrary_seqlen"] and is_training and ( device_compute_capability < (9, 0) @@ -1492,6 +1527,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False fused_attention_backend = None + fused_attention_backend_int = None # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 From f7564f9c1158f00e0e027099593143c58798976c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 15:14:25 +0200 Subject: [PATCH 3/9] Replace the FusedAttnBackend dict with a python-side IntEnum The pybind NVTE_Fused_Attn_Backend enum is not traceable by torch.compile: its C-implemented __eq__ cannot be traced, and a pybind enum instance baked through assume_constant_result produces guards dynamo cannot evaluate when compared against module-level enum values. FusedAttnBackend is now a plain python IntEnum generated at import time from tex.NVTE_Fused_Attn_Backend.__members__ (values always in sync with the C enum), and all remaining direct uses of the pybind enum on the python side are replaced with it. Name lookup (FusedAttnBackend["FP8"]) behaves the same as with the previous dict, and the backend value never crosses into a pybind call, so no boundary conversion is needed. This removes the previous int-based workaround in get_attention_backend (_fused_attn_backend_from_int and the precomputed int table): the assume_constant_result wrapper now simply returns the IntEnum and comparisons are traceable directly. Signed-off-by: Pawel Gadzinski --- .../dot_product_attention/backends.py | 13 ++--- .../dot_product_attention/context_parallel.py | 8 +-- .../attention/dot_product_attention/utils.py | 56 ++++++------------- .../pytorch/cpp_extensions/fused_attn.py | 27 ++++++--- 4 files changed, 45 insertions(+), 59 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8f42983553..977d94386c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1972,7 +1972,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, @@ -1994,7 +1994,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) @@ -2087,15 +2087,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!" @@ -2115,8 +2115,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 61a46a8652..75599d07c1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3144,7 +3144,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 @@ -3154,7 +3154,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:] @@ -3865,14 +3865,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 e5df286934..250c50f758 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -343,30 +343,18 @@ def warning(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend_int(*args): +def _get_fused_attn_backend(*args): """ - Wrapper for tex.get_fused_attn_backend returning the backend as a plain int. - The result only depends on the attention configuration (not on tensor - values), so it is marked with assume_constant_result to keep - get_attention_backend traceable by torch.compile without a graph break on - the C extension call. It returns an int rather than the pybind enum: an int - is a python literal that safely crosses the compile boundary, while - comparing a baked pybind enum object against module-level enum values - generates guards dynamo cannot evaluate. + Wrapper for tex.get_fused_attn_backend returning the backend as the + python-side FusedAttnBackend enum. The result only depends on the attention + configuration (not on tensor values), so it is marked with + assume_constant_result to keep get_attention_backend traceable by + torch.compile without a graph break on the C extension call. The pybind + enum is converted to FusedAttnBackend here: a python enum safely crosses + the compile boundary, while comparing a baked pybind enum object against + module-level enum values generates guards dynamo cannot evaluate. """ - return int(tex.get_fused_attn_backend(*args)) - - -@torch.compiler.assume_constant_result -def _fused_attn_backend_from_int(backend_int): - """Reconstruct the NVTE_Fused_Attn_Backend enum from its int value.""" - return tex.NVTE_Fused_Attn_Backend(backend_int) - - -# Plain-int values of the FusedAttnBackend enums, precomputed eagerly so that -# backend comparisons inside get_attention_backend are traceable (see -# _get_fused_attn_backend_int). -_FUSED_ATTN_BACKEND_INT = {name: int(backend) for name, backend in FusedAttnBackend.items()} + return FusedAttnBackend(int(tex.get_fused_attn_backend(*args))) def get_attention_backend( @@ -385,7 +373,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. @@ -1409,7 +1397,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False # Filter: cuDNN support fused_attention_backend = None - fused_attention_backend_int = None if use_fused_attention: # ``DType`` is implicitly convertible to ``transformer_engine::DType`` # on the C++ side, so pass it straight to the pybind function. @@ -1422,7 +1409,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # (dynamo's automatic dynamic on recompilation, or dynamic shapes); # assume_constant_result requires concrete values and currently graph # breaks on symbolic scalars. - fused_attention_backend_int = _get_fused_attn_backend_int( + fused_attention_backend = _get_fused_attn_backend( is_training, q_type, kv_type, @@ -1443,24 +1430,18 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cuda_graph, deterministic, ) - fused_attention_backend = _fused_attn_backend_from_int(fused_attention_backend_int) - if fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["No_Backend"]: + if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug("Disabling FusedAttention as no backend supports the provided input") use_fused_attention = False fused_attention_backend = None - fused_attention_backend_int = None - elif ( - has_score_mod - and fused_attention_backend_int != _FUSED_ATTN_BACKEND_INT["F16_arbitrary_seqlen"] - ): + elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " "F16/BF16 arbitrary-seqlen", - fused_attention_backend_int, + int(fused_attention_backend), ) use_fused_attention = False fused_attention_backend = None - fused_attention_backend_int = None # Filter: Determinism # backend | deterministic # --------------------------------------------- @@ -1502,9 +1483,8 @@ 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 - fused_attention_backend_int = None if ( - fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["FP8"] + fused_attention_backend == FusedAttnBackend["FP8"] and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1514,9 +1494,8 @@ 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 - fused_attention_backend_int = None if ( - fused_attention_backend_int == _FUSED_ATTN_BACKEND_INT["F16_arbitrary_seqlen"] + fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] and is_training and ( device_compute_capability < (9, 0) @@ -1527,7 +1506,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False fused_attention_backend = None - fused_attention_backend_int = None # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9c22c56bd1..d3f7b26f54 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,19 @@ "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, -} +# Python-side mirror of tex.NVTE_Fused_Attn_Backend, generated at import time +# so the values always match the C enum. 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. +FusedAttnBackend = IntEnum( + "FusedAttnBackend", + { + name[len("NVTE_") :]: int(value) + for name, value in NVTE_Fused_Attn_Backend.__members__.items() + }, +) BACKEND_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 @@ -124,7 +133,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 +185,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"; @@ -397,7 +406,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 +462,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] From 43b10847835c79f11a03dbc85ec597f4a8a16ae0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 15:23:27 +0200 Subject: [PATCH 4/9] Soften the FusedAttnBackend enum transition following the DType pattern Mirror how constants.DType migrated off the pybind enum: explicit IntEnum members pinned to the C values with an import-time sync assert, an __eq__ override comparing by integer value against NVTE_Fused_Attn_Backend (with matching __ne__/__hash__) so mixed comparisons stay equivalent regardless of the pybind11 version, and a cast() classmethod. fused_attn_fwd/bwd normalize their fused_attention_backend argument through cast(), so external callers still passing the pybind enum keep working. Signed-off-by: Pawel Gadzinski --- .../pytorch/cpp_extensions/fused_attn.py | 74 ++++++++++++++++--- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index d3f7b26f54..046019ee58 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -98,18 +98,64 @@ "learnable": NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, } -# Python-side mirror of tex.NVTE_Fused_Attn_Backend, generated at import time -# so the values always match the C enum. 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. -FusedAttnBackend = IntEnum( - "FusedAttnBackend", - { - name[len("NVTE_") :]: int(value) - for name, value in NVTE_Fused_Attn_Backend.__members__.items() - }, + +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 @@ -294,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:" @@ -547,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:" From c76304ba1b7ffe4617892f5079d57c3b6a85e5fe Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 15:26:35 +0200 Subject: [PATCH 5/9] Add FP8 backend-selection coverage to the torch.compile tests test_get_attention_backend_traceable_fp8 compiles the selection with fullgraph=True for AttentionParams(fp8=True) with a DelayedScaling(fp8_dpa) recipe, covering the FP8-only branch (run_config env reads, recipe filters, get_fp8_te_dtype) and checks that flipping NVTE_UnfusedDPA_Emulate_FP8 recompiles and keeps matching eager. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 50 +++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index d07d03bb4e..40768eb72c 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -525,3 +525,53 @@ def fn_no_backend(x, params): compiled_no_backend = torch.compile(fn_no_backend, fullgraph=True) torch.testing.assert_close(compiled_no_backend(x, params), x + 4.0) + + +def test_get_attention_backend_traceable_fp8(monkeypatch): + """Backend selection for FP8 attention (fp8_dpa recipes) must also trace + under torch.compile(fullgraph=True), including the FP8-only env var reads + and recipe filters, and the FP8 env vars must be guarded.""" + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + + 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) + + attention_params = dpa_utils.AttentionParams( + fp8=True, + fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}, + ) + + def fn(x): + ( + use_flash_attention, + _, + use_fused_attention, + _, + use_unfused_attention, + _, + ) = dpa_utils.get_attention_backend(attention_params) + return ( + x + + (1 if use_flash_attention else 0) + + (2 if use_fused_attention else 0) + + (4 if use_unfused_attention else 0) + ) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + x = torch.zeros(8, device="cuda") + torch.testing.assert_close(compiled(x), fn(x)) + + # FP8-only env var: allowing FP8 emulation enables UnfusedDotProductAttention, + # which must trigger recompilation (guard on os.environ) and match eager. + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") + torch.testing.assert_close(compiled(x), fn(x)) From ba39c933c973fba3e8bdd3b6af3fdf1747345648 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 15:33:46 +0200 Subject: [PATCH 6/9] Simplify the fused-attn probe wrapper and drop obsolete logging guards The is_compiling() guards around the available/selected-backend debug logs protected int() on the pybind enum, which crashed dynamo during tracing. With FusedAttnBackend now a python IntEnum, int() on it and str() on the flash-attn PkgVersion both trace cleanly (verified under fullgraph=True), so the logging blocks return to their upstream shape. The probe wrapper also reuses FusedAttnBackend.cast() and a shorter docstring. Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/utils.py | 61 ++++++++----------- 1 file changed, 25 insertions(+), 36 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 250c50f758..5639b968fa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -344,17 +344,10 @@ def warning(self, *args, **kwargs): @torch.compiler.assume_constant_result def _get_fused_attn_backend(*args): - """ - Wrapper for tex.get_fused_attn_backend returning the backend as the - python-side FusedAttnBackend enum. The result only depends on the attention - configuration (not on tensor values), so it is marked with - assume_constant_result to keep get_attention_backend traceable by - torch.compile without a graph break on the C extension call. The pybind - enum is converted to FusedAttnBackend here: a python enum safely crosses - the compile boundary, while comparing a baked pybind enum object against - module-level enum values generates guards dynamo cannot evaluate. - """ - return FusedAttnBackend(int(tex.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).""" + return FusedAttnBackend.cast(tex.get_fused_attn_backend(*args)) def get_attention_backend( @@ -1571,21 +1564,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_flash_attention_4: flash_attention_backend = FlashAttentionUtils.fa4_version - if not torch.compiler.is_compiling(): - # int()/str() on the backend objects are not traceable by dynamo - logger.debug( - "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s," - " UnfusedDotProductAttention=%s}", - bool(available_backends[0]), - (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), - bool(available_backends[1]), - ( - f" (sub-backend {int(fused_attention_backend)})" - if fused_attention_backend is not None - else "" - ), - bool(available_backends[2]), - ) + logger.debug( + "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s," + " UnfusedDotProductAttention=%s}", + bool(available_backends[0]), + (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), + bool(available_backends[1]), + ( + f" (sub-backend {int(fused_attention_backend)})" + if fused_attention_backend is not None + else "" + ), + bool(available_backends[2]), + ) # Select FusedAttention for performance if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0): @@ -1601,16 +1592,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_unfused_attention = False elif use_fused_attention: use_unfused_attention = False - if not torch.compiler.is_compiling(): - # int()/str() on the backend objects are not traceable by dynamo - selected_backend = "NoBackend" - if use_flash_attention: - selected_backend = f"FlashAttention ({str(flash_attention_backend)})" - elif use_fused_attention: - selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" - elif use_unfused_attention: - selected_backend = "UnfusedDotProductAttention" - logger.debug("Selected backend = %s.", selected_backend) + selected_backend = "NoBackend" + if use_flash_attention: + selected_backend = f"FlashAttention ({str(flash_attention_backend)})" + elif use_fused_attention: + selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" + elif use_unfused_attention: + selected_backend = "UnfusedDotProductAttention" + logger.debug("Selected backend = %s.", selected_backend) return ( use_flash_attention, From 17932268726832b05f311dd4ff11f0c3a03de15d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 16:05:36 +0200 Subject: [PATCH 7/9] Consolidate backend-selection compile tests; keep probe args literal Merge the three get_attention_backend tests into one covering: fullgraph tracing with the probe consulted at trace time only, env var flips (F16 and FP8), attention-param changes, and a forced No_Backend result driving the selection. The bitmask output now also encodes the fused sub-backend, which previously went unchecked. The probe wrapper takes layout/bias/mask/softmax as string keys and resolves the pybind enums internally, so every argument is a literal or a python enum - required for assume_constant_result(specialize_args=True) (pytorch#189042) to derive value guards once available. Scalars must stay concrete until then: the test pins specialize_int/float=True, because a symbolic scalar currently graph breaks at the probe and dynamo's resume then corrupts the returned fused backend (binds the wrapper function object instead of its result; surfaced by the sub-backend bits, minimal repro exists). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 185 ++++++------------ .../attention/dot_product_attention/utils.py | 46 +++-- 2 files changed, 95 insertions(+), 136 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 40768eb72c..0907c57584 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -391,86 +391,54 @@ def fn(inp): # --------------------------------------------------------------------------- +# Scalars in AttentionParams must stay concrete: assume_constant_result cannot +# convert symbolic scalars (dynamo's automatic dynamic would make changed ints +# symbolic on recompilation). Specializing them is the intended behavior and +# will be handled by assume_constant_result(specialize_args=True) +# (pytorch#189042); until then 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), - i.e. without any graph break, and NVTE_* env var reads must be guarded so - that changing an env var triggers recompilation instead of silently - reusing a stale backend selection.""" + """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 tex.get_fused_attn_backend probe must be + consulted at trace time only, with its baked result driving the + selection.""" from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils - attention_params = dpa_utils.AttentionParams() - - def fn(x): - ( - use_flash_attention, - _, - use_fused_attention, - _, - use_unfused_attention, - _, - ) = dpa_utils.get_attention_backend(attention_params) - return ( - x - + (1 if use_flash_attention else 0) - + (2 if use_fused_attention else 0) - + (4 if use_unfused_attention 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 in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): - monkeypatch.setenv(env_var, "1") - - torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) - - x = torch.zeros(8, device="cuda") - torch.testing.assert_close(compiled(x), fn(x)) - - # 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), fn(x)) - monkeypatch.setenv(env_var, "1") - - -def test_get_attention_backend_fused_backend_constant(monkeypatch): - """The tex.get_fused_attn_backend call inside get_attention_backend is - wrapped with assume_constant_result. Check that under torch.compile it is - (1) invoked at trace time only (not on every run), (2) still consulted - (with correct results) when attention params change and dynamo turns the - changed ints into symbolic scalars via automatic dynamic, and (3) its - return value actually drives the backend selection (verified by - monkeypatching it to report no available sub-backend).""" - from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils - - # Restrict candidates to FusedAttention vs UnfusedDotProductAttention so - # the selection outcome encodes the tex.get_fused_attn_backend result. - monkeypatch.setenv("NVTE_FLASH_ATTN", "0") - monkeypatch.setenv("NVTE_FUSED_ATTN", "1") - monkeypatch.setenv("NVTE_UNFUSED_ATTN", "1") - 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) ) - x = torch.zeros(8, device="cuda") - params = dpa_utils.AttentionParams() - if fn(x, params)[0].item() != 2.0: - pytest.skip("FusedAttention not available for the default attention params") + # 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) calls = [] real_get_backend = tex.get_fused_attn_backend @@ -483,9 +451,10 @@ def counting_get_backend(*args): torch._dynamo.reset() compiled = torch.compile(fn, fullgraph=True) + x = torch.zeros(8, device="cuda") + params = dpa_utils.AttentionParams() - out = compiled(x, params) - torch.testing.assert_close(out, x + 2.0) + torch.testing.assert_close(compiled(x, params), fn(x, params)) calls_after_trace = len(calls) assert calls_after_trace >= 1, "tex.get_fused_attn_backend not invoked during tracing" @@ -493,12 +462,26 @@ def counting_get_backend(*args): compiled(x, params) assert len(calls) == calls_after_trace - # Changing attention params: the changed ints (head_dim, seqlens) become - # symbolic scalars on recompilation (dynamo's automatic dynamic). - # assume_constant_result cannot convert symbolic scalars to constants and - # graph breaks on them, so this phase compiles without fullgraph and only - # requires that the selection stays correct (tex consulted again, eagerly). - compiled_dyn = torch.compile(fn) + # 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 consult the probe again, 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), @@ -506,14 +489,16 @@ def counting_get_backend(*args): dpa_utils.AttentionParams(qkv_dtype=torch.float16), ): num_calls = len(calls) - out = compiled_dyn(x, changed_params) + out = compiled(x, changed_params) assert len(calls) > num_calls, f"tex not consulted for {changed_params}" torch.testing.assert_close(out, fn(x, changed_params)) - # The baked result must drive the selection: report no fused sub-backend - # and expect UnfusedDotProductAttention instead of FusedAttention. Use a - # fresh frame: already-compiled frames keep the previously baked constant - # (assume_constant_result installs no guard on the wrapped function). + # 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", @@ -525,53 +510,3 @@ def fn_no_backend(x, params): compiled_no_backend = torch.compile(fn_no_backend, fullgraph=True) torch.testing.assert_close(compiled_no_backend(x, params), x + 4.0) - - -def test_get_attention_backend_traceable_fp8(monkeypatch): - """Backend selection for FP8 attention (fp8_dpa recipes) must also trace - under torch.compile(fullgraph=True), including the FP8-only env var reads - and recipe filters, and the FP8 env vars must be guarded.""" - from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils - - 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) - - attention_params = dpa_utils.AttentionParams( - fp8=True, - fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}, - ) - - def fn(x): - ( - use_flash_attention, - _, - use_fused_attention, - _, - use_unfused_attention, - _, - ) = dpa_utils.get_attention_backend(attention_params) - return ( - x - + (1 if use_flash_attention else 0) - + (2 if use_fused_attention else 0) - + (4 if use_unfused_attention else 0) - ) - - torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) - - x = torch.zeros(8, device="cuda") - torch.testing.assert_close(compiled(x), fn(x)) - - # FP8-only env var: allowing FP8 emulation enables UnfusedDotProductAttention, - # which must trigger recompilation (guard on os.environ) and match eager. - monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") - torch.testing.assert_close(compiled(x), fn(x)) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 5639b968fa..07ac9da53f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -343,11 +343,34 @@ def warning(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend(*args): +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).""" - return FusedAttnBackend.cast(tex.get_fused_attn_backend(*args)) + 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 literal or a python enum — required for + assume_constant_result(specialize_args=True) to derive value guards.""" + 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( @@ -1398,18 +1421,19 @@ 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 - # NOTE: under torch.compile the numeric args below may be symbolic - # (dynamo's automatic dynamic on recompilation, or dynamic shapes); - # assume_constant_result requires concrete values and currently graph - # breaks on symbolic scalars. + # 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. Once + # assume_constant_result(specialize_args=True) (pytorch#189042) is + # available, symbolic scalars will be specialized with a guard instead. 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, From 719b18436e80275a75c7ff55cab93b4f793da530 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 16:08:38 +0200 Subject: [PATCH 8/9] Drop references to unreleased dynamo features from comments The probe-argument and static-scalar comments referenced assume_constant_result(specialize_args=True), which is not part of any released PyTorch; describe the current behavior only. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 4 +--- .../pytorch/attention/dot_product_attention/utils.py | 7 ++----- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 0907c57584..4274008d60 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -393,9 +393,7 @@ def fn(inp): # Scalars in AttentionParams must stay concrete: assume_constant_result cannot # convert symbolic scalars (dynamo's automatic dynamic would make changed ints -# symbolic on recompilation). Specializing them is the intended behavior and -# will be handled by assume_constant_result(specialize_args=True) -# (pytorch#189042); until then pin them static explicitly. +# 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) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 07ac9da53f..4998e7f8ba 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -357,8 +357,7 @@ def _get_fused_attn_backend( 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 literal or a python enum — required for - assume_constant_result(specialize_args=True) to derive value guards.""" + that every argument is a python literal or a python enum.""" return FusedAttnBackend.cast( tex.get_fused_attn_backend( is_training, @@ -1423,9 +1422,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt kv_type = q_type # 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. Once - # assume_constant_result(specialize_args=True) (pytorch#189042) is - # available, symbolic scalars will be specialized with a guard instead. + # dynamic by automatic dynamic currently graph break here. fused_attention_backend = _get_fused_attn_backend( is_training, q_type, From 6386e2c8229a7c284e16a5e5c9bd0af86bb3c7f0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 8 Jul 2026 16:11:21 +0200 Subject: [PATCH 9/9] Assert selection results only in the compile test Drop the call-counting monkeypatch and its assertions; compiled-vs-eager output equality is what matters and already fails on stale selections. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 27 ++++----------------------- 1 file changed, 4 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 4274008d60..5447cfcc49 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -399,9 +399,8 @@ 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 tex.get_fused_attn_backend probe must be - consulted at trace time only, with its baked result driving the - selection.""" + 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): @@ -438,27 +437,12 @@ def fn(x, params): ): monkeypatch.setenv(env_var, value) - calls = [] - real_get_backend = tex.get_fused_attn_backend - - def counting_get_backend(*args): - calls.append(args) - return real_get_backend(*args) - - monkeypatch.setattr(dpa_utils.tex, "get_fused_attn_backend", counting_get_backend) - 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)) - calls_after_trace = len(calls) - assert calls_after_trace >= 1, "tex.get_fused_attn_backend not invoked during tracing" - - # Baked as a constant: running the compiled function again must not call it. - compiled(x, params) - assert len(calls) == calls_after_trace # Flip env vars one by one: the compiled function must recompile (guards # on os.environ) and keep matching eager. @@ -479,17 +463,14 @@ def counting_get_backend(*args): monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "0") # Changing attention params (ints, layout string, dtype) must recompile - # and consult the probe again, still with no graph break. + # 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), ): - num_calls = len(calls) - out = compiled(x, changed_params) - assert len(calls) > num_calls, f"tex not consulted for {changed_params}" - torch.testing.assert_close(out, fn(x, changed_params)) + 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