From 505f31589c623b4d912103076bef84209b2157ca Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 08:00:32 +0800 Subject: [PATCH] refactor(attention): use varlen operators for vision and audio --- lightllm/common/attention/varlen.py | 42 ++- .../basemodel/attention_vit/__init__.py | 0 .../basemodel/attention_vit/base_att.py | 38 -- .../basemodel/attention_vit/create_utils.py | 47 --- .../basemodel/attention_vit/fa3/__init__.py | 0 .../common/basemodel/attention_vit/fa3/fp.py | 60 ---- .../basemodel/attention_vit/sdpa/__init__.py | 0 .../common/basemodel/attention_vit/sdpa/fp.py | 47 --- .../attention_vit/triton/__init__.py | 0 .../basemodel/attention_vit/triton/fp.py | 24 -- .../attention_vit/xformers/__init__.py | 0 .../basemodel/attention_vit/xformers/fp.py | 45 --- .../models/deepseek_v4/deepseek_v4_visual.py | 15 +- lightllm/models/glm5_next/glm5_next_visual.py | 6 +- lightllm/models/qwen2_5_vl/qwen2_5_visual.py | 32 +- lightllm/models/qwen2_vl/qwen2_visual.py | 30 +- .../qwen3_omni_audio.py | 32 +- .../layer_infer/transformer_layer_infer.py | 12 +- .../vit/triton_kernel/flashattention_nopad.py | 229 ------------ lightllm/server/api_cli.py | 3 +- lightllm/server/core/objs/start_args_type.py | 4 +- lightllm/server/visualserver/__init__.py | 46 ++- lightllm/server/visualserver/manager.py | 36 +- .../visualserver/model_infer/model_rpc.py | 47 +-- .../visualserver/visual_only_manager.py | 44 +-- .../common/attention/test_varlen_selection.py | 131 +++++++ .../basemodel/attention/test_lazy_import.py | 45 +-- unit_tests/models/deepseek_v4/test_vision.py | 30 +- unit_tests/models/glm5_next/test_visual.py | 22 +- .../vit/test_flash_attention_forward.py | 335 +++++++++++------- 30 files changed, 586 insertions(+), 816 deletions(-) delete mode 100644 lightllm/common/basemodel/attention_vit/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/base_att.py delete mode 100644 lightllm/common/basemodel/attention_vit/create_utils.py delete mode 100644 lightllm/common/basemodel/attention_vit/fa3/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/fa3/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/sdpa/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/sdpa/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/triton/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/triton/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/xformers/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/xformers/fp.py delete mode 100644 lightllm/models/vit/triton_kernel/flashattention_nopad.py create mode 100644 unit_tests/common/attention/test_varlen_selection.py diff --git a/lightllm/common/attention/varlen.py b/lightllm/common/attention/varlen.py index aa758efdea..b6b5b068e7 100644 --- a/lightllm/common/attention/varlen.py +++ b/lightllm/common/attention/varlen.py @@ -1,35 +1,43 @@ from lightllm.platform import get_device_runtime -from lightllm.utils.log_utils import init_logger from .api import AttentionFamily, AttentionPhase, UnsupportedAttention from .selection import BackendSelectionError, get_varlen_attention_op_class -logger = init_logger(__name__) - class AutoVarlenAttentionOp: - """音频编码器的自动路径；失败的实现只在当前调用方禁用，并在本次调用中重算输出。""" + """One encoder owns this runner and shares stateless operators across its layers. + + Unsupported inputs can try the next provider for this call. Kernel and programming + errors propagate: retrying on a potentially failed device is not a safe fallback. + Environment rejections are cached by provider/device; input rejections are not. + """ def __init__(self, providers=("fa3", "triton")): self.providers = providers self.ops = {} - self.failures = {} + self.unavailable = {} def forward(self, inputs, *, out): + failures = {} + caps = get_device_runtime(inputs.query.device).caps for name in self.providers: - if name in self.failures: + key = (name, inputs.query.device) + if key in self.unavailable: + failures[name] = self.unavailable[key] continue - try: - if name not in self.ops: + if key not in self.ops: + try: op_cls = get_varlen_attention_op_class(name) - caps = get_device_runtime(inputs.query.device).caps op_cls.check_environment(caps) - op_cls.check_support(inputs, caps) - self.ops[name] = op_cls() - return self.ops[name].forward(inputs, out=out) + self.ops[key] = op_cls() + except (UnsupportedAttention, ImportError) as exc: + self.unavailable[key] = failures[name] = str(exc) + continue + op = self.ops[key] + try: + op.check_support(inputs, caps) except UnsupportedAttention as exc: - self.failures[name] = str(exc) - except Exception as exc: - self.failures[name] = str(exc) - logger.warning("Disabling %s varlen attention for this encoder: %s", name, exc) - raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, self.failures) + failures[name] = str(exc) + continue + return op.forward(inputs, out=out) + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) diff --git a/lightllm/common/basemodel/attention_vit/__init__.py b/lightllm/common/basemodel/attention_vit/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/base_att.py b/lightllm/common/basemodel/attention_vit/base_att.py deleted file mode 100644 index 49bf6ad74c..0000000000 --- a/lightllm/common/basemodel/attention_vit/base_att.py +++ /dev/null @@ -1,38 +0,0 @@ -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") diff --git a/lightllm/common/basemodel/attention_vit/create_utils.py b/lightllm/common/basemodel/attention_vit/create_utils.py deleted file mode 100644 index 21eb9a7af9..0000000000 --- a/lightllm/common/basemodel/attention_vit/create_utils.py +++ /dev/null @@ -1,47 +0,0 @@ -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 - -logger = init_logger(__name__) - - -vit_att_backend = { - "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: - 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: - 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 - 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" diff --git a/lightllm/common/basemodel/attention_vit/fa3/__init__.py b/lightllm/common/basemodel/attention_vit/fa3/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py deleted file mode 100644 index a8e56417c5..0000000000 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ /dev/null @@ -1,60 +0,0 @@ -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], - attention_chunk=0, - softcap=0.0, - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - - return o diff --git a/lightllm/common/basemodel/attention_vit/sdpa/__init__.py b/lightllm/common/basemodel/attention_vit/sdpa/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/sdpa/fp.py b/lightllm/common/basemodel/attention_vit/sdpa/fp.py deleted file mode 100644 index 6c6da2c2b1..0000000000 --- a/lightllm/common/basemodel/attention_vit/sdpa/fp.py +++ /dev/null @@ -1,47 +0,0 @@ -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 diff --git a/lightllm/common/basemodel/attention_vit/triton/__init__.py b/lightllm/common/basemodel/attention_vit/triton/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/triton/fp.py b/lightllm/common/basemodel/attention_vit/triton/fp.py deleted file mode 100644 index 88867102a4..0000000000 --- a/lightllm/common/basemodel/attention_vit/triton/fp.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.models.vit.triton_kernel.flashattention_nopad import _flash_attention_triton_fwd - - -class TritonVitAttBackend(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, - ): - _flash_attention_triton_fwd( - q, - k, - v, - o, - cu_seqlens, # q k v cu_seqlens, - max_seqlen, - ) - return o diff --git a/lightllm/common/basemodel/attention_vit/xformers/__init__.py b/lightllm/common/basemodel/attention_vit/xformers/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/xformers/fp.py b/lightllm/common/basemodel/attention_vit/xformers/fp.py deleted file mode 100644 index 361b5db050..0000000000 --- a/lightllm/common/basemodel/attention_vit/xformers/fp.py +++ /dev/null @@ -1,45 +0,0 @@ -import torch -import torch.nn.functional as F - -try: - from xformers import ops as xformers_ops - from xformers.ops import fmha -except ImportError: - xformers_ops = None - fmha = None - -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class XformersVitAttBackend(BaseVitAttBackend): - @torch.no_grad() - @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 - assert q.shape == k.shape == v.shape == o.shape - - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int64).tolist() - seqlens = [int(L) for L in seqlens if int(L) > 0] - - if len(seqlens) == 0: - return o - if max_seqlen: - assert max(seqlens) <= max_seqlen - - attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens, device=q.device) - - q_ = q.unsqueeze(0) # [1, T, H, D] - k_ = k.unsqueeze(0) # [1, T, H, D] - v_ = v.unsqueeze(0) # [1, T, H, D] - - out = xformers_ops.memory_efficient_attention(q_, k_, v_, attn_bias=attn_bias, p=0.0) - o.copy_(out.squeeze(0)) # [T, H, D] - return o diff --git a/lightllm/models/deepseek_v4/deepseek_v4_visual.py b/lightllm/models/deepseek_v4/deepseek_v4_visual.py index 3e730c5b2a..24985d08fd 100644 --- a/lightllm/models/deepseek_v4/deepseek_v4_visual.py +++ b/lightllm/models/deepseek_v4/deepseek_v4_visual.py @@ -13,14 +13,11 @@ from safetensors import safe_open from torch import nn -from lightllm.models.deepseek_v4.image_processor import ( - IMAGE, - build_canonical_image_block, - load_image, -) +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.models.deepseek_v4.image_processor import IMAGE, build_canonical_image_block, load_image from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op @lru_cache(8) @@ -56,7 +53,7 @@ class PatchEmbed(nn.Module): def __init__(self, args): super().__init__() self.proj = nn.Linear( - 3 * args.vision_patch_size ** 2, + 3 * args.vision_patch_size**2, args.vision_dim, dtype=torch.bfloat16, ) @@ -79,7 +76,7 @@ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cu_seql q = apply_rotary(q, cos, sin) k = apply_rotary(k, cos, sin) o = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, o, cu_seqlens, n) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, n), out=o) return self.wo(o.reshape(n, -1)) @@ -143,7 +140,7 @@ class Aligner(nn.Module): def __init__(self, args): super().__init__() self.downsample_ratio = args.vision_downsample_ratio - in_dim = args.vision_dim * self.downsample_ratio ** 2 + in_dim = args.vision_dim * self.downsample_ratio**2 self.w1 = nn.Linear(in_dim, args.hidden_size, dtype=torch.bfloat16) self.w2 = nn.Linear(args.hidden_size, args.hidden_size, dtype=torch.bfloat16) diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py index 43637c2cbb..c27e68aa76 100644 --- a/lightllm/models/glm5_next/glm5_next_visual.py +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -9,10 +9,12 @@ import torch.nn.functional as F from safetensors import safe_open +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op + from .vision_process import Glm5NextImageProcessor @@ -56,7 +58,7 @@ def forward(self, x, cu_seqlens, max_seqlen, rotary_cos, rotary_sin): q = apply_rotary_pos_emb_triton(q, rotary_cos, rotary_sin) k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) out = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return self.proj(out.reshape(x.shape[0], -1)) diff --git a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py index 1b3a5f0db7..fcfe0726c8 100644 --- a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py +++ b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py @@ -1,21 +1,24 @@ -import os import json +import os +from io import BytesIO +from typing import List, Optional + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List, Optional -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn -from transformers.activations import ACT2FN -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding -from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from transformers.activations import ACT2FN + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op class Qwen2RMSNorm(nn.Module): @@ -74,7 +77,7 @@ def forward( k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -122,7 +125,7 @@ def forward( class Qwen2_5_VLPatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: super().__init__() - self.hidden_size = context_dim * (spatial_merge_size ** 2) + self.hidden_size = context_dim * (spatial_merge_size**2) self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(self.hidden_size, self.hidden_size), @@ -350,7 +353,6 @@ def load_image(self, img: List[ImageItem]): return pixel_values.to(dtype=self.data_type), image_grid_thw def load_model(self, weight_dir): - bin_weight_files = [file_ for file_ in os.listdir(weight_dir) if file_.endswith(".bin")] if bin_weight_files: weight_dict = {} @@ -390,7 +392,7 @@ def encode(self, images: List[ImageItem]): raise Exception("Unsupported input types: {} for {}".format(type(img), img)) # must devide merge_length - cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size ** 2) + cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size**2) valid_ids.append([valid_id, valid_id + cur_num]) valid_id += cur_num diff --git a/lightllm/models/qwen2_vl/qwen2_visual.py b/lightllm/models/qwen2_vl/qwen2_visual.py index e02c3d9aa3..172e5aacb9 100644 --- a/lightllm/models/qwen2_vl/qwen2_visual.py +++ b/lightllm/models/qwen2_vl/qwen2_visual.py @@ -18,24 +18,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import json +import os +from io import BytesIO +from typing import List + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List -from torchvision import transforms as T -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn +from safetensors import safe_open from torch.nn import LayerNorm +from torchvision import transforms as T from transformers.activations import ACT2FN -from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op # adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -73,7 +76,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class PatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: super().__init__() - self.hidden_size = context_dim * (spatial_merge_size ** 2) + self.hidden_size = context_dim * (spatial_merge_size**2) self.ln_q = LayerNorm(context_dim, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(self.hidden_size, self.hidden_size), @@ -147,7 +150,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -239,7 +242,6 @@ def _init_datatype(self): return def load_model(self, weight_dir): - processor_config_path = os.path.join(weight_dir, "preprocessor_config.json") with open(processor_config_path, "r") as f: processor_config_dict = json.load(f) @@ -322,7 +324,7 @@ def encode(self, images: List[ImageItem]): raise Exception("Unsupported input types: {} for {}".format(type(img), img)) # must devide merge_length - cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size ** 2) + cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size**2) valid_ids.append([valid_id, valid_id + cur_num]) valid_id += cur_num diff --git a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py index b9ea66fb15..5bd945d000 100644 --- a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py +++ b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py @@ -1,19 +1,21 @@ -import os import json import math -import torch +import os +from typing import Callable, List, Optional, Union + import numpy as np -from torch import Tensor, nn +import torch from safetensors import safe_open +from torch import Tensor, nn from torch.nn import functional as F -from typing import Callable, Optional, Union, List from transformers.activations import ACT2FN -from lightllm.server.multimodal_params import AudioItem -from lightllm.utils.log_utils import init_logger +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.attention.varlen import AutoVarlenAttentionOp from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd from lightllm.models.qwen3_omni_moe_thinker.audio_process import WhisperFeatureExtractor +from lightllm.server.multimodal_params import AudioItem +from lightllm.utils.log_utils import init_logger QWEN3_OMNI_CONV_CHUNKSIZE = int(os.getenv("LIGHTLLM_QWEN3_OMNI_CONV_CHUNKSIZE", 200)) @@ -24,7 +26,6 @@ def _get_feat_extract_output_lengths(input_lengths): """ Computes the output length of the convolutional layers and the output length of the audio encoder """ - input_lengths_leave = input_lengths % 100 feat_lengths = (input_lengths_leave - 1) // 2 + 1 output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 @@ -41,10 +42,13 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=None, ): super().__init__() self.embed_dim = d_model - self.self_attn = Qwen3OmniMoeAudioAttention(d_model, encoder_attention_heads, attention_dropout) + self.self_attn = Qwen3OmniMoeAudioAttention( + d_model, encoder_attention_heads, attention_dropout, attention_op=attention_op + ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = dropout self.activation_fn = ACT2FN[activation_function] @@ -86,7 +90,7 @@ def forward( class Qwen3OmniMoeAudioAttention(nn.Module): - def __init__(self, d_model, encoder_attention_heads, attention_dropout): + def __init__(self, d_model, encoder_attention_heads, attention_dropout, *, attention_op=None): super().__init__() self.embed_dim = d_model self.num_heads = encoder_attention_heads @@ -99,10 +103,11 @@ def __init__(self, d_model, encoder_attention_heads, attention_dropout): f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" f" and `num_heads`: {self.num_heads})." ) - self.scaling = self.head_dim ** -0.5 + self.scaling = self.head_dim**-0.5 self.attention_dropout = 0.0 self.is_decoder = False self.is_causal = False + self.attention_op = attention_op if attention_op is not None else AutoVarlenAttentionOp() self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) @@ -116,7 +121,6 @@ def forward( **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" - seq_length, _ = hidden_states.size() q = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1) @@ -125,7 +129,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - flash_attention_fwd(q, k, v, attn_output, cu_seqlens, max_seqlen) + self.attention_op.forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.out_proj(attn_output) @@ -176,6 +180,7 @@ def __init__( self.embed_scale = math.sqrt(self.embed_dim) if scale_embedding else 1.0 self.n_window = n_window self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, self.embed_dim) + attention_op = AutoVarlenAttentionOp() self.layers = nn.ModuleList( [ Qwen3OmniMoeAudioEncoderLayer( @@ -186,6 +191,7 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=attention_op, ) for _ in range(encoder_layers) ] diff --git a/lightllm/models/vit/layer_infer/transformer_layer_infer.py b/lightllm/models/vit/layer_infer/transformer_layer_infer.py index cada55e58b..f8dbf2b7d4 100644 --- a/lightllm/models/vit/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/vit/layer_infer/transformer_layer_infer.py @@ -1,13 +1,15 @@ +from typing import Tuple + import torch import torch.distributed as dist -from typing import Tuple +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.vit.layer_weights.transformer_layer_weight import ViTTransformerLayerWeight -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.models.vit.triton_kernel.gelu_vit import gelu_fwd from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.server.visualserver import get_vit_attention_op +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size class ViTTransformerLayerInfer: @@ -105,7 +107,7 @@ def _context_attention_kernel(self, q, k, v) -> torch.Tensor: q, k, v, out = map(reshape, (q, k, v, out)) cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) * seq_len max_seqlen = seq_len - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return out.reshape(batch_size, seq_len, -1) def _get_o(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor: diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py deleted file mode 100644 index af062c2a5b..0000000000 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ /dev/null @@ -1,229 +0,0 @@ -import torch -import triton -import triton.language as tl -import math -import time -import torch.nn.functional as F -from typing import Optional, Tuple -from lightllm.utils.device_utils import is_hopper, is_4090 - -if triton.__version__ >= "2.1.0": - - @triton.jit - def _fwd_kernel( - Q, - K, - V, - sm_scale, - Out, - q_stride_s, - q_stride_h, - q_stride_d, - k_stride_s, - k_stride_h, - k_stride_d, - v_stride_s, - v_stride_h, - v_stride_d, - o_stride_s, - o_stride_h, - o_stride_d, - head_dim_act, - cu_seqlens, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - ): - cur_batch = tl.program_id(2) - cur_head = tl.program_id(1) - start_m = tl.program_id(0) - - seq_start = tl.load(cu_seqlens + cur_batch).to(tl.int32) - seq_end = tl.load(cu_seqlens + cur_batch + 1).to(tl.int32) - seq_len = seq_end - seq_start - - # initialize offsets - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - - mask_d = offs_d < head_dim_act - off_q = cur_head * q_stride_h + (seq_start + offs_m[:, None]) * q_stride_s + offs_d[None, :] * q_stride_d - q = tl.load(Q + off_q, mask=(offs_m[:, None] < seq_len) & mask_d[None, :], other=0.0) - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - - for start_n in range(0, seq_len, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - off_k = ( - (seq_start + start_n + offs_n[None, :]) * k_stride_s - + cur_head * k_stride_h - + offs_d[:, None] * k_stride_d - ) - k = tl.load(K + off_k, mask=((start_n + offs_n[None, :]) < seq_len) & mask_d[:, None], other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) - qk *= sm_scale - qk += tl.where((start_n + offs_n[None, :]) < seq_len, 0, float("-inf")) - - # -- compute m_ij, p, l_ij - m_ij = tl.maximum(tl.max(qk, 1), m_i) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, 1) - - acc_scale = tl.exp(m_i - m_ij) - acc = acc * acc_scale[:, None] - - # update acc - off_v = ( - (seq_start + start_n + offs_n[:, None]) * v_stride_s - + cur_head * v_stride_h - + offs_d[None, :] * v_stride_d - ) - v = tl.load(V + off_v, mask=((start_n + offs_n[:, None]) < seq_len) & mask_d[None, :], other=0.0) - p = p.to(v.dtype) - acc += tl.dot(p, v) - # update m_i and l_i - m_i = m_ij - l_i_new = tl.exp(l_i - m_ij) + l_ij - l_i = m_ij + tl.log(l_i_new) - - o_scale = tl.exp(m_i - l_i) - acc = acc * o_scale[:, None] - # initialize pointers to output - off_o = (seq_start + offs_m[:, None]) * o_stride_s + cur_head * o_stride_h + offs_d[None, :] * o_stride_d - out_ptrs = Out + off_o - tl.store(out_ptrs, acc, mask=(offs_m[:, None] < seq_len) & mask_d[None, :]) - return - - @torch.no_grad() - def _flash_attention_triton_fwd( - q, - k, - v, - o, - cu_seqlens, # q k v cu_seqlens, - max_seqlen, - ): - BLOCK = 64 - # shape constraints - assert q.ndim == k.ndim == v.ndim == o.ndim == 3, "q, k, v, o must be 3D tensors" - _, head_num, head_dim = q.shape - batch_size = cu_seqlens.numel() - 1 - - sm_scale = 1.0 / (head_dim ** 0.5) # 计算scale系数 - d_pad = triton.next_power_of_2(head_dim) - grid = (triton.cdiv(max_seqlen, BLOCK), head_num, batch_size) # batch, head, - num_warps = 4 - _fwd_kernel[grid]( - q, - k, - v, - sm_scale, - o, - q.stride(0), - q.stride(1), - q.stride(2), - k.stride(0), - k.stride(1), - k.stride(2), - v.stride(0), - v.stride(1), - v.stride(2), - o.stride(0), - o.stride(1), - o.stride(2), - head_dim, - cu_seqlens, - BLOCK_M=BLOCK, - BLOCK_DMODEL=d_pad, - BLOCK_N=BLOCK, - num_warps=num_warps, - num_stages=2, - ) - return - -else: - raise Exception("error triton version!") - -_flash_attn_v3_available = False -try: - from sgl_kernel.flash_attn import flash_attn_varlen_func - - _flash_attn_v3_available = True - - def flash_attention_v3_fwd( - q, - k, - v, - o, - cu_seqlens, - max_seqlen, - ): - 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, # attention_chunk - 0.0, # softcap - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - return o - -except ImportError: - print("Failed to import _flash_attn_forward from hopper.flash_attn_interface.") - _flash_attn_v3_available = False - - -def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): - """ - 统一的 Flash Attention 接口。如果 sgl_kernel 存在， - 则使用 sgl_kernel里的接口，否则使用 Triton 版本。 - """ - global _flash_attn_v3_available - if _flash_attn_v3_available and (is_hopper() or is_4090()): - try: - flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) - return - except Exception as e: - print(f"Failed to use sgl_kernel: {e}") - _flash_attn_v3_available = False - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 16a8735a7c..068fc207b1 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -536,9 +536,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--vit_att_backend", type=str, - nargs="+", choices=["auto", "triton", "fa3", "sdpa", "xformers"], - default=["auto"], + default="auto", help="""vit attention kernel used in vlm. auto: automatically select best backend based on GPU and available packages (priority: fa3 > xformers > sdpa > triton)""", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0609132e62..033ca9e68a 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -187,9 +187,7 @@ class StartArgs: default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]}, ) - vit_att_backend: List[str] = field( - default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]} - ) + vit_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]}) llm_kv_type: str = field( default="None", metadata={"choices": ["None", "int8kv", "int4kv", "fp8kv_sph", "fp8kv_spt", "fp8kv_dsa"]} ) diff --git a/lightllm/server/visualserver/__init__.py b/lightllm/server/visualserver/__init__.py index 6bc0923da0..3eaf6bf701 100644 --- a/lightllm/server/visualserver/__init__.py +++ b/lightllm/server/visualserver/__init__.py @@ -1,16 +1,42 @@ -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from lightllm.common.attention.api import VarlenAttentionOp +from lightllm.common.attention.selection import get_varlen_attention_op_class, select_varlen_attention_op +from lightllm.platform import get_device_runtime -VIT_ATTN_BACKEND: BaseVitAttBackend = None +VIT_ATTENTION_OP: VarlenAttentionOp = None -def set_vit_att_backend(backend_name: str): - global VIT_ATTN_BACKEND - VIT_ATTN_BACKEND = get_vit_att_backend_class(backend_name) +class _InitializingVarlenAttentionOp: + """用 worker 预热时的真实 Q/K/V 选择实现，之后通过固定的 forward 委托执行。 + + 部分视觉模型固定使用 BF16，与 llm 的 data_type 不同。manager 无法据启动 + 参数推断实际 dtype/head dim，因此选择发生在各设备的首次预热输入到达时。 + """ + + def __init__(self, requested): + self.requested = requested + self._op = None + + def forward(self, inputs, *, out): + if self._op is None: + runtime = get_device_runtime(inputs.query.device) + if runtime.is_capturing(): + raise RuntimeError("warm up vision attention before Graph capture") + name = select_varlen_attention_op( + self.requested, device=inputs.query.device, dtype=inputs.query.dtype, head_dim=inputs.query.shape[-1] + ) + self._op = get_varlen_attention_op_class(name)() + return self._op.forward(inputs, out=out) + + +def set_vit_attention_op(provider: str, *, defer_validation=False): + global VIT_ATTENTION_OP + VIT_ATTENTION_OP = ( + _InitializingVarlenAttentionOp(provider) if defer_validation else get_varlen_attention_op_class(provider)() + ) return -def get_vit_attn_backend(): - if VIT_ATTN_BACKEND is None: - raise RuntimeError("VIT_ATTN_BACKEND is not initialized. Call init_vit_att_backend() first.") - return VIT_ATTN_BACKEND._vit_att_fwd +def get_vit_attention_op(): + if VIT_ATTENTION_OP is None: + raise RuntimeError("VIT_ATTENTION_OP is not initialized. Call set_vit_attention_op() first.") + return VIT_ATTENTION_OP diff --git a/lightllm/server/visualserver/manager.py b/lightllm/server/visualserver/manager.py index 873eff8456..d6faad3ca0 100644 --- a/lightllm/server/visualserver/manager.py +++ b/lightllm/server/visualserver/manager.py @@ -1,29 +1,31 @@ -import zmq -import zmq.asyncio import asyncio -import uvloop -import rpyc -import socket -import pickle +import collections import inspect -import setproctitle +import pickle +import socket import threading -import collections from typing import List -from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes + +import rpyc +import setproctitle +import uvloop +import zmq +import zmq.asyncio + from lightllm.server.core.objs import ShmReqManager, StartArgs +from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -62,12 +64,10 @@ def __init__( self.lock = asyncio.Lock() async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..6d1146c0a2 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -1,39 +1,40 @@ import os -import rpyc -import torch -import socket -import torch.multiprocessing as mp import queue +import socket import threading import time +from typing import Deque, Dict, List, Optional, Tuple + +import rpyc +import torch import torch.distributed as dist -from typing import Dict, List, Tuple, Deque, Optional -from transformers.configuration_utils import PretrainedConfig +import torch.multiprocessing as mp from rpyc.utils.classic import obtain -from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer -from lightllm.models.llava.llava_visual import LlavaVisionModel +from transformers.configuration_utils import PretrainedConfig + +from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel from lightllm.models.gemma3.gemma3_visual import Gemma3VisionModel from lightllm.models.gemma4.gemma4_visual import Gemma4VisionModel -from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel -from lightllm.models.vit.model import VisionTransformer -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer +from lightllm.models.llava.llava_visual import LlavaVisionModel +from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel from lightllm.models.qwen2_5_vl.qwen2_5_visual import Qwen2_5_VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel +from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel -from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel -from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel -from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer -from lightllm.utils.infer_utils import set_random_seed +from lightllm.models.vit.model import VisionTransformer +from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.server.visualserver import set_vit_attention_op +from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_port_args import get_shm_port_args -from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient -from lightllm.server.visualserver import set_vit_att_backend -from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger -from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder - +from lightllm.utils.shm_port_args import get_shm_port_args logger = init_logger(__name__) @@ -67,7 +68,7 @@ def exposed_init_model(self, kvargs): self.is_visual_only_mode = get_env_start_args().run_mode == "visual_only" self.data_type = kvargs["data_type"] self.vit_attn_backend = kvargs["vit_attn_backend"] - set_vit_att_backend(self.vit_attn_backend) + set_vit_attention_op(self.vit_attn_backend, defer_validation=True) init_vision_distributed_env(kvargs) model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir) diff --git a/lightllm/server/visualserver/visual_only_manager.py b/lightllm/server/visualserver/visual_only_manager.py index 835a25fe3e..ba501eaa7e 100644 --- a/lightllm/server/visualserver/visual_only_manager.py +++ b/lightllm/server/visualserver/visual_only_manager.py @@ -1,35 +1,38 @@ import asyncio -import uvloop -import rpyc import inspect -import setproctitle -import threading -import uuid +import os import pickle -import websockets +import signal import socket import sys -import os -import signal +import threading import time -from lightllm.utils.net_utils import get_hostname_ip -from .objs import VIT_Obj +import uuid from typing import List + +import rpyc +import setproctitle +import uvloop +import websockets + from lightllm.server.core.objs import StartArgs +from lightllm.utils.net_utils import get_hostname_ip + +from .objs import VIT_Obj asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.embed_cache.utils import create_shm, free_shm, get_shm_name_data +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain -from lightllm.server.embed_cache.utils import create_shm, get_shm_name_data, free_shm -from .manager import VisualManager +from .manager import VisualManager +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -74,7 +77,6 @@ async def register_to_config_server_loop(self, args: StartArgs): try: uri = f"ws://{args.config_server_host}:{ports.config_server_port}/visual_register" async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket: - sock = websocket.transport.get_extra_info("socket") sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) @@ -94,12 +96,10 @@ async def register_to_config_server_loop(self, args: StartArgs): logger.info("reconnection to config_server") async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/unit_tests/common/attention/test_varlen_selection.py b/unit_tests/common/attention/test_varlen_selection.py new file mode 100644 index 0000000000..56f85213e5 --- /dev/null +++ b/unit_tests/common/attention/test_varlen_selection.py @@ -0,0 +1,131 @@ +import pytest +import torch + +from lightllm.common.attention import selection +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.utils import backend_validator + + +def test_auto_probes_actual_varlen_path_on_assigned_device(monkeypatch): + calls = [] + + def probe(name, **kwargs): + calls.append((name, kwargs)) + return name == "sdpa" + + monkeypatch.setattr(backend_validator, "_validate", probe) + assert selection.select_varlen_attention_op(device=3) == "sdpa" + assert calls == [ + (name, {"varlen": True, "device": torch.device("cuda:3"), "dtype": None, "head_dim": 80}) + for name in ("fa3", "xformers", "sdpa") + ] + + +@pytest.mark.parametrize("requested", ["auto", "fa3", "triton", "missing"]) +def test_rejected_ops_are_not_returned(monkeypatch, requested): + calls = [] + + def probe(name, **kwargs): + calls.append(name) + return False + + monkeypatch.setattr(backend_validator, "_validate", probe) + with pytest.raises(selection.BackendSelectionError) as error: + selection.select_varlen_attention_op(requested) + expected = list(selection.varlen_attention_ops) if requested == "auto" else [requested] + assert list(error.value.failures) == expected + assert calls == ([] if requested == "missing" else expected) + + +def test_varlen_probe_exercises_selected_op(monkeypatch): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + assert backend_validator._validate_varlen("sdpa") == (True, None) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda self, inputs, out: out.fill_(float("nan"))) + success, reason = backend_validator._validate_varlen("sdpa") + assert not success + assert "mismatch" in reason + + +@pytest.mark.parametrize("check", ["check_environment", "check_support"]) +def test_rejected_probe_does_not_execute_op(monkeypatch, check): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + + def reject(cls, *args): + raise UnsupportedAttention("layout") + + monkeypatch.setattr(SdpaVarlenAttentionOp, check, classmethod(reject)) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda *args, **kwargs: pytest.fail("rejected operator ran")) + assert backend_validator._validate_varlen("sdpa") == (False, "layout") + + +def test_worker_initialization_keeps_fixed_forward_and_uses_actual_inputs(monkeypatch): + import torch + + from lightllm.common.attention.api import VarlenAttentionInputs + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + from lightllm.server import visualserver + + calls = [] + + def select(requested, **kwargs): + calls.append((requested, kwargs)) + return "sdpa" + + monkeypatch.setattr(visualserver, "select_varlen_attention_op", select) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", None) + visualserver.set_vit_attention_op("auto", defer_validation=True) + first = visualserver.get_vit_attention_op() + q = torch.randn(4, 2, 64, dtype=torch.float16) + inputs = VarlenAttentionInputs(q, q, q, torch.tensor([0, 4], dtype=torch.int32), 4) + first.forward(inputs, out=torch.empty_like(q)) + assert calls == [("auto", {"device": q.device, "dtype": q.dtype, "head_dim": 64})] + assert visualserver.get_vit_attention_op() is first + assert isinstance(first._op, SdpaVarlenAttentionOp) + assert "forward" not in vars(first) + first.forward(inputs, out=torch.empty_like(q)) + assert len(calls) == 1 + + +def test_probe_cache_separates_device_dtype_and_layout(monkeypatch): + from dataclasses import replace + + import torch + + from unit_tests.common.attention.helpers import make_spec + + calls = [] + monkeypatch.setattr(backend_validator, "_validate", lambda *args, **kwargs: calls.append(kwargs) or True) + backend_validator.validate.cache_clear() + try: + spec = make_spec() + for device, dtype, shape in ( + ("cuda:0", torch.float16, spec), + ("cuda:1", torch.float16, spec), + ("cuda:0", torch.bfloat16, spec), + ("cuda:0", torch.float16, replace(spec, head_size=64)), + ("cuda:0", torch.float16, spec), + ): + assert backend_validator.validate("provider", device=device, dtype=dtype, spec=shape) + assert len(calls) == 4 + finally: + backend_validator.validate.cache_clear() + + +def test_empty_device_assignment_does_not_probe(monkeypatch): + monkeypatch.setattr(backend_validator, "_validate", lambda *a, **k: pytest.fail("empty assignment probed")) + with pytest.raises(ValueError, match="at least one device"): + selection.select_varlen_attention_op(devices=()) + + +def test_default_device_is_resolved_before_probe(monkeypatch): + calls = [] + monkeypatch.setattr( + selection, "get_device_runtime", lambda device: type("Runtime", (), {"device": torch.device("cpu")})() + ) + monkeypatch.setattr(backend_validator, "_validate", lambda *a, **kw: calls.append(kw["device"]) or True) + assert selection.select_varlen_attention_op("sdpa") == "sdpa" + assert calls == [torch.device("cpu")] diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index e7423d30f9..283d58c587 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -3,13 +3,12 @@ import importlib import importlib.abc import importlib.util -from pathlib import Path import sys -from types import ModuleType, SimpleNamespace import unittest +from pathlib import Path +from types import ModuleType, SimpleNamespace from unittest.mock import Mock, patch - ROOT = Path(__file__).resolve().parents[4] ATTENTION = "lightllm.common.basemodel.attention" @@ -212,37 +211,15 @@ def test_linear_selection_falls_back_to_triton(self): self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), 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() + from lightllm.common.attention import registry, selection - 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"]) + loaded = [] + sentinel = object() + def import_module(name): + loaded.append(name) + return SimpleNamespace(SdpaVarlenAttentionOp=sentinel) -if __name__ == "__main__": - unittest.main() + with patch.object(registry.importlib, "import_module", import_module): + assert selection.get_varlen_attention_op_class("sdpa") is sentinel + assert loaded == ["lightllm.common.attention.backends.dense.sdpa.varlen"] diff --git a/unit_tests/models/deepseek_v4/test_vision.py b/unit_tests/models/deepseek_v4/test_vision.py index ad931ac2b5..b037a2f84d 100644 --- a/unit_tests/models/deepseek_v4/test_vision.py +++ b/unit_tests/models/deepseek_v4/test_vision.py @@ -12,12 +12,10 @@ from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend -from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual -from lightllm.models.deepseek_v4 import image_processor as ours +from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp +from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual, image_processor as ours from lightllm.server import visualserver - REFERENCE_DIR = "/mtc/models/DeepSeek-V4-Flash-Vision-Exp/inference" REFERENCE_PROCESSOR_PATH = os.path.join(REFERENCE_DIR, "image_processor.py") REFERENCE_VISION_PATH = os.path.join(REFERENCE_DIR, "vision.py") @@ -105,7 +103,7 @@ def test_tokenizer_spans_match_reference_blocks(reference_processor, prefix_len, expected_cache = [] for index, (image, (height, width)) in enumerate(zip(params.images, grids)): canonical_types, _ = reference_processor.build_image_block(height, width, start_pos=0) - image.token_id = 2 ** 62 + (0 if reuse_image and index == 2 else index * 1024) + image.token_id = 2**62 + (0 if reuse_image and index == 2 else index * 1024) image.token_num = len(canonical_types) image.start_index_in_embed_cache = 100 + (0 if reuse_image and index == 2 else index * 1024) expected_cache.append((image.token_id, image.token_num, image.start_index_in_embed_cache)) @@ -201,7 +199,7 @@ def _tiny_vision_args(): @pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) def test_vit_matches_reference(reference_vision, n_vit_h, n_vit_w, monkeypatch): - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", SdpaVitAttBackend) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) torch.manual_seed(0) args = _tiny_vision_args() model = visual.ViT(args).eval() @@ -348,3 +346,23 @@ def fake_read_shm(name): assert torch.equal(types[:3], torch.full((3,), ours.IMAGE_PAD, dtype=torch.int64)) assert types[3] == ours.IMAGE_START assert types[-1] == ours.IMAGE_END + + +def test_attention_uses_sequence_boundaries(monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) + torch.manual_seed(17) + attention = visual.Attention(_tiny_vision_args()).eval() + x = torch.randn(7, 16, dtype=torch.bfloat16) + cos = torch.ones(7, 1, 2) + sin = torch.zeros_like(cos) + cu = torch.tensor([0, 3, 7], dtype=torch.int32) + with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH): + q, k, v = [t.reshape(7, 4, 4) for t in attention.wqkv(x).chunk(3, dim=-1)] + parts = [] + for start, end in ((0, 3), (3, 7)): + qi, ki, vi = [t[start:end].float().transpose(0, 1) for t in (q, k, v)] + weights = (qi @ ki.transpose(-1, -2) * 0.5).softmax(-1) + parts.append((weights @ vi).transpose(0, 1).to(x.dtype)) + expected = attention.wo(torch.cat(parts).reshape(7, 16)) + actual = attention(x, cos, sin, cu) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py index 6a25a8d4ea..a7b7dde2ce 100644 --- a/unit_tests/models/glm5_next/test_visual.py +++ b/unit_tests/models/glm5_next/test_visual.py @@ -25,7 +25,7 @@ def test_image_canvas_and_token_expansion(tmp_path, width, height, canvas): image.image_w, image.image_h = width, height image.token_id = 1000 image.token_num = tokenizer.get_image_token_length(image) - assert image.token_num == canvas[0] * canvas[1] // 28 ** 2 + assert image.token_num == canvas[0] * canvas[1] // 28**2 multi = MultimodalParams() multi.images = [image] ids = tokenizer.encode([1, 100, 102, 101, 2], multi) @@ -55,8 +55,8 @@ def test_image_patch_order_temporal_repeat_and_black_padding(): def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): import lightllm.models.glm5_next.glm5_next_visual as glm_visual import lightllm.server.visualserver as visualserver - from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.common.attention.backends.dense.fa3.varlen import Fa3VarlenAttentionOp + from lightllm.server.visualserver import set_vit_attention_op rope_inputs = [] norm_inputs = [] @@ -70,16 +70,18 @@ def fake_rope(x, cos, sin): rope_inputs.append((x, cos, sin)) return x - def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): + def fake_fa3(self, inputs, *, out): + q, k, v = inputs.query, inputs.key, inputs.value + o, cu_seqlens, max_seqlen = out, inputs.cu_seqlens, inputs.max_seqlen attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) o.copy_(q) return o monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) - monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) - set_vit_att_backend("fa3") + monkeypatch.setattr(Fa3VarlenAttentionOp, "forward", fake_fa3) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", visualserver.VIT_ATTENTION_OP) + set_vit_attention_op("fa3") attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) x = torch.randn(3, 16) @@ -149,10 +151,10 @@ def glm_rms_norm(x, weight, eps): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_vision_batch_preserves_independent_image_attention(): from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.server.visualserver import set_vit_attention_op torch.manual_seed(53) - set_vit_att_backend("sdpa") + set_vit_attention_op("sdpa") model = ( Glm5NextVisionTransformer( {"data_type": "float32"}, @@ -167,7 +169,7 @@ def test_vision_batch_preserves_independent_image_attention(): .eval() ) grid = torch.tensor([[1, 4, 6], [1, 6, 4]]) - x = torch.randn(48, 3 * 2 * 14 ** 2, device="cuda") + x = torch.randn(48, 3 * 2 * 14**2, device="cuda") with torch.inference_mode(): batch = model(x, grid) individual = torch.cat([model(x[:24], grid[:1]), model(x[24:], grid[1:])]) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index f736ec4cd6..6242037812 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -1,139 +1,228 @@ -import torch -import math -import time import pytest -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd -from lightllm.models.vit.triton_kernel import flashattention_nopad -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend +import torch +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs +from lightllm.common.attention.selection import get_varlen_attention_op_class +from lightllm.common.attention.varlen import AutoVarlenAttentionOp +from lightllm.platform import get_device_runtime + + +def make_inputs(lengths=(3, 0, 5), *, device="cpu", dtype=torch.float32, head_dim=80, strided=False): + tokens = sum(lengths) + if strided: + qkv = torch.randn(tokens, 3, 2, head_dim, dtype=dtype, device=device) + q, k, v = qkv.unbind(1) + else: + q, k, v = [torch.randn(tokens, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, *lengths], dtype=torch.int32, device=device).cumsum(0, dtype=torch.int32) + return VarlenAttentionInputs(q, k, v, cu, max(lengths)) + + +def reference_attention(inputs): + out = torch.empty_like(inputs.query) + offsets = inputs.cu_seqlens.cpu().tolist() + for start, end in zip(offsets, offsets[1:]): + if start == end: + continue + q, k, v = [x[start:end].float().transpose(0, 1) for x in (inputs.query, inputs.key, inputs.value)] + scores = q @ k.transpose(-1, -2) * inputs.query.shape[-1] ** -0.5 + out[start:end] = (scores.softmax(-1) @ v).transpose(0, 1).to(out.dtype) + return out -_fa3_mark = pytest.mark.skipif( - not flashattention_nopad._flash_attn_v3_available - or not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] < 9, - reason="FA3 requires sgl_kernel and a Hopper GPU", -) +@pytest.mark.parametrize("error_type", [None, UnsupportedAttention]) +def test_audio_fallback_only_rejects_current_input(monkeypatch, error_type): + from lightllm.common.attention import varlen -def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): calls = [] - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("fa3") - o[0].fill_(99) # A backend may fail after partially writing output. - raise RuntimeError("shape-dependent FA3 failure") - - def fallback(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("triton") - o.fill_(7) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) - monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - assert not flashattention_nopad._flash_attn_v3_available - - # Later calls should use the disabled-backend path without retrying FA3. - out.fill_(-123) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - -def test_fa3_success_does_not_run_fallback(monkeypatch): - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fa3(q, k, v, o, cu_seqlens, max_seqlen): - o.fill_(9) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) + class FakeOp: + @classmethod + def check_environment(cls, caps): + return None + + @classmethod + def check_support(cls, inputs, caps=None): + calls.append("fa3") + if error_type: + raise error_type("shape-dependent FA3 failure") + + def forward(self, inputs, *, out): + out.fill_(9) + return out + + class Fallback(FakeOp): + @classmethod + def check_support(cls, inputs, caps=None): + return None + + def forward(self, inputs, *, out): + calls.append("triton") + out.fill_(7) + return out + + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", lambda name: {"fa3": FakeOp, "triton": Fallback}[name]) + op = AutoVarlenAttentionOp() + inputs = make_inputs() + out = torch.full_like(inputs.query, -123) + assert op.forward(inputs, out=out) is out + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) + out.fill_(-123) + op.forward(inputs, out=out) + assert calls == (["fa3", "triton", "fa3", "triton"] if error_type else ["fa3", "fa3"]) + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) + + +@pytest.mark.parametrize("lengths", [(3, 0, 5), (1, 7), (0, 0)]) +def test_sdpa_preserves_sequence_boundaries_and_output_storage(lengths): + inputs = make_inputs(lengths) + inputs.query.zero_() + inputs.key.zero_() + start = 0 + for i, length in enumerate(lengths): + inputs.value[start : start + length].fill_(i + 1) + start += length + out = torch.empty_like(inputs.query) + result = get_varlen_attention_op_class("sdpa")().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, inputs.value) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) +@pytest.mark.parametrize("head_dim,strided", [(64, False), (80, True), (128, False)]) +def test_varlen(provider, dtype, atol, head_dim, strided): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + torch.manual_seed(0) + inputs = make_inputs((1, 0, 65, 129), device="cuda", dtype=dtype, head_dim=head_dim, strided=strided) + storage = torch.full((*inputs.query.shape[:2], head_dim * 2), -123, device="cuda", dtype=dtype) + out = storage[..., :head_dim] + before = [x.clone() for x in (inputs.query, inputs.key, inputs.value)] + result = op().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, reference_attention(inputs), atol=atol, rtol=0) + assert torch.equal(storage[..., head_dim:], torch.full_like(out, -123)) + for actual, expected in zip((inputs.query, inputs.key, inputs.value), before): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +def test_empty_batch(provider): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + inputs = make_inputs((0, 0), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + assert op().forward(inputs, out=out) is out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton"]) +def test_graph_replay_uses_new_inputs(provider): + op_cls = get_varlen_attention_op_class(provider) + try: + op_cls.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + op = op_cls() + assert op.supports_graph + inputs = make_inputs((3, 5), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op.forward(inputs, out=out) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op.forward(inputs, out=out) + inputs.value.add_(1) + inputs.cu_seqlens.copy_(torch.tensor([0, 5, 8], device="cuda", dtype=torch.int32)) + graph.replay() + torch.testing.assert_close(out, reference_attention(inputs), atol=2e-2, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_audio_attention_matches_projected_qkv(monkeypatch): + from lightllm.models.qwen3_omni_moe_thinker import qwen3_omni_audio as audio + + torch.manual_seed(19) + attention = audio.Qwen3OmniMoeAudioAttention(128, 2, 0).cuda().to(torch.bfloat16).eval() monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + audio.g_cache_manager, + "alloc_tensor", + lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device=device), ) - flash_attention_fwd(q, q, q, out, cu, 2) - assert torch.equal(out, torch.full_like(q, 9)) - assert flashattention_nopad._flash_attn_v3_available - - -def reference_attention_varlen(q, k, v, cu): - """ - q, k, v : (total_len, n_head, D) - cu_seqlen : prefix sums (batch+1,) - """ - total, n_head, d = q.shape - out = torch.empty_like(q) - scale = 1.0 / math.sqrt(d) + hidden = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda") + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device="cuda") + with torch.inference_mode(): + q, k, v = [proj(hidden).reshape(8, 2, 64) for proj in (attention.q_proj, attention.k_proj, attention.v_proj)] + expected = attention.out_proj(reference_attention(VarlenAttentionInputs(q, k, v, cu, 5)).reshape(8, 128)) + actual = attention(hidden, cu_seqlens=cu, max_seqlen=5) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) - for b in range(cu.numel() - 1): - s, e = cu[b].item(), cu[b + 1].item() - q_b, k_b, v_b = q[s:e], k[s:e], v[s:e] # (seq, head, D) - q_hsd = q_b.permute(1, 0, 2) # (head, seq, D) - k_hds = k_b.permute(1, 2, 0) # (head, D, seq) - v_hsd = v_b.permute(1, 0, 2) # (head, seq, D) +@pytest.mark.parametrize("error_type", [RuntimeError, TypeError, AssertionError]) +def test_audio_kernel_errors_propagate_without_retry(monkeypatch, error_type): + from lightllm.common.attention import varlen - scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) - probs = torch.softmax(scores.float(), dim=-1) + class Broken: + @classmethod + def check_environment(cls, caps): + pass - out_hsd = torch.matmul(probs, v_hsd.float()) # (head, seq, D) - out[s:e] = out_hsd.permute(1, 0, 2).to(q.dtype) # back to (seq, head, D) + @classmethod + def check_support(cls, inputs, caps=None): + pass - return out + def forward(self, inputs, *, out): + out.fill_(99) + raise error_type("kernel failed") + calls = [] -@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -@pytest.mark.parametrize( - "attention_fn", - [ - flash_attention_fwd, - pytest.param( - getattr(flashattention_nopad, "flash_attention_v3_fwd", None), - marks=_fa3_mark, - ), - pytest.param( - Fa3VitAttBackend._vit_att_fwd, - marks=_fa3_mark, - ), - ], -) -def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): - torch.manual_seed(0) - lengths = torch.randint(1, 257, (batch,)) - max_len = int(lengths.max().item()) - - cu = torch.zeros(batch + 1, dtype=torch.int32, device=device) - cu[1:] = torch.cumsum(lengths, 0) - tot = int(cu[-1]) - - q = torch.randn(tot, heads, d, dtype=dtype, device=device) - k = torch.randn_like(q) - v = torch.randn_like(q) - out_tri = torch.randn_like(q) - attention_fn(q, k, v, out_tri, cu, max_len) - a = time.time() - for _ in range(100): - attention_fn(q, k, v, out_tri, cu, max_len) - b = time.time() - print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") - out_ref = reference_attention_varlen(q, k, v, cu) - - max_err = (out_ref - out_tri).abs().max().item() - mean_err = (out_ref - out_tri).abs().mean().item() - print(f"{dtype}: max {max_err:.6f}, mean {mean_err:.6f}") - torch.testing.assert_close(out_tri, out_ref, atol=atol, rtol=0) - - -if __name__ == "__main__": - pytest.main() + def load(name): + calls.append(name) + return Broken + + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", load) + op = AutoVarlenAttentionOp() + inputs = make_inputs() + with pytest.raises(error_type, match="kernel failed"): + op.forward(inputs, out=torch.empty_like(inputs.query)) + assert calls == ["fa3"] + assert not op.unavailable + + +def test_audio_encoder_shares_runner_only_with_its_layers(): + from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_audio import Qwen3OmniMoeAudioEncoder + + def encoder(): + return Qwen3OmniMoeAudioEncoder( + {}, + d_model=16, + encoder_attention_heads=2, + encoder_layers=2, + encoder_ffn_dim=32, + downsample_hidden_size=4, + output_dim=16, + ) + + first, second = encoder(), encoder() + assert first.layers[0].self_attn.attention_op is first.layers[1].self_attn.attention_op + assert first.layers[0].self_attn.attention_op is not second.layers[0].self_attn.attention_op + + +@pytest.mark.parametrize("boundaries,maximum", [([1, 3, 8], 5), ([0, 7, 6, 8], 7), ([0, 3, 7], 5), ([0, 3, 8], 4)]) +def test_varlen_rejects_invalid_boundaries(boundaries, maximum): + inputs = make_inputs()._replace(cu_seqlens=torch.tensor(boundaries, dtype=torch.int32), max_seqlen=maximum) + with pytest.raises(ValueError, match="cu_seqlens"): + inputs.validate()