SDPA first on Blackwell, FA2 only for Ampere/Hopper, txt caption support

train_engine.py

@@ -541,37 +541,35 @@ def _ensure_acestep_imports():
 
 
 def _attn_candidates(device: str) -> List[str]:
-    """FA2 -> SDPA -> eager, filtered by availability.
+    """SDPA -> FA2 -> eager, filtered by availability.
 
-    On CUDA with flash_attn installed and compute capability >= 8.0,
-    flash_attention_2 is tried first.  On CPU, flash_attention_2 is
-    always skipped (it requires CUDA).
+    SDPA is preferred (faster on Blackwell SM12.0, native cuDNN).
+    FA2 is fallback for older GPUs where SDPA is slower.
+    On CPU, only SDPA and eager are tried.
     """
     candidates = []
     if device.startswith("cuda"):
+        candidates.append("sdpa")
         try:
             import flash_attn  # noqa: F401
             dev_idx = int(device.split(":")[1]) if ":" in device else 0
             props = torch.cuda.get_device_properties(dev_idx)
-            if props.major >= 8:
-                candidates.append("flash_attention_2")
-                logger.info(
-                    "flash_attention_2 available (compute %d.%d, flash_attn installed)",
-                    props.major, props.minor,
-                )
+            if props.major >= 8 and props.major < 12:
+                # FA2 is faster on Ampere/Hopper (SM 8.x-9.x), slower on Blackwell (SM 12.x)
+                candidates.insert(0, "flash_attention_2")
+                logger.info("FA2 prioritized (compute %d.%d, Ampere/Hopper)", props.major, props.minor)
             else:
-                logger.info(
-                    "flash_attention_2 skipped: compute %d.%d < 8.0 (need Ampere+)",
-                    props.major, props.minor,
-                )
+                logger.info("FA2 available but SDPA preferred (compute %d.%d)", props.major, props.minor)
         except ImportError:
             logger.info("flash_attention_2 skipped: flash_attn package not installed")
         except Exception as exc:
             logger.info("flash_attention_2 skipped: %s", exc)
     else:
+        candidates.append("sdpa")
         logger.info("flash_attention_2 skipped: device is %s (not CUDA)", device)
-    candidates.extend(["sdpa", "eager"])
-    return candidates
+    if "eager" not in candidates:
+        candidates.append("eager")
+    return list(dict.fromkeys(candidates))
 
 
 def load_model_for_training(