huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 12:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: improve condtional and error message

Browse Source
This commit is contained in:
drbh 2024-08-01 16:17:29 +00:00
parent cae28dcbf1
commit 5b649d67c4
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
server/text_generation_server/layers/attention/cuda.py
View File

@ -171,8 +171,9 @@ def paged_attention(
try:
if major <= 8:
raise ImportError("Flash Attention V2 requires CUDA 11.0 or higher")
is_ampere_or_newer = major >= 8 and minor >= 0
if not is_ampere_or_newer:
raise ImportError("FlashAttention only supports Ampere GPUs or newer.")
import flash_attn_2_cuda