Spaces:

Ex0bit
/

mLattice

Running

mLattice / src /mulattice /sparse_scan.py

μLattice: live demo — your own facts, store at scale

20a3a74 verified 28 days ago

6.26 kB

"""
Batched, differentiable μLattice lattice scan.
Same published math as mulattice.lattice.LatticeMemory, but the sequential
recurrence is executed with tensor ops over the whole batch instead of a Python
loop over (token x 9 cells), so the synthetic MQAR mixers can be trained.
Sparsity is preserved: cells are keyed by (sequence, z_x, z_y) via
torch.unique, and only the keys that actually occur are allocated.  No dense
grid over the address bounding box is ever materialised.  Cell count is
unbounded -- it is the number of distinct addresses written.
Semantics identical to the reference implementation:
    phi from the published bump kernel, exactly 9 cells per token,
    v_bar = alpha * W_c kappa ;  r = v - v_bar ;  W_c <- alpha W_c + beta r kappa^T
    read  = sum over the 3x3 neighbourhood of phi_c * (W_c q2), empty cells 0.
"""

from __future__ import annotations

import torch

from .kernel import MAX_CELLS_PER_TOKEN, phi_2d


def compact_ids(keys: torch.Tensor) -> tuple[torch.Tensor, int]:
    """Map arbitrary int64 keys to dense contiguous ids, via sort.
    Uses sort + run detection instead of torch.unique: on MPS, torch.unique corrupts int64
    values outside the int32 range (measured: 4959 distinct keys reported as 2789,
    with returned values like -30933022 instead of 549724880885).  That scrambled
    the cell ids, collapsed the lattice onto a few cells, and made training
    diverge to inf gradients on MPS while CPU was fine.  Sort + run detection is
    exact on every backend.
    """
    if keys.numel() == 0:
        return keys.clone(), 0
    order = torch.argsort(keys)
    ks = keys[order]
    first = torch.ones_like(ks, dtype=torch.bool)
    first[1:] = ks[1:] != ks[:-1]
    ids_sorted = torch.cumsum(first.to(torch.int64), 0) - 1
    ids = torch.empty_like(ids_sorted)
    ids[order] = ids_sorted
    return ids, int(first.sum())


def _encode_keys(coords: torch.Tensor, n_seq: int, size: int = 1 << 20) -> torch.Tensor:
    """coords (B,T,9,2) int64 -> flat int64 keys, collision-free for |z| < size/2."""
    z = coords + (size // 2)
    b = torch.arange(n_seq, device=coords.device, dtype=torch.int64).view(n_seq, 1, 1)
    return (b * (size * size)) + z[..., 0] * size + z[..., 1]


class SparseLatticeScan(torch.nn.Module):
    def __init__(self, d_k: int, d_v: int, n_heads: int, size: int = 1 << 20,
                 weighted_write: bool = True):
        super().__init__()
        self.d_k, self.d_v, self.n_heads, self.size = d_k, d_v, n_heads, size
        self.weighted_write = weighted_write
        self.last_n_cells = 0

    def forward(
        self,
        q_addr: torch.Tensor,     # (B,T,2) real address of the query
        k_addr: torch.Tensor,     # (B,T,2) real address of the key
        q2: torch.Tensor,         # (B,T,H,d_k) content query
        kappa: torch.Tensor,      # (B,T,H,d_k) content key
        v: torch.Tensor,          # (B,T,H,d_v) value
        alpha: torch.Tensor,      # (B,T,H,1) decay in (0,1)
        beta: torch.Tensor,       # (B,T,H,1) write gate in (0,1)
        read_mode: str = "last",  # "last" | "all" | "none"
        causal_read: bool = True,
    ):
        B, T, H = kappa.shape[0], kappa.shape[1], self.n_heads
        dev, dt = kappa.device, kappa.dtype

        wk, zk = phi_2d(k_addr)                                  # (B,T,3,3), (B,T,3,3,2)
        wq, zq = phi_2d(q_addr)
        wk = wk.reshape(B, T, MAX_CELLS_PER_TOKEN)
        wq = wq.reshape(B, T, MAX_CELLS_PER_TOKEN)
        zk = zk.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)
        zq = zq.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)

        need_q = read_mode != "none"
        keys_w = _encode_keys(zk, B, self.size).reshape(-1)
        keys_q = _encode_keys(zq, B, self.size).reshape(-1) if need_q else keys_w[:0]
        all_keys = torch.cat([keys_w, keys_q]) if need_q else keys_w
        inv, n_cells = compact_ids(all_keys)
        self.last_n_cells = int(n_cells)
        inv_w = inv[: keys_w.numel()].reshape(B, T, MAX_CELLS_PER_TOKEN)
        inv_q = (
            inv[keys_w.numel():].reshape(B, T, MAX_CELLS_PER_TOKEN) if need_q else None
        )

        # cells allocated lazily: only distinct keys that occur, zero-initialised
        cells = torch.zeros(n_cells, H, self.d_v, self.d_k, device=dev, dtype=dt)

        outs = []
        for t in range(T):
            idx = inv_w[:, t]                                     # (B,9)
            W = cells.index_select(0, idx.reshape(-1)).reshape(B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
            kap = kappa[:, t].unsqueeze(1)                        # (B,1,H,dk)
            a = alpha[:, t].unsqueeze(1)                          # (B,1,H,1)
            b = beta[:, t].unsqueeze(1)
            v_bar = a * (W @ kap.unsqueeze(-1)).squeeze(-1)       # (B,9,H,dv)
            r = v[:, t].unsqueeze(1) - v_bar
            gate = b * wk[:, t].unsqueeze(-1).unsqueeze(-1) if self.weighted_write else b
            W = a.unsqueeze(-1) * W + gate.unsqueeze(-1) * (r.unsqueeze(-1) * kap.unsqueeze(-2))
            cells.index_put_((idx.reshape(-1),), W.reshape(-1, H, self.d_v, self.d_k))

            if read_mode == "all":
                qi = inv_q[:, t] if causal_read else None
                if qi is None:                                    # non-causal: after full seq
                    outs.append(None)
                else:
                    Wq = cells.index_select(0, qi.reshape(-1)).reshape(
                        B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
                    rq = (Wq @ q2[:, t].unsqueeze(1).unsqueeze(-1)).squeeze(-1)   # (B,9,H,dv)
                    outs.append((wq[:, t].unsqueeze(-1).unsqueeze(-1) * rq).sum(1))

        if read_mode == "last":
            qi = inv_q[:, T - 1]                                  # (B,9)
            Wq = cells.index_select(0, qi.reshape(-1)).reshape(
                B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
            rq = (Wq @ q2[:, T - 1].unsqueeze(1).unsqueeze(-1)).squeeze(-1)
            out = (wq[:, T - 1].unsqueeze(-1).unsqueeze(-1) * rq).sum(1)   # (B,H,dv)
            return out, n_cells
        if read_mode == "all":
            return torch.stack(outs, dim=1), n_cells              # (B,T,H,dv)
        return None, n_cells

	"""
	Batched, differentiable μLattice lattice scan.

	Same published math as mulattice.lattice.LatticeMemory, but the sequential
	recurrence is executed with tensor ops over the whole batch instead of a Python
	loop over (token x 9 cells), so the synthetic MQAR mixers can be trained.

	Sparsity is preserved: cells are keyed by (sequence, z_x, z_y) via
	torch.unique, and only the keys that actually occur are allocated. No dense
	grid over the address bounding box is ever materialised. Cell count is
	unbounded -- it is the number of distinct addresses written.

	Semantics identical to the reference implementation:
	phi from the published bump kernel, exactly 9 cells per token,
	v_bar = alpha * W_c kappa ; r = v - v_bar ; W_c <- alpha W_c + beta r kappa^T
	read = sum over the 3x3 neighbourhood of phi_c * (W_c q2), empty cells 0.
	"""

	from __future__ import annotations

	import torch

	from .kernel import MAX_CELLS_PER_TOKEN, phi_2d


	def compact_ids(keys: torch.Tensor) -> tuple[torch.Tensor, int]:
	"""Map arbitrary int64 keys to dense contiguous ids, via sort.

	Uses sort + run detection instead of torch.unique: on MPS, torch.unique corrupts int64
	values outside the int32 range (measured: 4959 distinct keys reported as 2789,
	with returned values like -30933022 instead of 549724880885). That scrambled
	the cell ids, collapsed the lattice onto a few cells, and made training
	diverge to inf gradients on MPS while CPU was fine. Sort + run detection is
	exact on every backend.
	"""
	if keys.numel() == 0:
	return keys.clone(), 0
	order = torch.argsort(keys)
	ks = keys[order]
	first = torch.ones_like(ks, dtype=torch.bool)
	first[1:] = ks[1:] != ks[:-1]
	ids_sorted = torch.cumsum(first.to(torch.int64), 0) - 1
	ids = torch.empty_like(ids_sorted)
	ids[order] = ids_sorted
	return ids, int(first.sum())


	def _encode_keys(coords: torch.Tensor, n_seq: int, size: int = 1 << 20) -> torch.Tensor:
	"""coords (B,T,9,2) int64 -> flat int64 keys, collision-free for \|z\| < size/2."""
	z = coords + (size // 2)
	b = torch.arange(n_seq, device=coords.device, dtype=torch.int64).view(n_seq, 1, 1)
	return (b * (size * size)) + z[..., 0] * size + z[..., 1]


	class SparseLatticeScan(torch.nn.Module):
	def __init__(self, d_k: int, d_v: int, n_heads: int, size: int = 1 << 20,
	weighted_write: bool = True):
	super().__init__()
	self.d_k, self.d_v, self.n_heads, self.size = d_k, d_v, n_heads, size
	self.weighted_write = weighted_write
	self.last_n_cells = 0

	def forward(
	self,
	q_addr: torch.Tensor, # (B,T,2) real address of the query
	k_addr: torch.Tensor, # (B,T,2) real address of the key
	q2: torch.Tensor, # (B,T,H,d_k) content query
	kappa: torch.Tensor, # (B,T,H,d_k) content key
	v: torch.Tensor, # (B,T,H,d_v) value
	alpha: torch.Tensor, # (B,T,H,1) decay in (0,1)
	beta: torch.Tensor, # (B,T,H,1) write gate in (0,1)
	read_mode: str = "last", # "last" \| "all" \| "none"
	causal_read: bool = True,
	):
	B, T, H = kappa.shape[0], kappa.shape[1], self.n_heads
	dev, dt = kappa.device, kappa.dtype

	wk, zk = phi_2d(k_addr) # (B,T,3,3), (B,T,3,3,2)
	wq, zq = phi_2d(q_addr)
	wk = wk.reshape(B, T, MAX_CELLS_PER_TOKEN)
	wq = wq.reshape(B, T, MAX_CELLS_PER_TOKEN)
	zk = zk.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)
	zq = zq.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)

	need_q = read_mode != "none"
	keys_w = _encode_keys(zk, B, self.size).reshape(-1)
	keys_q = _encode_keys(zq, B, self.size).reshape(-1) if need_q else keys_w[:0]
	all_keys = torch.cat([keys_w, keys_q]) if need_q else keys_w
	inv, n_cells = compact_ids(all_keys)
	self.last_n_cells = int(n_cells)
	inv_w = inv[: keys_w.numel()].reshape(B, T, MAX_CELLS_PER_TOKEN)
	inv_q = (
	inv[keys_w.numel():].reshape(B, T, MAX_CELLS_PER_TOKEN) if need_q else None
	)

	# cells allocated lazily: only distinct keys that occur, zero-initialised
	cells = torch.zeros(n_cells, H, self.d_v, self.d_k, device=dev, dtype=dt)

	outs = []
	for t in range(T):
	idx = inv_w[:, t] # (B,9)
	W = cells.index_select(0, idx.reshape(-1)).reshape(B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
	kap = kappa[:, t].unsqueeze(1) # (B,1,H,dk)
	a = alpha[:, t].unsqueeze(1) # (B,1,H,1)
	b = beta[:, t].unsqueeze(1)
	v_bar = a * (W @ kap.unsqueeze(-1)).squeeze(-1) # (B,9,H,dv)
	r = v[:, t].unsqueeze(1) - v_bar
	gate = b * wk[:, t].unsqueeze(-1).unsqueeze(-1) if self.weighted_write else b
	W = a.unsqueeze(-1) * W + gate.unsqueeze(-1) * (r.unsqueeze(-1) * kap.unsqueeze(-2))
	cells.index_put_((idx.reshape(-1),), W.reshape(-1, H, self.d_v, self.d_k))

	if read_mode == "all":
	qi = inv_q[:, t] if causal_read else None
	if qi is None: # non-causal: after full seq
	outs.append(None)
	else:
	Wq = cells.index_select(0, qi.reshape(-1)).reshape(
	B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
	rq = (Wq @ q2[:, t].unsqueeze(1).unsqueeze(-1)).squeeze(-1) # (B,9,H,dv)
	outs.append((wq[:, t].unsqueeze(-1).unsqueeze(-1) * rq).sum(1))

	if read_mode == "last":
	qi = inv_q[:, T - 1] # (B,9)
	Wq = cells.index_select(0, qi.reshape(-1)).reshape(
	B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
	rq = (Wq @ q2[:, T - 1].unsqueeze(1).unsqueeze(-1)).squeeze(-1)
	out = (wq[:, T - 1].unsqueeze(-1).unsqueeze(-1) * rq).sum(1) # (B,H,dv)
	return out, n_cells
	if read_mode == "all":
	return torch.stack(outs, dim=1), n_cells # (B,T,H,dv)
	return None, n_cells