"""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)