0.8.1 — a read cache can cross a vmap boundary

README.md

@@ -116,7 +116,7 @@ it was written on.
 | [`lumen.gdn`](./docs/GDN.md) | Gated DeltaNet — a fixed-size associative memory with a delta-rule write. Linear in sequence length, constant-memory generation, configurable head layout, reads that do not write, opt-in `fla` kernel path verified against the fp64 oracle. [Design record](./docs/design/GATED_DELTANET.md) |
 | [`lumen.sdm`](./docs/SDM.md) | Sparse Delta Memory — a gated delta rule over a large, sparsely addressed table. State size decoupled from parameter count and per-token compute; a slot nobody writes to is frozen, decay included; an empty or a learned initial table; runs under `vmap`; an opt-in Triton backend. [Design record](./docs/design/SPARSE_DELTA_MEMORY.md) |
 | [`lumen.block`](./docs/BLOCK.md) | A residual `Block` and a `Stack` of them — the level a result is actually read off. Modality-free, sub-layer instances rather than a config, and the depth-scaled init that cannot live any lower. [Design record](./docs/design/BLOCK.md) |
-| [`lumen.pytree`](./docs/PYTREE.md) | Every streaming state is a pytree node, so `vmap`, `functional_call` and `torch.compile` can traverse one. The state is half the mixer contract; this is what keeps it from being where the abstraction stops |
+| [`lumen.pytree`](./docs/PYTREE.md) | Every streaming state is a pytree node, and so is the GDN read cache, so `vmap`, `functional_call` and `torch.compile` can traverse one. The state is half the mixer contract; this is what keeps it from being where the abstraction stops |
 
 A roadmap of things that do not exist yet is a promise, not a description, so
 there isn't one here.
@@ -224,7 +224,7 @@ suite checks against; they are deliberately not exported.
 
 ## Status
 
-`0.8.0` — alpha, and versioned honestly. The *interface* is the thing this
+`0.8.1` — 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/PYTREE.md

@@ -1,7 +1,9 @@
 # States as pytrees
 
 Every streaming state in this package is a registered pytree node, so the
-`torch` transforms that have to flatten a structure can walk one.
+`torch` transforms that have to flatten a structure can walk one. So is the
+gated delta net's read cache, which is not a state but meets the same
+boundaries — see [Read caches](#read-caches).
 
 `GatedDeltaNetState`, `UndertowState`, `SparseDeltaMemoryState`, `BlockState`
 and `StackState` are frozen dataclasses of tensors. Until they are declared, `torch` sees each of them as a
@@ -103,20 +105,75 @@ automatically, same outputs, and a cost that grows with the table. A learned
 initial table is a parameter like any other, so it is stacked per set with the
 rest.
 
+## Read caches
+
+`GatedDeltaNet.forward(..., return_cache=True)` returns a
+`GatedDeltaNetReadCache`, which `reread` reads again with a fresh query and no
+write. A cache made and consumed inside one vmapped function never crosses the
+transform boundary, and one cache shared by every set crosses it unbatched
+(`in_dims=None`). What the registration is for is the write and the re-read in
+*separate* batched calls — `G` writers returning `G` caches out of `vmap`, and
+`G` re-readers taking them back in:
+
+```python
+import torch
+from torch.func import functional_call, stack_module_state, vmap
+
+from lumen.gdn import GatedDeltaNet
+
+class Rereading(torch.nn.Module):
+    """`functional_call` only calls `forward`; this one's forward is `reread`."""
+
+    def __init__(self, layer: GatedDeltaNet) -> None:
+        super().__init__()
+        self.layer = layer
+
+    def forward(self, x, cache):
+        return self.layer.reread(x, cache)
+
+layers = [GatedDeltaNet(config) for _ in range(SETS)]
+params, buffers = stack_module_state(layers)
+base = GatedDeltaNet(config).to("meta")
+rereading = Rereading(base)
+
+def prefixed(tensors):
+    return {f"layer.{name}": tensor for name, tensor in tensors.items()}
+
+def write(parameters, buffered, x):
+    return functional_call(
+        base, (parameters, buffered), (x,), {"return_cache": True}
+    )[1]
+
+def reread(parameters, buffered, x, cache):
+    weights = (prefixed(parameters), prefixed(buffered))
+    return functional_call(rereading, weights, (x, cache))
+
+caches = vmap(write)(params, buffers, x)           # every leaf leads with SETS
+y = vmap(reread)(params, buffers, fresh, caches)   # (SETS, BATCH, SEQ, D)
+```
+
+`ReadCache.seq_len`, the unpadded length inside the cache, is
+`UndertowState.seen` again: an `int` that drives control flow, so it rides in
+the treespec context and every leaf is a tensor. A cache from a continued pass
+carries the short-conv context it started from, and that is batched like
+everything else.
+
 ## Checking what registered
 
 ```python
 import lumen
 
 lumen.REGISTERED
 # ('GatedDeltaNetState', 'BlockState', 'StackState', 'UndertowState',
-#  'SparseDeltaMemoryState')
+#  'SparseDeltaMemoryState', 'ReadCache', 'GatedDeltaNetReadCache')
 ```
 
-Five names is the whole set. A shorter tuple means `torch` on this machine did
-not expose an entry point one of the states needs — the library is unaffected
-and every path documented elsewhere behaves identically, but a `torch.func`
-transform will still refuse the states that are missing. Reading it is cheaper
+Seven names is the whole set: the five states, then the read cache's two layers
+— the kernel's `ReadCache` and the layer's `GatedDeltaNetReadCache` around it.
+A shorter tuple means `torch` on this machine did not expose an entry point one
+of them needs — the library is unaffected and every path documented elsewhere
+behaves identically, but a `torch.func` transform will still refuse the ones
+that are missing. Reading it is cheaper
 than rediscovering the same fact from a `vmap` error inside a recurrence.
 
 `register_state_pytrees()` returns the same tuple and can be called at any time;

docs/design/GATED_DELTANET.md

@@ -752,11 +752,12 @@ The cache's members are tensors autograd already saves for backward — each is
 a matmul or multiply operand — so under a training graph retaining it costs
 nothing; the cache is returned as a frozen value rather than stashed on the
 module, because a cache held on `self` leaks when a block raises and is a
-side effect `torch.compile` cannot trace. It is **not** registered as a pytree
-node: `seq_len` is an int that describes the tree rather than living in it —
-the `UndertowState.seen` case `lumen.pytree` documents — so the registration
-is a hand-rolled flatten with its own round-trip tests, and is its own change.
-Until then a cache is an ordinary value, and a leaf at a `torch.func` boundary.
+side effect `torch.compile` cannot trace. It is registered as a pytree node,
+and by hand rather than through `register_dataclass`: `seq_len` is an int that
+describes the tree rather than living in it — the `UndertowState.seen` case
+`lumen.pytree` documents — so it rides in the treespec context, and a batch of
+caches crosses a `vmap` boundary the way a batch of states does
+(`latticedynamics/lumen#20`).
 
 **The re-read query is formed by `_features`, through the short conv.** The
 memory was written under keys that went through `k_conv`; the address space is

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "lumen"
-version = "0.8.0"
+version = "0.8.1"
 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

@@ -55,9 +55,9 @@
 from lumen.undertow import UndertowAttention, UndertowConfig, UndertowState
 from lumen.stack import Stack, StackState
 
-# Imported for its side effect as much as its surface: the state containers are
-# declared as pytree nodes here, so `torch.func` can traverse them.  Additive
-# and numerically inert -- see :mod:`lumen.pytree`.
+# Imported for its side effect as much as its surface: the state containers and
+# the GDN read caches are declared as pytree nodes here, so `torch.func` can
+# traverse them.  Additive and numerically inert -- see :mod:`lumen.pytree`.
 from lumen.pytree import REGISTERED, register_state_pytrees
 from lumen.probe import (
     ArchFacts,
@@ -73,7 +73,7 @@
     probe_triton,
 )
 
-__version__ = "0.8.0"
+__version__ = "0.8.1"
 
 __all__ = [
     "__version__",

src/lumen/pytree.py

@@ -10,9 +10,9 @@
 because a pytree that has never been declared is a leaf, and a leaf that is not
 a tensor is rejected.
 
-Nothing about these five classes is hostile to registration.  They are frozen
-dataclasses holding tensors, tuples of tensors, ``None``, and one ``int``.  The
-declaration is the whole fix.  See `latticedynamics/lumen#12`.
+Nothing about these classes is hostile to registration.  They are frozen
+dataclasses holding tensors, tuples of tensors, ``None``, and in two places an
+``int``.  The declaration is the whole fix.  See `latticedynamics/lumen#12`.
 
 **This is additive and numerically inert.**  No tensor changes value, no
 checkpoint key moves, no existing call path behaves differently.  What is
@@ -30,45 +30,65 @@
 because the states are Lumen's own classes — no other library can be holding a
 registration for them to collide with.
 
-Why three registrations are public and one is not
--------------------------------------------------
+Why the read caches are here too
+--------------------------------
+
+A :class:`~lumen.gdn.GatedDeltaNetReadCache` is not a state — it continues
+nothing — but it meets the same boundary.  Made and consumed inside one
+transformed call it never crosses it, and one cache shared by every member of a
+batch crosses it unbatched, as a constant.  What needs the declaration is the
+write and the re-read happening in *separate* calls over ``G`` parameter sets:
+``G`` writers return ``G`` caches out of ``vmap``, and ``G`` re-readers take them
+back in at ``in_dims=0``.  Both are refused while the cache is a leaf.  See
+`latticedynamics/lumen#20`.
+
+Why two registrations are not public
+------------------------------------
 
 :func:`torch.export.register_dataclass` handles the case that looks hardest.
 ``GatedDeltaNetState.conv`` is ``None`` on a layer with no convolution and
 ``BlockState.local`` is ``None`` on a block with no local sub-layer; it drops
 absent fields from the children and records their absence in the context, which
 is exactly what a hand-rolled flatten gets wrong on the first attempt.
+``GatedDeltaNetReadCache.conv`` is ``None`` on a pass that began a sequence, and
+is the same case again.
 
 It cannot express the case that looks easy.  ``UndertowState.seen`` is present,
 permanent, and not a tensor, so it becomes a child and ``vmap`` refuses it for
 the same reason it refuses ``None``.  ``seen`` is not data: it drives Python
 control flow — the valid-position count, the front-padding width, the slice
 bounds — and it describes the tree rather than living in it.  That is what a
 treespec context is for, and ``register_dataclass`` has no way to put a field
-there.
+there.  ``ReadCache.seq_len`` is the same case a second time: it is where a
+re-read trims its padded output and what a query's length is checked against.
 
 ``drop_field_names`` is the public-looking escape and is a worse trap than the
 one it solves: flattening succeeds and unflattening raises, because a dropped
 field is reconstructed from nothing.  **The repair that suggests itself next is
 the one to refuse loudest** — giving ``seen`` a default makes every round-tripped
 state rebuild as ``seen=0``, and a window state that has forgotten its count
-masks the wrong slots and returns numbers instead of an error.
-
-So one class is registered against ``torch.utils._pytree``, which is private.
-The degradation discipline covers it on the same terms as everything else here:
-attempt it, skip it if the surface moved, and change nothing that works today.
-A skipped registration costs the ``torch.func`` path, which is unavailable
-anyway on a version that lacks the API.
+masks the wrong slots and returns numbers instead of an error.  ``seq_len``
+gets no default either, for the same reason.
+
+So two classes are registered against ``torch.utils._pytree``, which is private,
+through one helper: the tensor fields as children, the one field that is not a
+tensor as context.  The degradation discipline covers it on the same terms as
+everything else here: attempt it, skip it if the surface moved, and change
+nothing that works today.  A skipped registration costs the ``torch.func`` path,
+which is unavailable anyway on a version that lacks the API.
 """
 
 from __future__ import annotations
 
+from dataclasses import fields
+from functools import partial
 from typing import Any, Callable, Iterable
 
 import torch
 
 from lumen.block import BlockState
-from lumen.gdn import GatedDeltaNetState
+from lumen.gdn import GatedDeltaNetReadCache, GatedDeltaNetState
+from lumen.gdn.reference import ReadCache
 from lumen.sdm import SparseDeltaMemoryState
 from lumen.stack import StackState
 from lumen.undertow import UndertowState
@@ -83,29 +103,23 @@
 _REGISTERED: set[type] = set()
 
 
-def _flatten_undertow(
-    state: UndertowState,
-) -> tuple[list[torch.Tensor], int]:
-    """Two tensors as children, the count as context."""
-    return [state.keys, state.values], state.seen
-
-
-def _unflatten_undertow(
-    children: Iterable[torch.Tensor], seen: int
-) -> UndertowState:
-    keys, values = children
-    return UndertowState(keys=keys, values=values, seen=seen)
-
-
 def _register_dataclass(cls: type) -> None:
     from torch.export import register_dataclass
 
     register_dataclass(cls)
 
 
-def _register_undertow(cls: type) -> None:
-    """Register with paths if this ``torch`` has them, without if it does not.
+def _register_with_context(cls: type, context: str) -> None:
+    """Every field but ``context`` as a child; ``context`` in the treespec.
 
+    The children are read off the dataclass rather than listed here, so a field
+    added to the class later becomes a child without anyone remembering this
+    file.  Listed by hand, a new field would be dropped by the flatten and
+    either rebuilt from its default or refused by the constructor -- the first
+    silently.  Read off the class, the worst a new non-tensor field can do is
+    be refused by ``vmap``, which is loud.
+
+    Registered with paths if this ``torch`` has them, without if it does not.
     The keyed flatten is what makes ``tree_map_with_path`` treat every state
     alike; omitting it leaves the path-carrying traversals working on the
     dataclass-registered states and failing on this one, which is the asymmetry
@@ -120,53 +134,62 @@ def _register_undertow(cls: type) -> None:
     """
     from torch.utils._pytree import register_pytree_node
 
+    children = tuple(field.name for field in fields(cls) if field.name != context)
+
+    def flatten(node: Any) -> tuple[list[torch.Tensor], Any]:
+        return [getattr(node, name) for name in children], getattr(node, context)
+
+    def unflatten(leaves: Iterable[torch.Tensor], value: Any) -> Any:
+        return cls(**dict(zip(children, leaves, strict=True)), **{context: value})
+
     try:
         from torch.utils._pytree import GetAttrKey
     except ImportError:
-        register_pytree_node(cls, _flatten_undertow, _unflatten_undertow)
+        register_pytree_node(cls, flatten, unflatten)
         return
 
-    def flatten_with_keys(
-        state: UndertowState,
-    ) -> tuple[list[tuple[Any, torch.Tensor]], int]:
-        return [
-            (GetAttrKey("keys"), state.keys),
-            (GetAttrKey("values"), state.values),
-        ], state.seen
+    def flatten_with_keys(node: Any) -> tuple[list[tuple[Any, torch.Tensor]], Any]:
+        keyed = [(GetAttrKey(name), getattr(node, name)) for name in children]
+        return keyed, getattr(node, context)
 
     try:
         register_pytree_node(
-            cls,
-            _flatten_undertow,
-            _unflatten_undertow,
-            flatten_with_keys_fn=flatten_with_keys,
+            cls, flatten, unflatten, flatten_with_keys_fn=flatten_with_keys
         )
     except TypeError:
         # The keyword is not accepted on this version, so nothing was registered
         # -- binding fails before any work.  Retrying without it is safe.
-        register_pytree_node(cls, _flatten_undertow, _unflatten_undertow)
+        register_pytree_node(cls, flatten, unflatten)
 
 
-#: Each state and the entry point that can express its shape.  Enumerated in one
-#: place rather than declared beside each class: the set is no longer a uniform
-#: loop, and the asymmetry is the thing most worth keeping visible.
+#: Each container and the entry point that can express its shape.  Enumerated in
+#: one place rather than declared beside each class: the set is no longer a
+#: uniform loop, and the asymmetry is the thing most worth keeping visible.
+#:
+#: The read caches come last, so the five states keep the positions in
+#: :data:`REGISTERED` they have always had.
 _STATES: tuple[tuple[type, Callable[[type], None]], ...] = (
     (GatedDeltaNetState, _register_dataclass),
     (BlockState, _register_dataclass),
     (StackState, _register_dataclass),
-    (UndertowState, _register_undertow),
+    (UndertowState, partial(_register_with_context, context="seen")),
     (SparseDeltaMemoryState, _register_dataclass),
+    (ReadCache, partial(_register_with_context, context="seq_len")),
+    (GatedDeltaNetReadCache, _register_dataclass),
 )
 
 
 def register_state_pytrees() -> tuple[str, ...]:
-    """Declare the state containers as pytree nodes.  Idempotent.
+    """Declare the state containers and the read caches as pytree nodes.
+
+    Idempotent.
 
     Returns:
-        The names of every state that is registered once this call returns,
-        in declaration order.  A short tuple is a partial registration and a
-        working library: the states behave exactly as they did before, and only
-        the ``torch.func`` path on the missing ones is unavailable.
+        The names of every container that is registered once this call
+        returns, in declaration order.  A short tuple is a partial registration
+        and a working library: the containers behave exactly as they did
+        before, and only the ``torch.func`` path on the missing ones is
+        unavailable.
 
     Failures are captured rather than raised, matching :mod:`lumen.probe` and
     for the same reason -- one absent API in ``torch`` should be a shorter tuple
@@ -186,8 +209,8 @@ def register_state_pytrees() -> tuple[str, ...]:
 
 #: What registration achieved on this ``torch``, evaluated once at import.
 #:
-#: Five names is the whole set.  Anything shorter names the states a
-#: ``torch.func`` transform will still refuse, which is the only question a
-#: consumer needs answered and is cheaper to read than to rediscover from a
-#: ``vmap`` error.
+#: Seven names is the whole set: five states and the two read caches.  Anything
+#: shorter names the containers a ``torch.func`` transform will still refuse,
+#: which is the only question a consumer needs answered and is cheaper to read
+#: than to rediscover from a ``vmap`` error.
 REGISTERED: tuple[str, ...] = register_state_pytrees()