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 5, 11:00 UTC
[FlashAttention] Symlink FA4 instead of copying when using `VLLM_FLASH_ATTN_SRC_DIR` (#38814)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Apr 8, 2026 at 08:04 UTC
308cec5864890f5c0724e1d4531d9fe2ee0a8209
2 files changed
+47
-16
cmake/external_projects/vllm_flash_attn.cmake
+27
-15
@@ -87,18 +87,30 @@ endforeach()
87
#
88
add_custom_target(_vllm_fa4_cutedsl_C)
89
90
-# Copy flash_attn/cute directory (needed for FA4) and transform imports
91
-# The cute directory uses flash_attn.cute imports internally, which we replace
92
-# with vllm.vllm_flash_attn.cute to match our package structure.
93
-install(CODE "
94
- file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
95
- foreach(SRC_FILE \${CUTE_PY_FILES})
96
- file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
97
- set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
98
- get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
99
- file(MAKE_DIRECTORY \${DST_DIR})
100
- file(READ \${SRC_FILE} FILE_CONTENTS)
101
- string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
102
- file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
103
- endforeach()
104
-" COMPONENT _vllm_fa4_cutedsl_C)
90
+# Install flash_attn/cute directory (needed for FA4).
91
+# When using a local source dir (VLLM_FLASH_ATTN_SRC_DIR), create a symlink
92
+# so edits to cute-dsl Python files take effect immediately without rebuilding.
93
+# Otherwise, copy files and transform flash_attn.cute imports to
94
+# vllm.vllm_flash_attn.cute to match our package structure.
95
+if(VLLM_FLASH_ATTN_SRC_DIR)
96
+ install(CODE "
97
+ set(LINK_TARGET \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\")
98
+ set(LINK_NAME \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute\")
99
+ file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")
100
+ file(REMOVE_RECURSE \"\${LINK_NAME}\")
101
+ file(CREATE_LINK \"\${LINK_TARGET}\" \"\${LINK_NAME}\" SYMBOLIC)
102
+ " COMPONENT _vllm_fa4_cutedsl_C)
103
+else()
104
+ install(CODE "
105
+ file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
106
+ foreach(SRC_FILE \${CUTE_PY_FILES})
107
+ file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
108
+ set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
109
+ get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
110
+ file(MAKE_DIRECTORY \${DST_DIR})
111
+ file(READ \${SRC_FILE} FILE_CONTENTS)
112
+ string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
113
+ file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
114
+ endforeach()
115
+ " COMPONENT _vllm_fa4_cutedsl_C)
116
+endif()
vllm/vllm_flash_attn/__init__.py
+20
-1
@@ -1,7 +1,26 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from vllm.vllm_flash_attn.flash_attn_interface import (
4
+import importlib.machinery
5
+import os
6
+import sys
7
+import types
8
+
9
+# In symlink mode (VLLM_FLASH_ATTN_SRC_DIR), cute/ is a symlink to the real
10
+# source tree and its files use `flash_attn.cute.*` imports (not rewritten).
11
+# Register a virtual `flash_attn` package so those imports resolve.
12
+_cute_dir = os.path.join(os.path.dirname(__file__), "cute")
13
+if os.path.islink(_cute_dir) and "flash_attn" not in sys.modules:
14
+ _fa_mod = types.ModuleType("flash_attn")
15
+ _fa_mod.__path__ = [os.path.dirname(os.path.realpath(_cute_dir))]
16
+ _fa_mod.__package__ = "flash_attn"
17
+ _fa_mod.__spec__ = importlib.machinery.ModuleSpec(
18
+ "flash_attn", None, is_package=True
19
+ )
20
+ _fa_mod.__spec__.submodule_search_locations = _fa_mod.__path__
21
+ sys.modules["flash_attn"] = _fa_mod
22
+
23
+from vllm.vllm_flash_attn.flash_attn_interface import ( # noqa: E402
24
FA2_AVAILABLE,
25
FA3_AVAILABLE,
26
fa_version_unsupported_reason,