cortex/growing_memory.py

"""Experimental attention: shared KV per 2-layer set (CLA-2) + a growing,
causal KV-memory that deeper sets read with their own (dense) queries.

Idea (user proposal): every set of 2 layers computes K,V once and shares it
(CLA-2 -> ~2x smaller kv-cache). Each set also writes a *causal* summary of its
values into a growing memory; every deeper set's per-layer queries attend over
BOTH their local shared KV AND the accumulated memory. Queries stay dense and
their reach "grows" with depth.

Causality: the memory a set writes is a causal cumulative mean of its values, so
the summary at position t only depends on positions <= t. Deeper queries attend
to these per-position summaries along the depth axis (G slots, one per earlier
set) -> no future leak. At inference the cumulative mean is an O(1) running state
per set, so the memory is constant-size (cheap).

This is a standalone trunk for clean A/B (use_mem False == plain CLA-2 baseline,
True == CLA-2 + growing memory). It deliberately omits AttnRes/DiffAttn/MTP to
isolate the new structure.
"""

from __future__ import annotations

import math

import jax
import jax.numpy as jnp
from flax import nnx

from .config import ModelConfig
from .model import compute_rope, apply_rope, causal_bias, _linear, _rmsnorm, MLP


def _causal_cummean(v):                       # v: [B, T, H, d] -> causal running mean over T
    T = v.shape[1]
    cs = jnp.cumsum(v, axis=1)
    denom = jnp.arange(1, T + 1, dtype=v.dtype).reshape(1, T, 1, 1)
    return cs / denom


class _Sublayer(nnx.Module):
    """KV-agnostic transformer sublayer: receives processed (k, v); own Q + MLP.
    Optionally also reads a growing memory via a depth-axis attention."""

    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs, use_mem: bool):
        self.cfg = cfg
        self.use_mem = use_mem
        self.window = cfg.sliding_window               # None => full causal (default)
        self.n_sink = cfg.n_sink                       # global attention sinks
        self.scale = 1.0 / math.sqrt(cfg.head_dim)
        resid = 1.0 / math.sqrt(2 * cfg.n_layers)
        self.norm1 = _rmsnorm(cfg.d_model, cfg, rngs)
        self.wq = _linear(cfg.d_model, cfg.q_dim, cfg, rngs)
        self.wo = _linear(cfg.q_dim, cfg.d_model, cfg, rngs, scale=resid)
        self.norm2 = _rmsnorm(cfg.d_model, cfg, rngs)
        self.mlp = MLP(cfg, rngs)
        if cfg.use_qk_norm:
            self.q_norm = _rmsnorm(cfg.head_dim, cfg, rngs)
        if use_mem:
            self.mem_gate = nnx.Param(jnp.zeros((), cfg.param_dtype))   # init 0 -> starts as plain CLA-2

    def __call__(self, h, rope, k, v, memory):
        cfg = self.cfg
        B, T, _ = h.shape
        cos, sin = rope
        q = self.wq(self.norm1(h)).reshape(B, T, cfg.n_q_heads, cfg.head_dim)
        if cfg.use_qk_norm:
            q = self.q_norm(q)
        q_r = apply_rope(q, cos, sin)

        scores = jnp.einsum("bthd,bshd->bhts", q_r, k).astype(jnp.float32) * self.scale
        scores = scores + causal_bias(T, self.window, self.n_sink)[None, None]
        out = jnp.einsum("bhts,bshd->bthd", jax.nn.softmax(scores, -1).astype(h.dtype), v)

        if self.use_mem and memory is not None:
            mem_k, mem_v = memory                       # [B, T, G, Hq, hd]
            ms = jnp.einsum("bthd,btghd->bthg", q, mem_k).astype(jnp.float32) * self.scale
            ma = jax.nn.softmax(ms, axis=-1).astype(h.dtype)
            out_mem = jnp.einsum("bthg,btghd->bthd", ma, mem_v)
            out = out + self.mem_gate.value.astype(h.dtype) * out_mem

        h = h + self.wo(out.reshape(B, T, cfg.q_dim))
        h = h + self.mlp(self.norm2(h))
        return h


class MemSet(nnx.Module):
    """An N-layer set (N=cfg.kv_share_group) sharing one KV (CLA-N), writes an optional causal memory."""

    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs, use_mem: bool):
        self.cfg = cfg
        self.use_mem = use_mem
        self.wk = _linear(cfg.d_model, cfg.kv_dim, cfg, rngs)
        self.wv = _linear(cfg.d_model, cfg.kv_dim, cfg, rngs)
        if cfg.use_qk_norm:
            self.k_norm = _rmsnorm(cfg.head_dim, cfg, rngs)
        self.subs = nnx.List([_Sublayer(cfg, rngs, use_mem) for _ in range(cfg.kv_share_group)])
        if use_mem:
            self.wmk = _linear(cfg.kv_dim, cfg.kv_dim, cfg, rngs)   # pooled values -> memory keys

    def __call__(self, h, rope, memory):
        cfg = self.cfg
        B, T, _ = h.shape
        cos, sin = rope
        rep = cfg.n_q_heads // cfg.n_kv_heads
        k_raw = self.wk(h).reshape(B, T, cfg.n_kv_heads, cfg.head_dim)
        v_raw = self.wv(h).reshape(B, T, cfg.n_kv_heads, cfg.head_dim)
        k = self.k_norm(k_raw) if cfg.use_qk_norm else k_raw
        k = jnp.repeat(apply_rope(k, cos, sin), rep, axis=2)        # [B,T,Hq,hd]
        v = jnp.repeat(v_raw, rep, axis=2)

        for sub in self.subs:
            h = sub(h, rope, k, v, memory)

        new_mem = None
        if self.use_mem:
            cmv = _causal_cummean(v_raw)                            # [B,T,Hkv,hd], causal
            mk = self.wmk(cmv.reshape(B, T, cfg.kv_dim)).reshape(B, T, cfg.n_kv_heads, cfg.head_dim)
            new_mem = (mk, cmv)
        return h, new_mem


class GrowingMemoryLM(nnx.Module):
    def __init__(self, cfg: ModelConfig, rngs: nnx.Rngs, use_mem: bool):
        g = cfg.kv_share_group
        assert cfg.n_layers % g == 0, f"n_layers ({cfg.n_layers}) must be divisible by kv_share_group ({g})"
        self.cfg = cfg
        self.use_mem = use_mem
        self.n_sets = cfg.n_layers // g
        self.embed = nnx.Embed(cfg.vocab_size, cfg.d_model,
                               embedding_init=nnx.initializers.normal(stddev=0.02),
                               dtype=cfg.compute_dtype, param_dtype=cfg.param_dtype, rngs=rngs)
        self.sets = nnx.List([MemSet(cfg, rngs, use_mem) for _ in range(self.n_sets)])
        self.final_norm = _rmsnorm(cfg.d_model, cfg, rngs)

    def __call__(self, tokens):
        cfg = self.cfg
        T = tokens.shape[1]
        x = self.embed(tokens)
        rope = compute_rope(T, cfg.head_dim, cfg.rope_base, cfg, x.dtype)
        rep = cfg.n_q_heads // cfg.n_kv_heads

        mem_list = []
        for mset in self.sets:
            memory = None
            if self.use_mem and mem_list:
                mk = jnp.stack([m[0] for m in mem_list], axis=2)   # [B,T,G,Hkv,hd]
                mv = jnp.stack([m[1] for m in mem_list], axis=2)
                memory = (jnp.repeat(mk, rep, axis=3), jnp.repeat(mv, rep, axis=3))  # [B,T,G,Hq,hd]
            x, new_mem = mset(x, rope, memory)
            if new_mem is not None:
                mem_list.append(new_mem)

        x = self.final_norm(x)
        emb = self.embed.embedding.value.astype(x.dtype)
        return jnp.einsum("btd,vd->btv", x, emb).astype(jnp.float32)