71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
"""Attention layer."""
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.attention.backends.abstract import AttentionMetadata
|
|
from vllm.attention.selector import get_attn_backend
|
|
from vllm.config import CacheConfig
|
|
|
|
|
|
class Attention(nn.Module):
|
|
"""Attention layer.
|
|
|
|
This class takes query, key, and value tensors as input. The input tensors
|
|
can either contain prompt tokens or generation tokens.
|
|
The class does the following:
|
|
|
|
1. Store the input key and value tensors in the KV cache.
|
|
2. Perform (multi-head/multi-query/grouped-query) attention.
|
|
3. Return the output tensor.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_heads: int,
|
|
head_size: int,
|
|
scale: float,
|
|
num_kv_heads: Optional[int] = None,
|
|
alibi_slopes: Optional[List[float]] = None,
|
|
sliding_window: Optional[int] = None,
|
|
cache_config: Optional[CacheConfig] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
if cache_config is not None:
|
|
kv_cache_dtype = cache_config.cache_dtype
|
|
block_size = cache_config.block_size
|
|
else:
|
|
kv_cache_dtype = "auto"
|
|
block_size = 16
|
|
if num_kv_heads is None:
|
|
num_kv_heads = num_heads
|
|
# During model initialization, the default dtype is set as the model
|
|
# weight and activation dtype.
|
|
dtype = torch.get_default_dtype()
|
|
attn_backend = get_attn_backend(num_heads, head_size, num_kv_heads,
|
|
sliding_window, dtype, kv_cache_dtype,
|
|
block_size)
|
|
impl_cls = attn_backend.get_impl_cls()
|
|
self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
|
|
alibi_slopes, sliding_window, kv_cache_dtype)
|
|
|
|
def forward(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
kv_cache: Optional[torch.Tensor],
|
|
attn_metadata: AttentionMetadata,
|
|
kv_scale: float = 1.0,
|
|
) -> torch.Tensor:
|
|
return self.impl.forward(query, key, value, kv_cache, attn_metadata,
|
|
kv_scale)
|
|
|
|
def extra_repr(self) -> str:
|
|
s = f"head_size={self.impl.head_size}" # type: ignore
|
|
s += f", num_heads={self.impl.num_heads}" # type: ignore
|
|
s += f", num_kv_heads={self.impl.num_kv_heads}" # type: ignore
|
|
s += f", scale={self.impl.scale}" # type: ignore
|
|
return s
|