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 11, 08:45 UTC
[MM] Pass FA version in ViT Attn (#30756)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Nicolò Lucchesi committed
Dec 17, 2025 at 00:54 UTC
e087fbc393055fb69e9acf71fa124be0190498ec
2 files changed
+14
-1
vllm/attention/layers/mm_encoder_attention.py
+6
@@ -10,6 +10,7 @@ from vllm.attention.ops.vit_attn_wrappers import (
10
vit_flash_attn_wrapper,
11
vit_torch_sdpa_wrapper,
12
)
13
+from vllm.attention.utils.fa_utils import get_flash_attn_version
14
from vllm.config import MultiModalConfig
15
from vllm.logger import init_logger
16
from vllm.model_executor.custom_op import CustomOp
@@ -101,6 +102,10 @@ class MMEncoderAttention(CustomOp):
102
self.attn_backend,
103
)
104
105
+ if self.is_flash_attn_backend:
106
+ assert self.flash_attn_varlen_func is not None
107
+ self._fa_version = get_flash_attn_version()
108
+
109
logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")
110
111
@classmethod
@@ -204,6 +209,7 @@ class MMEncoderAttention(CustomOp):
209
max_seqlen=max_seqlen,
210
batch_size=bsz,
211
is_rocm_aiter=(self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA),
212
+ fa_version=self._fa_version,
213
)
214
return output
215
vllm/attention/ops/vit_attn_wrappers.py
+8
-1
@@ -28,11 +28,15 @@ def flash_attn_maxseqlen_wrapper(
28
max_seqlen: torch.Tensor,
29
batch_size: int,
30
is_rocm_aiter: bool,
31
+ fa_version: int,
32
) -> torch.Tensor:
33
+ kwargs = {}
34
if is_rocm_aiter:
35
from aiter import flash_attn_varlen_func
36
else:
37
from vllm.attention.utils.fa_utils import flash_attn_varlen_func
38
+
39
+ kwargs["fa_version"] = fa_version
40
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
41
output = flash_attn_varlen_func(
42
q,
@@ -44,6 +48,7 @@ def flash_attn_maxseqlen_wrapper(
48
max_seqlen_k=max_seqlen.item(),
49
dropout_p=0.0,
50
causal=False,
51
+ **kwargs,
52
)
53
context_layer = einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size)
54
return context_layer
@@ -57,6 +62,7 @@ def flash_attn_maxseqlen_wrapper_fake(
62
max_seqlen: torch.Tensor,
63
batch_size: int,
64
is_rocm_aiter: bool,
65
+ fa_version: int,
66
) -> torch.Tensor:
67
return torch.empty_like(q)
68
@@ -76,9 +82,10 @@ def vit_flash_attn_wrapper(
82
max_seqlen: torch.Tensor,
83
batch_size: int,
84
is_rocm_aiter: bool,
85
+ fa_version: int,
86
) -> torch.Tensor:
87
return torch.ops.vllm.flash_attn_maxseqlen_wrapper(
81
- q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter
88
+ q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter, fa_version
89
)
90
91