from __future__ import annotations import torch import torch.nn.functional as F def gdn_recurrent( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor, g: torch.Tensor, scale: float, initial_state: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: q, k, v, beta, g = ( t.transpose(1, 2).contiguous().to(torch.float32) for t in (q, k, v, beta, g) ) B, H, T, K = q.shape V = v.shape[-1] o = torch.zeros(B, H, T, V, dtype=torch.float32, device=q.device) if initial_state is None: h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) else: h = initial_state.to(torch.float32).clone() q = q * scale for i in range(T): b_q = q[:, :, i] b_k = k[:, :, i] b_v = v[:, :, i].clone() h = h * g[:, :, i].exp()[..., None, None] b_v = b_v - (h * b_k[..., None]).sum(-2) b_v = b_v * beta[:, :, i][..., None] h = h + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) o[:, :, i] = torch.einsum("bhd,bhdm->bhm", b_q, h) return (o.transpose(1, 2).contiguous(), h) def gdn_chunk_recurrent( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor, g: torch.Tensor, scale: float, initial_state: torch.Tensor | None = None, chunk_size: int = 64, ) -> tuple[torch.Tensor, torch.Tensor]: bt = chunk_size q, k, v, beta, g = ( t.transpose(1, 2).contiguous().to(torch.float32) for t in (q, k, v, beta, g) ) B, H, T, K = q.shape Vd = v.shape[-1] pad = (bt - T % bt) % bt if pad: q = F.pad(q, (0, 0, 0, pad)) k = F.pad(k, (0, 0, 0, pad)) v = F.pad(v, (0, 0, 0, pad)) beta = F.pad(beta, (0, pad)) g = F.pad(g, (0, pad)) L = q.shape[2] n = L // bt q = q * scale v = v * beta[..., None] k_beta = k * beta[..., None] def _chunks(x): return x.reshape(B, H, n, bt, x.shape[-1]) q, k, v, k_beta = (_chunks(q), _chunks(k), _chunks(v), _chunks(k_beta)) decay = g.reshape(B, H, n, bt).cumsum(-1) decay_exp = decay.exp()[..., None] l_mask = (decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().tril() eye_mask = torch.triu( torch.ones(bt, bt, dtype=torch.bool, device=q.device), diagonal=0 ) attn = -(k_beta @ k.transpose(-1, -2) * l_mask).masked_fill(eye_mask, 0) for i in range(1, bt): attn[..., i, :i] = attn[..., i, :i].clone() + ( attn[..., i, :i, None].clone() * attn[..., :i, :i].clone() ).sum(-2) attn = attn + torch.eye(bt, dtype=torch.float32, device=q.device) v = attn @ v k_cumdecay = attn @ (k_beta * decay_exp) S = q.new_zeros(B, H, K, Vd) if initial_state is not None: S = initial_state.to(torch.float32).clone() o = torch.zeros_like(v) causal = torch.triu( torch.ones(bt, bt, dtype=torch.bool, device=q.device), diagonal=1 ) for i in range(n): q_i, k_i, v_i = (q[:, :, i], k[:, :, i], v[:, :, i]) a_i = (q_i @ k_i.transpose(-1, -2) * l_mask[:, :, i]).masked_fill(causal, 0) v_new = v_i - k_cumdecay[:, :, i] @ S o_inter = q_i * decay[:, :, i, :, None].exp() @ S o[:, :, i] = o_inter + a_i @ v_new S = ( S * decay[:, :, i, -1, None, None].exp() + ( k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()[..., None] ).transpose(-1, -2) @ v_new ) o = o.reshape(B, H, L, Vd)[:, :, :T] return (o.transpose(1, 2).contiguous(), S) def causal_depthwise_conv1d( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None, activation: str | None = "silu", ) -> torch.Tensor: B, T, C = x.shape W = weight.shape[-1] xt = x.transpose(1, 2) xp = F.pad(xt, (W - 1, 0)) y = F.conv1d(xp, weight, bias=bias, groups=C)[..., :T].transpose(1, 2) if activation == "silu": y = F.silu(y) elif activation is not None: raise ValueError(f"unsupported activation {activation!r}") return y def _conv_chunk( p: torch.Tensor, weight: torch.Tensor, conv_state: torch.Tensor | None, bias: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: B, L, C = p.shape W = weight.shape[-1] pt = p.transpose(1, 2) if conv_state is None: conv_state = pt.new_zeros(B, C, W - 1) full = torch.cat([conv_state, pt], dim=-1) y = F.silu(F.conv1d(full, weight, bias=bias, groups=C)).transpose(1, 2) return (y, full[..., -(W - 1) :]) def gdn_gate( g: torch.Tensor, A_log: torch.Tensor, dt_bias: torch.Tensor | None = None ) -> torch.Tensor: g = g.float() if dt_bias is not None: g = g + dt_bias.float() return -A_log.float().exp() * F.softplus(g) def gated_rmsnorm( o: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float = 1e-05 ) -> torch.Tensor: o = o.float() gate = gate.float() rstd = torch.rsqrt(o.pow(2).mean(-1, keepdim=True) + eps) return o * rstd * weight.float() * (gate * torch.sigmoid(gate)) class CPUGatedDeltaNet: def __init__(self, gdn): self.gdn = gdn self.H = gdn.num_heads self.Dk = gdn.head_k_dim self.Dv = gdn.head_v_dim self.C = gdn.value_dim self.W = gdn.conv_size self.scale = self.Dk ** (-0.5) self.eps = getattr(gdn.o_norm, "eps", 1e-05) self.chunk_size = 64 assert not gdn.allow_neg_eigval, "allow_neg_eigval not supported" assert gdn.use_gate and gdn.use_short_conv def _project(self, x, conv_state): g = self.gdn q, sq = _conv_chunk( g.q_proj(x), g.q_conv1d.weight, None if conv_state is None else conv_state[0], ) k, sk = _conv_chunk( g.k_proj(x), g.k_conv1d.weight, None if conv_state is None else conv_state[1], ) v, sv = _conv_chunk( g.v_proj(x), g.v_conv1d.weight, None if conv_state is None else conv_state[2], ) B, L = (x.shape[0], x.shape[1]) q = F.normalize(q.view(B, L, self.H, self.Dk), p=2, dim=-1) k = F.normalize(k.view(B, L, self.H, self.Dk), p=2, dim=-1) v = v.view(B, L, self.H, self.Dv) beta = torch.sigmoid(g.b_proj(x)) gate_g = gdn_gate(g.a_proj(x), g.A_log, g.dt_bias) return (q, k, v, beta, gate_g, (sq, sk, sv)) def _output(self, o, x): g = self.gdn B, L = (x.shape[0], x.shape[1]) gate = g.g_proj(x).view(B, L, self.H, self.Dv) o = gated_rmsnorm(o, gate, g.o_norm.weight, self.eps) o = o.reshape(B, L, self.C) return g.o_proj(o) def init_state(self, batch_size: int = 1, device=None): if device is None: device = self.gdn.q_proj.weight.device return { "recurrent_state": torch.zeros( batch_size, self.H, self.Dk, self.Dv, device=device ), "conv": tuple( ( torch.zeros(batch_size, self.C, self.W - 1, device=device) for _ in range(3) ) ), } def _recurrence(self, q, k, v, beta, g, initial_state): if q.shape[1] >= self.chunk_size: return gdn_chunk_recurrent( q, k, v, beta, g, self.scale, initial_state=initial_state, chunk_size=self.chunk_size, ) return gdn_recurrent(q, k, v, beta, g, self.scale, initial_state=initial_state) def forward(self, x: torch.Tensor) -> torch.Tensor: q, k, v, beta, gate_g, _ = self._project(x, None) o, _ = self._recurrence(q, k, v, beta, gate_g, initial_state=None) return self._output(o, x).to(x.dtype) def forward_chunk(self, x: torch.Tensor, state): conv_state = None if state is None else state["conv"] rec = None if state is None else state["recurrent_state"] q, k, v, beta, gate_g, new_conv = self._project(x, conv_state) o, rec = self._recurrence(q, k, v, beta, gate_g, initial_state=rec) return ( self._output(o, x).to(x.dtype), {"recurrent_state": rec, "conv": new_conv}, ) 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)