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
[Chore] added stubs for `vllm_flash_attn` during development mode (#17228)
Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
Aaron Pham committed
Apr 26, 2025 at 10:45 UTC
e782e0a170a652c879593234c8723a7350c4784a
4 files changed
+269
-2
pyproject.toml
+2
-1
@@ -58,7 +58,8 @@ ignore_patterns = [
58
line-length = 80
59
exclude = [
60
# External file, leaving license intact
61
- "examples/other/fp8/quantizer/quantize.py"
61
+ "examples/other/fp8/quantizer/quantize.py",
62
+ "vllm/vllm_flash_attn/flash_attn_interface.pyi"
63
]
64
65
[tool.ruff.lint.per-file-ignores]
setup.py
-1
@@ -378,7 +378,6 @@ class repackage_wheel(build_ext):
378
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
379
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
380
"vllm/vllm_flash_attn/flash_attn_interface.py",
381
- "vllm/vllm_flash_attn/__init__.py",
381
"vllm/cumem_allocator.abi3.so",
382
# "vllm/_version.py", # not available in nightly wheels yet
383
]
vllm/vllm_flash_attn/__init__.py
+22
@@ -0,0 +1,22 @@
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
new
+245
@@ -0,0 +1,245 @@
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: ...