keys = repeat_kv(keys, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)values = repeat_kv(values, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)
xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)values = values.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)
scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)

self.freqs_cis = precompute_freqs_cis(    self.params.dim // self.params.n_heads, self.params.max_seq_len * 2)


def apply_rotary_emb(    xq: torch.Tensor,    xk: torch.Tensor,    freqs_cis: torch.Tensor,) -> Tuple[torch.Tensor, torch.Tensor]:    """    Apply rotary embeddings to input tensors using the given frequency tensor.
    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are    returned as real tensors.
    Args:        xq (torch.Tensor): Query tensor to apply rotary embeddings.        xk (torch.Tensor): Key tensor to apply rotary embeddings.        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.
    Returns:        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.    "        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)    return xq_out.type_as(xq), xk_out.type_as(xk)

def apply_rotary_emb(    xq: torch.Tensor,    xk: torch.Tensor,    freqs_cis: torch.Tensor,) -> Tuple[torch.Tensor, torch.Tensor]:    """    Apply rotary embeddings to input tensors using the given frequency tensor.
    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are    returned as real tensors.
    Args:        xq (torch.Tensor): Query tensor to apply rotary embeddings.        xk (torch.Tensor): Key tensor to apply rotary embeddings.        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.
    Returns:        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.    """        xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)    return xq_out.type_as(xq), xk_out.type_as(xk)