[4/5] refactor(attention): use varlen operators for vision and audio

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)

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)
 

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))