feat: disable flash attn if not supported CUDA version or device capability

modeling_clip.py

@@ -144,6 +144,25 @@ def _resolve_attention_libs(config: JinaCLIPConfig):
                     'for installation instructions, disabling'
                 )
                 return False
+            major, minor, *_ = torch.version.cuda.split('.')
+            major, minor = int(major), int(minor)
+            if major < 11 or (major == 11 and minor < 7):
+                warnings.warn(
+                    'Flash attention requires CUDA>=11.7. Found version '
+                    f'{major}.{minor}, disabling'
+                )
+                return False
+            capability = torch.cuda.get_device_capability()
+            major, *_ = capability
+            major = int(major)
+            if major < 8:
+                device_name = torch.cuda.get_device_properties(0).name
+                warnings.warn(
+                    'Flash attention requires device capability>=8.0 (NVIDIA Ampere, '
+                    f'Hopper or ADA). Found device {device_name} with capability '
+                    f'{capability}, disabling'
+                )
+                return False
             return True
         return False