Spaces:

Ex0bit
/

mLattice

Running

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

20a3a74 verified 28 days ago

8.73 kB

"""
Cell-parallel μLattice scan -- the same math, ordered by CELL instead of
by time.
The delta recurrence is sequential *within* a cell but completely independent
*across* cells, so a token-by-token loop pays T sequential steps when only
max_writes_per_cell steps are needed.  This module:
  1. computes the 3x3 neighbourhood (published bump weights phi) for every token;
  2. keys each write by (sequence, z_x, z_y) and assigns it a RANK = how many
     earlier writes landed on the same cell;
  3. runs max_rank+1 rounds; in round r every cell receives exactly one write,
     so the round is one batched gather/matmul/scatter;
  4. schedules each read into the round where its cell reaches the required
     rank, accumulating its phi-weighted contribution.
Cell count is the number of distinct addresses (no dense grid, no cap); the
number of dispatch rounds is max writes per cell, not the sequence length.
Equivalence with mulattice.sparse_scan.SparseLatticeScan is asserted in
tests/test_kernel.py (max abs diff 8.9e-16 in float64, 16 configurations).
NOTE: the rounds counter is max writes per (sequence, cell), which is <= T; if you
see rounds > T the host's argsort/cummax is not honouring `stable=True` (observed
on one MPS run) and the result must not be trusted.
"""

from __future__ import annotations

import torch

from .kernel import MAX_CELLS_PER_TOKEN, phi_2d
from .sparse_scan import compact_ids as _compact_ids


def _encode(coords: torch.Tensor, n_seq: int, size: int = 1 << 20) -> torch.Tensor:
    """coords (B,T,9,2) int64 -> (B,T,9) collision-free int64 keys."""
    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 CellParallelScan(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_rounds = 0
        self.last_n_cells = 0

    def forward(
        self,
        q_addr: torch.Tensor,        # (B,T,2) query address (read)  -- same order
        k_addr: torch.Tensor,        # (B,T,2) key address  (write)  -- as SparseLatticeScan
        q2: torch.Tensor,            # (B,T,H,d_k)
        kappa: torch.Tensor,         # (B,T,H,d_k)
        v: torch.Tensor,             # (B,T,H,d_v)
        alpha: torch.Tensor,         # (B,T,H,1)
        beta: torch.Tensor,          # (B,T,H,1)
        read_mode: str = "all",      # "all" | "none"
    ):
        """For every token t: write into its 3x3 cells, then read them.
        Semantics are identical to the token-sequential reference; only the
        execution order changes.
        """
        B, T, H = kappa.shape[0], kappa.shape[1], self.n_heads
        dev, dt = kappa.device, kappa.dtype
        M = MAX_CELLS_PER_TOKEN

        wk, zk = phi_2d(k_addr.reshape(-1, 2))
        wq, zq = phi_2d(q_addr.reshape(-1, 2))
        wk = wk.reshape(B, T, M)
        wq = wq.reshape(B, T, M)
        zk = zk.reshape(B, T, M, 2)
        zq = zq.reshape(B, T, M, 2)

        keys_w = _encode(zk, B, self.size).reshape(-1)              # (B*T*M,)
        keys_q = _encode(zq, B, self.size).reshape(-1)
        inv, n_cells = _compact_ids(torch.cat([keys_w, keys_q]))
        self.last_n_cells = int(n_cells)
        inv_w = inv[: keys_w.numel()].reshape(B, T, M)
        inv_q = inv[keys_w.numel():].reshape(B, T, M)
        cells = torch.zeros(n_cells, H, self.d_v, self.d_k, device=dev, dtype=dt)

        # ---- write events, ordered by (cell, time) -------------------------
        cw = inv_w.reshape(-1)                                       # cell id
        tw = (torch.arange(B * T, device=dev).view(B, T, 1)
              .expand(B, T, M).reshape(-1) % T)                      # time within sequence
        order_w = torch.argsort(cw * (T + 1) + tw, stable=True)
        cw_s, tw_s = cw[order_w], tw[order_w]
        first = torch.ones_like(cw_s, dtype=torch.bool)
        first[1:] = cw_s[1:] != cw_s[:-1]
        ar = torch.arange(cw_s.numel(), device=dev)
        run_start = torch.cummax(torch.where(first, ar, torch.zeros_like(ar)), 0).values
        rank = ar - run_start
        max_rank = int(rank.max()) if rank.numel() else -1
        self.last_rounds = max_rank + 1
        # index of the first write of each cell in the sorted write array
        cell_start = torch.full((n_cells,), cw_s.numel(), device=dev, dtype=torch.int64)
        cell_start[cw_s[first]] = ar[first]

        row_w = (order_w // M)                                       # (b,t) row of the write
        kf = kappa.reshape(B * T, H, self.d_k)[row_w]
        vf = v.reshape(B * T, H, self.d_v)[row_w]
        af = alpha.reshape(B * T, H, 1)[row_w]
        bf = beta.reshape(B * T, H, 1)[row_w]
        # per-write bump weight, in the same (cell, time) order as the writes
        wphi = wk.reshape(-1)[order_w].unsqueeze(-1) if self.weighted_write else None

        # ---- read events: rank of the cell at the moment of the read --------
        acc = None
        if read_mode == "all":
            cq = inv_q.reshape(-1)
            tq = tw
            wkeys_sorted = cw * (T + 1) + tw
            wkeys_sorted = wkeys_sorted[order_w]                     # sorted write keys
            rkeys = cq * (T + 1) + tq
            # writes to THIS cell with key <= read key == writes to this cell with
            # time <= t.  Restricting to the cell's contiguous run is essential:
            # a plain searchsorted would also count writes to lower-id cells.
            need = torch.searchsorted(wkeys_sorted, rkeys, right=True) - cell_start[cq]
            need = need.clamp_(min=0)
            phi_q = wq.reshape(-1)
            row_q = torch.arange(B * T, device=dev).repeat_interleave(M)
            # cell index WITHIN a token (0..M-1).  Each (row, cell) pair is unique, so
            # accumulating into a (B*T, M, H, dv) tensor with index_put_ is deterministic
            # -- unlike index_add_ over duplicate rows, whose CUDA accumulation order is
            # unspecified and makes training non-reproducible across runs.
            cell_i = torch.arange(M, device=dev).repeat(B * T)
            acc = torch.zeros(B * T, M, H, self.d_v, device=dev, dtype=dt)
            qflat = q2.reshape(B * T, H, self.d_k)

        # ---- rounds --------------------------------------------------------
        for r in range(max_rank + 1):
            sel = rank == r
            idx = cw_s[sel]
            W = cells.index_select(0, idx)
            a, b, kap = af[sel], bf[sel], kf[sel]
            if wphi is not None:
                b = b * wphi[sel].unsqueeze(-1)
            v_bar = a * (W @ kap.unsqueeze(-1)).squeeze(-1)
            r_vec = vf[sel] - v_bar
            W = a.unsqueeze(-1) * W + b.unsqueeze(-1) * (
                r_vec.unsqueeze(-1) * kap.unsqueeze(-2))
            # FUNCTIONAL (out-of-place) update.  The in-place index_put_ chained across
            # rounds with interleaved index_select reads corrupted the autograd graph:
            # forward matched the token-loop to ~1e-16 but backward diverged by O(1),
            # which broke every fastscan-trained model (caught by
            # tests/test_kernel.py::test_fastscan_gradient_matches_sparse).  Out-of-place
            # index_put creates a proper graph node per round, so the sequential-recurrence
            # gradient flows correctly.  idx is unique within a round (one write per cell
            # per rank), so the scatter semantics are well defined.
            cells = cells.index_put((idx,), W)
            if read_mode == "all":
                take = need == (r + 1)          # state after exactly r+1 writes
                if take.any():
                    c_idx = cq[take]
                    rows = row_q[take]
                    ci = cell_i[take]
                    Wr = cells.index_select(0, c_idx)
                    out = (Wr @ qflat[rows].unsqueeze(-1)).squeeze(-1)
                    acc.index_put_((rows, ci), phi_q[take].unsqueeze(-1).unsqueeze(-1) * out)

        if read_mode == "all":
            # A read whose cells were never written contributes nothing, which is
            # correct -- but if NO read in the batch hits a written cell, `acc`
            # would stay a plain zeros tensor with no graph, and backward() would
            # fail.  This zero-valued term keeps the read path connected.
            acc = acc.sum(dim=1) + qflat.sum() * 0.0
            return acc.reshape(B, T, H, self.d_v), n_cells
        return None, n_cells

	"""
	Cell-parallel μLattice scan -- the same math, ordered by CELL instead of
	by time.

	The delta recurrence is sequential within a cell but completely independent
	across cells, so a token-by-token loop pays T sequential steps when only
	max_writes_per_cell steps are needed. This module:

	1. computes the 3x3 neighbourhood (published bump weights phi) for every token;
	2. keys each write by (sequence, z_x, z_y) and assigns it a RANK = how many
	earlier writes landed on the same cell;
	3. runs max_rank+1 rounds; in round r every cell receives exactly one write,
	so the round is one batched gather/matmul/scatter;
	4. schedules each read into the round where its cell reaches the required
	rank, accumulating its phi-weighted contribution.

	Cell count is the number of distinct addresses (no dense grid, no cap); the
	number of dispatch rounds is max writes per cell, not the sequence length.
	Equivalence with mulattice.sparse_scan.SparseLatticeScan is asserted in
	tests/test_kernel.py (max abs diff 8.9e-16 in float64, 16 configurations).
	NOTE: the rounds counter is max writes per (sequence, cell), which is <= T; if you
	see rounds > T the host's argsort/cummax is not honouring `stable=True` (observed
	on one MPS run) and the result must not be trusted.
	"""

	from __future__ import annotations

	import torch

	from .kernel import MAX_CELLS_PER_TOKEN, phi_2d
	from .sparse_scan import compact_ids as _compact_ids


	def _encode(coords: torch.Tensor, n_seq: int, size: int = 1 << 20) -> torch.Tensor:
	"""coords (B,T,9,2) int64 -> (B,T,9) collision-free int64 keys."""
	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 CellParallelScan(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_rounds = 0
	self.last_n_cells = 0

	def forward(
	self,
	q_addr: torch.Tensor, # (B,T,2) query address (read) -- same order
	k_addr: torch.Tensor, # (B,T,2) key address (write) -- as SparseLatticeScan
	q2: torch.Tensor, # (B,T,H,d_k)
	kappa: torch.Tensor, # (B,T,H,d_k)
	v: torch.Tensor, # (B,T,H,d_v)
	alpha: torch.Tensor, # (B,T,H,1)
	beta: torch.Tensor, # (B,T,H,1)
	read_mode: str = "all", # "all" \| "none"
	):
	"""For every token t: write into its 3x3 cells, then read them.

	Semantics are identical to the token-sequential reference; only the
	execution order changes.
	"""
	B, T, H = kappa.shape[0], kappa.shape[1], self.n_heads
	dev, dt = kappa.device, kappa.dtype
	M = MAX_CELLS_PER_TOKEN

	wk, zk = phi_2d(k_addr.reshape(-1, 2))
	wq, zq = phi_2d(q_addr.reshape(-1, 2))
	wk = wk.reshape(B, T, M)
	wq = wq.reshape(B, T, M)
	zk = zk.reshape(B, T, M, 2)
	zq = zq.reshape(B, T, M, 2)

	keys_w = _encode(zk, B, self.size).reshape(-1) # (BTM,)
	keys_q = _encode(zq, B, self.size).reshape(-1)
	inv, n_cells = _compact_ids(torch.cat([keys_w, keys_q]))
	self.last_n_cells = int(n_cells)
	inv_w = inv[: keys_w.numel()].reshape(B, T, M)
	inv_q = inv[keys_w.numel():].reshape(B, T, M)
	cells = torch.zeros(n_cells, H, self.d_v, self.d_k, device=dev, dtype=dt)

	# ---- write events, ordered by (cell, time) -------------------------
	cw = inv_w.reshape(-1) # cell id
	tw = (torch.arange(B * T, device=dev).view(B, T, 1)
	.expand(B, T, M).reshape(-1) % T) # time within sequence
	order_w = torch.argsort(cw * (T + 1) + tw, stable=True)
	cw_s, tw_s = cw[order_w], tw[order_w]
	first = torch.ones_like(cw_s, dtype=torch.bool)
	first[1:] = cw_s[1:] != cw_s[:-1]
	ar = torch.arange(cw_s.numel(), device=dev)
	run_start = torch.cummax(torch.where(first, ar, torch.zeros_like(ar)), 0).values
	rank = ar - run_start
	max_rank = int(rank.max()) if rank.numel() else -1
	self.last_rounds = max_rank + 1
	# index of the first write of each cell in the sorted write array
	cell_start = torch.full((n_cells,), cw_s.numel(), device=dev, dtype=torch.int64)
	cell_start[cw_s[first]] = ar[first]

	row_w = (order_w // M) # (b,t) row of the write
	kf = kappa.reshape(B * T, H, self.d_k)[row_w]
	vf = v.reshape(B * T, H, self.d_v)[row_w]
	af = alpha.reshape(B * T, H, 1)[row_w]
	bf = beta.reshape(B * T, H, 1)[row_w]
	# per-write bump weight, in the same (cell, time) order as the writes
	wphi = wk.reshape(-1)[order_w].unsqueeze(-1) if self.weighted_write else None

	# ---- read events: rank of the cell at the moment of the read --------
	acc = None
	if read_mode == "all":
	cq = inv_q.reshape(-1)
	tq = tw
	wkeys_sorted = cw * (T + 1) + tw
	wkeys_sorted = wkeys_sorted[order_w] # sorted write keys
	rkeys = cq * (T + 1) + tq
	# writes to THIS cell with key <= read key == writes to this cell with
	# time <= t. Restricting to the cell's contiguous run is essential:
	# a plain searchsorted would also count writes to lower-id cells.
	need = torch.searchsorted(wkeys_sorted, rkeys, right=True) - cell_start[cq]
	need = need.clamp_(min=0)
	phi_q = wq.reshape(-1)
	row_q = torch.arange(B * T, device=dev).repeat_interleave(M)
	# cell index WITHIN a token (0..M-1). Each (row, cell) pair is unique, so
	# accumulating into a (B*T, M, H, dv) tensor with index_put_ is deterministic
	# -- unlike index_add_ over duplicate rows, whose CUDA accumulation order is
	# unspecified and makes training non-reproducible across runs.
	cell_i = torch.arange(M, device=dev).repeat(B * T)
	acc = torch.zeros(B * T, M, H, self.d_v, device=dev, dtype=dt)
	qflat = q2.reshape(B * T, H, self.d_k)

	# ---- rounds --------------------------------------------------------
	for r in range(max_rank + 1):
	sel = rank == r
	idx = cw_s[sel]
	W = cells.index_select(0, idx)
	a, b, kap = af[sel], bf[sel], kf[sel]
	if wphi is not None:
	b = b * wphi[sel].unsqueeze(-1)
	v_bar = a * (W @ kap.unsqueeze(-1)).squeeze(-1)
	r_vec = vf[sel] - v_bar
	W = a.unsqueeze(-1) * W + b.unsqueeze(-1) * (
	r_vec.unsqueeze(-1) * kap.unsqueeze(-2))
	# FUNCTIONAL (out-of-place) update. The in-place index_put_ chained across
	# rounds with interleaved index_select reads corrupted the autograd graph:
	# forward matched the token-loop to ~1e-16 but backward diverged by O(1),
	# which broke every fastscan-trained model (caught by
	# tests/test_kernel.py::test_fastscan_gradient_matches_sparse). Out-of-place
	# index_put creates a proper graph node per round, so the sequential-recurrence
	# gradient flows correctly. idx is unique within a round (one write per cell
	# per rank), so the scatter semantics are well defined.
	cells = cells.index_put((idx,), W)
	if read_mode == "all":
	take = need == (r + 1) # state after exactly r+1 writes
	if take.any():
	c_idx = cq[take]
	rows = row_q[take]
	ci = cell_i[take]
	Wr = cells.index_select(0, c_idx)
	out = (Wr @ qflat[rows].unsqueeze(-1)).squeeze(-1)
	acc.index_put_((rows, ci), phi_q[take].unsqueeze(-1).unsqueeze(-1) * out)

	if read_mode == "all":
	# A read whose cells were never written contributes nothing, which is
	# correct -- but if NO read in the batch hits a written cell, `acc`
	# would stay a plain zeros tensor with no graph, and backward() would
	# fail. This zero-valued term keeps the read path connected.
	acc = acc.sum(dim=1) + qflat.sum() * 0.0
	return acc.reshape(B, T, H, self.d_v), n_cells
	return None, n_cells