diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py index 76d8c380ad..bbe6f7b7dd 100644 --- a/lightllm/common/basemodel/attention/__init__.py +++ b/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, diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index b3405d268d..6a40a05183 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -1,5 +1,4 @@ -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 @@ -7,12 +6,12 @@ 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", } @@ -20,38 +19,38 @@ def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashql 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): diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py index 272c500319..a4e2641a9f 100644 --- a/lightllm/common/basemodel/attention/create_utils.py +++ b/lightllm/common/basemodel/attention/create_utils.py @@ -3,87 +3,61 @@ 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. @@ -91,7 +65,7 @@ def _auto_select_backend( 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: @@ -99,13 +73,13 @@ def _auto_select_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,9 +95,8 @@ 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: @@ -131,10 +104,9 @@ def get_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", " 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: @@ -142,9 +114,8 @@ def get_mla_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "fl 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: @@ -152,10 +123,9 @@ def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer 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) diff --git a/lightllm/common/basemodel/attention/import_utils.py b/lightllm/common/basemodel/attention/import_utils.py new file mode 100644 index 0000000000..14cb0b076d --- /dev/null +++ b/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) diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/basemodel/attention/linear/__init__.py index bfc8445390..e69de29bb2 100644 --- a/lightllm/common/basemodel/attention/linear/__init__.py +++ b/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", -] diff --git a/lightllm/common/basemodel/attention/nsa/__init__.py b/lightllm/common/basemodel/attention/nsa/__init__.py index f9db52dc2b..e69de29bb2 100644 --- a/lightllm/common/basemodel/attention/nsa/__init__.py +++ b/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", -] diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py index 9d62346555..536db3d4fe 100644 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ b/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", + } +} diff --git a/lightllm/common/basemodel/attention_vit/create_utils.py b/lightllm/common/basemodel/attention_vit/create_utils.py index 67f830ba0d..21eb9a7af9 100644 --- a/lightllm/common/basemodel/attention_vit/create_utils.py +++ b/lightllm/common/basemodel/attention_vit/create_utils.py @@ -1,27 +1,24 @@ -import torch +from importlib import import_module + from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.backend_validator import _validate from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend -from lightllm.common.basemodel.attention_vit.triton.fp import TritonVitAttBackend -from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend -from lightllm.common.basemodel.attention_vit.xformers.fp import XformersVitAttBackend logger = init_logger(__name__) vit_att_backend = { - "triton": TritonVitAttBackend, - "sdpa": SdpaVitAttBackend, - "fa3": Fa3VitAttBackend, - "xformers": XformersVitAttBackend, + "triton": "triton.fp:TritonVitAttBackend", + "sdpa": "sdpa.fp:SdpaVitAttBackend", + "fa3": "fa3.fp:Fa3VitAttBackend", + "xformers": "xformers.fp:XformersVitAttBackend", } def get_vit_att_backend_class(backend_name: str) -> BaseVitAttBackend: - vit_att_backend_class = vit_att_backend[backend_name] - return vit_att_backend_class + module_name, class_name = vit_att_backend[backend_name].split(":") + return getattr(import_module(f"{__package__}.{module_name}"), class_name) def init_vit_att_backend(index=0, priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: @@ -30,8 +27,7 @@ def init_vit_att_backend(index=0, priority_list: list = ["fa3", "xformers", "sdp if backend_name != "auto": logger.info(f"Selected {backend_name} backend for VIT") return backend_name - else: - return _select_vit_backend(priority_list=priority_list) + return _select_vit_backend(priority_list=priority_list) def _select_vit_backend(priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 9a1811a50a..9c4a42fffc 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -1,9 +1,7 @@ import torch -from lightllm.common.basemodel.attention import ( - Fa3AttBackend, - Fp8Fa3AttBackend, -) +from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend +from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.models.llama.model import LlamaTpPartModel diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index b21c897d10..eeb2fb1438 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -313,11 +313,10 @@ def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backen def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_args): validate_calls = [] - optimized_backend = type("OptimizedLinearDecodeAttBackend", (), {}) monkeypatch.setattr( linear_create_utils, "linear_decode_att_backend_classes", - {"optimized": optimized_backend, "triton": TritonLinearAttBackend}, + {"optimized": "linear.flashqla:FlashQlaLinearAttBackend", "triton": "linear.triton:TritonLinearAttBackend"}, ) monkeypatch.setattr( linear_create_utils, @@ -329,7 +328,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ index=1, priority_list=("optimized", "triton") ) - assert backend_class is optimized_backend + assert backend_class is FlashQlaLinearAttBackend assert validate_calls == ["optimized"] @@ -343,12 +342,10 @@ def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypat def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_backend_args): validate_calls = [] - flashqla2_backend = type("FlashQla2LinearAttBackend", (), {}) - flashqla3_backend = type("FlashQla3LinearAttBackend", (), {}) monkeypatch.setattr( linear_create_utils, "linear_prefill_att_backend_classes", - {"flashqla2": flashqla2_backend, "flashqla3": flashqla3_backend}, + {"flashqla2": "linear.triton:TritonLinearAttBackend", "flashqla3": "linear.flashqla:FlashQlaLinearAttBackend"}, ) monkeypatch.setattr( linear_create_utils, @@ -360,5 +357,5 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ index=1, priority_list=("flashqla2", "flashqla3") ) - assert backend_class is flashqla3_backend + assert backend_class is FlashQlaLinearAttBackend assert validate_calls == ["flashqla2", "flashqla3"] diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py new file mode 100644 index 0000000000..1efed8423e --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -0,0 +1,240 @@ +"""Check import isolation and selection without loading optional GPU dependencies.""" + +import importlib +import importlib.abc +import importlib.util +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +import unittest +from unittest.mock import Mock, patch + + +ROOT = Path(__file__).resolve().parents[4] +ATTENTION = "lightllm.common.basemodel.attention" + + +class BackendImports(importlib.abc.MetaPathFinder, importlib.abc.Loader): + """Reject unselected implementations, and supply lightweight selected backends.""" + + def __init__(self, package=ATTENTION): + self.package = package + self.allowed = {} + self.imported = [] + + def allow(self, suffix, class_name, **attributes): + module = ModuleType(f"{self.package}.{suffix}") + setattr(module, class_name, type(class_name, (), {})) + module.__dict__.update(attributes) + self.allowed[module.__name__] = module + return getattr(module, class_name) + + def find_spec(self, fullname, path=None, target=None): + suffix = fullname.removeprefix(f"{self.package}.") + if not fullname.startswith(f"{self.package}.") or "." not in suffix: + return None + if suffix.split(".")[0] not in {"triton", "fa3", "flashinfer", "nsa", "linear", "sdpa", "xformers"}: + return None + self.imported.append(fullname) + if fullname not in self.allowed: + raise AssertionError(f"Unexpected attention backend import: {fullname}") + return importlib.util.spec_from_loader(fullname, self) + + def create_module(self, spec): + return self.allowed[spec.name] + + def exec_module(self, module): + pass + + +class TestLazyAttentionImports(unittest.TestCase): + def setUp(self): + # Isolate attention from basemodel's unrelated GPU initialization. + modules_patch = patch.dict(sys.modules) + modules_patch.start() + self.addCleanup(modules_patch.stop) + for name in list(sys.modules): + if name == ATTENTION or name.startswith(f"{ATTENTION}."): + del sys.modules[name] + for name in ("lightllm", "lightllm.common", "lightllm.common.basemodel", "lightllm.utils"): + module = ModuleType(name) + module.__path__ = [str(ROOT.joinpath(*name.split(".")))] + sys.modules[name] = module + + self.args = SimpleNamespace( + llm_kv_type="None", + llm_prefill_att_backend=["triton"], + llm_decode_att_backend=["triton"], + enable_ep_moe=False, + mtp_step=0, + ) + self.validate = Mock(return_value=True) + dependencies = { + f"{ATTENTION}.base_att": { + name: type(name, (), {}) + for name in ("BaseAttBackend", "BasePrefillAttState", "BaseDecodeAttState", "AttControl") + }, + "lightllm.utils.envs_utils": {"get_env_start_args": lambda: self.args}, + "lightllm.utils.backend_validator": {"validate": self.validate}, + "lightllm.utils.log_utils": {"init_logger": lambda name: Mock()}, + } + for name, attributes in dependencies.items(): + module = ModuleType(name) + module.__dict__.update(attributes) + sys.modules[name] = module + + self.imports = BackendImports() + sys.meta_path.insert(0, self.imports) + self.addCleanup(sys.meta_path.remove, self.imports) + self.attention = importlib.import_module(ATTENTION) + + def test_importing_selection_apis_loads_no_implementations(self): + importlib.import_module(f"{ATTENTION}.create_linear_utils") + importlib.import_module(f"{ATTENTION}.linear") + importlib.import_module(f"{ATTENTION}.nsa") + self.assertEqual(self.imports.imported, []) + + def test_explicit_regular_selection_skips_neo(self): + backend = self.imports.allow("fa3.fp", "Fa3AttBackend") + self.args.llm_prefill_att_backend = ["fa3"] + + self.assertIs(self.attention.get_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.fp"]) + self.validate.assert_not_called() + + def test_auto_selection_does_not_import_failed_candidates(self): + backend = self.imports.allow("flashinfer.fp", "FlashInferAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + self.validate.side_effect = lambda name: name == "flashinfer" + + self.assertIs(self.attention.get_prefill_att_backend_class(), backend) + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "flashinfer"]) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.flashinfer.fp"]) + + def test_mtp_decode_preserves_fa3_priority(self): + backend = self.imports.allow("fa3.fp", "Fa3AttBackend") + self.args.llm_decode_att_backend = ["auto"] + self.args.mtp_step = 2 + + self.assertIs(self.attention.get_decode_att_backend_class(), backend) + self.validate.assert_called_once_with("fa3") + + def test_expert_parallelism_still_excludes_flashinfer(self): + backend = self.imports.allow("triton.fp", "TritonAttBackend") + self.args.llm_decode_att_backend = ["auto"] + self.args.enable_ep_moe = True + self.validate.side_effect = lambda name: name == "triton" + + self.assertIs(self.attention.get_decode_att_backend_class(), backend) + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "triton"]) + + def test_auto_selection_falls_back_to_triton(self): + backend = self.imports.allow("triton.mla", "MlaTritonAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + self.validate.return_value = False + + self.assertIs(self.attention.get_mla_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.mla"]) + + def test_nsa_selection_does_not_import_other_kv_types(self): + backend = self.imports.allow("nsa.fp8_flashmla_sparse", "NsaFlashMlaFp8SparseAttBackend") + self.args.llm_kv_type = "fp8kv_dsa" + self.args.llm_decode_att_backend = ["flashmla_sparse"] + + self.assertIs(self.attention.get_nsa_decode_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.fp8_flashmla_sparse"]) + + def test_custom_backend_map_uses_import_paths(self): + backend = self.imports.allow("nsa.custom", "CustomBackend") + backend_map = {"None": {"custom": "nsa.custom:CustomBackend"}} + self.args.llm_prefill_att_backend = ["custom"] + self.args.llm_decode_att_backend = ["auto"] + + self.assertIs(self.attention.get_nsa_prefill_att_backend_class(backend_map=backend_map), backend) + self.assertIs( + self.attention.get_nsa_decode_att_backend_class(priority_list=["custom"], backend_map=backend_map), + backend, + ) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.custom"]) + + def test_neo_triton_selection_does_not_probe_fa3(self): + backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.args.llm_prefill_att_backend = ["auto"] + self.assertIs(self.attention.get_neo_prefill_att_backend_class(priority_list=["triton"]), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.neo"]) + + def test_neo_auto_selection_still_falls_back_without_image_token_interface(self): + self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=False) + backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.validate.assert_called_once_with("triton") + self.args.llm_prefill_att_backend = ["fa3"] + with self.assertRaises(KeyError): + self.attention.get_neo_prefill_att_backend_class() + + def test_neo_auto_selection_uses_supported_fa3(self): + backend = self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=True) + self.args.llm_prefill_att_backend = ["auto"] + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) + self.validate.assert_called_once_with("fa3") + + def test_linear_selection_imports_only_validated_backend(self): + backend = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + + self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.flashqla"]) + self.validate.assert_called_once_with("flashqla") + + def test_linear_selection_falls_back_without_importing_flashqla(self): + backend = self.imports.allow("linear.triton", "TritonLinearAttBackend") + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + self.validate.return_value = False + + self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) + self.assertIs(selection.get_qwen35_linear_decode_att_backend_class(), backend) + self.args.llm_prefill_att_backend.append("triton") + self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.triton"]) + + def test_vit_selection_loads_only_requested_implementation(self): + package = "lightllm.common.basemodel.attention_vit" + for name in list(sys.modules): + if name == package or name.startswith(f"{package}."): + del sys.modules[name] + base = ModuleType(f"{package}.base_att") + base.BaseVitAttBackend = type("BaseVitAttBackend", (), {}) + sys.modules[base.__name__] = base + sys.modules["lightllm.utils.backend_validator"]._validate = self.validate + imports = BackendImports(package) + sys.meta_path.insert(0, imports) + self.addCleanup(sys.meta_path.remove, imports) + selection = importlib.import_module(f"{package}.create_utils") + self.assertEqual(imports.imported, []) + + for name, class_name in ( + ("fa3", "Fa3VitAttBackend"), + ("triton", "TritonVitAttBackend"), + ("sdpa", "SdpaVitAttBackend"), + ("xformers", "XformersVitAttBackend"), + ): + backend = imports.allow(f"{name}.fp", class_name) + self.args.vit_att_backend = [name] + self.assertEqual(selection.init_vit_att_backend(), name) + self.assertIs(selection.get_vit_att_backend_class(name), backend) + self.validate.assert_not_called() + + self.args.vit_att_backend = ["auto"] + self.validate.side_effect = lambda name: name == "sdpa" + self.assertEqual(selection.init_vit_att_backend(), "sdpa") + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "xformers", "sdpa"]) + + +if __name__ == "__main__": + unittest.main()