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, 22:35 UTC
[FEAT][ROCm] Enable running Flash Attention as ViT attn backend for Qwen-VL models on ROCm platform. (#22069)
Signed-off-by: tjtanaavllm <tunjian.tan@amd.com> Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: tjtanaavllm <tunjian.tan@amd.com>
vllmellm committed
Aug 2, 2025 at 14:53 UTC
d3a6f2120bb6b67fc58a3f1000d624cfb351eb05
6 files changed
+64
-39
vllm/model_executor/models/qwen2_5_vl.py
+13
-5
@@ -246,11 +246,15 @@ class Qwen2_5_VisionAttention(nn.Module):
246
# Detect attention implementation.
247
self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
248
if self.attn_backend not in {
249
- _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS
249
+ _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
250
+ _Backend.ROCM_AITER_FA
251
}:
252
raise RuntimeError(
253
f"Qwen2.5-VL does not support {self.attn_backend} backend now."
254
)
255
+ self.is_flash_attn_backend = self.attn_backend in {
256
+ _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA
257
+ }
258
259
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
260
# [s, b, 3 * head * head_dim]
@@ -297,10 +301,13 @@ class Qwen2_5_VisionAttention(nn.Module):
301
q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
302
k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
303
300
- if self.attn_backend == _Backend.FLASH_ATTN:
304
+ if self.is_flash_attn_backend:
305
# from vllm_flash_attn.flash_attn_interface import (
306
# flash_attn_varlen_func)
303
- from flash_attn import flash_attn_varlen_func
307
+ if self.attn_backend == _Backend.ROCM_AITER_FA:
308
+ from aiter import flash_attn_varlen_func
309
+ else:
310
+ from flash_attn import flash_attn_varlen_func
311
312
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
313
@@ -311,7 +318,7 @@ class Qwen2_5_VisionAttention(nn.Module):
318
cu_seqlens_k=cu_seqlens,
319
max_seqlen_q=max_seqlen,
320
max_seqlen_k=max_seqlen,
314
- dropout_p=0,
321
+ dropout_p=0.0,
322
causal=False)
323
324
context_layer = rearrange(output,
@@ -635,7 +642,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
642
cu_seqlens: torch.Tensor,
643
) -> tuple[Optional[int], Optional[list[int]]]:
644
max_seqlen, seqlens = None, None
638
- if self.attn_backend == _Backend.FLASH_ATTN:
645
+ if (self.attn_backend == _Backend.FLASH_ATTN
646
+ or self.attn_backend == _Backend.ROCM_AITER_FA):
647
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
648
elif self.attn_backend == _Backend.XFORMERS:
649
seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
vllm/model_executor/models/qwen2_vl.py
+13
-5
@@ -274,10 +274,14 @@ class Qwen2VisionAttention(nn.Module):
274
# Detect attention implementation.
275
self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
276
if self.attn_backend not in {
277
- _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS
277
+ _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
278
+ _Backend.ROCM_AITER_FA
279
}:
280
raise RuntimeError(
281
f"Qwen2-VL does not support {self.attn_backend} backend now.")
282
+ self.is_flash_attn_backend = self.attn_backend in {
283
+ _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA
284
+ }
285
286
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
287
# [s, b, 3 * head * head_dim]
@@ -324,10 +328,13 @@ class Qwen2VisionAttention(nn.Module):
328
q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
329
k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)
330
327
- if self.attn_backend == _Backend.FLASH_ATTN:
331
+ if self.is_flash_attn_backend:
332
# from vllm_flash_attn.flash_attn_interface import (
333
# flash_attn_varlen_func)
330
- from flash_attn import flash_attn_varlen_func
334
+ if self.attn_backend == _Backend.ROCM_AITER_FA:
335
+ from aiter import flash_attn_varlen_func
336
+ else:
337
+ from flash_attn import flash_attn_varlen_func
338
339
q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
340
@@ -338,7 +345,7 @@ class Qwen2VisionAttention(nn.Module):
345
cu_seqlens_k=cu_seqlens,
346
max_seqlen_q=max_seqlen,
347
max_seqlen_k=max_seqlen,
341
- dropout_p=0,
348
+ dropout_p=0.0,
349
causal=False)
350
351
context_layer = rearrange(output,
@@ -620,7 +627,8 @@ class Qwen2VisionTransformer(nn.Module):
627
self, cu_seqlens: torch.Tensor
628
) -> tuple[Optional[int], Optional[list[int]]]:
629
max_seqlen, seqlens = None, None
623
- if self.attn_backend == _Backend.FLASH_ATTN:
630
+ if (self.attn_backend == _Backend.FLASH_ATTN
631
+ or self.attn_backend == _Backend.ROCM_AITER_FA):
632
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
633
elif self.attn_backend == _Backend.XFORMERS:
634
seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
vllm/model_executor/models/vision.py
+7
-29
@@ -7,9 +7,7 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union
7
import torch
8
from transformers import PretrainedConfig
9
10
-import vllm.envs as envs
11
-from vllm.attention.selector import (backend_name_to_enum,
12
- get_global_forced_attn_backend)
10
+from vllm.attention.selector import get_env_variable_attn_backend
11
from vllm.logger import init_logger
12
from vllm.platforms import _Backend, current_platform
13
@@ -75,32 +73,12 @@ def get_vit_attn_backend(support_fa: bool = False) -> _Backend:
73
Get the available attention backend for Vision Transformer.
74
"""
75
# TODO(Isotr0py): Remove `support_fa` after support FA for all ViTs attn.
78
- selected_backend: Optional[_Backend] = get_global_forced_attn_backend()
79
- if selected_backend is None:
80
- backend_by_env_var: Optional[str] = envs.VLLM_ATTENTION_BACKEND
81
- if backend_by_env_var is not None:
82
- selected_backend = backend_name_to_enum(backend_by_env_var)
83
- if selected_backend is None:
84
- if current_platform.is_cuda():
85
- device_available = current_platform.has_device_capability(80)
86
- if device_available and support_fa:
87
- from transformers.utils import is_flash_attn_2_available
88
- if is_flash_attn_2_available():
89
- selected_backend = _Backend.FLASH_ATTN
90
- else:
91
- logger.warning_once(
92
- "Current `vllm-flash-attn` has a bug inside vision "
93
- "module, so we use xformers backend instead. You can "
94
- "run `pip install flash-attn` to use flash-attention "
95
- "backend.")
96
- selected_backend = _Backend.XFORMERS
97
- else:
98
- # For Volta and Turing GPUs, use xformers instead.
99
- selected_backend = _Backend.XFORMERS
100
- else:
101
- # Default to torch SDPA for other non-GPU platforms.
102
- selected_backend = _Backend.TORCH_SDPA
103
- return selected_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)
82
83
84
def resolve_visual_encoder_outputs(
vllm/platforms/cuda.py
+14
@@ -206,6 +206,20 @@ class CudaPlatformBase(Platform):
206
torch.cuda.reset_peak_memory_stats(device)
207
return torch.cuda.max_memory_allocated(device)
208
209
+ @classmethod
210
+ def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
211
+ if cls.has_device_capability(80) and support_fa:
212
+ from transformers.utils import is_flash_attn_2_available
213
+ if is_flash_attn_2_available():
214
+ return _Backend.FLASH_ATTN
215
+ logger.warning_once(
216
+ "Current `vllm-flash-attn` has a bug inside vision "
217
+ "module, so we use xformers backend instead. You can "
218
+ "run `pip install flash-attn` to use flash-attention "
219
+ "backend.")
220
+ # Fallback for Volta/Turing GPUs or FA not supported
221
+ return _Backend.XFORMERS
222
+
223
@classmethod
224
def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
225
kv_cache_dtype, block_size, use_v1,
vllm/platforms/interface.py
+5
@@ -46,6 +46,7 @@ class _Backend(enum.Enum):
46
ROCM_FLASH = enum.auto()
47
ROCM_AITER_MLA = enum.auto() # Supported by V1
48
ROCM_AITER_MLA_VLLM_V1 = enum.auto()
49
+ ROCM_AITER_FA = enum.auto() # used for ViT attn backend
50
TORCH_SDPA = enum.auto()
51
FLASHINFER = enum.auto()
52
FLASHINFER_VLLM_V1 = enum.auto()
@@ -186,6 +187,10 @@ class Platform:
187
else:
188
return device_id
189
190
+ @classmethod
191
+ def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
192
+ return _Backend.TORCH_SDPA
193
+
194
@classmethod
195
def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
196
dtype: torch.dtype, kv_cache_dtype: Optional[str],
vllm/platforms/rocm.py
+12
@@ -173,6 +173,18 @@ class RocmPlatform(Platform):
173
"quark", "ptpc_fp8"
174
]
175
176
+ @classmethod
177
+ def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
178
+ if support_fa:
179
+ if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA
180
+ and on_gfx9()):
181
+ # Note: AITER FA is only supported for Qwen-VL models.
182
+ # TODO: Add support for other VL models in their model class.
183
+ return _Backend.ROCM_AITER_FA
184
+ if on_gfx9():
185
+ return _Backend.FLASH_ATTN
186
+ return _Backend.TORCH_SDPA
187
+
188
@classmethod
189
def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
190
kv_cache_dtype, block_size, use_v1,