0.8.0 — a state holds only itself, and a step can give it up

README.md

@@ -224,7 +224,7 @@ suite checks against; they are deliberately not exported.
 
 ## Status
 
-`0.7.0` — alpha, and versioned honestly. The *interface* is the thing this
+`0.8.0` — alpha, and versioned honestly. The *interface* is the thing this
 library exists to protect, so it changes slowly and visibly; a merged component
 is required to reproduce the outputs of what it replaces, to a stated tolerance,
 before anything switches to it. Internals are fair game.

docs/BLOCK.md

@@ -111,6 +111,19 @@ y, state = trunk.step(x, state)                      # x is (256, 1, d_model)
 sharing a buffer, which matters wherever a prefix is continued down more than one
 future.
 
+A caller that never forks can waive that:
+
+```python
+y, state = trunk.step(x, state, donate=True)         # the old state is spent
+```
+
+`donate=True` is passed to every stateful sub-layer, and each may write its
+successor into the buffers it was given. The donated state must not be read
+again. Sparse Delta Memory uses it to skip copying its table per stream per
+step, which at a large batch is most of the step (see [SDM.md](SDM.md)). The other
+mixers accept it and ignore it. A block passes it on only when it is set, so a
+stateful sub-layer of your own that does not take the keyword keeps working.
+
 ---
 
 ## Initialisation

docs/SDM.md

@@ -2,9 +2,11 @@
 
 A gated delta rule over a large, sparsely addressed table. Each head holds a
 table of `N` slots; each position writes `W` of them and reads `R`, chosen by
-product keys. Per-token work is `O((W+R)·d_v)` whatever `N` is, and a slot
-nobody writes to is left exactly as it was — decay included — so **state size
-is decoupled from parameter count and from compute.**
+product keys. Per-token arithmetic is `O((W+R)·d_v)` whatever `N` is, and a
+slot nobody writes to is left exactly as it was — decay included — so **state
+size is decoupled from parameter count and from compute.** One cost does grow
+with `N`: by default `step` returns a fresh table per stream, which at a large
+batch is most of a decode step. [Generation](#generation) says how to waive it.
 
 The design rationale — the diagonal-decay derivation, which defaults are earned
 and which are not, what is deliberately excluded — is in
@@ -94,10 +96,32 @@ through its state carries a gradient back to the table.
 **`step` returns a successor and never mutates the state it was given**, the
 same guarantee every mixer here makes, so branching a stream cannot leave two
 branches sharing a buffer. For this layer that guarantee is not free: the
-successor is a full table, `O(n_slots · d_v)` per head per step against
-`O((W+R) · d_v)` of actual work. At small tables it does not matter; at very
-large ones it dominates decode. An in-place decode path is an open question,
-not a feature (design record §3.6).
+successor is a full table per stream, `O(B · n_slots · d_v)` per head per step
+against `O(B · (W+R) · d_v)` of actual work. The batch matters as much as the
+table. Measured on one machine at `d_model = 512`, a step at `B = 1` takes about
+1.2 ms at every table size, because it is launch-bound. At `B = 128` it rises
+from 4.1 ms at 32² slots to 37.6 ms at 128².
+
+A caller that will not read a state again can say so:
+
+```python
+for _ in range(max_new_tokens):
+    y, state = mixer.step(embed(token), state, donate=True)
+    token = sample(y[:, -1])
+```
+
+`donate=True` lets the layer write the successor into the table it was given.
+On the same machine, at `B = 128` and 128² slots, that is 2.3 ms against 37.6.
+The outputs are the same bits either way. The donated state then shares its
+buffer with the successor, so do not read it or step it again, and do not donate
+a state you have forked.
+
+Donating is a permission, not a demand. The first donated step after
+`init_state` still copies, because the initial table is a shared view, and
+under `"learned"` a view of the parameter. A step under autograd or a
+`torch.func` transform also copies. `Block.step` and `Stack.step` take the same
+keyword and pass it to every sub-layer. Design record §3.6 has the whole
+measurement.
 
 ## Reading without writing
 

docs/design/SPARSE_DELTA_MEMORY.md

@@ -4,8 +4,9 @@
 place (§3.11) and partners are found by search (§3.1); in 0.6.2 only the
 table-touching work stays in the chunk loop, and the pairwise terms can be
 recomputed (§3.12); in 0.7.0 an opt-in Triton backend computes the same
-function as kernels (§3.13). The implementation is expected to match this
-record; where the two disagree, one of them is a bug.
+function as kernels (§3.13); in 0.8.0 a decode step given `donate=True`
+writes the table in place (§3.6). The implementation is expected to match
+this record; where the two disagree, one of them is a bug.
 
 Unlike Gated DeltaNet and Undertow, this is **not a consolidation.** No
 existing implementation was merged, so there is no port to verify and no
@@ -304,12 +305,58 @@ Decisions:
 successor rather than mutating — the house guarantee that branching a stream
 cannot leave two branches sharing a buffer.
 
-For GDN that guarantee is free. **Here it costs a table copy per step:**
-`O(N·d_v)` per head per position, against `O((W+R)·d_v)` of actual work. For
-small tables that is tens of MB per step and acceptable; at `N = 2¹⁸,
-d_v = 512` it is half a gigabyte per step per layer. The reference keeps the
-guarantee. An in-place decode path — explicit about aliasing, opt-in — is a
-later decision with its own measurement, not a refactor of this one.
+For GDN that guarantee is free. **Here it costs a table copy per stream per
+step:** `O(B·N·d_v)` per head per position, against `O(B·(W+R)·d_v)` of actual
+work. It is easy to read that as a cost of large tables, and it is not only
+that: the copy scales with the batch exactly as it scales with `N`. At `B = 1` a
+step is launch-bound and flat in `N`, which hides it. Measured on one machine (a
+Pascal card, fp32, one layer at `d_model = 512`, 2 heads, `W = R = 64`), against
+the same arithmetic written into the table in place:
+
+| `N` | `B = 1` | `B = 32` | `B = 128` |
+|---|---|---|---|
+| 32² | 1.22 → 1.20 ms | 1.34 → 1.16 ms | 4.10 → 1.89 ms |
+| 64² | 1.20 → 1.18 ms | 3.06 → 1.14 ms | 11.07 → 2.22 ms |
+| 128² | 1.21 → 1.18 ms | 9.70 → 1.15 ms | 37.62 → 2.27 ms |
+
+In a five-block trunk with two such layers at 64² or 128², the copy was
+nothing at `B = 8`, and 61% and 86% of the whole step at `B = 128`. At
+`N = 512²` it is most of a step already at `B = 8`: 36 ms against 0.2.
+
+**So the guarantee is the default, and the copy can be waived:**
+`step(x, state, donate=True)`. Donating a state hands its buffers to the layer,
+which may write the successor into the table it was given. The caller must not
+read a donated state again; the successor is the stream. A caller who does not
+opt in keeps the functional path unchanged. The two paths are one function,
+`recurrent_sparse_delta`, with `scatter` or `scatter_` as its last write, so
+they agree bit for bit, and a test holds it.
+
+Why a keyword and not a separate `step_`. The path that needs it is a trunk
+stepping many streams, so whatever is chosen has to pass through `Block.step`
+and `Stack.step`. A method on this layer alone would not reach that path, and a
+fourth method on every component would break the surface they share. Gated
+DeltaNet and Undertow accept the keyword and ignore it, because their successors
+cost no more than their arithmetic. `Block` forwards it only when it is set, so a
+stateful sub-layer from outside the library that predates it keeps stepping.
+
+Why not copy-on-write at the granularity of pages of rows. Product keys spread a
+position's writes across the table on purpose. With 64-row pages, 64 writes touch
+essentially every page of a 32² table and up to a quarter of a 128² one. So
+paging saves little at the sizes where the copy bites, and every kernel would
+then have to read a paged table.
+
+**Donating is a permission, not a demand.** The layer copies anyway when the
+table is not its own to write:
+
+- any view. This includes the broadcast that `init_state` hands out, which
+  under `"learned"` shares the parameter's storage. At `B = 1` that view is not
+  even overlapping, so an in-place write would silently change the learned
+  table;
+- a table autograd is recording;
+- a table under a `torch.func` transform.
+
+The first donated step from `init_state` is therefore a copy, the same one the
+successor rule makes anyway, and every step after it is in place.
 
 `init_state` hands out the initial table as a **broadcast view**, `(B, H, N,
 d_v)` over `(H, N, d_v)`: no stream costs a copy until its first write, which is
@@ -896,8 +943,3 @@ Each of these is a standing test in `tests/test_sdm.py`.
   kept), and, for one or two sequences, more parallelism than splitting value
   columns gives. Unmeasured on any card with tensor cores, or with more shared
   memory per block than 48 KB -- where the cap on chunk length may not bind.
-- **An in-place decode** (§3.6). Measured on one machine, the successor copy is
-  most of a step once the table is large: at `B = 8`, 36 ms per step at
-  `N = 512²` against 0.2 ms for the same arithmetic in place, which also costs
-  0.2 ms at every smaller table. Whether that buys an opt-in `step_` against the
-  house's successor guarantee is an API decision, not an optimisation.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "lumen"
-version = "0.7.0"
+version = "0.8.0"
 description = "One place for the components worth keeping: consolidated architectures, measured hardware capability, no assumptions."
 readme = "README.md"
 license = { text = "MIT" }

src/lumen/__init__.py

@@ -73,7 +73,7 @@
     probe_triton,
 )
 
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 
 __all__ = [
     "__version__",

src/lumen/block.py

@@ -255,25 +255,35 @@ def init_state(
         )
 
     def step(
-        self, x: torch.Tensor, state: BlockState
+        self, x: torch.Tensor, state: BlockState, *, donate: bool = False
     ) -> tuple[torch.Tensor, BlockState]:
-        """One position — ``(B, 1, d_model)`` → output and the successor state."""
+        """One position — ``(B, 1, d_model)`` → output and the successor state.
+
+        ``donate=True`` says the caller will not read ``state`` again, and is
+        passed to every stateful sub-layer so each may reuse its buffers.
+        It is passed only when set: a stateful sub-layer from outside this
+        library that does not take the keyword steps as it always did, and
+        fails loudly only when someone asks it to donate.
+        """
         if x.shape[1] != 1:
             raise ValueError(
                 f"step() consumes one position at a time, got {x.shape[1]}; "
                 f"use forward(x, state=..., return_state=True) for a chunk"
             )
 
+        donation = {"donate": True} if donate else {}
         s_local = None
         if self.local is not None:
             if _is_stateful(self.local):
-                y, s_local = self.local.step(self.norm_local(x), state.local)
+                y, s_local = self.local.step(
+                    self.norm_local(x), state.local, **donation
+                )
             else:
                 y = self.local(self.norm_local(x))
             x = x + y
 
         if _is_stateful(self.mixer):
-            y, s_mixer = self.mixer.step(self.norm(x), state.mixer)
+            y, s_mixer = self.mixer.step(self.norm(x), state.mixer, **donation)
         else:
             y, s_mixer = self.mixer(self.norm(x)), None
         x = x + y

src/lumen/gdn/layer.py

@@ -772,9 +772,14 @@ def init_state(
         return GatedDeltaNetState(memory=memory, conv=conv)
 
     def step(
-        self, x: torch.Tensor, state: GatedDeltaNetState
+        self, x: torch.Tensor, state: GatedDeltaNetState, *, donate: bool = False
     ) -> tuple[torch.Tensor, GatedDeltaNetState]:
-        """One position — `(B, 1, d_model)` → output and the successor state."""
+        """One position — `(B, 1, d_model)` → output and the successor state.
+
+        ``donate`` is accepted so a trunk can pass it to every sub-layer, and
+        ignored: the successor's memory is this step's own arithmetic, so there
+        is no copy to save.
+        """
         if x.shape[1] != 1:
             raise ValueError(
                 f"step() consumes one position at a time, got {x.shape[1]}; "

src/lumen/sdm/__init__.py

@@ -1,9 +1,11 @@
 """Sparse Delta Memory — a gated delta rule over a large, sparsely addressed table.
 
 Per head, a table of `N` slots.  Each position writes `W` of them and reads `R`,
-chosen by product keys, so per-token work is `O((W+R)·d_v)` whatever `N` is,
-and a slot nobody writes to is frozen — decay included.  State size is
-decoupled from parameter count and from compute.
+chosen by product keys, so per-token arithmetic is `O((W+R)·d_v)` whatever
+`N` is, and a slot nobody writes to is frozen — decay included.  State size is
+decoupled from parameter count and from compute — except that ``step`` copies
+the table per stream unless its state is donated; see
+:class:`SparseDeltaMemoryState`.
 
 Public surface is the layer, its config and its state::
 

src/lumen/sdm/layer.py

@@ -6,10 +6,12 @@
 
 where `k_t` and `q_t` are nonzero only on the `W` slots position `t` writes and
 the `R` it reads, chosen by product keys, and `Λ_t` decays the written slots and
-nothing else.  Per-token work is `O((W+R)·d_v)` whatever `N` is, and the only
-parameters that grow with `N` are the two address projections, as `√N`.  **State
-size is decoupled from parameter count and from compute** — which is the reason
-to have this layer.
+nothing else.  Per-token arithmetic is `O((W+R)·d_v)` whatever `N` is, and the
+only parameters that grow with `N` are the two address projections, as `√N`.
+**State size is decoupled from parameter count and from compute** — which is the
+reason to have this layer.  One cost is not: :meth:`SparseDeltaMemory.step`
+returns a new table per stream, `O(B·N·d_v)` a step, unless its state is
+donated — see :class:`SparseDeltaMemoryState`.
 
 The design record is drafted ahead of promotion.  Points from it worth
 repeating where the code lives:
@@ -50,6 +52,7 @@
 from lumen.nn import rms_norm
 from lumen.sdm import triton_kernels
 from lumen.sdm.reference import (
+    _transforms_active,
     chunk_sparse_delta,
     read_sparse_delta,
     recurrent_sparse_delta,
@@ -73,14 +76,37 @@ class SparseDeltaMemoryState:
     Frozen — :meth:`SparseDeltaMemory.step` returns a successor rather than
     mutating in place, so branching a stream cannot leave two branches quietly
     sharing a buffer.  **Here that guarantee has a price** that it does not have
-    for Gated DeltaNet: the successor is a full table, `O(N·d_v)` per head per
-    step against `O((W+R)·d_v)` of actual work.  Paid on purpose; an in-place
-    decode is a separate, measured decision.
+    for Gated DeltaNet: the successor is a full table per stream, `O(B·N·d_v)`
+    per head per step against `O(B·(W+R)·d_v)` of actual work.  At `B = 1` the
+    step is launch-bound and the copy hides; at a large batch it is most of the
+    step, at tables nobody would call large.
+
+    ``step(x, state, donate=True)`` waives it: the successor is written into
+    the donated table, and the donated state must not be read again.  The
+    default keeps the guarantee.  Design record §3.6 has the measurements.
     """
 
     memory: torch.Tensor
 
 
+def _writable_in_place(memory: torch.Tensor) -> bool:
+    """May a donated table be written where it stands?
+
+    Donation says the caller is done with the state, not that its table is the
+    stream's own.  It is not when it is a view: ``init_state`` hands out a
+    broadcast, and under ``"learned"`` that broadcast shares the parameter's
+    storage -- at `B = 1` it does not even overlap, so a write would go through
+    and change the learned table.  Nor under a ``torch.func`` transform, or
+    while autograd records the table, where an in-place write is refused or
+    breaks the gradient.  Each of those steps by copy instead.
+    """
+    if _transforms_active():
+        return False
+    if torch.is_grad_enabled() and memory.requires_grad:
+        return False
+    return memory._base is None
+
+
 @dataclass(frozen=True)
 class SparseDeltaMemoryConfig:
     """Configuration for :class:`SparseDeltaMemory`, validated on construction.
@@ -91,8 +117,10 @@ class SparseDeltaMemoryConfig:
                   ``d_v = d_model / n_heads``.
         n_slots:  `N`, slots **per head**, a perfect square — product keys
                   address it as `√N × √N`.  The knob that grows the state
-                  without growing compute: per-token work does not depend on
-                  it, and the address projections grow as `√N`.
+                  without growing compute: per-token arithmetic does not
+                  depend on it, and the address projections grow as `√N`.
+                  Decode's successor copy does, `O(B·N·d_v)` a step, unless
+                  the state is donated — see :class:`SparseDeltaMemoryState`.
         initial_memory: ``"zero"`` or ``"learned"``, and **required** — see the
                   module docstring.  ``"learned"`` adds `H · N · d_v`
                   parameters, reported like any others; the paper this layer
@@ -468,16 +496,29 @@ def init_state(
         return SparseDeltaMemoryState(memory=self._initial_table(batch, device, dtype))
 
     def step(
-        self, x: torch.Tensor, state: SparseDeltaMemoryState
+        self, x: torch.Tensor, state: SparseDeltaMemoryState, *, donate: bool = False
     ) -> tuple[torch.Tensor, SparseDeltaMemoryState]:
-        """One position — `(B, 1, d_model)` → output and the successor state."""
+        """One position — `(B, 1, d_model)` → output and the successor state.
+
+        Args:
+            donate: the caller will not read ``state`` again, so the successor
+                may be written into its table rather than a copy of it -- at a
+                large batch, most of the step.  A permission, not a demand: a
+                table that is not the stream's own to write is copied anyway,
+                which makes the first donated step from :meth:`init_state` a
+                copy and every one after it in place.
+        """
         if x.shape[1] != 1:
             raise ValueError(
                 f"step() consumes one position at a time, got {x.shape[1]}; "
                 f"use forward(x, state=..., return_state=True) for a chunk"
             )
         features = tuple(t[:, :, 0] for t in self._features(x))
-        o, memory = recurrent_sparse_delta(state.memory, *features)
+        o, memory = recurrent_sparse_delta(
+            state.memory,
+            *features,
+            in_place=donate and _writable_in_place(state.memory),
+        )
         return self._out(o.unsqueeze(2), x), SparseDeltaMemoryState(memory=memory)
 
     def read(self, x: torch.Tensor, state: SparseDeltaMemoryState) -> torch.Tensor: