Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
b6dcb4d4
Unverified
Commit
b6dcb4d4
authored
Apr 26, 2024
by
Roy
Committed by
GitHub
Apr 25, 2024
Browse files
[Misc] Fix flash attention backend log (#4368)
parent
b5b4a398
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
5 additions
and
5 deletions
+5
-5
vllm/attention/selector.py
vllm/attention/selector.py
+5
-5
No files found.
vllm/attention/selector.py
View file @
b6dcb4d4
...
...
@@ -25,7 +25,7 @@ class _Backend(enum.Enum):
def
get_attn_backend
(
dtype
:
torch
.
dtype
)
->
Type
[
AttentionBackend
]:
backend
=
_which_attn_to_use
(
dtype
)
if
backend
==
_Backend
.
FLASH_ATTN
:
logger
.
info
(
"Using FlashAttention backend."
)
logger
.
info
(
"Using FlashAttention
-2
backend."
)
from
vllm.attention.backends.flash_attn
import
(
# noqa: F401
FlashAttentionBackend
)
return
FlashAttentionBackend
...
...
@@ -62,12 +62,12 @@ def _which_attn_to_use(dtype: torch.dtype) -> _Backend:
# NVIDIA GPUs.
if
torch
.
cuda
.
get_device_capability
()[
0
]
<
8
:
# Volta and Turing NVIDIA GPUs.
logger
.
info
(
"Cannot use FlashAttention backend for Volta and Turing "
logger
.
info
(
"Cannot use FlashAttention
-2
backend for Volta and Turing "
"GPUs."
)
return
_Backend
.
XFORMERS
if
dtype
not
in
(
torch
.
float16
,
torch
.
bfloat16
):
logger
.
info
(
"Cannot use FlashAttention backend for dtype other than "
logger
.
info
(
"Cannot use FlashAttention
-2
backend for dtype other than "
"torch.float16 or torch.bfloat16."
)
return
_Backend
.
XFORMERS
...
...
@@ -75,8 +75,8 @@ def _which_attn_to_use(dtype: torch.dtype) -> _Backend:
import
flash_attn
# noqa: F401
except
ImportError
:
logger
.
info
(
"Cannot use FlashAttention backend because the flash_attn
package
"
"is not found. Please install it for better performance."
)
"Cannot use FlashAttention
-2
backend because the flash_attn "
"
package
is not found. Please install it for better performance."
)
return
_Backend
.
XFORMERS
backend_by_env_var
=
os
.
getenv
(
VLLM_ATTENTION_BACKEND
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment