A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Model] GritLM supports other attention backends (#18109)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
May 14, 2025 at 18:33 UTC
d62a076e8467ddc9bba13849911275bf66b8ed6c
4 files changed
+84
-107
tests/models/language/pooling/test_gritlm.py
+31
-46
@@ -11,7 +11,6 @@ from scipy.spatial.distance import cosine
11
12
from vllm import LLM, SamplingParams
13
from vllm.config import ModelConfig
14
-from vllm.utils import STR_BACKEND_ENV_VAR
14
15
from ....utils import RemoteOpenAIServer
16
@@ -117,44 +116,37 @@ def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
116
assert math.isclose(cosine_sim_q1_d1, 0.534, abs_tol=0.001)
117
118
120
-def test_gritlm_offline_embedding(monkeypatch: pytest.MonkeyPatch,
121
- vllm_runner):
122
- # GritLM embedding implementation is only supported by XFormers backend.
123
- with monkeypatch.context() as m:
124
- m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
125
-
126
- queries, q_instruction, documents, d_instruction = get_test_data()
119
+def test_gritlm_offline_embedding(vllm_runner):
120
+ queries, q_instruction, documents, d_instruction = get_test_data()
121
128
- with vllm_runner(
129
- MODEL_NAME,
130
- task="embed",
131
- max_model_len=MAX_MODEL_LEN,
132
- ) as vllm_model:
133
- llm = vllm_model.model
122
+ with vllm_runner(
123
+ MODEL_NAME,
124
+ task="embed",
125
+ max_model_len=MAX_MODEL_LEN,
126
+ ) as vllm_model:
127
+ llm = vllm_model.model
128
135
- d_rep = run_llm_encode(
136
- llm,
137
- documents,
138
- d_instruction,
139
- )
140
- q_rep = run_llm_encode(
141
- llm,
142
- queries,
143
- q_instruction,
144
- )
129
+ d_rep = run_llm_encode(
130
+ llm,
131
+ documents,
132
+ d_instruction,
133
+ )
134
+ q_rep = run_llm_encode(
135
+ llm,
136
+ queries,
137
+ q_instruction,
138
+ )
139
146
- validate_embed_output(q_rep, d_rep)
140
+ validate_embed_output(q_rep, d_rep)
141
142
143
@pytest.mark.asyncio
144
async def test_gritlm_api_server_embedding():
145
queries, q_instruction, documents, d_instruction = get_test_data()
146
153
- # GritLM embedding implementation is only supported by XFormers backend.
147
args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
155
- env_dict = {STR_BACKEND_ENV_VAR: "XFORMERS"}
148
157
- with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
149
+ with RemoteOpenAIServer(MODEL_NAME, args) as server:
150
client_embedding = server.get_async_client()
151
152
d_rep = await run_client_embeddings(
@@ -172,35 +164,28 @@ async def test_gritlm_api_server_embedding():
164
165
166
def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
175
- # GritLM embedding implementation is only supported by XFormers backend.
176
- with monkeypatch.context() as m:
177
- m.setenv("VLLM_USE_V1", "0")
178
- m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
179
-
180
- input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
167
+ input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
168
182
- with vllm_runner(
183
- MODEL_NAME,
184
- task="generate",
185
- max_model_len=MAX_MODEL_LEN,
186
- ) as vllm_model:
187
- llm = vllm_model.model
169
+ with vllm_runner(
170
+ MODEL_NAME,
171
+ task="generate",
172
+ max_model_len=MAX_MODEL_LEN,
173
+ ) as vllm_model:
174
+ llm = vllm_model.model
175
189
- sampling_params = SamplingParams(temperature=0.0, max_tokens=256)
190
- outputs = llm.generate(input, sampling_params=sampling_params)
176
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=256)
177
+ outputs = llm.generate(input, sampling_params=sampling_params)
178
192
- assert outputs[0].outputs[0].text == "The capital of France is Paris."
179
+ assert outputs[0].outputs[0].text == "The capital of France is Paris."
180
181
182
@pytest.mark.asyncio
183
async def test_gritlm_api_server_generate():
184
input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
185
199
- # GritLM embedding implementation is only supported by XFormers backend.
186
args = ["--task", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
201
- env_dict = {"VLLM_USE_V1": "0", STR_BACKEND_ENV_VAR: "XFORMERS"}
187
203
- with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as server:
188
+ with RemoteOpenAIServer(MODEL_NAME, args) as server:
189
client_generate = server.get_async_client()
190
191
outputs = await client_generate.completions.create(
vllm/model_executor/models/gritlm.py
+12
-34
@@ -1,22 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
from array import array
4
-from typing import Optional, Union
4
+from typing import Optional
5
6
import torch
7
import torch.nn as nn
8
-from xformers.ops.fmha.attn_bias import BlockDiagonalMask
8
10
-from vllm.attention.backends.xformers import XFormersImpl
9
from vllm.config import ModelConfig, VllmConfig
12
-from vllm.forward_context import get_forward_context
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.pooler import PoolerHead
12
from vllm.model_executor.models.llama import LlamaForCausalLM
13
from vllm.model_executor.pooling_metadata import (PoolingMetadata,
14
PoolingTensors)
18
-from vllm.sequence import (IntermediateTensors, PoolerOutput,
19
- PoolingSequenceGroupOutput)
15
+from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
16
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
17
18
from .interfaces import SupportsV0Only
@@ -204,38 +200,20 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
200
prefix: str = "",
201
**kwargs,
202
) -> None:
207
- super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
203
+ # Use full attention for pooling
204
+ if vllm_config.model_config.runner_type == "pooling":
205
+ hf_config = vllm_config.model_config.hf_config
206
+ hf_config.is_causal = False
207
209
- self.runner_type = vllm_config.model_config.runner_type
208
+ vllm_config.cache_config.sliding_window = None
209
211
- self._pooler = GritLMPooler(vllm_config.model_config)
210
+ for attr in ("sliding_window", "interleaved_sliding_window"):
211
+ if hasattr(hf_config, attr):
212
+ delattr(hf_config, attr)
213
213
- for layer in self.model.layers:
214
- if self.runner_type == "pooling" and hasattr(layer, "self_attn"):
215
- assert isinstance(layer.self_attn.attn.impl, XFormersImpl), (
216
- "GritLM embedding is only supported by XFormers backend, "
217
- "which can be forced by VLLM_ATTENTION_BACKEND=XFORMERS")
214
+ super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
215
219
- def forward(
220
- self,
221
- input_ids: torch.Tensor,
222
- positions: torch.Tensor,
223
- **kwargs,
224
- ) -> Union[torch.Tensor, IntermediateTensors]:
225
-
226
- # Change attention to non-causal for pooling tasks.
227
- if self.runner_type == "pooling":
228
- attn_metadata = get_forward_context().attn_metadata
229
- assert attn_metadata.prefill_metadata.attn_bias is None
230
- attn_metadata.prefill_metadata.attn_bias = [
231
- BlockDiagonalMask.from_seqlens(attn_metadata.seq_lens)
232
- ]
233
-
234
- return super().forward(
235
- input_ids=input_ids,
236
- positions=positions,
237
- **kwargs,
238
- )
216
+ self._pooler = GritLMPooler(vllm_config.model_config)
217
218
def pooler(
219
self,
vllm/model_executor/models/llama.py
+28
-14
@@ -28,7 +28,7 @@ import torch
28
from torch import nn
29
from transformers import LlamaConfig
30
31
-from vllm.attention import Attention
31
+from vllm.attention import Attention, AttentionType
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -96,19 +96,22 @@ class LlamaMLP(nn.Module):
96
97
class LlamaAttention(nn.Module):
98
99
- def __init__(self,
100
- config: LlamaConfig,
101
- hidden_size: int,
102
- num_heads: int,
103
- num_kv_heads: int,
104
- rope_theta: float = 10000,
105
- rope_scaling: Optional[Dict[str, Any]] = None,
106
- max_position_embeddings: int = 8192,
107
- quant_config: Optional[QuantizationConfig] = None,
108
- bias: bool = False,
109
- bias_o_proj: bool = False,
110
- cache_config: Optional[CacheConfig] = None,
111
- prefix: str = "") -> None:
99
+ def __init__(
100
+ self,
101
+ config: LlamaConfig,
102
+ hidden_size: int,
103
+ num_heads: int,
104
+ num_kv_heads: int,
105
+ rope_theta: float = 10000,
106
+ rope_scaling: Optional[Dict[str, Any]] = None,
107
+ max_position_embeddings: int = 8192,
108
+ quant_config: Optional[QuantizationConfig] = None,
109
+ bias: bool = False,
110
+ bias_o_proj: bool = False,
111
+ cache_config: Optional[CacheConfig] = None,
112
+ prefix: str = "",
113
+ attn_type: str = AttentionType.DECODER,
114
+ ) -> None:
115
super().__init__()
116
layer_idx = extract_layer_index(prefix)
117
self.hidden_size = hidden_size
@@ -194,6 +197,7 @@ class LlamaAttention(nn.Module):
197
cache_config=cache_config,
198
quant_config=quant_config,
199
per_layer_sliding_window=sliding_window,
200
+ attn_type=attn_type,
201
prefix=f"{prefix}.attn",
202
)
203
@@ -238,6 +242,15 @@ class LlamaDecoderLayer(nn.Module):
242
if hasattr(config, 'qkv_bias'):
243
attention_bias = config.qkv_bias
244
245
+ # By default, Llama uses causal attention as it is a decoder-only model.
246
+ # You can override the HF config with `is_causal=False` to enable
247
+ # bidirectional attention, which is used in some embedding models
248
+ # (e.g. parasail-ai/GritLM-7B-vllm)
249
+ if getattr(config, "is_causal", True):
250
+ attn_type = AttentionType.DECODER
251
+ else:
252
+ attn_type = AttentionType.ENCODER_ONLY
253
+
254
self.self_attn = LlamaAttention(
255
config=config,
256
hidden_size=self.hidden_size,
@@ -252,6 +265,7 @@ class LlamaDecoderLayer(nn.Module):
265
bias_o_proj=bias_o_proj,
266
cache_config=cache_config,
267
prefix=f"{prefix}.self_attn",
268
+ attn_type=attn_type,
269
)
270
self.mlp = LlamaMLP(
271
hidden_size=self.hidden_size,
vllm/model_executor/models/qwen2.py
+13
-13
@@ -100,19 +100,19 @@ class Qwen2MLP(nn.Module):
100
class Qwen2Attention(nn.Module):
101
102
def __init__(
103
- self,
104
- hidden_size: int,
105
- num_heads: int,
106
- num_kv_heads: int,
107
- max_position: int = 4096 * 32,
108
- rope_theta: float = 10000,
109
- cache_config: Optional[CacheConfig] = None,
110
- quant_config: Optional[QuantizationConfig] = None,
111
- rope_scaling: Optional[Tuple] = None,
112
- prefix: str = "",
113
- attn_type: str = AttentionType.DECODER,
114
- dual_chunk_attention_config: Optional[dict[str,
115
- Any]] = None) -> None:
103
+ self,
104
+ hidden_size: int,
105
+ num_heads: int,
106
+ num_kv_heads: int,
107
+ max_position: int = 4096 * 32,
108
+ rope_theta: float = 10000,
109
+ cache_config: Optional[CacheConfig] = None,
110
+ quant_config: Optional[QuantizationConfig] = None,
111
+ rope_scaling: Optional[Tuple] = None,
112
+ prefix: str = "",
113
+ attn_type: str = AttentionType.DECODER,
114
+ dual_chunk_attention_config: Optional[dict[str, Any]] = None,
115
+ ) -> None:
116
super().__init__()
117
self.hidden_size = hidden_size
118
tp_size = get_tensor_model_parallel_world_size()