[Diffusion] Improve HuggingFace hub kernel loading in flash_attn_hub (#4977)

Signed-off-by: samithuang <285365963@qq.com>

vllm_omni/diffusion/attention/backends/flash_attn_hub.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import threading
 from functools import partial
 
 import torch
@@ -18,6 +19,33 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
 logger = init_logger(__name__)
 
 
+_hub_modules: dict[str, object] = {}
+_hub_lock = threading.Lock()
+
+
+def _load_hub_module(repo_id: str):
+    from kernels import get_kernel
+
+    logger.info("Loading %s kernel from HuggingFace Hub...", repo_id)
+    last_error = None
+    for version in (1, 2, None):
+        try:
+            if version is not None:
+                return get_kernel(repo_id, version=version)
+            return get_kernel(repo_id)
+        except Exception as exc:
+            logger.info("Failed to load %s version %s: %s", repo_id, version, exc)
+            last_error = exc
+    raise RuntimeError(f"Failed to load HuggingFace Hub kernel {repo_id!r}") from last_error
+
+
+def _get_hub_module(repo_id: str):
+    with _hub_lock:
+        if repo_id not in _hub_modules:
+            _hub_modules[repo_id] = _load_hub_module(repo_id)
+        return _hub_modules[repo_id]
+
+
 def _run_varlen_dense(
     varlen_func,
     query: torch.Tensor,
@@ -88,23 +116,9 @@ class FlashAttentionHubImpl(AttentionImpl):
         if backend_kwargs:
             logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
-        # Lazily get the kernel from the Hub to avoid network/disk overhead on import
-        from kernels import get_kernel
-
-        logger.info("Loading flash-attn2 kernel from HuggingFace Hub...")
-        try:
-            hub_module = get_kernel("kernels-community/flash-attn2", version=1)
-        except Exception as e:
-            try:
-                logger.info("Failed to load version 1, attempting version 2: %s", e)
-                hub_module = get_kernel("kernels-community/flash-attn2", version=2)
-            except Exception as e2:
-                logger.info("Failed to load version 2, attempting default: %s", e2)
-                hub_module = get_kernel("kernels-community/flash-attn2")
-
+        hub_module = _get_hub_module("kernels-community/flash-attn2")
         self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
         self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
-
         if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
             raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found")
 
@@ -267,23 +281,9 @@ class FlashAttention3HubImpl(AttentionImpl):
         if backend_kwargs:
             logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
-        # Lazily get the kernel from the Hub to avoid network/disk overhead on import
-        from kernels import get_kernel
-
-        logger.info("Loading flash-attn3 kernel from HuggingFace Hub...")
-        try:
-            hub_module = get_kernel("kernels-community/flash-attn3", version=1)
-        except Exception as e:
-            try:
-                logger.info("Failed to load version 1, attempting version 2: %s", e)
-                hub_module = get_kernel("kernels-community/flash-attn3", version=2)
-            except Exception as e2:
-                logger.info("Failed to load version 2, attempting default: %s", e2)
-                hub_module = get_kernel("kernels-community/flash-attn3")
-
+        hub_module = _get_hub_module("kernels-community/flash-attn3")
         self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
         self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
-
         if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
             raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found")