cached_attention

cached_attention

Attention with generation and conditioning cache support.

CachedAttention

CachedAttention(embedding_dimension, number_of_heads, number_of_key_value_heads=None, head_dimension=None, dropout=0.0, bias=True, attention_type=value)

Bases: Module

Base attention module with KV cache support.

Supports both Multi-Head Attention (MHA) and Grouped Query Attention (GQA). Can be used for self-attention or cross-attention.

Initialize cached attention module.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Model embedding dimension	required
number_of_heads	int	Number of query heads	required
number_of_key_value_heads	int | None	Number of key/value heads (for GQA)	None
head_dimension	int | None	Per-head dimension. Defaults to embedding_dimension // number_of_heads. Override for architectures where hidden_size != number_of_heads * head_dim (e.g. Gemma2).	None
dropout	float	Dropout probability for attention weights	0.0
bias	bool	Whether to include bias in projections	True
attention_type	str	Type of attention (use AttentionType enum values)	value

Raises:

Type	Description
ValueError	If dimensions don't match or invalid attention type

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    number_of_key_value_heads: int | None = None,
    head_dimension: int | None = None,
    dropout: float = 0.0,
    bias: bool = True,
    attention_type: str = AttentionType.MULTI_HEAD.value,
):
    """Initialize cached attention module.

    Args:
        embedding_dimension: Model embedding dimension
        number_of_heads: Number of query heads
        number_of_key_value_heads: Number of key/value heads (for GQA)
        head_dimension: Per-head dimension. Defaults to embedding_dimension // number_of_heads.
            Override for architectures where hidden_size != number_of_heads * head_dim (e.g. Gemma2).
        dropout: Dropout probability for attention weights
        bias: Whether to include bias in projections
        attention_type: Type of attention (use AttentionType enum values)

    Raises:
        ValueError: If dimensions don't match or invalid attention type
    """
    super().__init__()
    if number_of_heads <= 0:
        raise ValueError(
            f"number_of_heads must be positive, got {number_of_heads}."
        )
    if head_dimension is not None and head_dimension <= 0:
        raise ValueError(f"head_dimension must be positive, got {head_dimension}.")
    if head_dimension is None and embedding_dimension % number_of_heads != 0:
        raise ValueError(
            f"embedding_dimension ({embedding_dimension}) must be divisible "
            f"by number_of_heads ({number_of_heads})."
        )
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    self.head_dimension: int = (
        head_dimension
        if head_dimension is not None
        else embedding_dimension // number_of_heads
    )
    self.dropout = dropout
    self.attention_type = attention_type
    if attention_type == AttentionType.GROUPED_QUERY.value:
        if number_of_key_value_heads is None:
            raise ValueError("number_of_key_value_heads required for GQA")
        if number_of_key_value_heads <= 0:
            raise ValueError(
                "number_of_key_value_heads must be positive, "
                f"got {number_of_key_value_heads}."
            )
        if number_of_heads % number_of_key_value_heads != 0:
            raise ValueError(
                f"number_of_heads ({number_of_heads}) must be divisible "
                f"by number_of_key_value_heads ({number_of_key_value_heads})."
            )
        self.number_of_key_value_heads = number_of_key_value_heads
        self.group_size = number_of_heads // number_of_key_value_heads
    elif attention_type == AttentionType.MULTI_HEAD.value:
        if (
            number_of_key_value_heads is not None
            and number_of_key_value_heads != number_of_heads
        ):
            raise ValueError(
                "number_of_key_value_heads must be None or equal to "
                "number_of_heads for multi-head attention, got "
                f"{number_of_key_value_heads}."
            )
        self.number_of_key_value_heads = number_of_heads
        self.group_size = 1
    else:
        raise ValueError(
            f"Unsupported attention type: {attention_type}. "
            f"Must be one of {[e.value for e in AttentionType]}."
        )
    self.query_projection = nn.Linear(
        embedding_dimension,
        number_of_heads * self.head_dimension,
        bias=bias,
    )
    self.key_projection = nn.Linear(
        embedding_dimension,
        self.number_of_key_value_heads * self.head_dimension,
        bias=bias,
    )
    self.value_projection = nn.Linear(
        embedding_dimension,
        self.number_of_key_value_heads * self.head_dimension,
        bias=bias,
    )
    self.output_projection = nn.Linear(
        number_of_heads * self.head_dimension,
        embedding_dimension,
        bias=bias,
    )
    self.output_projection.SQUARE_ROOT_WEIGHT = (
        True  # Flag for initialization (GPT2 style)
    )

compute_query

compute_query(query_input)

Project and reshape query input.

Parameters:

Name	Type	Description	Default
query_input	Tensor	(B, query_len, D)	required

Returns:

Type	Description
Tensor	Projected queries (B, number_of_heads, query_len, head_dim).

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def compute_query(self, query_input: torch.Tensor) -> torch.Tensor:
    """Project and reshape query input.

    Args:
        query_input: (B, query_len, D)

    Returns:
        Projected queries (B, number_of_heads, query_len, head_dim).
    """
    batch_size, query_length, _ = query_input.shape
    projected = self.query_projection(query_input)
    # (B, L, number_of_heads * head_dim) -> (B, number_of_heads, L, head_dim)
    return projected.view(
        batch_size, query_length, self.number_of_heads, self.head_dimension
    ).transpose(1, 2)

compute_key

compute_key(key_input)

Project and reshape key input.

Parameters:

Name	Type	Description	Default
key_input	Tensor	(B, key_len, D)	required

Returns:

Type	Description
Tensor	Projected keys (B, kv_heads, key_len, head_dim).

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def compute_key(self, key_input: torch.Tensor) -> torch.Tensor:
    """Project and reshape key input.

    Args:
        key_input: (B, key_len, D)

    Returns:
        Projected keys (B, kv_heads, key_len, head_dim).
    """
    batch_size, key_length, _ = key_input.shape
    projected = self.key_projection(key_input)
    # (B, L, kv_heads * head_dim) -> (B, kv_heads, L, head_dim)
    return projected.view(
        batch_size, key_length, self.number_of_key_value_heads, self.head_dimension
    ).transpose(1, 2)

compute_value

compute_value(value_input)

Project and reshape value input.

Parameters:

Name	Type	Description	Default
value_input	Tensor	(B, value_len, D)	required

Returns:

Type	Description
Tensor	Projected values (B, kv_heads, value_len, head_dim).

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def compute_value(self, value_input: torch.Tensor) -> torch.Tensor:
    """Project and reshape value input.

    Args:
        value_input: (B, value_len, D)

    Returns:
        Projected values (B, kv_heads, value_len, head_dim).
    """
    batch_size, value_length, _ = value_input.shape
    projected = self.value_projection(value_input)
    # (B, L, kv_heads * head_dim) -> (B, kv_heads, L, head_dim)
    return projected.view(
        batch_size,
        value_length,
        self.number_of_key_value_heads,
        self.head_dimension,
    ).transpose(1, 2)

compute_query_key_value

compute_query_key_value(query_input, key_input, value_input)

Project inputs to query, key, value.

Parameters:

Name	Type	Description	Default
query_input	Tensor	Query input (B, query_len, D)	required
key_input	Tensor	Key input (B, key_len, D)	required
value_input	Tensor	Value input (B, value_len, D)	required

Returns:

Type	Description
Tensor	Tuple of (queries, keys, values). Queries: (B, number_of_heads, query_len, head_dim).
Tensor	Keys/values: (B, kv_heads, key_len, head_dim).

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def compute_query_key_value(
    self,
    query_input: torch.Tensor,
    key_input: torch.Tensor,
    value_input: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Project inputs to query, key, value.

    Args:
        query_input: Query input (B, query_len, D)
        key_input: Key input (B, key_len, D)
        value_input: Value input (B, value_len, D)

    Returns:
        Tuple of (queries, keys, values). Queries: (B, number_of_heads, query_len, head_dim).
        Keys/values: (B, kv_heads, key_len, head_dim).
    """
    return (
        self.compute_query(query_input),
        self.compute_key(key_input),
        self.compute_value(value_input),
    )

compute_attention

compute_attention(queries, keys, values, attention_mask=None)

Compute scaled dot-product attention.

Parameters:

Name	Type	Description	Default
queries	Tensor	Query tensor (B, number_of_heads, query_len, head_dim)	required
keys	Tensor	Key tensor (B, kv_heads, key_len, head_dim) - compact for GQA	required
values	Tensor	Value tensor (B, kv_heads, value_len, head_dim) - compact for GQA	required
attention_mask	Tensor | None	Optional bool mask (B, 1, query_len, key_len) where True means masked.	None

Returns:

Type	Description
Tensor	Attention output (B, query_len, embedding_dimension)

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def compute_attention(
    self,
    queries: torch.Tensor,
    keys: torch.Tensor,
    values: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Compute scaled dot-product attention.

    Args:
        queries: Query tensor (B, number_of_heads, query_len, head_dim)
        keys: Key tensor (B, kv_heads, key_len, head_dim) - compact for GQA
        values: Value tensor (B, kv_heads, value_len, head_dim) - compact for GQA
        attention_mask: Optional bool mask (B, 1, query_len, key_len) where True means masked.

    Returns:
        Attention output (B, query_len, embedding_dimension)
    """
    batch_size = queries.shape[0]
    query_length = queries.shape[2]
    if self.group_size > 1:  # For GQ attention
        keys = torch.repeat_interleave(
            keys, self.group_size, dim=1
        )  # (B, number_of_heads, kv_length, head_dim)
        values = torch.repeat_interleave(
            values, self.group_size, dim=1
        )  # (B, number_of_heads, kv_length, head_dim)

    sdpa_mask = None
    if attention_mask is not None:
        sdpa_mask = (
            ~attention_mask if attention_mask is not None else None
        )  # False means don't attend/padded
        # cf. https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

    attended_values = F.scaled_dot_product_attention(
        queries,
        keys,
        values,
        attn_mask=sdpa_mask,
        dropout_p=self.dropout if self.training else 0.0,
        scale=self.head_dimension**-0.5,
    )
    attended_values = attended_values.transpose(
        1, 2
    ).contiguous()  # (B, query_len, number_of_heads, head_dim)
    attended_values = attended_values.view(
        batch_size,
        query_length,
        self.number_of_heads
        * self.head_dimension,  # (B, query_len, embedding_dimension)
    )
    output = self.output_projection(attended_values)
    return output

forward

forward(query_input, key_input=None, value_input=None, attention_mask=None, generation_cache=None, positional_encoding=None, conditioning_cache=None)

Forward pass with optional generation and conditioning caches.

Parameters:

Name	Type	Description	Default
query_input	Tensor	Query input (B, query_len, D).	required
key_input	Tensor | None	Key input (B, key_len, D). None when using conditioning_cache.	None
value_input	Tensor | None	Value input (B, value_len, D). None when using conditioning_cache.	None
attention_mask	Tensor | None	Bool mask (B, 1, query_len, key_len), True = masked.	None
generation_cache	GenerationLayerCache | None	Cached K/V from the main sequence. When provided, an updated cache is returned.	None
positional_encoding	Module | None	Optional RoPE module.	None
conditioning_cache	ConditioningLayerCache | None	Precomputed K/V for static conditioning. When present, key_input/value_input are ignored and cached K/V is used directly.	None

Returns:

Type	Description
tuple[Tensor, GenerationLayerCache | None]	Tuple of (output (B, query_len, D), updated GenerationLayerCache or None).

Source code in src/versatil/models/layers/transformer/attention/cached_attention.py

def forward(
    self,
    query_input: torch.Tensor,
    key_input: torch.Tensor | None = None,
    value_input: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    generation_cache: GenerationLayerCache | None = None,
    positional_encoding: nn.Module | None = None,
    conditioning_cache: ConditioningLayerCache | None = None,
) -> tuple[torch.Tensor, GenerationLayerCache | None]:
    """Forward pass with optional generation and conditioning caches.

    Args:
        query_input: Query input (B, query_len, D).
        key_input: Key input (B, key_len, D). None when using conditioning_cache.
        value_input: Value input (B, value_len, D). None when using conditioning_cache.
        attention_mask: Bool mask (B, 1, query_len, key_len), True = masked.
        generation_cache: Cached K/V from the main sequence. When provided,
            an updated cache is returned.
        positional_encoding: Optional RoPE module.
        conditioning_cache: Precomputed K/V for static conditioning. When present,
            key_input/value_input are ignored and cached K/V is used directly.

    Returns:
        Tuple of (output (B, query_len, D), updated GenerationLayerCache or None).
    """
    if conditioning_cache is not None:
        queries = self.compute_query(query_input)
        keys = conditioning_cache.keys
        values = conditioning_cache.values
    else:
        if key_input is None or value_input is None:
            raise ValueError(
                "key_input and value_input required when conditioning_cache is not provided"
            )

        queries, keys, values = self.compute_query_key_value(
            query_input, key_input, value_input
        )
        cache_position = 0
        if generation_cache is not None and generation_cache.keys.numel() > 0:
            cache_position = generation_cache.get_length()
        # Apply RoPE before concatenation so cached keys retain original rotations
        if positional_encoding is not None:
            queries, keys = apply_rope_positional_encoding(
                queries=queries,
                keys=keys,
                positional_encoding=positional_encoding,
                cache_position=cache_position,
            )
        if generation_cache is not None and generation_cache.keys.numel() > 0:
            keys = torch.cat([generation_cache.keys, keys], dim=2)
            values = torch.cat([generation_cache.values, values], dim=2)

    output = self.compute_attention(queries, keys, values, attention_mask)

    new_cache = None
    if generation_cache is not None and conditioning_cache is None:
        new_cache = GenerationLayerCache(keys=keys, values=values)

    return output, new_cache