fix(pytorch): disable FlashAttention 2 for softcap with dropout

flash-attn rejects a nonzero softcap combined with nonzero dropout at
dispatch: "Softcapping does not support dropout for now" in
csrc/flash_attn/flash_api.cpp, present in mha_fwd and mha_varlen_fwd
from v2.6.0 (the earliest version TE allows softcap on) onwards. Backend
selection did not model this, so a softcap + attention-dropout config
passed selection, routed to FA2, and crashed inside flash-attn.

Dropout only reaches the kernel while training, since backends.py passes
`self.attention_dropout if self.training else 0.0`, so the gate is on
`attention_dropout != 0.0 and is_training` to avoid blocking valid
inference configs. UnfusedDotProductAttention supports both softcap and
dropout and stays available, so this steers rather than hard-fails.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

Author: nvegesna-netizen

transformer_engine/pytorch/attention/dot_product_attention/utils.py

@@ -811,6 +811,12 @@ def _disable_all_flash_attention() -> None:
         if use_flash_attention_2 and not FlashAttentionUtils.v2_6_0_plus:
             logger.debug("Disabling FlashAttention 2 for softcap (requires flash-attn >= 2.6.0)")
             use_flash_attention_2 = False
+        if use_flash_attention_2 and attention_dropout != 0.0 and is_training:
+            # FA2 hard-rejects a nonzero softcap combined with nonzero dropout at dispatch
+            # ("Softcapping does not support dropout for now", flash_api.cpp). Dropout only reaches
+            # the kernel while training -- backends.py passes 0.0 in eval -- hence the is_training.
+            logger.debug("Disabling FlashAttention 2 for softcap with dropout")
+            use_flash_attention_2 = False
 
     # Filter: score_mod
     if has_score_mod_bprop and not has_score_mod: