add-choose-vit-backend (#1191)

Co-authored-by: sangchengmeng <sangchengmeng@sensetime.com>
Co-authored-by: wangzaijun <wangzaijun@sensetime.com>
Co-authored-by: shihaobai <1798930569@qq.com>
Co-authored-by: hiworldwzj <30762946+hiworldwzj@users.noreply.github.com>

Author: 5 people

docs/CN/source/tutorial/api_server_args.rst

@@ -284,6 +284,17 @@ PD 分离模式参数
 
     为 ViT 构建分布式环境的 NCCL 端口列表，例如 29500 29501 29502，默认为 [29500]
 
+.. option:: --vit_att_backend
+
+    设置 ViT 使用的注意力后端。可选值为：
+
+    * ``auto``: 自动选择最佳后端（默认值），优先级为 fa3 > xformers > sdpa > triton
+    * ``fa3``: 使用 Flash-Attention 3 后端
+    * ``xformers``: 使用 xformers 后端
+    * ``sdpa``: 使用 sdpa 后端
+    * ``triton``: 使用 Triton 后端
+
+
 性能优化参数
 ------------
 

docs/EN/source/tutorial/api_server_args.rst

@@ -282,6 +282,17 @@ Multimodal Parameters
 
     List of NCCL ports for ViT, e.g., 29500 29501 29502, default is [29500]
 
+.. option:: --vit_att_backend
+
+    Set the attention backend for ViT. Available options:
+
+    * ``auto``: Automatically select the best backend (default), with priority fa3 > xformers > sdpa > triton
+    * ``fa3``: Use Flash-Attention 3 backend
+    * ``xformers``: Use xformers backend
+    * ``sdpa``: Use sdpa backend
+    * ``triton``: Use Triton backend
+
+
 Performance Optimization Parameters
 -----------------------------------
 

lightllm/common/basemodel/attention_vit/__init__.py

@@ -0,0 +1,38 @@
+import torch
+from abc import ABC, abstractmethod
+
+
+class BaseVitAttBackend(ABC):
+    """
+    用于创建支持各种不同的AttBackend, 如 fa3, sdpa, triton 实现等，
+    这个是单列模式, 每种backend只有一个实例
+    """
+
+    _instances = {}
+
+    def __new__(cls, *args, **kwargs):
+        """
+        重写__new__方法实现单例模式
+        """
+        # 检查是否已经有该类的实例
+        if cls not in cls._instances:
+            # 创建新实例并存储
+            instance = super().__new__(cls)
+            cls._instances[cls] = instance
+        # 返回已有的实例
+        return cls._instances[cls]
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def _vit_att_fwd(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        raise NotImplementedError("not impl")

lightllm/common/basemodel/attention_vit/base_att.py

@@ -0,0 +1,51 @@
+import torch
+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,
+}
+
+
+def get_vit_att_backend_class(backend_name: str) -> BaseVitAttBackend:
+    vit_att_backend_class = vit_att_backend[backend_name]
+    return vit_att_backend_class
+
+
+def init_vit_att_backend(index=0, priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str:
+    args = get_env_start_args()
+    backend_name = args.vit_att_backend[index]
+    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)
+
+
+def _select_vit_backend(priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str:
+    """Auto-select the best available backend with validation for VIT.
+
+    Priority: FA3 > Xformers > Sdpa > Triton
+    Each backend is validated in a subprocess with ground truth checks.
+    """
+
+    for backend_name in priority_list:
+        if _validate(backend_name):
+            logger.info(f"Auto-selected {backend_name} backend (validated) for VIT")
+            return backend_name
+
+    # Fallback to triton without validation (should not happen)
+    logger.warning("No backend validation succeeded, falling back to triton")
+    return "triton"

lightllm/common/basemodel/attention_vit/create_utils.py

@@ -0,0 +1,57 @@
+import dataclasses
+import torch
+from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend
+
+
+class Fa3VitAttBackend(BaseVitAttBackend):
+    @staticmethod
+    def _vit_att_fwd(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> None:
+
+        head_dim = q.shape[-1]
+        softmax_scale = head_dim ** -0.5
+        window_size = (-1, -1)
+        torch.ops.sgl_kernel.fwd.default(
+            q,
+            k,
+            v,
+            None,  # k_new
+            None,  # v_new
+            None,  # qv
+            o,  # out
+            cu_seqlens,
+            cu_seqlens,
+            None,  # cu_seqlens_k_new
+            None,
+            None,
+            max_seqlen,
+            max_seqlen,
+            None,  # page_table,
+            None,  # kv_batch_idx
+            None,  # leftpad_k
+            None,  # rotary cos
+            None,  # rotary sin
+            None,  # seqlens_rotary
+            None,
+            None,
+            None,
+            softmax_scale,
+            False,
+            window_size[0],
+            window_size[1],
+            0.0,
+            is_rotary_interleaved=False,
+            scheduler_metadata=None,
+            num_splits=1,
+            pack_gqa=None,
+            sm_margin=0,
+            sinks=None,
+        )
+
+        return o

lightllm/common/basemodel/attention_vit/fa3/__init__.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn.functional as F
+from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend
+
+
+class SdpaVitAttBackend(BaseVitAttBackend):
+    @staticmethod
+    def _vit_att_fwd(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        assert q.ndim == k.ndim == v.ndim == o.ndim == 3
+        assert cu_seqlens is not None and cu_seqlens.ndim == 1
+        cu_seqlens = cu_seqlens.detach().to("cpu")
+        B = cu_seqlens.numel() - 1
+
+        with torch.no_grad():
+            for b in range(B):
+                s = int(cu_seqlens[b])
+                e = int(cu_seqlens[b + 1])
+                L = e - s
+                if L <= 0:
+                    continue
+                if max_seqlen:
+                    assert L <= max_seqlen
+
+                # [L, H, D] -> [1, H, L, D]
+                q_ = q[s:e].permute(1, 0, 2).unsqueeze(0)
+                k_ = k[s:e].permute(1, 0, 2).unsqueeze(0)
+                v_ = v[s:e].permute(1, 0, 2).unsqueeze(0)
+
+                out = F.scaled_dot_product_attention(
+                    q_,
+                    k_,
+                    v_,
+                    attn_mask=None,
+                    dropout_p=0.0,
+                    is_causal=False,
+                )
+                # [1, H, L, D] -> [L, H, D]
+                o[s:e].copy_(out.squeeze(0).permute(1, 0, 2))
+
+        return o