xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 06:35:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Misc] Minor type annotation fix (#3716)

Browse Source
This commit is contained in:
Woosuk Kwon 2024-03-28 21:12:24 -07:00 committed by GitHub
parent 26422e477b
commit 395aa823ea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/attention/selector.py
View File

@ -1,4 +1,5 @@
from functools import lru_cache
from typing import Type
import torch
@ -10,7 +11,7 @@ logger = init_logger(__name__)
@lru_cache(maxsize=None)
def get_attn_backend(dtype: torch.dtype) -> AttentionBackend:
def get_attn_backend(dtype: torch.dtype) -> Type[AttentionBackend]:
if _can_use_flash_attn(dtype):
logger.info("Using FlashAttention backend.")
from vllm.attention.backends.flash_attn import ( # noqa: F401