code/yunmo_next/reference_recurrent.py

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def causal_depthwise_conv1d(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    """Torch reference for FLA's causal depthwise short convolution."""

    if x.ndim != 3 or weight.ndim != 3 or weight.shape[1] != 1:
        raise ValueError("expected x=[B,T,D] and depthwise weight=[D,1,W]")
    sequence = x.shape[1]
    channels = x.shape[2]
    if weight.shape[0] != channels:
        raise ValueError("convolution channels do not match the hidden projection")
    convolved = F.conv1d(
        x.transpose(1, 2),
        weight,
        bias,
        padding=weight.shape[-1] - 1,
        groups=channels,
    )
    return convolved[:, :, :sequence].transpose(1, 2)


def l2_normalize_sum(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """Match FLA's L2 normalization (sum, not mean, over head dim)."""

    normalized = x.float() * torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True) + eps)
    return normalized.to(x.dtype)


def recurrent_kda_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    log_decay: torch.Tensor,
    beta: torch.Tensor,
    *,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Differentiable token-loop KDA oracle following the published delta rule."""

    if q.ndim != 4 or k.shape != q.shape or v.ndim != 4:
        raise ValueError("q/k/v must be [B,T,H,D] tensors")
    if v.shape[:3] != q.shape[:3] or log_decay.shape != q.shape:
        raise ValueError("KDA q/k/v/gate dimensions must match")
    if beta.shape != q.shape[:3]:
        raise ValueError("KDA beta must be [B,T,H]")
    batch, sequence, heads, key_dim = q.shape
    value_dim = v.shape[-1]
    state = q.new_zeros(batch, heads, key_dim, value_dim, dtype=torch.float32)
    if initial_state is not None:
        if initial_state.shape != state.shape:
            raise ValueError("KDA initial state has the wrong shape")
        state = state + initial_state.float()
    attention_scale = key_dim**-0.5 if scale is None else scale
    qf, kf, vf = q.float(), k.float(), v.float()
    gf, bf = log_decay.float(), beta.float()
    outputs: list[torch.Tensor] = []
    for token in range(sequence):
        key = kf[:, token]
        state = state * gf[:, token].exp().unsqueeze(-1)
        prediction = torch.einsum("bhk,bhkv->bhv", key, state)
        residual = vf[:, token] - prediction
        state = state + torch.einsum(
            "bhk,bhv->bhkv", bf[:, token].unsqueeze(-1) * key, residual
        )
        outputs.append(
            torch.einsum("bhk,bhkv->bhv", qf[:, token] * attention_scale, state)
        )
    return torch.stack(outputs, dim=1).to(v.dtype), state


class ReferenceRMSNormGated(nn.Module):
    def __init__(self, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
        normalized = x.float() * torch.rsqrt(
            x.float().square().mean(dim=-1, keepdim=True) + self.eps
        )
        return (
            normalized
            * self.weight.float()
            * torch.sigmoid(gate.float())
        ).to(x.dtype)


class ReferenceKimiDeltaAttention(nn.Module):
    """Correctness-first KDA layer for CPU/MPS probes.

    Parameter names and shapes intentionally match FLA's
    ``KimiDeltaAttention`` so a CUDA checkpoint can be loaded strictly. CUDA
    production must continue to use FLA; this token loop is an oracle, not an
    efficiency backend.
    """

    def __init__(
        self,
        *,
        hidden_size: int,
        head_dim: int,
        num_heads: int,
        expand_v: float = 1.0,
        use_short_conv: bool = True,
        conv_size: int = 4,
        conv_bias: bool = False,
        layer_idx: int | None = None,
        norm_eps: float = 1e-5,
        safe_gate: bool = False,
        lower_bound: float | None = None,
        **_: object,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.head_dim = head_dim
        self.num_heads = num_heads
        self.head_v_dim = int(head_dim * expand_v)
        self.key_dim = num_heads * head_dim
        self.value_dim = num_heads * self.head_v_dim
        self.gate_dim = num_heads * head_dim
        self.use_short_conv = use_short_conv
        self.layer_idx = layer_idx
        self.safe_gate = safe_gate
        self.lower_bound = lower_bound
        if safe_gate and (lower_bound is None or not -5.0 <= lower_bound < 0.0):
            raise ValueError("safe KDA gate requires lower_bound in [-5, 0)")
        if self.head_v_dim != head_dim:
            raise ValueError("reference KDA currently requires expand_v=1")

        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
        if use_short_conv:
            self.q_conv1d = nn.Conv1d(
                self.key_dim, self.key_dim, conv_size, groups=self.key_dim, bias=conv_bias
            )
            self.k_conv1d = nn.Conv1d(
                self.key_dim, self.key_dim, conv_size, groups=self.key_dim, bias=conv_bias
            )
            self.v_conv1d = nn.Conv1d(
                self.value_dim,
                self.value_dim,
                conv_size,
                groups=self.value_dim,
                bias=conv_bias,
            )
        self.f_proj = nn.Sequential(
            nn.Linear(hidden_size, self.head_v_dim, bias=False),
            nn.Linear(self.head_v_dim, self.gate_dim, bias=False),
        )
        self.b_proj = nn.Linear(hidden_size, num_heads, bias=False)
        self.A_log = nn.Parameter(
            torch.zeros(num_heads)
            if safe_gate
            else torch.log(torch.empty(num_heads).uniform_(1, 16))
        )
        self.A_log._no_weight_decay = True  # type: ignore[attr-defined]
        dt = torch.exp(
            torch.rand(self.gate_dim) * (math.log(0.1) - math.log(0.001))
            + math.log(0.001)
        ).clamp(min=1e-4)
        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
        self.dt_bias._no_weight_decay = True  # type: ignore[attr-defined]
        self.g_proj = nn.Sequential(
            nn.Linear(hidden_size, self.head_v_dim, bias=False),
            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
        )
        self.o_norm = ReferenceRMSNormGated(self.head_v_dim, norm_eps)
        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)

    def _project(self, projection: nn.Linear, convolution: nn.Conv1d, x: torch.Tensor) -> torch.Tensor:
        return F.silu(causal_depthwise_conv1d(projection(x), convolution.weight, convolution.bias))

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        past_key_values: object | None = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        cu_seqlens: torch.Tensor | None = None,
        **_: object,
    ) -> tuple[torch.Tensor, None, None]:
        if attention_mask is not None or past_key_values is not None or use_cache:
            raise NotImplementedError("reference KDA currently supports training without cache/padding")
        if output_attentions:
            raise NotImplementedError("KDA does not materialize attention weights")
        if cu_seqlens is not None:
            if hidden_states.shape[0] != 1:
                raise ValueError("packed reference KDA requires batch size one")
            boundaries = [int(value) for value in cu_seqlens.detach().cpu().tolist()]
            pieces = [
                self.forward(hidden_states[:, start:end])[0]
                for start, end in zip(boundaries[:-1], boundaries[1:], strict=True)
                if end > start
            ]
            return torch.cat(pieces, dim=1), None, None

        if self.use_short_conv:
            q = self._project(self.q_proj, self.q_conv1d, hidden_states)
            k = self._project(self.k_proj, self.k_conv1d, hidden_states)
            v = self._project(self.v_proj, self.v_conv1d, hidden_states)
        else:
            q, k, v = (
                F.silu(projection(hidden_states))
                for projection in (self.q_proj, self.k_proj, self.v_proj)
            )
        batch, sequence, _ = hidden_states.shape
        shape = (batch, sequence, self.num_heads, self.head_dim)
        q = l2_normalize_sum(q.view(shape))
        k = l2_normalize_sum(k.view(shape))
        v = v.view(batch, sequence, self.num_heads, self.head_v_dim)
        raw_gate = self.f_proj(hidden_states).view(shape)
        gate_input = raw_gate.float() + self.dt_bias.float().view(
            1, 1, self.num_heads, self.head_dim
        )
        if self.safe_gate:
            log_decay = self.lower_bound * torch.sigmoid(
                self.A_log.float().exp().view(1, 1, self.num_heads, 1)
                * gate_input
            )
        else:
            log_decay = -self.A_log.float().exp().view(
                1, 1, self.num_heads, 1
            ) * F.softplus(gate_input)
        beta = torch.sigmoid(self.b_proj(hidden_states))
        output, _ = recurrent_kda_reference(q, k, v, log_decay, beta)
        gate = self.g_proj(hidden_states).view(
            batch, sequence, self.num_heads, self.head_v_dim
        )
        output = self.o_norm(output, gate).reshape(batch, sequence, self.value_dim)
        return self.o_proj(output), None, None