Recurrent State Manager#

class RecurrentStateManager#

Dense per-slot store for the recurrent axis: mixer recurrent and convolution states plus any auxiliary convolution state a layer carries. Rows are resident slots, the index space state_indices carries into the engine. Each state is one contiguous slab whose row stride is its own per-slot byte size, which is what makes slot * bytesPerSlot addressing valid without consulting tensor strides.

State order is the order the specs were given in and never changes, because capture and restore match a checkpoint to the pool positionally.

The store carries no mixer-specific behaviour: it knows a state’s geometry, not how the state evolves or how speculative decoding commits it.

Public Functions

RecurrentStateManager() noexcept = default#

RecurrentStateManager( std::vector<StateSpec> specs, int32_t slotCount, std::string const &name, cudaStream_t stream )#

Allocate one zeroed [slotCount] + shape slab per spec. No spec or no slot allocates nothing.

Parameters:

specs – Per-slot states in the order checkpoints will match them.
slotCount – Resident slot rows; may be 0, which leaves every slot index invalid.
name – Prefix for the slab tensor names, which surface in allocation diagnostics.
stream – CUDA stream for the zeroing memsets.

Throws:

std::runtime_error – if a spec has non-positive per-slot volume, names a negative layer, or repeats a {kind, absLayerIdx} already claimed.

RecurrentStateManager(RecurrentStateManager const&) = delete#

RecurrentStateManager &operator=( RecurrentStateManager const& ) = delete#

RecurrentStateManager(RecurrentStateManager&&) noexcept = default#

RecurrentStateManager &operator=( RecurrentStateManager&& ) noexcept = default#

int32_t slotCount() const noexcept#: Number of resident slot rows.

int32_t numStates() const noexcept#: Number of per-slot states held.

StateSpec const &spec(int32_t stateIdx) const#

The spec one state was allocated from.

Parameters:: stateIdx – Index in [0, numStates()).

std::vector<StateSpec> const &specs() const noexcept#: Every spec, in pool order.

Tensor &state(int32_t stateIdx)#

Storage for one state, shaped [slotCount] + spec.shape.

Parameters:: stateIdx – Index in [0, numStates()).

Tensor const &state(int32_t stateIdx) const#

int32_t findState(StateKind kind, int32_t absLayerIdx) const#: Index of the state with the given lineage, or -1 when the layer declares none.

size_t stateBytesPerSlot(int32_t stateIdx) const#

Bytes one slot occupies in one state.

Parameters:: stateIdx – Index in [0, numStates()).

size_t bytesPerSlot() const#: Bytes one slot occupies across every state held.

void clear(cudaStream_t stream)#: Zero every slot of every state.

void clearSlot(int32_t slot, cudaStream_t stream)#: Zero one slot of every state, leaving every other slot untouched.

std::vector<Tensor> capture(int32_t slot, cudaStream_t stream)#: Copy one slot out into freshly allocated [1] + shape tensors, one per state in pool order.

void restore( std::vector<Tensor> const &saved, int32_t slot, cudaStream_t stream )#: Write a checkpoint back into one slot. States beyond saved.size() are zeroed rather than left alone, so a checkpoint taken before a state existed resumes that state from zero instead of from the previous occupant’s data. Nothing is written until every entry has been validated against its spec.

void copySlotFrom( RecurrentStateManager const &source, int32_t destinationSlot, int32_t sourceSlot, cudaStream_t stream )#: Copy one slot of every state from a pool holding the same states in the same order. Pairs a live pool with a snapshot pool of a different slot count.

Public Static Functions

static size_t bytesPerSlot(std::vector<StateSpec> const &specs)#: Bytes one slot occupies across specs. Pool sizing runs against a byte budget before any allocation exists, so this is a property of the specs alone.