From c65146e75e71f72bdb93542efdc40c87caccb42e Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 5 Jun 2024 09:18:59 -0700 Subject: [PATCH] [Misc] Fix docstring of get_attn_backend (#5271) --- vllm/attention/selector.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/attention/selector.py b/vllm/attention/selector.py index 9ceda343..7253483f 100644 --- a/vllm/attention/selector.py +++ b/vllm/attention/selector.py @@ -31,15 +31,14 @@ def get_attn_backend( block_size: int, is_blocksparse: bool = False, ) -> Type[AttentionBackend]: + """Selects which attention backend to use and lazily imports it.""" if is_blocksparse: logger.info("Using BlocksparseFlashAttention backend.") from vllm.attention.backends.blocksparse_attn import ( 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, sliding_window, dtype, kv_cache_dtype, block_size)