squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Misc] Fix docstring of get_attn_backend (#5271)

Browse Source
This commit is contained in:
Woosuk Kwon 2024-06-05 09:18:59 -07:00 committed by GitHub
parent 41ca62cf03
commit c65146e75e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/attention/selector.py
View File

@ -31,15 +31,14 @@ def get_attn_backend(
block_size: int, block_size: int,
is_blocksparse: bool = False, is_blocksparse: bool = False,
) -> Type[AttentionBackend]: ) -> Type[AttentionBackend]:
"""Selects which attention backend to use and lazily imports it."""
if is_blocksparse: if is_blocksparse:
logger.info("Using BlocksparseFlashAttention backend.") logger.info("Using BlocksparseFlashAttention backend.")
from vllm.attention.backends.blocksparse_attn import ( from vllm.attention.backends.blocksparse_attn import (
BlocksparseFlashAttentionBackend) BlocksparseFlashAttentionBackend)
return BlocksparseFlashAttentionBackend return BlocksparseFlashAttentionBackend
"""Determine which attention backend to use and only import
the selected backend module.
"""
backend = which_attn_to_use(num_heads, head_size, num_kv_heads, backend = which_attn_to_use(num_heads, head_size, num_kv_heads,
sliding_window, dtype, kv_cache_dtype, sliding_window, dtype, kv_cache_dtype,
block_size) block_size)