[Model] Add MiniMax H3 diffusion support (#5691)

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -22,6 +22,15 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+
+def _is_blackwell() -> bool:
+    try:
+        capability = current_omni_platform.get_device_capability()
+    except Exception:
+        return False
+    return capability is not None and capability.major >= 10
+
+
 # Flash Attention function detection with fallback chain
 flash_attn_func = None
 flash_attn_varlen_func = None
@@ -48,12 +57,27 @@ elif current_omni_platform.is_musa():
     except (ImportError, ModuleNotFoundError):
         pass
 else:
-    # CUDA: try FA3 -> FA2 fallback chain
+    # CUDA: prefer the CuTe-based FA4 package when present.  This matters on
+    # Blackwell, where an importable Hopper-only FA3 wheel may otherwise be
+    # selected and fail at launch with "no kernel image is available".
+    if _is_blackwell():
+        try:
+            from flash_attn.cute import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+
+            logger.info("Using CuTe FlashAttention-4 on Blackwell")
+        except Exception as exc:
+            # Optional FA4 dependencies may be present but ABI-incompatible
+            # (for example, mismatched CUTLASS DSL and Quack builds). Treat
+            # that as unavailable so backend discovery can continue.
+            logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+
+    # Try FA3 -> FA2 fallback chain.
     # Try FA3 from fa3-fwd PyPI package
-    try:
-        from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
-    except (ImportError, ModuleNotFoundError):
-        pass
+    if flash_attn_func is None:
+        try:
+            from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        except (ImportError, ModuleNotFoundError):
+            pass
 
     # Fallback: Try FA3 from flash-attention source build
     if flash_attn_func is None:
@@ -100,7 +124,16 @@ def is_flash_attn_installed() -> bool:
     Shared by CUDA/ROCm/MUSA platforms.
     """
     try:
-        # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
+        # Check for any FA backend: FA4 (flash_attn.cute), FA3
+        # (fa3_fwd_interface, flash_attn_interface), or FA2 (flash_attn).
+        if _is_blackwell():
+            try:
+                from flash_attn.cute import flash_attn_varlen_func  # noqa: F401
+
+                return True
+            except Exception as exc:
+                logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+
         # Try FA3 from fa3-fwd PyPI package
         try:
             import fa3_fwd_interface  # noqa: F401