[Misc] Minor type annotation fix (#3716)

395aa823 · Woosuk Kwon · GitHub · 26422e47 · 395aa823
Unverified Commit 395aa823 authored Mar 28, 2024 by Woosuk Kwon Committed by GitHub Mar 28, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

vllm/attention/selector.py vllm/attention/selector.py +2 -1

No files found.
--- a/vllm/attention/selector.py
+++ b/vllm/attention/selector.py
 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