Aluode/PerceptionLabPortable

1from typing import Optional2 3import torch4from torch import nn5 6 7def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:8    """9    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,10    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)11    """12    batch, num_key_value_heads, slen, head_dim = hidden_states.shape13    if n_rep == 1:14        return hidden_states15    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)16    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)17 18 19def eager_paged_attention_forward(20    module: nn.Module,21    query: torch.Tensor,22    key: torch.Tensor,23    value: torch.Tensor,24    attention_mask: Optional[torch.Tensor],  # shape [seqlen_q, seqlen_k]25    scaling: float,26    **kwargs,27):28    # Add KV cache to the key and value tensors29    cache = kwargs.pop("cache", None)30    if cache is not None:31        # This changes the shape of k and v from [1, num_kv_heads, seqlen_kv, head_dim] to [-1, num_kv_heads, head_dim]32        key, value = cache.update(key, value, module.layer_idx, **kwargs)33        key = key.transpose(0, 1).unsqueeze(0)34        value = value.transpose(0, 1).unsqueeze(0)35 36    # Repeat the key and value tensors for each group of key-value heads37    if hasattr(module, "num_key_value_groups"):38        key = repeat_kv(key, module.num_key_value_groups)39        value = repeat_kv(value, module.num_key_value_groups)40 41    # Get the right causal mask for the current layer42    if isinstance(attention_mask, dict):43        sliding_window = getattr(module, "sliding_window", 1)44        layer_type = "full_attention" if sliding_window == 1 or sliding_window is None else "sliding_attention"45        causal_mask = attention_mask[layer_type]46    else:47        causal_mask = attention_mask48 49    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling50    if causal_mask is not None:51        attn_weights = attn_weights + causal_mask52 53    # Handle attention sinks if the model has them54    if hasattr(module, "sinks"):55        # Retrieve the sink and add it to the attention weights56        sinks = module.sinks.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)57        attn_weights = torch.cat([attn_weights, sinks], dim=-1)58        # Normalize the attention weights for better numerical stability59        attn_weights = attn_weights - attn_weights.max(dim=-1, keepdim=True).values60        # Apply softmax and drop the sink. Not exactly the same code as eager w/ sink, but the same code does not produce the same results.61        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)62        attn_weights = attn_weights[..., :-1]63    else:64        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)65 66    attn_output = torch.matmul(attn_weights, value)67    attn_output = attn_output.transpose(1, 2).contiguous()68 69    return attn_output, attn_weights70