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 14, 05:30 UTC
Support encoder_only attention for FlexAttention (#22273)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
Maximilien de Bayser committed
Aug 6, 2025 at 22:37 UTC
f825c6bd22133a8b2242457069f59654a2ae401b
2 files changed
+138
-47
tests/kernels/test_flex_attention.py
+68
-20
@@ -9,7 +9,9 @@ import pytest
9
import torch
10
from packaging import version
11
12
-from vllm import LLM, SamplingParams
12
+from vllm import SamplingParams
13
+
14
+from ..models.utils import check_embeddings_close
15
16
TORCH_VERSION = version.parse(torch.__version__)
17
MINIMUM_TORCH_VERSION = version.parse("2.7.0")
@@ -28,7 +30,7 @@ def set_seed(seed):
30
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
31
reason="CUDA not available or PyTorch version < 2.7",
32
)
31
-def test_flex_attention_vs_default_backend(monkeypatch):
33
+def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
34
"""Test that FlexAttention produces the same outputs as the default backend.
35
36
This test compares the outputs from the FlexAttention backend with
@@ -36,7 +38,7 @@ def test_flex_attention_vs_default_backend(monkeypatch):
38
"""
39
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
40
seed = 42
39
- max_tokens = 32
41
+ max_tokens = 24
42
prompts = [
43
"Hello, my name is",
44
"The president of the United States is",
@@ -54,33 +56,30 @@ def test_flex_attention_vs_default_backend(monkeypatch):
56
m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
57
58
set_seed(seed)
57
-
58
- llm_flex = LLM(
59
- model_name,
60
- tensor_parallel_size=1,
61
- num_gpu_blocks_override=128,
62
- enforce_eager=True,
63
- )
64
- output_flex = llm_flex.generate(prompts, sampling_params)
59
+ with vllm_runner(model_name,
60
+ runner="generate",
61
+ tensor_parallel_size=1,
62
+ num_gpu_blocks_override=128,
63
+ enforce_eager=True) as llm_flex:
64
+ output_flex = llm_flex.generate(prompts, sampling_params)
65
66
# Run with default backend
67
with monkeypatch.context() as m:
68
m.setenv("VLLM_USE_V1", "1")
69
set_seed(seed)
70
- llm_default = LLM(
71
- model_name,
72
- tensor_parallel_size=1,
73
- num_gpu_blocks_override=128,
74
- enforce_eager=True,
75
- )
76
- output_default = llm_default.generate(prompts, sampling_params)
70
+ with vllm_runner(model_name,
71
+ runner="generate",
72
+ tensor_parallel_size=1,
73
+ num_gpu_blocks_override=128,
74
+ enforce_eager=True) as llm_default:
75
+ output_default = llm_default.generate(prompts, sampling_params)
76
77
# Compare outputs from both backends
78
for i, (flex_result,
79
default_result) in enumerate(zip(output_flex, output_default)):
80
prompt = prompts[i]
82
- flex_text = flex_result.outputs[0].text
83
- default_text = default_result.outputs[0].text
81
+ flex_text = flex_result[1][0]
82
+ default_text = default_result[1][0]
83
84
assert flex_text == default_text, (
85
f"FlexAttention output doesn't match default for: {prompt!r}\n"
@@ -88,5 +87,54 @@ def test_flex_attention_vs_default_backend(monkeypatch):
87
f"Default: {default_text!r}")
88
89
90
+@pytest.mark.skipif(
91
+ not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
92
+ reason="CUDA not available or PyTorch version < 2.7",
93
+)
94
+def test_encoder_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
95
+ """Test that FlexAttention produces the same outputs as the default backend.
96
+
97
+ This test compares the outputs from the FlexAttention backend with
98
+ the default backend for encoder models.
99
+ """
100
+ model_name = "BAAI/bge-base-en-v1.5"
101
+ prompts = [
102
+ "Hello, my name is",
103
+ "The president of the United States is",
104
+ "The capital of France is",
105
+ ]
106
+
107
+ # Run with flex attention
108
+ with monkeypatch.context() as m:
109
+ m.setenv("VLLM_USE_V1", "1")
110
+ m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
111
+ with vllm_runner(model_name,
112
+ runner="pooling",
113
+ dtype=torch.bfloat16,
114
+ tensor_parallel_size=1,
115
+ max_model_len=100,
116
+ enforce_eager=True) as llm_flex:
117
+ flex_outputs = llm_flex.embed(prompts)
118
+
119
+ # Run with default backend
120
+ with monkeypatch.context() as m:
121
+ m.setenv("VLLM_USE_V1", "1")
122
+ with vllm_runner(model_name,
123
+ runner="pooling",
124
+ dtype=torch.bfloat16,
125
+ tensor_parallel_size=1,
126
+ max_model_len=100,
127
+ enforce_eager=True) as llm_default:
128
+ default_outputs = llm_default.embed(prompts)
129
+
130
+ check_embeddings_close(
131
+ embeddings_0_lst=flex_outputs,
132
+ embeddings_1_lst=default_outputs,
133
+ name_0="flex",
134
+ name_1="default",
135
+ tol=1e-2,
136
+ )
137
+
138
+
139
if __name__ == "__main__":
140
pytest.main([__file__])
vllm/v1/attention/backends/flex_attention.py
+70
-27
@@ -148,6 +148,7 @@ def causal_mask_mod(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor,
148
149
@dataclass
150
class FlexAttentionMetadata:
151
+ causal: bool
152
num_actual_tokens: int # Number of tokens excluding padding.
153
max_query_len: int
154
query_start_loc: torch.Tensor
@@ -177,10 +178,9 @@ class FlexAttentionMetadata:
178
num_blocks = 0
179
block_mask: Optional[BlockMask] = None
180
score_mod: Optional[_score_mod_signature] = None
180
- mask_mod: Optional[_mask_mod_signature] = None
181
logical_mask_mod: _mask_mod_signature = causal_mask_mod
182
183
- def get_mask_mod(self) -> _mask_mod_signature:
183
+ def get_causal_mask_mod(self) -> _mask_mod_signature:
184
"""Creates the mask_mod function for FlexAttention.
185
186
This function creates the combined mask mod function that handles:
@@ -233,14 +233,39 @@ class FlexAttentionMetadata:
233
234
return final_mask_mod
235
236
+ def get_bidirectional_mask_mod(self) -> _mask_mod_signature:
237
+ """Creates the encoder mask_mod function for FlexAttention.
238
+
239
+ Since the encoder bidirectional attention doesn't run with
240
+ KV cache, this function creates a mask based on the
241
+ packed query sequences.
242
+ """
243
+ # Create a lookup mapping from query indices -> request number
244
+ request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)
245
+
246
+ def final_mask_mod(
247
+ b: torch.Tensor,
248
+ h: torch.Tensor,
249
+ q_idx: torch.Tensor,
250
+ kv_idx: torch.Tensor,
251
+ ) -> torch.Tensor:
252
+ return request_lookup[q_idx] == request_lookup[kv_idx]
253
+
254
+ return final_mask_mod
255
+
256
def build_block_mask(self) -> BlockMask:
237
- assert self.mask_mod is not None
257
+ if self.causal:
258
+ mask_mod = self.get_causal_mask_mod()
259
+ kv_len = self.total_cache_tokens
260
+ else:
261
+ mask_mod = self.get_bidirectional_mask_mod()
262
+ kv_len = self.num_actual_tokens
263
return create_block_mask_compiled(
239
- self.mask_mod,
264
+ mask_mod,
265
None,
266
None,
267
self.num_actual_tokens,
243
- self.total_cache_tokens,
268
+ kv_len,
269
device=self.block_table.device,
270
)
271
@@ -251,7 +276,6 @@ class FlexAttentionMetadata:
276
assert self.prefix_kv_lens is None, "Not implemented yet."
277
assert self.suffix_kv_lens is None, "Not implemented yet."
278
self.num_blocks = self.total_cache_tokens // self.block_size
254
- self.mask_mod = self.get_mask_mod()
279
self.block_mask = self.build_block_mask()
280
281
@@ -306,6 +330,7 @@ class FlexAttentionMetadataBuilder(
330
self.device, non_blocking=True)
331
332
out = FlexAttentionMetadata(
333
+ causal=common_attn_metadata.causal,
334
num_actual_tokens=num_actual_tokens,
335
max_query_len=max_query_len,
336
query_start_loc=query_start_loc,
@@ -350,6 +375,12 @@ class FlexAttentionImpl(AttentionImpl):
375
self.head_size = head_size
376
self.scale = float(scale)
377
self.num_kv_heads = num_kv_heads
378
+ self.attn_type = attn_type
379
+
380
+ if attn_type not in (AttentionType.ENCODER_ONLY,
381
+ AttentionType.DECODER):
382
+ raise NotImplementedError(
383
+ f"FlexAttention does not support {attn_type} attention")
384
385
if alibi_slopes is not None:
386
raise NotImplementedError(
@@ -425,26 +456,38 @@ class FlexAttentionImpl(AttentionImpl):
456
457
num_actual_tokens = attn_metadata.num_actual_tokens
458
428
- key_cache, value_cache = kv_cache.unbind(0)
429
-
430
- torch.ops._C_cache_ops.reshape_and_cache_flash(
431
- key,
432
- value,
433
- key_cache,
434
- value_cache,
435
- attn_metadata.slot_mapping,
436
- self.kv_cache_dtype,
437
- layer._k_scale,
438
- layer._v_scale,
439
- )
459
+ if not attn_metadata.causal:
460
+ assert self.attn_type == AttentionType.ENCODER_ONLY
461
+
462
+ query, key_tensor, value_tensor = map(
463
+ lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
464
+ (query, key, value),
465
+ )
466
+
467
+ else:
468
+ assert self.attn_type == AttentionType.DECODER
469
+ key_cache, value_cache = kv_cache.unbind(0)
470
+
471
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
472
+ key,
473
+ value,
474
+ key_cache,
475
+ value_cache,
476
+ attn_metadata.slot_mapping,
477
+ self.kv_cache_dtype,
478
+ layer._k_scale,
479
+ layer._v_scale,
480
+ )
481
+
482
+ # View out the block_size dim
483
+ key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
484
+ value_cache = value_cache.view(-1, self.num_kv_heads,
485
+ self.head_size)
486
+ query, key_tensor, value_tensor = map(
487
+ lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
488
+ (query, key_cache, value_cache),
489
+ )
490
441
- # View out the block_size dim
442
- key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
443
- value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
444
- query, key_cache, value_cache = map(
445
- lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
446
- (query, key_cache, value_cache),
447
- )
491
query = query[:, :, :num_actual_tokens, :]
492
# Doesn't work for now -> constraint violation
493
# torch._dynamo.try_mark_dynamic(query, 2)
@@ -465,8 +508,8 @@ class FlexAttentionImpl(AttentionImpl):
508
509
out = flex_attention_compiled(
510
query,
468
- key_cache,
469
- value_cache,
511
+ key_tensor,
512
+ value_tensor,
513
attn_metadata.score_mod,
514
attn_metadata.block_mask,
515
self.scale,