[Attention] Lazy-load LLM and ViT attention backends

lightllm/common/basemodel/attention/__init__.py

@@ -1,21 +1,4 @@
 from .base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl
-from .triton.fp import TritonAttBackend
-from .triton.int4kv import Int4kvTritonAttBackend
-from .triton.int8kv import Int8kvTritonAttBackend
-from .triton.mla import MlaTritonAttBackend
-from .triton.neo import NeoTritonAttBackend
-from .fa3.fp import Fa3AttBackend
-from .fa3.fp8 import Fp8Fa3AttBackend
-from .fa3.mla import MlaFa3AttBackend
-from .fa3.neo import NeoFa3AttBackend
-from .flashinfer.fp8 import Fp8FlashInferAttBackend
-from .flashinfer.fp import FlashInferAttBackend
-from .flashinfer.mla import MlaFlashInferAttBackend
-
-# NSA backend
-from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend
-from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend
-
 from .create_utils import (
     get_prefill_att_backend_class,
     get_decode_att_backend_class,

lightllm/common/basemodel/attention/create_linear_utils.py

@@ -1,57 +1,56 @@
-from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend
-from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend
+from .import_utils import load_backend_class
 from lightllm.utils.backend_validator import validate
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
 
 linear_prefill_att_backend_classes = {
-    "flashqla": FlashQlaLinearAttBackend,
-    "triton": TritonLinearAttBackend,
+    "flashqla": "linear.flashqla:FlashQlaLinearAttBackend",
+    "triton": "linear.triton:TritonLinearAttBackend",
 }
 
 linear_decode_att_backend_classes = {
-    "triton": TritonLinearAttBackend,
+    "triton": "linear.triton:TritonLinearAttBackend",
 }
 
 
 def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")):
     args = get_env_start_args()
     backend_str = _get_backend_str(args.llm_prefill_att_backend, index)
     if backend_str != "auto":
-        return linear_prefill_att_backend_classes[backend_str]
+        return load_backend_class(linear_prefill_att_backend_classes[backend_str])
 
     for backend_name in priority_list:
         backend_class = linear_prefill_att_backend_classes[backend_name]
         if backend_name == "triton":
             logger.info("Linear prefill attention backend: triton.")
-            return backend_class
+            return load_backend_class(backend_class)
         if validate(backend_name):
             logger.info(f"Linear prefill attention backend: {backend_name} (validated).")
-            return backend_class
+            return load_backend_class(backend_class)
 
     logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.")
-    return TritonLinearAttBackend
+    return load_backend_class("linear.triton:TritonLinearAttBackend")
 
 
 def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)):
     args = get_env_start_args()
     backend_str = _get_backend_str(args.llm_decode_att_backend, index)
     if backend_str != "auto":
-        return linear_decode_att_backend_classes[backend_str]
+        return load_backend_class(linear_decode_att_backend_classes[backend_str])
 
     for backend_name in priority_list:
         backend_class = linear_decode_att_backend_classes[backend_name]
         if backend_name == "triton":
             logger.info("Linear decode attention backend: triton.")
-            return backend_class
+            return load_backend_class(backend_class)
         if validate(backend_name):
             logger.info(f"Linear decode attention backend: {backend_name} (validated).")
-            return backend_class
+            return load_backend_class(backend_class)
 
     logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.")
-    return TritonLinearAttBackend
+    return load_backend_class("linear.triton:TritonLinearAttBackend")
 
 
 def _get_backend_str(backend_args, index):

lightllm/common/basemodel/attention/create_utils.py

@@ -3,109 +3,83 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.backend_validator import validate
-from typing import Dict
 from .base_att import BaseAttBackend
-from .triton.fp import TritonAttBackend
-from .triton.int4kv import Int4kvTritonAttBackend
-from .triton.int8kv import Int8kvTritonAttBackend
-from .triton.mla import MlaTritonAttBackend
-from .triton.neo import NeoTritonAttBackend
-from .fa3.fp import Fa3AttBackend
-from .fa3.fp8 import Fp8Fa3AttBackend
-from .fa3.mla import MlaFa3AttBackend
-from .fa3.neo import NeoFa3AttBackend, HAS_FLASH_ATTN_INTERFACE
-from .flashinfer.fp8 import Fp8FlashInferAttBackend
-from .flashinfer.fp import FlashInferAttBackend
-from .flashinfer.mla import MlaFlashInferAttBackend
-from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend
-from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend
+from .import_utils import load_backend_class
 
 logger = init_logger(__name__)
 
-# Backend class mappings by data type
+# Lazy backend class references by data type.
 data_type_to_backend = {
     "None": {
-        "triton": TritonAttBackend,
-        "fa3": Fa3AttBackend,
-        "flashinfer": FlashInferAttBackend,
+        "triton": "triton.fp:TritonAttBackend",
+        "fa3": "fa3.fp:Fa3AttBackend",
+        "flashinfer": "flashinfer.fp:FlashInferAttBackend",
     },
     "int4kv": {
-        "triton": Int4kvTritonAttBackend,
-        # "fa3": Fp8Fa3AttBackend,
-        # "flashinfer": Fp8FlashInferAttBackend,
+        "triton": "triton.int4kv:Int4kvTritonAttBackend",
     },
     "int8kv": {
-        "triton": Int8kvTritonAttBackend,
-        # "fa3": Fp8Fa3AttBackend,
-        # "flashinfer": Fp8FlashInferAttBackend,
+        "triton": "triton.int8kv:Int8kvTritonAttBackend",
     },
     "fp8kv_sph": {
-        "fa3": Fp8Fa3AttBackend,
+        "fa3": "fa3.fp8:Fp8Fa3AttBackend",
     },
     "fp8kv_spt": {
-        "flashinfer": Fp8FlashInferAttBackend,
+        "flashinfer": "flashinfer.fp8:Fp8FlashInferAttBackend",
     },
 }
 
 mla_data_type_to_backend = {
     "None": {
-        "triton": MlaTritonAttBackend,
-        "fa3": MlaFa3AttBackend,
-        "flashinfer": MlaFlashInferAttBackend,
+        "triton": "triton.mla:MlaTritonAttBackend",
+        "fa3": "fa3.mla:MlaFa3AttBackend",
+        "flashinfer": "flashinfer.mla:MlaFlashInferAttBackend",
     },
 }
 
 nsa_data_type_to_backend = {
     "None": {
-        "flashmla_sparse": NsaFlashMlaSparseAttBackend,
+        "flashmla_sparse": "nsa.flashmla_sparse:NsaFlashMlaSparseAttBackend",
         # Future backends: "fa3", "tilelang", "aiter"
     },
     "fp8kv_dsa": {
-        "flashmla_sparse": NsaFlashMlaFp8SparseAttBackend,
+        "flashmla_sparse": "nsa.fp8_flashmla_sparse:NsaFlashMlaFp8SparseAttBackend",
     },
 }
 
-neo_data_type_to_backend = (
-    {
-        "None": {
-            "triton": NeoTritonAttBackend,
-            "fa3": NeoFa3AttBackend,
-        }
-    }
-    if HAS_FLASH_ATTN_INTERFACE
-    else {
-        "None": {
-            "triton": NeoTritonAttBackend,
-        }
+neo_data_type_to_backend = {
+    "None": {
+        "triton": "triton.neo:NeoTritonAttBackend",
+        "fa3": "fa3.neo:NeoFa3AttBackend",
     }
-)
+}
 
 
 def _auto_select_backend(
     llm_dtype: str,
-    kv_type_to_backend: Dict[str, Dict[str, BaseAttBackend]],
+    kv_type_to_backend: dict,
     priority_list: list = ["fa3", "flashinfer", "triton"],
 ) -> type:
     """Auto-select the best available backend with validation.
 
     Priority follows the provided priority_list.
     Each backend is validated in a subprocess with ground truth checks.
     """
-    backend_map = kv_type_to_backend
+    backend_map = kv_type_to_backend[llm_dtype]
 
     args = get_env_start_args()
     if args.enable_ep_moe:
         logger.info("Expert parallelism with MoE enabled, excluding flashinfer attention backend")
         priority_list = [name for name in priority_list if name != "flashinfer"]
 
     for backend_name in priority_list:
-        if backend_name in backend_map[llm_dtype] and validate(backend_name):
+        if backend_name in backend_map and validate(backend_name):
             logger.info(f"Auto-selected {backend_name} backend (validated)")
-            return backend_map[llm_dtype][backend_name]
+            return load_backend_class(backend_map[backend_name])
 
     # Fallback to triton without validation (should not happen)
     logger.warning("No backend validation succeeded, falling back to triton")
-    return backend_map[llm_dtype]["triton"]
+    return load_backend_class(backend_map["triton"])
 
 
 def _get_decode_backend_priority(priority_list: list, mtp_step: int) -> list:
@@ -121,41 +95,37 @@ def get_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashi
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_prefill_att_backend[index]
     if backend_str != "auto":
-        return data_type_to_backend[llm_dtype][backend_str]
-    else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list)
+        return load_backend_class(data_type_to_backend[llm_dtype][backend_str])
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list)
 
 
 def get_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_decode_att_backend[index]
     if backend_str != "auto":
-        return data_type_to_backend[llm_dtype][backend_str]
-    else:
-        priority_list = _get_decode_backend_priority(priority_list, args.mtp_step)
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list)
+        return load_backend_class(data_type_to_backend[llm_dtype][backend_str])
+    priority_list = _get_decode_backend_priority(priority_list, args.mtp_step)
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list)
 
 
 def get_mla_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashinfer", "triton"]) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_prefill_att_backend[index]
     if backend_str != "auto":
-        return mla_data_type_to_backend[llm_dtype][backend_str]
-    else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
+        return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str])
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
 
 
 def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_decode_att_backend[index]
     if backend_str != "auto":
-        return mla_data_type_to_backend[llm_dtype][backend_str]
-    else:
-        priority_list = _get_decode_backend_priority(priority_list, args.mtp_step)
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
+        return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str])
+    priority_list = _get_decode_backend_priority(priority_list, args.mtp_step)
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list)
 
 
 def get_nsa_prefill_att_backend_class(
@@ -165,9 +135,8 @@ def get_nsa_prefill_att_backend_class(
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_prefill_att_backend[index]
     if backend_str != "auto":
-        return backend_map[llm_dtype][backend_str]
-    else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
+        return load_backend_class(backend_map[llm_dtype][backend_str])
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
 
 
 def get_nsa_decode_att_backend_class(
@@ -177,16 +146,21 @@ def get_nsa_decode_att_backend_class(
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_decode_att_backend[index]
     if backend_str != "auto":
-        return backend_map[llm_dtype][backend_str]
-    else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
+        return load_backend_class(backend_map[llm_dtype][backend_str])
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)
 
 
 def get_neo_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "triton"]) -> BaseAttBackend:
     args = get_env_start_args()
     llm_dtype = args.llm_kv_type
     backend_str = args.llm_prefill_att_backend[index]
+    backend_map = neo_data_type_to_backend
+    if backend_str == "fa3" or (backend_str == "auto" and "fa3" in priority_list):
+        # Only Neo FA3 selection needs the optional image-token attention interface.
+        from .fa3.neo import HAS_FLASH_ATTN_INTERFACE
+
+        if not HAS_FLASH_ATTN_INTERFACE:
+            backend_map = {"None": {"triton": neo_data_type_to_backend["None"]["triton"]}}
     if backend_str != "auto":
-        return neo_data_type_to_backend[llm_dtype][backend_str]
-    else:
-        return _auto_select_backend(llm_dtype, kv_type_to_backend=neo_data_type_to_backend, priority_list=priority_list)
+        return load_backend_class(backend_map[llm_dtype][backend_str])
+    return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list)

lightllm/common/basemodel/attention/import_utils.py

@@ -0,0 +1,6 @@
+from importlib import import_module
+
+
+def load_backend_class(backend):
+    module_name, class_name = backend.split(":")
+    return getattr(import_module(f"{__package__}.{module_name}"), class_name)

lightllm/common/basemodel/attention/linear/__init__.py

@@ -1,15 +0,0 @@
-from .gdn import (
-    LinearAttBackend,
-    LinearAttPrefillAttState,
-    LinearAttDecodeAttState,
-)
-from .flashqla import FlashQlaLinearAttBackend
-from .triton import TritonLinearAttBackend
-
-__all__ = [
-    "LinearAttBackend",
-    "LinearAttPrefillAttState",
-    "LinearAttDecodeAttState",
-    "FlashQlaLinearAttBackend",
-    "TritonLinearAttBackend",
-]

lightllm/common/basemodel/attention/nsa/__init__.py

@@ -1,21 +0,0 @@
-"""NSA (Native Sparse Attention) backend implementations."""
-
-from .flashmla_sparse import (
-    NsaFlashMlaSparseAttBackend,
-    NsaFlashMlaSparsePrefillAttState,
-    NsaFlashMlaSparseDecodeAttState,
-)
-from .fp8_flashmla_sparse import (
-    NsaFlashMlaFp8SparseAttBackend,
-    NsaFlashMlaFp8SparsePrefillAttState,
-    NsaFlashMlaFp8SparseDecodeAttState,
-)
-
-__all__ = [
-    "NsaFlashMlaSparseAttBackend",
-    "NsaFlashMlaSparsePrefillAttState",
-    "NsaFlashMlaSparseDecodeAttState",
-    "NsaFlashMlaFp8SparseAttBackend",
-    "NsaFlashMlaFp8SparsePrefillAttState",
-    "NsaFlashMlaFp8SparseDecodeAttState",
-]

lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py

@@ -196,4 +196,8 @@ def decode_att(
         return padded_out[:, : self.backend.tp_q_head_num]
 
 
-DSV4_NSA_BACKENDS = {"fp8kv_dsa": {"flashmla_sparse": DeepseekV4FlashMlaFp8SparseAttBackend}}
+DSV4_NSA_BACKENDS = {
+    "fp8kv_dsa": {
+        "flashmla_sparse": "nsa.dsv4_fp8_flashmla_sparse:DeepseekV4FlashMlaFp8SparseAttBackend",
+    }
+}