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, 13:30 UTC
[BugFix] Fix vllm_flash_attn install issues (#17267)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Aaron Pham <contact@aarnphm.xyz>
Lucas Wilkinson committed
Apr 27, 2025 at 20:27 UTC
d8bccde686348b38fc4590dffe8069c68627ab67
11 files changed
+28
-284
.github/CODEOWNERS
+1
@@ -12,6 +12,7 @@
12
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth
13
/vllm/model_executor/guided_decoding @mgoin @russellb
14
/vllm/multimodal @DarkLight1337 @ywang96
15
+/vllm/vllm_flash_attn @LucasWilkinson
16
CMakeLists.txt @tlrmchlsmth
17
18
# vLLM V1
.gitignore
-2
@@ -3,8 +3,6 @@
3
4
# vllm-flash-attn built from source
5
vllm/vllm_flash_attn/*
6
-!vllm/vllm_flash_attn/__init__.py
7
-!vllm/vllm_flash_attn/fa_utils.py
6
7
# Byte-compiled / optimized / DLL files
8
__pycache__/
setup.py
+19
-7
@@ -269,15 +269,17 @@ class cmake_build_ext(build_ext):
269
# First, run the standard build_ext command to compile the extensions
270
super().run()
271
272
- # copy vllm/vllm_flash_attn/*.py from self.build_lib to current
272
+ # copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current
273
# directory so that they can be included in the editable build
274
import glob
275
- files = glob.glob(
276
- os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "*.py"))
275
+ files = glob.glob(os.path.join(self.build_lib, "vllm",
276
+ "vllm_flash_attn", "**", "*.py"),
277
+ recursive=True)
278
for file in files:
279
dst_file = os.path.join("vllm/vllm_flash_attn",
279
- os.path.basename(file))
280
+ file.split("vllm/vllm_flash_attn/")[-1])
281
print(f"Copying {file} to {dst_file}")
282
+ os.makedirs(os.path.dirname(dst_file), exist_ok=True)
283
self.copy_file(file, dst_file)
284
285
@@ -377,12 +379,22 @@ class repackage_wheel(build_ext):
379
"vllm/_flashmla_C.abi3.so",
380
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
381
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
380
- "vllm/vllm_flash_attn/flash_attn_interface.py",
382
"vllm/cumem_allocator.abi3.so",
383
# "vllm/_version.py", # not available in nightly wheels yet
384
]
384
- file_members = filter(lambda x: x.filename in files_to_copy,
385
- wheel.filelist)
385
+
386
+ file_members = list(
387
+ filter(lambda x: x.filename in files_to_copy, wheel.filelist))
388
+
389
+ # vllm_flash_attn python code:
390
+ # Regex from
391
+ # `glob.translate('vllm/vllm_flash_attn/**/*.py', recursive=True)`
392
+ import re
393
+ compiled_regex = re.compile(
394
+ r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
395
+ file_members += list(
396
+ filter(lambda x: compiled_regex.match(x.filename),
397
+ wheel.filelist))
398
399
for file in file_members:
400
print(f"Extracting and including {file.filename} "
vllm/attention/backends/flash_attn.py
+3
-3
@@ -22,13 +22,13 @@ from vllm.attention.backends.utils import (
22
compute_slot_mapping_start_idx, get_num_prefill_decode_query_kv_tokens,
23
get_seq_len_block_table_args, is_all_cross_attn_metadata_set,
24
is_all_encoder_attn_metadata_set, is_block_tables_empty)
25
+from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
26
+ get_flash_attn_version)
27
from vllm.logger import init_logger
28
from vllm.multimodal import MultiModalPlaceholderMap
29
from vllm.utils import async_tensor_h2d, make_tensor_with_pad
30
from vllm.vllm_flash_attn import (flash_attn_varlen_func,
31
flash_attn_with_kvcache)
30
-from vllm.vllm_flash_attn.fa_utils import (flash_attn_supports_fp8,
31
- get_flash_attn_version)
32
33
if TYPE_CHECKING:
34
from vllm.worker.model_runner import (ModelInputForGPUBuilder,
@@ -689,7 +689,7 @@ class FlashAttentionImpl(AttentionImpl):
689
assert output is not None, "Output tensor must be provided."
690
691
# NOTE(woosuk): FlashAttention2 does not support FP8 KV cache.
692
- if self.vllm_flash_attn_version < 3 or output.dtype != torch.bfloat16:
692
+ if not flash_attn_supports_fp8() or output.dtype != torch.bfloat16:
693
assert (
694
layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0), (
695
"key/v_scale is only supported in FlashAttention 3 with "
vllm/attention/backends/mla/common.py
+1
-1
@@ -205,6 +205,7 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
205
compute_slot_mapping_start_idx,
206
is_block_tables_empty)
207
from vllm.attention.ops.merge_attn_states import merge_attn_states
208
+from vllm.attention.utils.fa_utils import get_flash_attn_version
209
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
210
LinearBase, RowParallelLinear,
211
UnquantizedLinearMethod)
@@ -214,7 +215,6 @@ from vllm.multimodal import MultiModalPlaceholderMap
215
from vllm.platforms import current_platform
216
from vllm.triton_utils import HAS_TRITON
217
from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down
217
-from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version
218
219
if HAS_TRITON:
220
from vllm.attention.ops.triton_flash_attention import triton_attention
vllm/attention/utils/fa_utils.py
renamed
vllm/engine/arg_utils.py
+1
-1
@@ -1377,7 +1377,7 @@ class EngineArgs:
1377
) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
1378
supported = False
1379
if fp8_attention and will_use_fa:
1380
- from vllm.vllm_flash_attn.fa_utils import (
1380
+ from vllm.attention.utils.fa_utils import (
1381
flash_attn_supports_fp8)
1382
supported = flash_attn_supports_fp8()
1383
if not supported:
vllm/v1/attention/backends/flash_attn.py
+2
-2
@@ -11,11 +11,11 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionMetadata, AttentionType,
12
is_quantized_kv_cache)
13
from vllm.attention.ops.merge_attn_states import merge_attn_states
14
+from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
15
+ get_flash_attn_version)
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
18
from vllm.utils import cdiv
17
-from vllm.vllm_flash_attn.fa_utils import (flash_attn_supports_fp8,
18
- get_flash_attn_version)
19
20
if TYPE_CHECKING:
21
from vllm.v1.core.sched.output import SchedulerOutput
vllm/v1/attention/backends/mla/common.py
+1
-1
@@ -197,6 +197,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
197
MLAAttentionImpl)
198
from vllm.attention.backends.utils import get_mla_dims
199
from vllm.attention.ops.merge_attn_states import merge_attn_states
200
+from vllm.attention.utils.fa_utils import get_flash_attn_version
201
from vllm.logger import init_logger
202
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
203
LinearBase, RowParallelLinear,
@@ -204,7 +205,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
205
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
206
from vllm.platforms import current_platform
207
from vllm.utils import cdiv, round_down
207
-from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version
208
209
try:
210
from vllm.vllm_flash_attn import flash_attn_varlen_func
vllm/vllm_flash_attn/__init__.py
deleted
-22
@@ -1,22 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import importlib.metadata
4
-
5
-try:
6
- __version__ = importlib.metadata.version("vllm-flash-attn")
7
-except importlib.metadata.PackageNotFoundError:
8
- # in this case, vllm-flash-attn is built from installing vllm editable
9
- __version__ = "0.0.0.dev0"
10
-
11
-from .flash_attn_interface import (fa_version_unsupported_reason,
12
- flash_attn_varlen_func,
13
- flash_attn_with_kvcache,
14
- get_scheduler_metadata,
15
- is_fa_version_supported, sparse_attn_func,
16
- sparse_attn_varlen_func)
17
-
18
-__all__ = [
19
- 'flash_attn_varlen_func', 'flash_attn_with_kvcache',
20
- 'get_scheduler_metadata', 'sparse_attn_func', 'sparse_attn_varlen_func',
21
- 'is_fa_version_supported', 'fa_version_unsupported_reason'
22
-]
vllm/vllm_flash_attn/flash_attn_interface.pyi
deleted
-245
@@ -1,245 +0,0 @@
1
-# ruff: ignore
2
-# SPDX-License-Identifier: Apache-2.0
3
-
4
-from __future__ import annotations
5
-
6
-from typing import Any, Literal, overload
7
-
8
-import torch
9
-
10
-def get_scheduler_metadata(
11
- batch_size: int,
12
- max_seqlen_q: int,
13
- max_seqlen_k: int,
14
- num_heads_q: int,
15
- num_heads_kv: int,
16
- headdim: int,
17
- cache_seqlens: torch.Tensor,
18
- qkv_dtype: torch.dtype = ...,
19
- headdim_v: int | None = ...,
20
- cu_seqlens_q: torch.Tensor | None = ...,
21
- cu_seqlens_k_new: torch.Tensor | None = ...,
22
- cache_leftpad: torch.Tensor | None = ...,
23
- page_size: int = ...,
24
- max_seqlen_k_new: int = ...,
25
- causal: bool = ...,
26
- window_size: tuple[int, int] = ...,
27
- has_softcap: bool = ...,
28
- num_splits: int = ...,
29
- pack_gqa: Any | None = ...,
30
- sm_margin: int = ...,
31
-): ...
32
-@overload
33
-def flash_attn_varlen_func(
34
- q: tuple[int, int, int],
35
- k: tuple[int, int, int],
36
- v: tuple[int, int, int],
37
- max_seqlen_q: int,
38
- cu_seqlens_q: torch.Tensor | None,
39
- max_seqlen_k: int,
40
- cu_seqlens_k: torch.Tensor | None = ...,
41
- seqused_k: Any | None = ...,
42
- q_v: Any | None = ...,
43
- dropout_p: float = ...,
44
- causal: bool = ...,
45
- window_size: list[int] | None = ...,
46
- softmax_scale: float = ...,
47
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
48
- deterministic: bool = ...,
49
- return_attn_probs: bool = ...,
50
- block_table: Any | None = ...,
51
- return_softmax_lse: Literal[False] = ...,
52
- out: Any = ...,
53
- # FA3 Only
54
- scheduler_metadata: Any | None = ...,
55
- q_descale: Any | None = ...,
56
- k_descale: Any | None = ...,
57
- v_descale: Any | None = ...,
58
- # Version selector
59
- fa_version: int = ...,
60
-) -> tuple[int, int, int]: ...
61
-@overload
62
-def flash_attn_varlen_func(
63
- q: tuple[int, int, int],
64
- k: tuple[int, int, int],
65
- v: tuple[int, int, int],
66
- max_seqlen_q: int,
67
- cu_seqlens_q: torch.Tensor | None,
68
- max_seqlen_k: int,
69
- cu_seqlens_k: torch.Tensor | None = ...,
70
- seqused_k: Any | None = ...,
71
- q_v: Any | None = ...,
72
- dropout_p: float = ...,
73
- causal: bool = ...,
74
- window_size: list[int] | None = ...,
75
- softmax_scale: float = ...,
76
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
77
- deterministic: bool = ...,
78
- return_attn_probs: bool = ...,
79
- block_table: Any | None = ...,
80
- return_softmax_lse: Literal[True] = ...,
81
- out: Any = ...,
82
- # FA3 Only
83
- scheduler_metadata: Any | None = ...,
84
- q_descale: Any | None = ...,
85
- k_descale: Any | None = ...,
86
- v_descale: Any | None = ...,
87
- # Version selector
88
- fa_version: int = ...,
89
-) -> tuple[tuple[int, int, int], tuple[int, int]]: ...
90
-@overload
91
-def flash_attn_with_kvcache(
92
- q: tuple[int, int, int, int],
93
- k_cache: tuple[int, int, int, int],
94
- v_cache: tuple[int, int, int, int],
95
- k: tuple[int, int, int, int] | None = ...,
96
- v: tuple[int, int, int, int] | None = ...,
97
- rotary_cos: tuple[int, int] | None = ...,
98
- rotary_sin: tuple[int, int] | None = ...,
99
- cache_seqlens: int | torch.Tensor | None = None,
100
- cache_batch_idx: torch.Tensor | None = None,
101
- cache_leftpad: torch.Tensor | None = ...,
102
- block_table: torch.Tensor | None = ...,
103
- softmax_scale: float = ...,
104
- causal: bool = ...,
105
- window_size: tuple[int, int] = ..., # -1 means infinite context window
106
- softcap: float = ...,
107
- rotary_interleaved: bool = ...,
108
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
109
- num_splits: int = ...,
110
- return_softmax_lse: Literal[False] = ...,
111
- *,
112
- out: Any = ...,
113
- # FA3 Only
114
- scheduler_metadata: Any | None = ...,
115
- q_descale: Any | None = ...,
116
- k_descale: Any | None = ...,
117
- v_descale: Any | None = ...,
118
- # Version selector
119
- fa_version: int = ...,
120
-) -> tuple[int, int, int, int]: ...
121
-@overload
122
-def flash_attn_with_kvcache(
123
- q: tuple[int, int, int, int],
124
- k_cache: tuple[int, int, int, int],
125
- v_cache: tuple[int, int, int, int],
126
- k: tuple[int, int, int, int] | None = ...,
127
- v: tuple[int, int, int, int] | None = ...,
128
- rotary_cos: tuple[int, int] | None = ...,
129
- rotary_sin: tuple[int, int] | None = ...,
130
- cache_seqlens: int | torch.Tensor | None = None,
131
- cache_batch_idx: torch.Tensor | None = None,
132
- cache_leftpad: torch.Tensor | None = ...,
133
- block_table: torch.Tensor | None = ...,
134
- softmax_scale: float = ...,
135
- causal: bool = ...,
136
- window_size: tuple[int, int] = ..., # -1 means infinite context window
137
- softcap: float = ...,
138
- rotary_interleaved: bool = ...,
139
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
140
- num_splits: int = ...,
141
- return_softmax_lse: Literal[True] = ...,
142
- *,
143
- out: Any = ...,
144
- # FA3 Only
145
- scheduler_metadata: Any | None = ...,
146
- q_descale: Any | None = ...,
147
- k_descale: Any | None = ...,
148
- v_descale: Any | None = ...,
149
- # Version selector
150
- fa_version: int = ...,
151
-) -> tuple[tuple[int, int, int], tuple[int, int]]: ...
152
-@overload
153
-def sparse_attn_func(
154
- q: tuple[int, int, int, int],
155
- k: tuple[int, int, int, int],
156
- v: tuple[int, int, int, int],
157
- block_count: tuple[int, int, float],
158
- block_offset: tuple[int, int, float, int],
159
- column_count: tuple[int, int, float],
160
- column_index: tuple[int, int, float, int],
161
- dropout_p: float = ...,
162
- softmax_scale: float = ...,
163
- causal: bool = ...,
164
- softcap: float = ...,
165
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
166
- deterministic: bool = ...,
167
- return_attn_probs: bool = ...,
168
- *,
169
- return_softmax_lse: Literal[False] = ...,
170
- out: Any = ...,
171
-) -> tuple[int, int, int]: ...
172
-@overload
173
-def sparse_attn_func(
174
- q: tuple[int, int, int, int],
175
- k: tuple[int, int, int, int],
176
- v: tuple[int, int, int, int],
177
- block_count: tuple[int, int, float],
178
- block_offset: tuple[int, int, float, int],
179
- column_count: tuple[int, int, float],
180
- column_index: tuple[int, int, float, int],
181
- dropout_p: float = ...,
182
- softmax_scale: float = ...,
183
- causal: bool = ...,
184
- softcap: float = ...,
185
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
186
- deterministic: bool = ...,
187
- return_attn_probs: bool = ...,
188
- *,
189
- return_softmax_lse: Literal[True] = ...,
190
- out: Any = ...,
191
-) -> tuple[tuple[int, int, int], tuple[int, int]]: ...
192
-@overload
193
-def sparse_attn_varlen_func(
194
- q: tuple[int, int, int],
195
- k: tuple[int, int, int],
196
- v: tuple[int, int, int],
197
- block_count: tuple[int, int, float],
198
- block_offset: tuple[int, int, float, int],
199
- column_count: tuple[int, int, float],
200
- column_index: tuple[int, int, float, int],
201
- cu_seqlens_q: torch.Tensor | None,
202
- cu_seqlens_k: torch.Tensor | None,
203
- max_seqlen_q: int,
204
- max_seqlen_k: int,
205
- dropout_p: float = ...,
206
- softmax_scale: float = ...,
207
- causal: bool = ...,
208
- softcap: float = ...,
209
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
210
- deterministic: bool = ...,
211
- return_attn_probs: bool = ...,
212
- *,
213
- return_softmax_lse: Literal[False] = ...,
214
- out: Any = ...,
215
-) -> tuple[int, int, int]: ...
216
-@overload
217
-def sparse_attn_varlen_func(
218
- q: tuple[int, int, int],
219
- k: tuple[int, int, int],
220
- v: tuple[int, int, int],
221
- block_count: tuple[int, int, float],
222
- block_offset: tuple[int, int, float, int],
223
- column_count: tuple[int, int, float],
224
- column_index: tuple[int, int, float, int],
225
- cu_seqlens_q: torch.Tensor | None,
226
- cu_seqlens_k: torch.Tensor | None,
227
- max_seqlen_q: int,
228
- max_seqlen_k: int,
229
- dropout_p: float = ...,
230
- softmax_scale: float = ...,
231
- causal: bool = ...,
232
- softcap: float = ...,
233
- alibi_slopes: tuple[int] | tuple[int, int] | None = ...,
234
- deterministic: bool = ...,
235
- return_attn_probs: bool = ...,
236
- *,
237
- return_softmax_lse: Literal[True] = ...,
238
- out: Any = ...,
239
-) -> tuple[tuple[int, int, int], tuple[int, int]]: ...
240
-def is_fa_version_supported(
241
- fa_version: int, device: torch.device | None = None
242
-) -> bool: ...
243
-def fa_version_unsupported_reason(
244
- fa_version: int, device: torch.device | None = None
245
-) -> str | None: ...