| """Exact PyTorch Diagonal KDN recurrences that carry state across calls. |
| |
| The Triton kernels cover stateless fixed-length training. These functions |
| take and return ``(memory, precision)`` so prefill, decode, and padded batches |
| continue a sequence exactly; the sequential form is also the test oracle. |
| |
| Per head, with one covariance coordinate ``p`` per key channel:: |
| |
| p_hat_t = alpha_t**2 * p_{t-1} + omega_t |
| kappa_t = p_hat_t * k_t / (r_t + sum_i p_hat_ti * k_ti**2) |
| 1 / p_t = 1 / p_hat_t + info_scale * k_t**2 / r_t |
| S_t = (I - kappa_t k_t^T) diag(alpha_t) S_{t-1} + kappa_t v_t^T |
| o_t = scale * S_t^T q_t |
| |
| Shapes: ``q, k, log_alpha, omega`` are ``[B, T, H, K]``, ``v`` is |
| ``[B, T, H, V]``, ``r`` is ``[B, T, H]``, ``memory`` is ``[B, H, K, V]`` and |
| ``precision`` (``1 / p``) is ``[B, H, K]``. All math runs in float32. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| OMEGA_MIN = 0.0 |
| R_MIN = 0.01 |
| INITIAL_PRECISION = 1.0 |
| INITIAL_PRECISION_FLOOR = 0.1 |
| LOG_MEMORY_ALPHA_FLOOR = math.log(1e-6) |
| NORM_EPS = 1e-12 |
| DEFAULT_CHUNK_SIZE = 64 |
|
|
|
|
| def initial_precision_logit() -> float: |
| """Raw parameter value whose floored softplus equals ``INITIAL_PRECISION``.""" |
|
|
| return math.log(math.expm1(INITIAL_PRECISION - INITIAL_PRECISION_FLOOR)) |
|
|
|
|
| def _work(tensor: torch.Tensor) -> torch.Tensor: |
| return tensor.to(torch.promote_types(tensor.dtype, torch.float32)) |
|
|
|
|
| def kalman_noise( |
| retention_raw: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| A_log: torch.Tensor, |
| dt_bias: torch.Tensor, |
| initial_precision_param: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Return ``(log_alpha, omega, r, initial_precision)`` like the Triton frontend.""" |
|
|
| heads, key_dim = retention_raw.shape[-2:] |
| rate = _work(A_log).exp().view(heads, 1) |
| log_alpha = -rate * F.softplus( |
| _work(retention_raw) + _work(dt_bias).view(heads, key_dim) |
| ) |
| omega = OMEGA_MIN + F.softplus(_work(omega_raw)) |
| r = R_MIN + F.softplus(_work(r_raw)) |
| initial_precision = INITIAL_PRECISION_FLOOR + F.softplus( |
| _work(initial_precision_param) |
| ) |
| return log_alpha, omega, r, initial_precision |
|
|
|
|
| def normalize_query_key( |
| q: torch.Tensor, k: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| return ( |
| F.normalize(_work(q), dim=-1, eps=NORM_EPS), |
| F.normalize(_work(k), dim=-1, eps=NORM_EPS), |
| ) |
|
|
|
|
| def initial_state( |
| initial_precision: torch.Tensor, |
| *, |
| batch_size: int, |
| key_dim: int, |
| value_dim: int, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Zero memory and the learned per-head prior precision for new sequences.""" |
|
|
| heads = initial_precision.shape[0] |
| prior = _work(initial_precision) |
| memory = prior.new_zeros((batch_size, heads, key_dim, value_dim)) |
| precision = prior.view(1, heads, 1).expand(batch_size, heads, key_dim).contiguous() |
| return memory, precision |
|
|
|
|
| def mask_padding( |
| token_mask: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| log_alpha: torch.Tensor, |
| omega: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Turn padded positions into identity steps that leave the state untouched.""" |
|
|
| keep = token_mask[:, :, None, None] |
| return ( |
| torch.where(keep, k, 0.0), |
| torch.where(keep, v, 0.0), |
| torch.where(keep, log_alpha, 0.0), |
| torch.where(keep, omega, 0.0), |
| ) |
|
|
|
|
| def diag_kdn_step( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| log_alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| memory: torch.Tensor, |
| precision: torch.Tensor, |
| *, |
| scale: float, |
| info_scale: float, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Advance one token: ``q, k, log_alpha, omega`` are ``[B, H, K]``.""" |
|
|
| k_squared = k.square() |
| predicted = (2.0 * log_alpha).exp() / precision + omega |
| noise = r[..., None] |
| gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True)) |
| precision = predicted.reciprocal() + info_scale * k_squared / noise |
| decayed = memory * log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR).exp()[..., None] |
| innovation = v - torch.einsum("bhk,bhkv->bhv", k, decayed) |
| memory = decayed + gain[..., None] * innovation[..., None, :] |
| output = scale * torch.einsum("bhk,bhkv->bhv", q, memory) |
| return output, memory, precision |
|
|
|
|
| def diag_kdn_recurrent( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| log_alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| memory: torch.Tensor, |
| precision: torch.Tensor, |
| *, |
| scale: float, |
| info_scale: float, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Token-by-token reference; returns ``(output [B,T,H,V], memory, precision)``.""" |
|
|
| outputs = [] |
| for token in range(q.shape[1]): |
| output, memory, precision = diag_kdn_step( |
| q[:, token], |
| k[:, token], |
| v[:, token], |
| log_alpha[:, token], |
| omega[:, token], |
| r[:, token], |
| memory, |
| precision, |
| scale=scale, |
| info_scale=info_scale, |
| ) |
| outputs.append(output) |
| return torch.stack(outputs, dim=1), memory, precision |
|
|
|
|
| def _compose( |
| later: tuple[torch.Tensor, ...], earlier: tuple[torch.Tensor, ...] |
| ) -> tuple[torch.Tensor, ...]: |
| """Product ``later @ earlier`` of 2x2 maps stored as (m00, m01, m10, m11).""" |
|
|
| a00, a01, a10, a11 = later |
| b00, b01, b10, b11 = earlier |
| product = ( |
| a00 * b00 + a01 * b10, |
| a00 * b01 + a01 * b11, |
| a10 * b00 + a11 * b10, |
| a10 * b01 + a11 * b11, |
| ) |
| |
| magnitude = torch.stack(product).abs().amax(dim=0).clamp_min( |
| torch.finfo(product[0].dtype).tiny |
| ) |
| return tuple(entry / magnitude for entry in product) |
|
|
|
|
| def _inclusive_scan(maps: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: |
| """Hillis-Steele prefix products ``M_t @ ... @ M_1`` along dim 1.""" |
|
|
| length = maps[0].shape[1] |
| offset = 1 |
| while offset < length: |
| combined = _compose( |
| tuple(entry[:, offset:] for entry in maps), |
| tuple(entry[:, :-offset] for entry in maps), |
| ) |
| maps = tuple( |
| torch.cat((entry[:, :offset], update), dim=1) |
| for entry, update in zip(maps, combined) |
| ) |
| offset *= 2 |
| return maps |
|
|
|
|
| def kalman_gain( |
| k: torch.Tensor, |
| log_alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| precision: torch.Tensor, |
| *, |
| info_scale: float, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Vector gains for every token via a parallel Mobius scan, plus final precision.""" |
|
|
| k_squared = k.square() |
| noise = r[..., None] |
| alpha_squared = (2.0 * log_alpha).exp() |
| information = info_scale * k_squared / noise |
| maps = _inclusive_scan( |
| ( |
| alpha_squared, |
| omega, |
| information * alpha_squared, |
| 1.0 + information * omega, |
| ) |
| ) |
| start = precision.reciprocal()[:, None] |
| covariance = (maps[0] * start + maps[1]) / (maps[2] * start + maps[3]) |
| previous = torch.cat((start, covariance[:, :-1]), dim=1) |
| predicted = alpha_squared * previous + omega |
| gain = predicted * k / (noise + (predicted * k_squared).sum(-1, keepdim=True)) |
| return gain, covariance[:, -1].reciprocal() |
|
|
|
|
| def chunked_memory( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| log_alpha: torch.Tensor, |
| gain: torch.Tensor, |
| memory: torch.Tensor, |
| *, |
| scale: float, |
| chunk_size: int = DEFAULT_CHUNK_SIZE, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Two-key delta-rule memory in WY form, one triangular solve per chunk.""" |
|
|
| if chunk_size <= 0: |
| raise ValueError("chunk_size must be positive") |
| log_alpha = log_alpha.clamp_min(LOG_MEMORY_ALPHA_FLOOR) |
| length = q.shape[1] |
| outputs = [] |
| for start in range(0, length, chunk_size): |
| stop = min(start + chunk_size, length) |
| size = stop - start |
| q_chunk, k_chunk, v_chunk, gain_chunk = ( |
| tensor[:, start:stop].transpose(1, 2) for tensor in (q, k, v, gain) |
| ) |
| cumulative = log_alpha[:, start:stop].transpose(1, 2).cumsum(dim=2) |
| causal = torch.ones(size, size, dtype=torch.bool, device=q.device).tril() |
| differences = cumulative[:, :, :, None, :] - cumulative[:, :, None, :, :] |
| weights = ( |
| torch.where(causal[:, :, None], differences, -math.inf).exp() |
| * gain_chunk[:, :, None, :, :] |
| ) |
| interaction = torch.einsum("bhtk,bhtsk->bhts", k_chunk, weights) |
| readout = torch.einsum("bhtk,bhtsk->bhts", q_chunk, weights) |
| decay = cumulative.exp() |
| corrected = torch.linalg.solve_triangular( |
| interaction, |
| v_chunk - (decay * k_chunk) @ memory, |
| upper=False, |
| unitriangular=True, |
| ) |
| outputs.append( |
| (scale * ((q_chunk * decay) @ memory + readout @ corrected)).transpose(1, 2) |
| ) |
| to_end = (cumulative[:, :, -1:] - cumulative).exp() * gain_chunk |
| memory = memory * decay[:, :, -1, :, None] + to_end.transpose(-1, -2) @ corrected |
| return torch.cat(outputs, dim=1), memory |
|
|
|
|
| def diag_kdn_chunked( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| log_alpha: torch.Tensor, |
| omega: torch.Tensor, |
| r: torch.Tensor, |
| memory: torch.Tensor, |
| precision: torch.Tensor, |
| *, |
| scale: float, |
| info_scale: float, |
| chunk_size: int = DEFAULT_CHUNK_SIZE, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Parallel form of :func:`diag_kdn_recurrent` with the same signature.""" |
|
|
| gain, precision = kalman_gain( |
| k, log_alpha, omega, r, precision, info_scale=info_scale |
| ) |
| output, memory = chunked_memory( |
| q, k, v, log_alpha, gain, memory, scale=scale, chunk_size=chunk_size |
| ) |
| return output, memory, precision |
|
|
