Feature/vit attention unification# 23880

vllm/attention/layer.py

@@ -350,22 +350,23 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
 
         dtype = torch.get_default_dtype()
-        attn_backend = get_attn_backend(head_size,
-                                        dtype,
-                                        kv_cache_dtype=None,
-                                        block_size=16,
-                                        is_attention_free=False)
-        backend = backend_name_to_enum(attn_backend.get_name())
+        try:
+            attn_backend = get_attn_backend(head_size,
+                                            dtype,
+                                            kv_cache_dtype=None,
+                                            block_size=16,
+                                            is_attention_free=False)
+            backend = backend_name_to_enum(attn_backend.get_name())
+        except (ValueError, AttributeError):
+            # Fallback to TORCH_SDPA if backend detection fails
+            backend = _Backend.TORCH_SDPA
         if current_platform.is_rocm():
             # currently, only torch_sdpa is supported on rocm
             self.attn_backend = _Backend.TORCH_SDPA
         else:
-            if backend in (_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1,
-                           _Backend.FLEX_ATTENTION):
-                backend = _Backend.XFORMERS
-
             self.attn_backend = backend if backend in {
-                _Backend.TORCH_SDPA, _Backend.XFORMERS, _Backend.PALLAS_VLLM_V1
+                _Backend.TORCH_SDPA, _Backend.XFORMERS, _Backend.PALLAS_VLLM_V1,
+                _Backend.ROCM_AITER_FA, _Backend.FLEX_ATTENTION
             } else _Backend.TORCH_SDPA
 
         if (self.attn_backend == _Backend.XFORMERS
@@ -413,6 +414,30 @@ def forward(
             from torch_xla.experimental.custom_kernel import flash_attention
             out = flash_attention(query, key, value, sm_scale=self.scale)
             out = out.transpose(1, 2)
+        elif self.attn_backend in (_Backend.FLASH_ATTN, 
+                                  _Backend.FLASH_ATTN_VLLM_V1):
+            # Flash Attention for ViT is not essential, fallback to PyTorch SDPA
+            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+            out = F.scaled_dot_product_attention(
+                query, key, value, scale=self.scale)
+            out = out.transpose(1, 2)
+        elif self.attn_backend == _Backend.ROCM_AITER_FA:
+            from aiter import flash_attn_varlen_func
+            out = flash_attn_varlen_func(
+                query, key, value, softmax_scale=self.scale)
+        elif self.attn_backend == _Backend.FLEX_ATTENTION:
+            # FlexAttention requires specific tensor format
+            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+            # Use a simple fallback to torch SDPA for now
+            out = F.scaled_dot_product_attention(
+                query, key, value, scale=self.scale)
+            out = out.transpose(1, 2)
+        else:
+            # Fallback to torch SDPA for unsupported backends
+            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+            out = F.scaled_dot_product_attention(
+                query, key, value, scale=self.scale)
+            out = out.transpose(1, 2)
 
         return out.reshape(bsz, q_len, -1)
 

vllm/model_executor/models/idefics2_vision_model.py

@@ -170,6 +170,7 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.out_proj",
             )
+        # Use unified MultiHeadAttention with Flash Attention support
         self.attn = MultiHeadAttention(self.num_heads_per_partition,
                                        self.head_dim, self.scale)
 
@@ -181,8 +182,11 @@ def forward(
             hidden_states
         )  # batch_size, q_len, 3 * num_heads_per_partition * head_dim
         query_states, key_states, value_states = qkv.chunk(3, dim=-1)
-        out = self.attn(query_states, key_states, value_states)
-        attn_output, _ = self.out_proj(out)
+
+        # Use unified MultiHeadAttention implementation
+        attn_output = self.attn(query_states, key_states, value_states)
+
+        attn_output, _ = self.out_proj(attn_output)
         return attn_output
 
 

vllm/model_executor/models/intern_vit.py

@@ -29,6 +29,7 @@
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+
 NORM2FN = {
     'rms_norm': RMSNorm,
     'layer_norm': nn.LayerNorm,
@@ -254,6 +255,10 @@ def __init__(
                                   var_hidden_size=self.embed_dim)
 
         self.proj = nn.Linear(self.dummy_dim, self.embed_dim)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        self.attn = MultiHeadAttention(
+            self.num_heads, self.head_dim, self.scale)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, N, C = x.shape
@@ -268,12 +273,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             B_, N_, H_, D_ = q.shape
             q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)
             k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        x = F.scaled_dot_product_attention(q, k, v, scale=self.scale)
-        x = x.transpose(1, 2).reshape(B, N, -1)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        x = self.attn(q, k, v)
+        x = x.reshape(B, N, -1)
 
         x = self.proj(x)
         return x

vllm/model_executor/models/interns1_vit.py

@@ -12,7 +12,6 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from transformers import PretrainedConfig
 from transformers.utils import torch_int
 
@@ -22,6 +21,7 @@
                                                RowParallelLinear)
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.attention.layer import MultiHeadAttention
 
 NORM2FN = {
     'rms_norm': RMSNorm,
@@ -205,6 +205,10 @@ def __init__(
                                   var_hidden_size=self.embed_dim)
 
         self.projection_layer = nn.Linear(self.dummy_dim, self.embed_dim)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        self.attn = MultiHeadAttention(
+            self.num_heads, self.head_dim, self.scale)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, N, C = x.shape
@@ -213,20 +217,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         k = self.k_proj(x)
         v = self.v_proj(x)
 
-        q = q.view(B, N, self.num_heads, self.head_dim)
-        k = k.view(B, N, self.num_heads, self.head_dim)
-        v = v.view(B, N, self.num_heads, self.head_dim)
-
         if self.qk_normalization:
             B_, N_, H_, D_ = q.shape
             q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)
             k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        x = F.scaled_dot_product_attention(q, k, v, scale=self.scale)
-        x = x.transpose(1, 2).reshape(B, N, -1)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        x = self.attn(q, k, v)
+        x = x.reshape(B, N, -1)
 
         x = self.projection_layer(x)
         return x

vllm/model_executor/models/mllama.py

@@ -53,6 +53,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader, maybe_remap_kv_scale_name)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
+
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalEncDecInputs,
@@ -70,6 +71,7 @@
 from .interfaces import SupportsMultiModal, SupportsV0Only
 from .llama import LlamaDecoderLayer, LlamaMLP
 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
+from vllm.attention.layer import MultiHeadAttention
 
 logger = init_logger(__name__)
 
@@ -516,6 +518,10 @@ def __init__(self,
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        self.attn = MultiHeadAttention(
+            self.num_local_heads, self.head_dim, 1.0 / math.sqrt(self.head_dim))
 
     def forward(
         self,
@@ -524,21 +530,11 @@ def forward(
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_state)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q = q.view(q.shape[0], q.shape[1], self.num_local_heads,
-                   self.head_dim).transpose(1, 2)
-        k = k.view(k.shape[0], k.shape[1], self.num_local_heads,
-                   self.head_dim).transpose(1, 2)
-        v = v.view(v.shape[0], v.shape[1], self.num_local_heads,
-                   self.head_dim).transpose(1, 2)
-
-        # TODO: remove padding in image encoder
-        attn_output = F.scaled_dot_product_attention(q,
-                                                     k,
-                                                     v,
-                                                     attn_mask=attention_mask,
-                                                     dropout_p=0.0)
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        attn_output = self.attn(q, k, v)
+
+        attn_output = attn_output.contiguous()
         attn_output = attn_output.reshape(attn_output.shape[0],
                                           attn_output.shape[1], -1)
         output, _ = self.o_proj(attn_output)

vllm/model_executor/models/pixtral.py

@@ -32,7 +32,9 @@
                                                RowParallelLinear)
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.attention.layer import MultiHeadAttention
 from vllm.model_executor.sampling_metadata import SamplingMetadata
+
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
                                     NestedTensors)
@@ -1079,6 +1081,10 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        self.attn = MultiHeadAttention(
+            self.n_heads, self.head_dim, 1.0 / math.sqrt(self.head_dim))
 
     def forward(
         self,
@@ -1098,19 +1104,13 @@ def forward(
         cos, sin = position_embeddings
         q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0)
 
-        if USE_XFORMERS_OPS:
-            # Transpose q and k back for attention
-            q = q.transpose(1, 2).contiguous()
-            k = k.transpose(1, 2).contiguous()
-            out = xops.memory_efficient_attention(q,
-                                                  k,
-                                                  v,
-                                                  attn_bias=attention_mask)
-        else:
-            v = v.transpose(1, 2)
-            out = nn.functional.scaled_dot_product_attention(
-                q, k, v, attn_mask=attention_mask)
-            out = out.transpose(1, 2)
+        # Reshape for MultiHeadAttention: (batch, seq, hidden_size)
+        q_reshaped = q.transpose(1, 2).contiguous().view(batch, patches, -1)
+        k_reshaped = k.transpose(1, 2).contiguous().view(batch, patches, -1)
+        v_reshaped = v.transpose(1, 2).contiguous().view(batch, patches, -1)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        out = self.attn(q_reshaped, k_reshaped, v_reshaped)
 
         out = out.view(batch, patches, self.n_heads * self.head_dim)
         attn_output, _ = self.o_proj(out)

vllm/model_executor/models/step3_vl.py

@@ -25,6 +25,7 @@
                                                RowParallelLinear)
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
+from vllm.attention.layer import MultiHeadAttention
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
@@ -696,10 +697,10 @@ def __init__(self,
                                               bias=True,
                                               quant_config=quant_config,
                                               prefix=prefix)
-
-    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
-        return tensor.view(bsz, seq_len, self.num_heads,
-                           self.head_dim).transpose(1, 2).contiguous()
+        
+        # Use unified MultiHeadAttention with automatic backend selection
+        self.attn = MultiHeadAttention(
+            self.num_heads, self.head_dim, self.scale)
 
     def forward(
         self,
@@ -711,19 +712,14 @@ def forward(
         # get query proj
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.chunk(chunks=3, dim=-1)
-        q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)
-        k = k.view(bsz, tgt_len, self.num_heads, self.head_dim)
-        v = v.view(bsz, tgt_len, self.num_heads, self.head_dim)
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-        attn_output = F.scaled_dot_product_attention(q,
-                                                     k,
-                                                     v,
-                                                     scale=self.scale,
-                                                     is_causal=False)
-        attn_output = attn_output.transpose(1, 2).reshape(
-            bsz, tgt_len, self.num_heads * self.head_dim)
+
+        # Reshape for MultiHeadAttention: (batch, seq, hidden_size)
+        q_reshaped = q.view(bsz, tgt_len, -1)
+        k_reshaped = k.view(bsz, tgt_len, -1)
+        v_reshaped = v.view(bsz, tgt_len, -1)
+
+        # Use unified MultiHeadAttention with automatic backend selection
+        attn_output = self.attn(q_reshaped, k_reshaped, v_reshaped)
 
         attn_output, _ = self.out_proj(attn_output)
 

vllm/model_executor/models/vision.py

@@ -122,4 +122,4 @@ def resolve_visual_encoder_outputs(
     uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)
     if post_layer_norm is not None and uses_last_layer:
         hs_pool[-1] = post_layer_norm(encoder_outputs)
-    return torch.cat(hs_pool, dim=-1)
+    return torch.cat(hs_pool, dim=-1)