cpu_delta.py

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)