[Model] GritLM supports other attention backends (#18109)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

tests/models/language/pooling/test_gritlm.py

@@ -11,7 +11,6 @@ from scipy.spatial.distance import cosine
 
 from vllm import LLM, SamplingParams
 from vllm.config import ModelConfig
-from vllm.utils import STR_BACKEND_ENV_VAR
 
 from ....utils import RemoteOpenAIServer
 
@@ -117,12 +116,7 @@ def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
     assert math.isclose(cosine_sim_q1_d1, 0.534, abs_tol=0.001)
 
 
-def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch,
-                                  vllm_runner):
-    # GritLM embedding implementation is only supported by XFormers backend.
-    with monkeypatch.context() as m:
-        m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
-
+def test_gritlm_offline_embedding(vllm_runner):
     queries, q_instruction, documents, d_instruction = get_test_data()
 
     with vllm_runner(
@@ -150,11 +144,9 @@ def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch,
 async def test_gritlm_api_server_embedding():
     queries, q_instruction, documents, d_instruction = get_test_data()
 
-    # GritLM embedding implementation is only supported by XFormers backend.
     args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
-    env_dict = {STR_BACKEND_ENV_VAR: "XFORMERS"}
 
-    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
+    with RemoteOpenAIServer(MODEL_NAME, args) as server:
         client_embedding = server.get_async_client()
 
         d_rep = await run_client_embeddings(
@@ -172,11 +164,6 @@ async def test_gritlm_api_server_embedding():
 
 
 def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
-    # GritLM embedding implementation is only supported by XFormers backend.
-    with monkeypatch.context() as m:
-        m.setenv("VLLM_USE_V1", "0")
-        m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
-
     input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
 
     with vllm_runner(
@@ -196,11 +183,9 @@ def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
 async def test_gritlm_api_server_generate():
     input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
 
-    # GritLM embedding implementation is only supported by XFormers backend.
     args = ["--task", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
-    env_dict = {"VLLM_USE_V1": "0", STR_BACKEND_ENV_VAR: "XFORMERS"}
 
-    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
+    with RemoteOpenAIServer(MODEL_NAME, args) as server:
         client_generate = server.get_async_client()
 
         outputs = await client_generate.completions.create(

vllm/model_executor/models/gritlm.py

 # SPDX-License-Identifier: Apache-2.0
 
 from array import array
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 import torch.nn as nn
-from xformers.ops.fmha.attn_bias import BlockDiagonalMask
 
-from vllm.attention.backends.xformers import XFormersImpl
 from vllm.config import ModelConfig, VllmConfig
-from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.pooler import PoolerHead
 from vllm.model_executor.models.llama import LlamaForCausalLM
 from vllm.model_executor.pooling_metadata import (PoolingMetadata,
                                                   PoolingTensors)
-from vllm.sequence import (IntermediateTensors, PoolerOutput,
-                           PoolingSequenceGroupOutput)
+from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
 from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
 
 from .interfaces import SupportsV0Only
@@ -204,38 +200,20 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
         prefix: str = "",
         **kwargs,
     ) -> None:
-        super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
-
-        self.runner_type = vllm_config.model_config.runner_type
+        # Use full attention for pooling
+        if vllm_config.model_config.runner_type == "pooling":
+            hf_config = vllm_config.model_config.hf_config
+            hf_config.is_causal = False
 
-        self._pooler = GritLMPooler(vllm_config.model_config)
+            vllm_config.cache_config.sliding_window = None
 
-        for layer in self.model.layers:
-            if self.runner_type == "pooling" and hasattr(layer, "self_attn"):
-                assert isinstance(layer.self_attn.attn.impl, XFormersImpl), (
-                    "GritLM embedding is only supported by XFormers backend, "
-                    "which can be forced by VLLM_ATTENTION_BACKEND=XFORMERS")
+            for attr in ("sliding_window", "interleaved_sliding_window"):
+                if hasattr(hf_config, attr):
+                    delattr(hf_config, attr)
 
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        **kwargs,
-    ) -> Union[torch.Tensor, IntermediateTensors]:
-
-        # Change attention to non-causal for pooling tasks.
-        if self.runner_type == "pooling":
-            attn_metadata = get_forward_context().attn_metadata
-            assert attn_metadata.prefill_metadata.attn_bias is None
-            attn_metadata.prefill_metadata.attn_bias = [
-                BlockDiagonalMask.from_seqlens(attn_metadata.seq_lens)
-            ]
+        super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
 
-        return super().forward(
-            input_ids=input_ids,
-            positions=positions,
-            **kwargs,
-        )
+        self._pooler = GritLMPooler(vllm_config.model_config)
 
     def pooler(
         self,

vllm/model_executor/models/llama.py

@@ -28,7 +28,7 @@ import torch
 from torch import nn
 from transformers import LlamaConfig
 
-from vllm.attention import Attention
+from vllm.attention import Attention, AttentionType
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -96,7 +96,8 @@ class LlamaMLP(nn.Module):
 
 class LlamaAttention(nn.Module):
 
-    def __init__(self,
+    def __init__(
+        self,
         config: LlamaConfig,
         hidden_size: int,
         num_heads: int,
@@ -108,7 +109,9 @@ class LlamaAttention(nn.Module):
         bias: bool = False,
         bias_o_proj: bool = False,
         cache_config: Optional[CacheConfig] = None,
-                 prefix: str = "") -> None:
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
         super().__init__()
         layer_idx = extract_layer_index(prefix)
         self.hidden_size = hidden_size
@@ -194,6 +197,7 @@ class LlamaAttention(nn.Module):
             cache_config=cache_config,
             quant_config=quant_config,
             per_layer_sliding_window=sliding_window,
+            attn_type=attn_type,
             prefix=f"{prefix}.attn",
         )
 
@@ -238,6 +242,15 @@ class LlamaDecoderLayer(nn.Module):
         if hasattr(config, 'qkv_bias'):
             attention_bias = config.qkv_bias
 
+        # By default, Llama uses causal attention as it is a decoder-only model.
+        # You can override the HF config with `is_causal=False` to enable
+        # bidirectional attention, which is used in some embedding models
+        # (e.g. parasail-ai/GritLM-7B-vllm)
+        if getattr(config, "is_causal", True):
+            attn_type = AttentionType.DECODER
+        else:
+            attn_type = AttentionType.ENCODER_ONLY
+
         self.self_attn = LlamaAttention(
             config=config,
             hidden_size=self.hidden_size,
@@ -252,6 +265,7 @@ class LlamaDecoderLayer(nn.Module):
             bias_o_proj=bias_o_proj,
             cache_config=cache_config,
             prefix=f"{prefix}.self_attn",
+            attn_type=attn_type,
         )
         self.mlp = LlamaMLP(
             hidden_size=self.hidden_size,

vllm/model_executor/models/qwen2.py

@@ -111,8 +111,8 @@ class Qwen2Attention(nn.Module):
         rope_scaling: Optional[Tuple] = None,
         prefix: str = "",
         attn_type: str = AttentionType.DECODER,
-            dual_chunk_attention_config: Optional[dict[str,
-                                                       Any]] = None) -> None:
+        dual_chunk_attention_config: Optional[dict[str, Any]] = None,
+    ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
         tp_size = get_tensor_model_parallel_world_size()