[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: switching os.getenv calls to os.environ.get so dynamo installs value guards on NVTE_* env vars; wrapping the opaque pybind tex.get_fused_attn_backend call in an assume_constant_result helper; replacing the FusedAttnBackend string→pybind dict with a proper Python IntEnum that is dynamo-traceable; and substituting a _NoOpLogger for the real logger during compilation to avoid logging-induced graph breaks.

• FusedAttnBackend is converted from a plain dict to a Python IntEnum with import-time sync assertion against the C++ NVTE_Fused_Attn_Backend enum, a cast() classmethod for backward-compat, and a custom __eq__/__ne__/__hash__ so mixed comparisons with the pybind enum remain transparent.
• get_device_compute_capability and get_cudnn_version are decorated with assume_constant_result; _get_cudnn_version retains its lru_cache and is now the private implementation backing the public wrapper.
• A new test test_get_attention_backend_traceable compiles the selection function with fullgraph=True, flips env vars, exercises the FP8 branch, and patches the probe to verify the baked constant drives selection.

Confidence Score: 5/5

Safe to merge. The backend-selection logic is correct in all paths; the FusedAttnBackend IntEnum mirrors the C++ enum exactly with an import-time assert as a safety net; os.environ.get guards and assume_constant_result are applied consistently; and backward compatibility is preserved through cast() normalisation in fused_attn_fwd/bwd.

All changes are in selection/dispatch logic with no tensor computation affected. The only findings are quality observations: warning/error logs are silently dropped during compilation (users lose fallback notices under torch.compile), a minor type annotation inaccuracy in ne, and three per-version env-var keys left out of the test pre-set. None of these affect correctness or introduce data-path regressions.

transformer_engine/pytorch/attention/dot_product_attention/utils.py deserves a second look regarding the _NoOpLogger swallowing warning/error log levels. tests/pytorch/test_torch_compile.py is missing NVTE_FLASH_ATTN_V2/V3/V4 pre-sets which limits dynamo guard coverage for those keys.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Replaces FusedAttnBackend dict with a properly-defined IntEnum; adds cast(), custom __eq__/__ne__/__hash__, and import-time sync assert; fused_attn_fwd/bwd now normalize via cast() for backward compatibility.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Core change: _get_fused_attn_backend wraps the pybind call with assume_constant_result; _NoOpLogger avoids logging graph breaks; all os.getenv → os.environ.get for dynamo guards. Warning/error logs are silently dropped during compilation.
transformer_engine/pytorch/utils.py	Adds assume_constant_result to get_device_compute_capability; renames old get_cudnn_version to _get_cudnn_version (retains lru_cache) and adds a thin public get_cudnn_version wrapper with assume_constant_result.
tests/pytorch/test_torch_compile.py	Adds test_get_attention_backend_traceable: tests fullgraph=True tracing, env-var guard firing, FP8 branch coverage, and parameter-change recompilation. Missing pre-set for NVTE_FLASH_ATTN_V2/V3/V4 variants means dynamo won't guard those keys.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Mechanical replacement of all tex.NVTE_Fused_Attn_Backend.* references with FusedAttnBackend[...] equivalents; default parameter type updated to FusedAttnBackend.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Four mechanical replacements of tex.NVTE_Fused_Attn_Backend.NVTE_FP8 / NVTE_F16_arbitrary_seqlen with FusedAttnBackend["FP8"] / FusedAttnBackend["F16_arbitrary_seqlen"].

_{Reviews (8): Last reviewed commit: "[PyTorch] Address review: add info()/err..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch

[pggPL]

/te-ci pytorch