Skip to content

Commit

Permalink
[Misc] Fix docstring of get_attn_backend (vllm-project#5271)
Browse files Browse the repository at this point in the history
  • Loading branch information
WoosukKwon authored Jun 5, 2024
1 parent e78010a commit 0805160
Showing 1 changed file with 2 additions and 3 deletions.
5 changes: 2 additions & 3 deletions vllm/attention/selector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 0805160

Please sign in to comment.