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, 13:45 UTC
[Bugfix] Pass FA version in `MultiHeadAttention` (#30575)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Dec 12, 2025 at 19:02 UTC
86a3261525858bca5d4a234691fae0496d6fed99
1 file changed
+10
vllm/attention/layer.py
+10
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer."""
4
5
+import functools
6
from collections.abc import Callable
7
from typing import cast
8
@@ -17,6 +18,7 @@ from vllm.attention.backends.abstract import (
18
)
19
from vllm.attention.backends.registry import AttentionBackendEnum
20
from vllm.attention.selector import get_attn_backend
21
+from vllm.attention.utils.fa_utils import get_flash_attn_version
22
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
23
from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
24
from vllm.config import CacheConfig, get_current_vllm_config
@@ -524,6 +526,14 @@ class MultiHeadAttention(nn.Module):
526
AttentionBackendEnum.ROCM_AITER_FA,
527
}
528
529
+ self.fa_version = None
530
+ if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:
531
+ self.fa_version = get_flash_attn_version()
532
+ assert self._flash_attn_varlen_func is not None
533
+ self._flash_attn_varlen_func = functools.partial(
534
+ self._flash_attn_varlen_func, fa_version=self.fa_version
535
+ )
536
+
537
logger.info_once(
538
f"Using {self.attn_backend} for MultiHeadAttention in multimodal encoder."
539
)