(self, _q, _k)
| 68 | return self.cos_cached, self.sin_cached |
| 69 | |
| 70 | def forward(self, _q, _k): |
| 71 | batch, seq_len, num_heads, head_dim = _q.shape |
| 72 | q = _q.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim) |
| 73 | k = _k.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim) |
| 74 | cos, sin = self.cos_sin(seq_len, q.device, q.dtype) |
| 75 | return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) |
| 76 | |
| 77 | |
| 78 | class FalconAttentionFused(nn.Module): |
nothing calls this directly
no test coverage detected