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, 06:00 UTC
[Bugfix][V1] Allow manual FlashAttention for Blackwell (#19492)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 12, 2025 at 06:40 UTC
af09b3f0a05479e36bab12ae4e05cbdd3fed9d27
1 file changed
+13
-4
vllm/platforms/cuda.py
+13
-4
@@ -226,15 +226,21 @@ class CudaPlatformBase(Platform):
226
if selected_backend == _Backend.FLASHINFER:
227
logger.info_once("Using FlashInfer backend on V1 engine.")
228
return "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
229
- if selected_backend == _Backend.FLEX_ATTENTION:
229
+ elif selected_backend == _Backend.FLEX_ATTENTION:
230
logger.info("Using FlexAttenion backend on V1 engine.")
231
return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
232
- if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
232
+ elif selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
233
logger.info_once("Using Triton backend on V1 engine.")
234
return ("vllm.v1.attention.backends."
235
"triton_attn.TritonAttentionBackend")
236
+ elif selected_backend == _Backend.FLASH_ATTN:
237
+ logger.info_once("Using Flash Attention backend on V1 engine.")
238
+ return ("vllm.v1.attention.backends."
239
+ "flash_attn.FlashAttentionBackend")
240
+
241
+ # Default backends for V1 engine
242
+ # Prefer FlashInfer for Blackwell GPUs if installed
243
if cls.is_device_capability(100):
237
- # Prefer FlashInfer for V1 on Blackwell GPUs if installed
244
try:
245
import flashinfer # noqa: F401
246
logger.info_once(
@@ -248,10 +254,13 @@ class CudaPlatformBase(Platform):
254
"Blackwell (SM 10.0) GPUs; it is recommended to "
255
"install FlashInfer for better performance.")
256
pass
251
- if cls.has_device_capability(80):
257
+ # FlashAttention is the default for SM 8.0+ GPUs
258
+ elif cls.has_device_capability(80):
259
logger.info_once("Using Flash Attention backend on V1 engine.")
260
return ("vllm.v1.attention.backends."
261
"flash_attn.FlashAttentionBackend")
262
+
263
+ # Backends for V0 engine
264
if selected_backend == _Backend.FLASHINFER:
265
logger.info("Using FlashInfer backend.")
266
return "vllm.attention.backends.flashinfer.FlashInferBackend"