(
xq: torch.Tensor,
xk: torch.Tensor,
freqs_cis: torch.Tensor,
)
| 52 | |
| 53 | |
| 54 | def apply_rotary_emb( |
| 55 | xq: torch.Tensor, |
| 56 | xk: torch.Tensor, |
| 57 | freqs_cis: torch.Tensor, |
| 58 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 59 | # xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) |
| 60 | # k_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) |
| 61 | xq_ = torch.view_as_complex( |
| 62 | xq.float().reshape(*xq.shape[:-1], 2, -1).transpose(-2, -1).contiguous() |
| 63 | ) |
| 64 | xk_ = torch.view_as_complex( |
| 65 | xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous() |
| 66 | ) |
| 67 | freqs_cis = reshape_for_broadcast(freqs_cis, xq_) |
| 68 | xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3) |
| 69 | xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3) |
| 70 | return xq_out.type_as(xq), xk_out.type_as(xk) |
| 71 | |
| 72 | |
| 73 | class LlamaAttentionFused(nn.Module): |
no test coverage detected