Unverified Commit c65146e7 authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[Misc] Fix docstring of get_attn_backend (#5271)

parent 41ca62cf
...@@ -31,15 +31,14 @@ def get_attn_backend( ...@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment