[PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2…#3204
Open
wdykas wants to merge 1 commit into
Open
[PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2…#3204wdykas wants to merge 1 commit into
wdykas wants to merge 1 commit into
Conversation
…/V3/V4) NVTE_FLASH_ATTN enables or disables the whole FlashAttention family, but the choice between FlashAttention 2, 3, and 4 is automatic (package presence and compute capability) with no user override. Some workloads need to pin the FlashAttention generation, e.g. RL training that must produce bitwise-identical logprobs to an inference engine running a specific FlashAttention version: different generations use different tile sizes and online-softmax accumulation orders, so mixed versions between training and inference break batch-invariant / train-inference parity guarantees. Add NVTE_FLASH_ATTN_V2, NVTE_FLASH_ATTN_V3, and NVTE_FLASH_ATTN_V4 (default 1) that disable a specific FlashAttention version even when it is installed, following the existing NVTE_FLASH_ATTN filter pattern. Behavior is unchanged when the variables are unset. Signed-off-by: wdykas <wdykas@nvidia.com>
github-actions
Bot
added
the
community-contribution
Contributor
Greptile SummaryThis PR adds per-version FlashAttention controls for PyTorch attention backend selection. The main changes are:
Confidence Score: 5/5This looks safe to merge after a small hardening cleanup.
transformer_engine/pytorch/attention/dot_product_attention/utils.py Important Files Changed
Reviews (1): Last reviewed commit: "[PyTorch] Add per-version FlashAttention..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+458
to
+460
| use_flash_attention_2 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V2", "1")) | ||
| use_flash_attention_3 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V3", "1")) | ||
| use_flash_attention_4 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) |
Contributor
There was a problem hiding this comment.
When any new per-version variable is set to a non-integer value such as true, yes, or an empty string, int(os.getenv(...)) raises before backend fallback runs. A single malformed NVTE_FLASH_ATTN_V2, NVTE_FLASH_ATTN_V3, or NVTE_FLASH_ATTN_V4 setting can make get_attention_backend fail instead of selecting another available attention backend.
Collaborator
|
/te-ci pytorch L1 |
cyanguwa
approved these changes
Jul 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
set fa version
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: