[PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2… (#3204)

[PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2/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>

Author: wdykas

docs/envvars.rst

@@ -141,6 +141,24 @@ backend-selection overview.
    :Default: ``1``
    :Description: Enable or disable FlashAttention backend for DotProductAttention. When set to ``0``, FlashAttention will not be used.
 
+.. envvar:: NVTE_FLASH_ATTN_V2
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 2 (the ``flash-attn`` package) for DotProductAttention, without affecting FlashAttention 3 or 4. When set to ``0``, FlashAttention 2 will not be used even if it is installed. Useful for pinning the FlashAttention version, e.g. so training-side attention runs the same kernel generation as an inference engine.
+
+.. envvar:: NVTE_FLASH_ATTN_V3
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 3 (the ``flash-attn-3`` package) for DotProductAttention, without affecting FlashAttention 2 or 4. When set to ``0``, FlashAttention 3 will not be used even if it is installed.
+
+.. envvar:: NVTE_FLASH_ATTN_V4
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FlashAttention 4 (the ``flash-attn-4`` package) for DotProductAttention, without affecting FlashAttention 2 or 3. When set to ``0``, FlashAttention 4 will not be used even if it is installed.
+
 .. envvar:: NVTE_FUSED_ATTN
 
    :Type: ``int`` (0 or 1)

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -455,18 +455,18 @@ def get_attention_backend(
 
     # Filter: Environment variables
     use_flash_attention = int(os.getenv("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
+    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"))
     flash_attention_backend = None
     use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1"))
     use_unfused_attention = int(os.getenv("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")
+        logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V2=0")
     if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
-        logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0")
+        logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V3=0")
     if not use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
-        logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0")
+        logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V4=0")
     if not use_fused_attention:
         logger.debug("Disabling FusedAttention due to NVTE_FUSED_ATTN=0")
     if not use_unfused_attention: