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
[Multi Modal] Add FA3 in VIT (#24347)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>
Wenlong Wang committed
Sep 12, 2025 at 06:27 UTC
72fc8aa41255a21716cc0f4d8eb81b9572f71bc2
13 files changed
+247
-66
tests/entrypoints/openai/test_vision.py
+2
-2
@@ -34,11 +34,11 @@ EXPECTED_MM_BEAM_SEARCH_RES = [
34
],
35
[
36
"The image shows a Venn diagram with three over",
37
- "The image shows a Venn diagram with three intersect",
37
+ "This image shows a Venn diagram with three over",
38
],
39
[
40
"This image displays a gradient of colors ranging from",
41
- "The image displays a gradient of colors ranging from",
41
+ "This image displays a gradient of colors forming a spectrum",
42
],
43
]
44
tests/kernels/attention/test_mha_attn.py
+35
-14
@@ -36,31 +36,52 @@ def test_mha_attn_platform(device: str):
36
torch.set_default_dtype(torch.float16)
37
38
if device == "cpu":
39
- with patch("vllm.attention.selector.current_platform",
40
- CpuPlatform()), \
41
- patch("vllm.platforms.current_platform", CpuPlatform()):
39
+ with patch("vllm.attention.layer.current_platform", CpuPlatform()), \
40
+ patch("vllm.model_executor.models.vision.current_platform",
41
+ CpuPlatform()):
42
attn = MultiHeadAttention(16, 64, scale=1)
43
- assert attn.attn_backend == _Backend.TORCH_SDPA_VLLM_V1
43
+ assert attn.attn_backend == _Backend.TORCH_SDPA
44
elif device == "hip":
45
- with patch("vllm.attention.selector.current_platform",
46
- RocmPlatform()), \
47
- patch("vllm.platforms.current_platform", RocmPlatform()), \
48
- patch("vllm.attention.layer.current_platform", RocmPlatform()):
45
+ with patch("vllm.attention.layer.current_platform", RocmPlatform()), \
46
+ patch("vllm.model_executor.models.vision.current_platform",
47
+ RocmPlatform()):
48
attn = MultiHeadAttention(16, 64, scale=1)
49
assert attn.attn_backend == _Backend.TORCH_SDPA
50
else:
52
- with patch("vllm.attention.selector.current_platform",
53
- CudaPlatform()), \
54
- patch("vllm.platforms.current_platform", CudaPlatform()):
51
+ # Test CUDA with head_size=64 (divisible by 32)
52
+ # - should use vLLM's FlashAttention
53
+ with patch("vllm.attention.layer.current_platform", CudaPlatform()), \
54
+ patch("vllm.model_executor.models.vision.current_platform",
55
+ CudaPlatform()):
56
attn = MultiHeadAttention(16, 64, scale=1)
56
- assert attn.attn_backend == _Backend.XFORMERS
57
+ assert attn.attn_backend == _Backend.FLASH_ATTN
58
58
- with patch("vllm.attention.selector.current_platform",
59
+ # Test CUDA with head_size=72 (not divisible by 32)
60
+ # - with upstream FA not available
61
+ # - should use xformers
62
+ with patch("vllm.attention.layer.current_platform", CudaPlatform()), \
63
+ patch("vllm.model_executor.models.vision.current_platform",
64
CudaPlatform()), \
60
- patch("vllm.platforms.current_platform", CudaPlatform()):
65
+ patch("vllm.attention.layer.check_upstream_fa_availability",
66
+ return_value=False):
67
attn = MultiHeadAttention(16, 72, scale=1)
68
assert attn.attn_backend == _Backend.XFORMERS
69
70
+ # Test CUDA with head_size=72 (not divisible by 32)
71
+ # - with upstream FA available
72
+ # - should use upstream FA
73
+ with patch("vllm.attention.layer.current_platform", CudaPlatform()), \
74
+ patch("vllm.model_executor.models.vision.current_platform",
75
+ CudaPlatform()), \
76
+ patch("vllm.attention.layer.check_upstream_fa_availability",
77
+ return_value=True), \
78
+ patch.dict('sys.modules', {'flash_attn': type('MockFlashAttn', (),
79
+ {
80
+ 'flash_attn_varlen_func': lambda *args, **kwargs: None
81
+ })()}):
82
+ attn = MultiHeadAttention(16, 72, scale=1)
83
+ assert attn.attn_backend == _Backend.FLASH_ATTN
84
+
85
86
def ref_attention(
87
query: torch.Tensor,
vllm/attention/layer.py
+67
-8
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
23
from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig)
25
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
26
+from vllm.model_executor.models.vision import get_vit_attn_backend
27
from vllm.platforms import _Backend, current_platform
28
from vllm.utils import direct_register_custom_op
29
@@ -55,6 +56,14 @@ def check_xformers_availability():
56
return USE_XFORMERS_OPS
57
58
59
+def check_upstream_fa_availability(dtype: torch.dtype):
60
+ if dtype in (torch.float16, torch.bfloat16) and current_platform.is_cuda(
61
+ ) and current_platform.has_device_capability(80):
62
+ from transformers.utils import is_flash_attn_2_available
63
+ return is_flash_attn_2_available()
64
+ return False
65
+
66
+
67
class Attention(nn.Module, AttentionLayerBase):
68
"""Attention layer.
69
@@ -349,29 +358,55 @@ class MultiHeadAttention(nn.Module):
358
f"divisible by num_kv_heads ({self.num_kv_heads})"
359
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
360
361
+ # During model initialization, the default dtype is set as the model
362
+ # weight and activation dtype.
363
dtype = torch.get_default_dtype()
353
- attn_backend = get_attn_backend(head_size,
354
- dtype,
355
- kv_cache_dtype=None,
356
- block_size=16,
357
- is_attention_free=False)
358
- backend = backend_name_to_enum(attn_backend.get_name())
364
+
365
+ # Determine the attention backend
366
+ backend = get_vit_attn_backend(head_size=head_size, dtype=dtype)
367
+
368
+ # Some auto-selected backends can be upgraded
369
+ # to upstream flash attention if available.
370
+ # If vllm native fa is selected, we use it directly.
371
+ use_upstream_fa = False
372
+ if backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
373
+ dtype):
374
+ backend = _Backend.FLASH_ATTN
375
+ use_upstream_fa = True
376
+
377
if current_platform.is_rocm():
378
# currently, only torch_sdpa is supported on rocm
379
self.attn_backend = _Backend.TORCH_SDPA
380
else:
381
+
382
self.attn_backend = backend if backend in {
383
_Backend.TORCH_SDPA,
384
_Backend.TORCH_SDPA_VLLM_V1,
385
_Backend.XFORMERS,
386
_Backend.PALLAS_VLLM_V1,
387
_Backend.ROCM_AITER_FA,
369
- } else current_platform.get_vit_attn_backend()
388
+ _Backend.FLASH_ATTN,
389
+ _Backend.FLASH_ATTN_VLLM_V1,
390
+ } else _Backend.TORCH_SDPA
391
392
if (self.attn_backend == _Backend.XFORMERS
393
and not check_xformers_availability()):
394
self.attn_backend = _Backend.TORCH_SDPA
395
396
+ if self.attn_backend in {
397
+ _Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1
398
+ }:
399
+ if use_upstream_fa:
400
+ from flash_attn import flash_attn_varlen_func
401
+ self._flash_attn_varlen_func = flash_attn_varlen_func
402
+ else:
403
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
404
+ self._flash_attn_varlen_func = flash_attn_varlen_func
405
+
406
+ logger.info_once(
407
+ f"MultiHeadAttention attn_backend: {self.attn_backend}, "
408
+ f"use_upstream_fa: {use_upstream_fa}")
409
+
410
def forward(
411
self,
412
query: torch.Tensor,
@@ -392,7 +427,31 @@ class MultiHeadAttention(nn.Module):
427
key = torch.repeat_interleave(key, num_repeat, dim=2)
428
value = torch.repeat_interleave(value, num_repeat, dim=2)
429
395
- if self.attn_backend == _Backend.XFORMERS:
430
+ if self.attn_backend in {
431
+ _Backend.FLASH_ATTN,
432
+ _Backend.FLASH_ATTN_VLLM_V1,
433
+ }:
434
+
435
+ cu_seqlens_q = torch.arange(0, (bsz + 1) * q_len,
436
+ step=q_len,
437
+ dtype=torch.int32,
438
+ device=query.device)
439
+ cu_seqlens_k = torch.arange(0, (bsz + 1) * kv_len,
440
+ step=kv_len,
441
+ dtype=torch.int32,
442
+ device=key.device)
443
+
444
+ out = self._flash_attn_varlen_func(
445
+ query.flatten(0, 1),
446
+ key.flatten(0, 1),
447
+ value.flatten(0, 1),
448
+ cu_seqlens_q=cu_seqlens_q,
449
+ cu_seqlens_k=cu_seqlens_k,
450
+ max_seqlen_q=q_len,
451
+ max_seqlen_k=kv_len,
452
+ softmax_scale=self.scale,
453
+ )
454
+ elif self.attn_backend == _Backend.XFORMERS:
455
from xformers import ops as xops
456
457
out = xops.memory_efficient_attention_forward(query,
vllm/model_executor/models/ernie45_vl.py
+20
-3
@@ -34,6 +34,7 @@ import torch.nn.functional as F
34
from einops import rearrange, repeat
35
from transformers import BatchFeature
36
37
+from vllm.attention.layer import check_upstream_fa_availability
38
from vllm.config import VllmConfig
39
from vllm.distributed import parallel_state
40
from vllm.distributed import utils as dist_utils
@@ -170,7 +171,16 @@ class Ernie4_5_VisionAttention(nn.Module):
171
prefix=f"{prefix}.proj")
172
173
# Detect attention implementation.
173
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
174
+ self.attn_backend = get_vit_attn_backend(
175
+ head_size=self.hidden_size_per_attention_head,
176
+ dtype=torch.get_default_dtype())
177
+
178
+ self.use_upstream_fa = False
179
+ if self.attn_backend != _Backend.FLASH_ATTN and \
180
+ check_upstream_fa_availability(torch.get_default_dtype()):
181
+ self.attn_backend = _Backend.FLASH_ATTN
182
+ self.use_upstream_fa = True
183
+
184
if self.attn_backend not in {
185
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
186
_Backend.ROCM_AITER_FA
@@ -233,7 +243,10 @@ class Ernie4_5_VisionAttention(nn.Module):
243
if self.attn_backend == _Backend.ROCM_AITER_FA:
244
from aiter import flash_attn_varlen_func
245
else:
236
- from flash_attn import flash_attn_varlen_func
246
+ if self.use_upstream_fa:
247
+ from flash_attn import flash_attn_varlen_func
248
+ else:
249
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
250
251
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
252
@@ -457,7 +470,11 @@ class Ernie4_5_VisionTransformer(nn.Module):
470
), "vit's config.hidden must be equal to config.embed_dim"
471
self.ln = nn.LayerNorm(hidden_size, eps=1e-6)
472
460
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
473
+ self.attn_backend = get_vit_attn_backend(
474
+ head_size=head_dim, dtype=torch.get_default_dtype())
475
+ if self.attn_backend != _Backend.FLASH_ATTN and \
476
+ check_upstream_fa_availability(torch.get_default_dtype()):
477
+ self.attn_backend = _Backend.FLASH_ATTN
478
479
@property
480
def dtype(self) -> torch.dtype:
vllm/model_executor/models/glm4_1v.py
+19
-3
@@ -44,6 +44,7 @@ from transformers.models.glm4v.video_processing_glm4v import (
44
Glm4vVideoProcessor)
45
from transformers.video_utils import VideoMetadata
46
47
+from vllm.attention.layer import check_upstream_fa_availability
48
from vllm.config import VllmConfig
49
from vllm.distributed import (get_tensor_model_parallel_world_size,
50
parallel_state)
@@ -260,7 +261,15 @@ class Glm4vVisionAttention(nn.Module):
261
)
262
263
# Detect attention implementation.
263
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
264
+ self.attn_backend = get_vit_attn_backend(
265
+ head_size=self.hidden_size_per_attention_head,
266
+ dtype=torch.get_default_dtype())
267
+ self.use_upstream_fa = False
268
+ if self.attn_backend != _Backend.FLASH_ATTN and \
269
+ check_upstream_fa_availability(torch.get_default_dtype()):
270
+ self.attn_backend = _Backend.FLASH_ATTN
271
+ self.use_upstream_fa = True
272
+
273
if self.attn_backend not in {
274
_Backend.FLASH_ATTN,
275
_Backend.TORCH_SDPA,
@@ -310,7 +319,10 @@ class Glm4vVisionAttention(nn.Module):
319
if self.attn_backend == _Backend.FLASH_ATTN:
320
# from vllm_flash_attn.flash_attn_interface import (
321
# flash_attn_varlen_func)
313
- from flash_attn import flash_attn_varlen_func
322
+ if self.use_upstream_fa:
323
+ from flash_attn import flash_attn_varlen_func
324
+ else:
325
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
326
327
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
328
@@ -715,7 +727,11 @@ class Glm4vVisionTransformer(nn.Module):
727
self.post_layernorm = RMSNorm(vision_config.hidden_size,
728
eps=vision_config.rms_norm_eps)
729
718
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
730
+ self.attn_backend = get_vit_attn_backend(
731
+ head_size=head_dim, dtype=torch.get_default_dtype())
732
+ if self.attn_backend != _Backend.FLASH_ATTN and \
733
+ check_upstream_fa_availability(torch.get_default_dtype()):
734
+ self.attn_backend = _Backend.FLASH_ATTN
735
736
@property
737
def dtype(self) -> torch.dtype:
vllm/model_executor/models/keye.py
+15
-2
@@ -17,6 +17,7 @@ from transformers.modeling_outputs import (BaseModelOutput,
17
BaseModelOutputWithPooling)
18
from transformers.utils import torch_int
19
20
+from vllm.attention.layer import check_upstream_fa_availability
21
from vllm.config import VllmConfig
22
from vllm.distributed import get_tensor_model_parallel_world_size
23
from vllm.logger import init_logger
@@ -374,7 +375,16 @@ class KeyeSiglipAttention(nn.Module):
375
)
376
377
# Detect attention implementation.
377
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
378
+ self.attn_backend = get_vit_attn_backend(
379
+ head_size=self.head_dim, dtype=torch.get_default_dtype())
380
+
381
+ self.use_upstream_fa = False
382
+ if self.attn_backend != _Backend.FLASH_ATTN and \
383
+ check_upstream_fa_availability(
384
+ torch.get_default_dtype()):
385
+ self.attn_backend = _Backend.FLASH_ATTN
386
+ self.use_upstream_fa = True
387
+
388
if self.attn_backend not in {_Backend.FLASH_ATTN, _Backend.XFORMERS}:
389
raise RuntimeError(
390
f"Keye-VL does not support {self.attn_backend} backend now.")
@@ -428,7 +438,10 @@ class KeyeSiglipAttention(nn.Module):
438
)
439
440
if self.attn_backend == _Backend.FLASH_ATTN:
431
- from flash_attn import flash_attn_varlen_func
441
+ if self.use_upstream_fa:
442
+ from flash_attn import flash_attn_varlen_func
443
+ else:
444
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
445
446
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
447
vllm/model_executor/models/qwen2_5_vl.py
+21
-3
@@ -38,6 +38,7 @@ from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor
38
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
39
Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig)
40
41
+from vllm.attention.layer import check_upstream_fa_availability
42
from vllm.config import VllmConfig
43
from vllm.distributed import parallel_state
44
from vllm.distributed import utils as dist_utils
@@ -298,7 +299,16 @@ class Qwen2_5_VisionAttention(nn.Module):
299
disable_tp=use_data_parallel)
300
301
# Detect attention implementation.
301
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
302
+ self.attn_backend = get_vit_attn_backend(
303
+ head_size=self.hidden_size_per_attention_head,
304
+ dtype=torch.get_default_dtype())
305
+ self.use_upstream_fa = False
306
+ if self.attn_backend != _Backend.FLASH_ATTN and \
307
+ check_upstream_fa_availability(
308
+ torch.get_default_dtype()):
309
+ self.attn_backend = _Backend.FLASH_ATTN
310
+ self.use_upstream_fa = True
311
+
312
if self.attn_backend not in {
313
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
314
_Backend.ROCM_AITER_FA
@@ -359,7 +369,10 @@ class Qwen2_5_VisionAttention(nn.Module):
369
if self.attn_backend == _Backend.ROCM_AITER_FA:
370
from aiter import flash_attn_varlen_func
371
else:
362
- from flash_attn import flash_attn_varlen_func
372
+ if self.use_upstream_fa:
373
+ from flash_attn import flash_attn_varlen_func
374
+ else:
375
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
376
377
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
378
@@ -628,7 +641,12 @@ class Qwen2_5_VisionTransformer(nn.Module):
641
prefix=f"{prefix}.merger",
642
use_data_parallel=use_data_parallel,
643
)
631
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
644
+ self.attn_backend = get_vit_attn_backend(
645
+ head_size=head_dim, dtype=torch.get_default_dtype())
646
+ if self.attn_backend != _Backend.FLASH_ATTN and \
647
+ check_upstream_fa_availability(
648
+ torch.get_default_dtype()):
649
+ self.attn_backend = _Backend.FLASH_ATTN
650
651
@property
652
def dtype(self) -> torch.dtype:
vllm/model_executor/models/qwen2_vl.py
+21
-3
@@ -41,6 +41,7 @@ from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
41
from transformers.models.qwen2_vl.video_processing_qwen2_vl import (
42
Qwen2VLVideoProcessor)
43
44
+from vllm.attention.layer import check_upstream_fa_availability
45
from vllm.config import VllmConfig
46
from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
47
from vllm.distributed import utils as dist_utils
@@ -314,7 +315,16 @@ class Qwen2VisionAttention(nn.Module):
315
prefix=f"{prefix}.proj")
316
317
# Detect attention implementation.
317
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
318
+ self.attn_backend = get_vit_attn_backend(
319
+ head_size=self.hidden_size_per_attention_head,
320
+ dtype=torch.get_default_dtype())
321
+ self.use_upstream_fa = False
322
+ if self.attn_backend != _Backend.FLASH_ATTN and \
323
+ check_upstream_fa_availability(
324
+ torch.get_default_dtype()):
325
+ self.attn_backend = _Backend.FLASH_ATTN
326
+ self.use_upstream_fa = True
327
+
328
if self.attn_backend not in {
329
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
330
_Backend.ROCM_AITER_FA
@@ -374,7 +384,10 @@ class Qwen2VisionAttention(nn.Module):
384
if self.attn_backend == _Backend.ROCM_AITER_FA:
385
from aiter import flash_attn_varlen_func
386
else:
377
- from flash_attn import flash_attn_varlen_func
387
+ if self.use_upstream_fa:
388
+ from flash_attn import flash_attn_varlen_func
389
+ else:
390
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
391
392
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
393
@@ -628,7 +641,12 @@ class Qwen2VisionTransformer(nn.Module):
641
quant_config=quant_config,
642
prefix=f"{prefix}.merger",
643
)
631
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
644
+ self.attn_backend = get_vit_attn_backend(
645
+ head_size=head_dim, dtype=torch.get_default_dtype())
646
+ if self.attn_backend != _Backend.FLASH_ATTN and \
647
+ check_upstream_fa_availability(
648
+ torch.get_default_dtype()):
649
+ self.attn_backend = _Backend.FLASH_ATTN
650
651
@property
652
def dtype(self) -> torch.dtype:
vllm/model_executor/models/siglip2navit.py
+14
-2
@@ -13,6 +13,7 @@ from torch.nn import functional as F
13
from transformers import Siglip2VisionConfig
14
from transformers.configuration_utils import PretrainedConfig
15
16
+from vllm.attention.layer import check_upstream_fa_availability
17
from vllm.config import QuantizationConfig
18
from vllm.distributed import divide, get_tensor_model_parallel_world_size
19
from vllm.model_executor.layers.activation import get_act_fn
@@ -236,7 +237,15 @@ class Siglip2Attention(nn.Module):
237
self.use_rope = config.use_rope
238
239
# Detect attention implementation.
239
- self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
240
+ self.attn_backend = get_vit_attn_backend(
241
+ head_size=self.head_dim, dtype=torch.get_default_dtype())
242
+ self.use_upstream_fa = False
243
+ if self.attn_backend != _Backend.FLASH_ATTN and \
244
+ check_upstream_fa_availability(
245
+ torch.get_default_dtype()):
246
+ self.attn_backend = _Backend.FLASH_ATTN
247
+ self.use_upstream_fa = True
248
+
249
if self.attn_backend not in {
250
_Backend.FLASH_ATTN, _Backend.TORCH_SDPA,
251
_Backend.ROCM_AITER_FA
@@ -280,7 +289,10 @@ class Siglip2Attention(nn.Module):
289
if self.attn_backend == _Backend.ROCM_AITER_FA:
290
from aiter import flash_attn_varlen_func
291
else:
283
- from flash_attn import flash_attn_varlen_func
292
+ if self.use_upstream_fa:
293
+ from flash_attn import flash_attn_varlen_func
294
+ else:
295
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
296
attn_output = flash_attn_varlen_func(
297
queries, keys, values, cu_seqlens, cu_seqlens, max_seqlen,
298
max_seqlen).reshape(seq_length, -1)
vllm/model_executor/models/vision.py
+5
-5
@@ -7,7 +7,6 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union
7
import torch
8
from transformers import PretrainedConfig
9
10
-from vllm.attention.selector import get_env_variable_attn_backend
10
from vllm.logger import init_logger
11
from vllm.platforms import _Backend, current_platform
12
@@ -68,17 +67,18 @@ def get_vision_encoder_info(
67
raise NotImplementedError(msg)
68
69
71
-def get_vit_attn_backend(support_fa: bool = False) -> _Backend:
70
+def get_vit_attn_backend(head_size: int, dtype: torch.dtype) -> _Backend:
71
"""
72
Get the available attention backend for Vision Transformer.
73
"""
75
- # TODO(Isotr0py): Remove `support_fa` after support FA for all ViTs attn.
74
+ # Lazy import to avoid circular dependency
75
+ from vllm.attention.selector import get_env_variable_attn_backend
76
77
selected_backend: Optional[_Backend] = get_env_variable_attn_backend()
78
if selected_backend is not None:
79
return selected_backend
80
81
- return current_platform.get_vit_attn_backend(support_fa)
81
+ return current_platform.get_vit_attn_backend(head_size, dtype)
82
83
84
def resolve_visual_encoder_outputs(
@@ -122,4 +122,4 @@ def resolve_visual_encoder_outputs(
122
uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)
123
if post_layer_norm is not None and uses_last_layer:
124
hs_pool[-1] = post_layer_norm(encoder_outputs)
125
- return torch.cat(hs_pool, dim=-1)
\ No newline at end of file
125
+ return torch.cat(hs_pool, dim=-1)
vllm/platforms/cuda.py
+17
-11
@@ -209,18 +209,24 @@ class CudaPlatformBase(Platform):
209
return torch.cuda.max_memory_allocated(device)
210
211
@classmethod
212
- def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
213
- if cls.has_device_capability(80) and support_fa:
214
- from transformers.utils import is_flash_attn_2_available
215
- if is_flash_attn_2_available():
212
+ def get_vit_attn_backend(cls, head_size: int,
213
+ dtype: torch.dtype) -> _Backend:
214
+ if dtype not in (torch.float16, torch.bfloat16):
215
+ return _Backend.XFORMERS
216
+
217
+ if cls.has_device_capability(80):
218
+ FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
219
+ from vllm.attention.selector import is_attn_backend_supported
220
+ is_default_fa_supported = is_attn_backend_supported(
221
+ FLASH_ATTN_V1, head_size, dtype, allow_import_error=False)
222
+ if is_default_fa_supported:
223
return _Backend.FLASH_ATTN
217
- logger.warning_once(
218
- "Current `vllm-flash-attn` has a bug inside vision "
219
- "module, so we use xformers backend instead. You can "
220
- "run `pip install flash-attn` to use flash-attention "
221
- "backend.")
222
- # Fallback for Volta/Turing GPUs or FA not supported
223
- return _Backend.XFORMERS
224
+ else:
225
+ # Fallback to XFORMERS
226
+ return _Backend.XFORMERS
227
+ else:
228
+ # Fallback for Volta/Turing GPUs or FA not supported
229
+ return _Backend.XFORMERS
230
231
@classmethod
232
def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
vllm/platforms/interface.py
+2
-1
@@ -192,7 +192,8 @@ class Platform:
192
return device_id
193
194
@classmethod
195
- def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
195
+ def get_vit_attn_backend(cls, head_size: int,
196
+ dtype: torch.dtype) -> _Backend:
197
return _Backend.TORCH_SDPA
198
199
@classmethod
vllm/platforms/rocm.py
+9
-9
@@ -175,15 +175,15 @@ class RocmPlatform(Platform):
175
]
176
177
@classmethod
178
- def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
179
- if support_fa:
180
- if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA
181
- and on_gfx9()):
182
- # Note: AITER FA is only supported for Qwen-VL models.
183
- # TODO: Add support for other VL models in their model class.
184
- return _Backend.ROCM_AITER_FA
185
- if on_gfx9():
186
- return _Backend.FLASH_ATTN
178
+ def get_vit_attn_backend(cls, head_size: int,
179
+ dtype: torch.dtype) -> _Backend:
180
+ if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA
181
+ and on_gfx9()):
182
+ # Note: AITER FA is only supported for Qwen-VL models.
183
+ # TODO: Add support for other VL models in their model class.
184
+ return _Backend.ROCM_AITER_FA
185
+ if on_gfx9():
186
+ return _Backend.FLASH_ATTN
187
return _Backend.TORCH_SDPA
188
189
@classmethod