Import flash_dmattn_varlen_func alongside flash_dmattn_func for backend availability

Author: algo-home

flash_dmattn/__init__.py

@@ -7,11 +7,11 @@
 
 # Import CUDA functions when available
 try:
-    from flash_dmattn.flash_dmattn_interface import flash_dmattn_func
+    from flash_dmattn.flash_dmattn_interface import flash_dmattn_func, flash_dmattn_varlen_func
     CUDA_AVAILABLE = True
 except ImportError:
     CUDA_AVAILABLE = False
-    flash_dmattn_func = None
+    flash_dmattn_func, flash_dmattn_varlen_func = None, None
 
 # Import Triton functions when available
 try:
@@ -89,6 +89,7 @@ def flash_dmattn_func_auto(backend: Optional[str] = None, **kwargs):
     "TRITON_AVAILABLE",
     "FLEX_AVAILABLE",
     "flash_dmattn_func",
+    "flash_dmattn_varlen_func",
     "triton_dmattn_func",
     "flex_dmattn_func",
     "get_available_backends",