File size: 3,276 Bytes
309d916 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | from __future__ import annotations
import torch
import torch.nn.functional as F
def banded_softmax_attention(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, w: int
) -> torch.Tensor:
T = q.shape[-2]
i = torch.arange(T, device=q.device)[:, None]
j = torch.arange(T, device=q.device)[None, :]
mask = (j <= i) & (j > i - w)
return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
def swa_fused_forward(mixer, x: torch.Tensor) -> torch.Tensor:
B, T, C = x.shape
H, w = (mixer.n_head, mixer.swa_fused_window)
D = C // H
q = mixer.swa_fused_q(x).reshape(B, T, H, D).transpose(1, 2)
k = mixer.swa_fused_k(x).reshape(B, T, H, D).transpose(1, 2)
v = mixer.swa_fused_v(x).reshape(B, T, H, D).transpose(1, 2)
o = banded_softmax_attention(q, k, v, w)
o = o.transpose(1, 2).reshape(B, T, C)
return x + mixer.swa_fused_o(o)
class CPUSlidingWindow:
def __init__(self, mixer):
self.mixer = mixer
self.H = mixer.n_head
self.C = mixer.swa_fused_q.out_features
self.D = self.C // self.H
self.w = mixer.swa_fused_window
self.scale = self.D ** (-0.5)
def init_state(self, batch_size: int = 1, device=None):
if device is None:
device = self.mixer.swa_fused_q.weight.device
W1 = self.w - 1
return {
"k": torch.zeros(batch_size, W1, self.H, self.D, device=device),
"v": torch.zeros(batch_size, W1, self.H, self.D, device=device),
"n": 0,
}
def _project(self, x):
B, L, _ = x.shape
q = self.mixer.swa_fused_q(x).reshape(B, L, self.H, self.D)
k = self.mixer.swa_fused_k(x).reshape(B, L, self.H, self.D)
v = self.mixer.swa_fused_v(x).reshape(B, L, self.H, self.D)
return (q, k, v)
def forward_chunk(self, x: torch.Tensor, state):
B, L, C = x.shape
w = self.w
if state is None:
state = self.init_state(B, device=x.device)
n = state["n"]
W1 = w - 1
ck = state["k"][:, W1 - n :, :, :] if n > 0 else state["k"][:, :0]
cv = state["v"][:, W1 - n :, :, :] if n > 0 else state["v"][:, :0]
q, k, v = self._project(x)
k_all = torch.cat([ck, k], dim=1)
v_all = torch.cat([cv, v], dim=1)
Tk = n + L
iq = torch.arange(L, device=x.device)[:, None] + n
jk = torch.arange(Tk, device=x.device)[None, :]
mask = (jk <= iq) & (jk > iq - w)
qh = q.transpose(1, 2)
kh = k_all.transpose(1, 2)
vh = v_all.transpose(1, 2)
o = F.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask)
o = o.transpose(1, 2).reshape(B, L, C)
out = x + self.mixer.swa_fused_o(o)
keep = min(W1, Tk)
new_k = state["k"].clone()
new_v = state["v"].clone()
if keep > 0:
new_k[:, W1 - keep :, :, :] = k_all[:, Tk - keep :, :, :]
new_v[:, W1 - keep :, :, :] = v_all[:, Tk - keep :, :, :]
new_state = {"k": new_k, "v": new_v, "n": keep}
return (out, new_state)
def step(self, x: torch.Tensor, state):
if x.dim() == 2:
x = x.unsqueeze(1)
o, state = self.forward_chunk(x, state)
return (o.squeeze(1), state)
|