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 15, 21:00 UTC
[Hardware] Update the flash attn tag to support Blackwell (#14244)
Pavani Majety committed
Mar 5, 2025 at 22:01 UTC
ed6ea06577ec06f0b3a9ac921b55ef254f19d923
2 files changed
+8
-3
cmake/external_projects/vllm_flash_attn.cmake
+2
-2
@@ -38,7 +38,7 @@ else()
38
FetchContent_Declare(
39
vllm-flash-attn
40
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41
- GIT_TAG 720c94869cf2e0ff5a706e9c7f1dce0939686ade
41
+ GIT_TAG 9bfa9869829d8c593527eb34c5271d0090f7ccc9
42
GIT_PROGRESS TRUE
43
# Don't share the vllm-flash-attn build between build types
44
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -64,4 +64,4 @@ install(
64
DESTINATION vllm_flash_attn
65
COMPONENT _vllm_fa3_C
66
FILES_MATCHING PATTERN "*.py"
67
-)
\ No newline at end of file
67
+)
vllm/attention/backends/utils.py
+6
-1
@@ -595,7 +595,7 @@ def get_flash_attn_version():
595
# if hopper default to FA3, otherwise stick to FA2 for now
596
# TODO(lucas): profile FA3 on ampere to see if it makes sense to
597
# use FA3 as default for both
598
- if current_platform.get_device_capability()[0] >= 9:
598
+ if current_platform.get_device_capability()[0] == 9:
599
fa_version = 3 if is_fa_version_supported(3) else 2
600
else:
601
fa_version = 2
@@ -603,6 +603,11 @@ def get_flash_attn_version():
603
if envs.VLLM_FLASH_ATTN_VERSION is not None:
604
assert envs.VLLM_FLASH_ATTN_VERSION in [2, 3]
605
fa_version = envs.VLLM_FLASH_ATTN_VERSION
606
+ if (current_platform.get_device_capability()[0] == 10
607
+ and envs.VLLM_FLASH_ATTN_VERSION == 3):
608
+ logger.warning("Cannot use FA version 3 on Blackwell platform",
609
+ "defaulting to FA version 2.")
610
+ fa_version = 2
611
612
if not is_fa_version_supported(fa_version):
613
logger.error("Cannot use FA version %d is not supported due to %s",