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

[pggPL]

Description

This PR is part of an effort to enable torch.compile support for attention: it makes get_attention_backend traceable.

It works on current PyTorch as-is — there is no need to wait for the two related PRs I opened in PyTorch. Those only improve the remaining rough edge: pytorch/pytorch#189042 (bigger and much more important; among other things it handles symbolic scalar arguments, which with stock dynamo must stay non-symbolic in the backend-selection probe — with static scalars, the common case, everything in this PR already works) and pytorch/pytorch#189027 (small, and possibly not needed anymore since this PR reads env vars via os.environ.get, which stock dynamo already guards).

get_attention_backend currently graph-breaks under torch.compile, forcing backend selection to run as an eager island. This PR makes it fully traceable with fullgraph=True, with dynamo guards on the NVTE_* environment variables so that changing an env var triggers recompilation instead of silently reusing a stale backend selection.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (see the FusedAttnBackend note below)
• Infra/Build change
• Code refactoring

Changes

• Read NVTE_* env vars in get_attention_backend via os.environ.get instead of os.getenv: dynamo installs value guards on os.environ accesses, while os.getenv reads are unguarded (the stdlib-global source is skipped) and would bake a stale backend selection into the compiled graph.

• Wrap the tex.get_fused_attn_backend pybind call in a torch.compiler.assume_constant_result helper — the result depends only on the attention configuration, not on tensor values.

• Mark get_device_compute_capability and get_cudnn_version with assume_constant_result (pybind/CUDA property calls are not traceable, and their results are constant for a process).

• Use a no-op logger when compiling (logging.Logger methods graph-break in dynamo) and skip the debug-log blocks that evaluate int()/str() on pybind enum values during tracing. Backend-selection debug logs are unchanged in eager mode.

• Add test_get_attention_backend_traceable to tests/pytorch/test_torch_compile.py: compiles a function calling get_attention_backend with fullgraph=True (any graph break fails the test) and flips NVTE_FUSED_ATTN/NVTE_UNFUSED_ATTN/NVTE_FLASH_ATTN to verify the guards trigger recompilation and the result keeps matching eager.

• FusedAttnBackend (in pytorch/cpp_extensions/fused_attn.py) is now a python IntEnum generated at import time from tex.NVTE_Fused_Attn_Backend.__members__, replacing the previous str->pybind-enum dict, and all direct python-side uses of the pybind enum are replaced with it. Rationale: the pybind enum is not traceable by dynamo (C __eq__), and a pybind enum baked through assume_constant_result produces guards dynamo cannot evaluate when compared against module-level enum values (hard crash when cuDNN rejects a config). The transition follows the pattern used for constants.DType: explicit members pinned to the C values with an import-time sync assert, an __eq__ override so mixed comparisons with tex.NVTE_Fused_Attn_Backend stay equivalent regardless of the pybind11 version, and a cast() classmethod; fused_attn_fwd/bwd normalize their backend argument through cast(), so callers still passing the pybind enum keep working. Name lookup FusedAttnBackend["FP8"] and int(...) behave as before. The remaining (minor) breaking surface is isinstance checks against the pybind enum and dict-API such as .items()/iteration; no such uses exist inside TE and none were found in Megatron-LM/NeMo.

Note: dynamo only guards os.environ keys that exist at trace time; reads of absent keys are not guarded yet (upstream PyTorch limitation).

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR makes get_attention_backend fully traceable under torch.compile(fullgraph=True) by replacing unguarded os.getenv calls with guarded os.environ.get, wrapping pybind11 calls in assume_constant_result, and converting the FusedAttnBackend dict to a Python IntEnum that dynamo can constant-fold without crashes.

• os.getenv → os.environ.get: dynamo installs value guards on os.environ accesses, triggering recompilation when NVTE_* env vars change, which os.getenv did not provide.
• FusedAttnBackend IntEnum: replaces the previous str→pybind-enum dict; members are pinned to the C++ integer values with an import-time sync assertion, an __eq__ override for mixed comparisons with the pybind enum, and a cast() classmethod for backward compatibility with callers still passing tex.NVTE_Fused_Attn_Backend.
• assume_constant_result applied to _get_fused_attn_backend, get_device_compute_capability, and get_cudnn_version; a _NoOpLogger is substituted for logging.Logger during compilation to avoid graph-breaking logger method calls.

Confidence Score: 5/5

Safe to merge; all changes are scoped to traceability improvements without altering attention computation logic, and backward compatibility is maintained.

The os.getenv to os.environ.get swap is correctly applied uniformly within get_attention_backend, the FusedAttnBackend IntEnum includes an import-time sync assertion that catches C++ drift, and assume_constant_result decorators are applied only to truly static configuration queries. No incorrect logic, data loss, or security concerns were found.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Replaces the FusedAttnBackend dict with an IntEnum; import-time sync assert, eq/ne/hash overrides for pybind interop, and cast() for backward compat.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Core of the PR: os.getenv to os.environ.get, _NoOpLogger during compilation, _get_fused_attn_backend with assume_constant_result, string args resolved inside the wrapper.
transformer_engine/pytorch/utils.py	Adds assume_constant_result to get_device_compute_capability and wraps _get_cudnn_version (still lru_cached) with a new assume_constant_result get_cudnn_version.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Replaces four tex.NVTE_Fused_Attn_Backend.* references with FusedAttnBackend[...]; purely mechanical.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Replaces four tex.NVTE_Fused_Attn_Backend.* references with FusedAttnBackend[...]; purely mechanical.
tests/pytorch/test_torch_compile.py	Adds test_get_attention_backend_traceable compiling with fullgraph=True, flips NVTE_* vars, covers FP8 branch and param changes.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["get_attention_backend(params)"] --> B{torch.compiler.is_compiling?}
    B -- Yes --> C[_no_op_logger]
    B -- No --> D[logging.Logger]
    A --> E["os.environ.get NVTE_FLASH_ATTN\nos.environ.get NVTE_FUSED_ATTN\n[dynamo guards installed]"]
    E --> F{use_fused_attention?}
    F -- Yes --> G["_get_fused_attn_backend\n@assume_constant_result"]
    G --> H["tex.get_fused_attn_backend\n[pybind11 called eagerly at compile time]"]
    H --> I["FusedAttnBackend.cast result\n[Python IntEnum constant-foldable]"]
    I --> J{backend == No_Backend?}
    J -- Yes --> K[use_fused_attention = False]
    J -- No --> L[fused_attention_backend set]
    F -- No --> M[skip cuDNN probe]
    A --> N["get_device_compute_capability\n@assume_constant_result"]
    A --> O["get_cudnn_version\n@assume_constant_result"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["get_attention_backend(params)"] --> B{torch.compiler.is_compiling?}
    B -- Yes --> C[_no_op_logger]
    B -- No --> D[logging.Logger]
    A --> E["os.environ.get NVTE_FLASH_ATTN\nos.environ.get NVTE_FUSED_ATTN\n[dynamo guards installed]"]
    E --> F{use_fused_attention?}
    F -- Yes --> G["_get_fused_attn_backend\n@assume_constant_result"]
    G --> H["tex.get_fused_attn_backend\n[pybind11 called eagerly at compile time]"]
    H --> I["FusedAttnBackend.cast result\n[Python IntEnum constant-foldable]"]
    I --> J{backend == No_Backend?}
    J -- Yes --> K[use_fused_attention = False]
    J -- No --> L[fused_attention_backend set]
    F -- No --> M[skip cuDNN probe]
    A --> N["get_device_compute_capability\n@assume_constant_result"]
    A --> O["get_cudnn_version\n@assume_constant_result"]

Loading

_{Reviews (6): Last reviewed commit: "Merge remote-tracking branch 'upstream/m..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch