Aluode/PerceptionLabPortable

1from typing import Optional2 3import torch4 5 6def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:7    """8    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,9    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)10    """11    batch, num_key_value_heads, slen, head_dim = hidden_states.shape12    if n_rep == 1:13        return hidden_states14    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)15    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)16 17 18def sdpa_attention_paged_forward(19    module: torch.nn.Module,20    query: torch.Tensor,21    key: torch.Tensor,22    value: torch.Tensor,23    attention_mask: Optional[torch.Tensor],24    dropout: float = 0.0,25    scaling: Optional[float] = None,26    **kwargs,27) -> tuple[torch.Tensor, None]: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    causal_mask = attention_mask43 44    # Run the actual attention45    query = query.contiguous()46    key = key.contiguous()47    value = value.contiguous()48    attn_output = torch.nn.functional.scaled_dot_product_attention(49        query,50        key,51        value,52        attn_mask=causal_mask,53        dropout_p=dropout,54        scale=scaling,55        # Packed sequence format is used for input, so that it can never be causal.56        is_causal=False,57    )58    attn_output = attn_output.transpose(1, 2).contiguous()59 60    return attn_output, None61