def METHOD_NAME(
    self,
    positions: torch.Tensor,  # [num_tokens]
    query: torch.Tensor,  # [num_tokens, num_heads, head_size]
    key: torch.Tensor,  # [num_tokens, num_heads, head_size]
) -> Tuple[torch.Tensor, torch.Tensor]:
    query_rot = query[..., :self.rotary_dim]
    query_pass = query[..., self.rotary_dim:]
    key_rot = key[..., :self.rotary_dim]
    key_pass = key[..., self.rotary_dim:]
    query_rot = query_rot.transpose(0, 1)
    key_rot = key_rot.transpose(0, 1)
    cos = F.embedding(positions, self.cos_cached)
    sin = F.embedding(positions, self.sin_cached)
    query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin)
    query_rot = query_rot.transpose(0, 1).contiguous()
    key_rot = key_rot.transpose(0, 1).contiguous()
    query = torch.cat((query_rot, query_pass), dim=-1)
    key = torch.cat((key_rot, key_pass), dim=-1)
    # Output query/key shape: [num_tokens, num_tokens, head_size]
    return query, key