Recurrent Commit Policy#

class RecurrentCommitPolicy#

How one recurrent-mixer kind commits its speculative-verify state. A layer’s mixer lineage fixes both the side buffers the verify pass writes and the math that folds the accepted path back into the persistent slab, so the commit path is a property of the kind, not a runtime flag:

  • mamba2 stashes the per-token replay inputs (dA/x/B/dt) and reconstructs the accepted state by replaying the SSD scan from the committed state;

  • gdn snapshots the per-token full recurrent state and commits the accepted endpoint by scatter, except for tree sizes the chunk-form verify handles, which commit by replaying the stash. Both kinds snapshot the convolution state per token and commit it by scatter.

The per-kind commit lives in its own translation unit (mamba2CommitPolicy / gdnConvCommitPolicy); this file holds only the shared base and the convolution side-buffer machinery both reuse, so a future mixer or attention-side commit adds a sibling file without touching the others.

The policy owns the spec-verify side buffers and the device pointer bundles the commit kernels read; it does not own the persistent recurrent/conv slabs, which it reaches through references supplied at construction and whose lifetime the caller guarantees for the policy’s lifetime.

Subclassed by trt_edgellm::rt::GdnConvCommitPolicy, trt_edgellm::rt::Mamba2CommitPolicy

Public Functions

virtual ~RecurrentCommitPolicy() = default#
RecurrentCommitPolicy(RecurrentCommitPolicy const&) = delete#
RecurrentCommitPolicy &operator=(
RecurrentCommitPolicy const&
) = delete#
RecurrentCommitPolicy(RecurrentCommitPolicy&&) = delete#
RecurrentCommitPolicy &operator=(RecurrentCommitPolicy&&) = delete#
int32_t numLayers() const noexcept#

Number of recurrent layers this policy commits.

virtual bool hasIntermediateRecurrentStates() const noexcept = 0#

Whether verify-pass recurrent side buffers were allocated (spec-verify enabled).

bool hasIntermediateConvStates() const noexcept#

Whether verify-pass convolution side buffers were allocated (spec-verify enabled).

virtual bool usesReplay() const noexcept = 0#

Whether the recurrent commit reconstructs by replay (mamba2) rather than scatter (gdn).

void reshapeIntermediate(int32_t activeBatchSize, int32_t seqLen)#

Reshape the verify-pass side buffers from their [maxBatch, maxIntermediateSeqLen, …] allocation to the [activeBatchSize, seqLen, …] the engine wrote, so stride-based reads address correctly. No-op when side buffers are not allocated.

rt::Tensor &intermediateConvState(int32_t recurrentLayerIdx)#

Verify-pass convolution side buffer for one layer, shaped by the last reshapeIntermediate.

virtual void commitLinear(
rt::Tensor const &acceptLengths,
rt::Tensor const &stateIndices,
cudaStream_t stream
) = 0#

Commit the single accepted recurrent/conv state per sequence after linear (chain) verification.

Parameters:

  • acceptLengths – [activeBatchSize] INT32 GPU accepted-token count per sequence.

  • stateIndices – [activeBatchSize] INT32 GPU active-row to resident-slot map.

virtual void commitTree(
rt::Tensor const &acceptedNodeIds,
rt::Tensor const &acceptLengths,
rt::Tensor const &stateIndices,
int32_t verifySize,
cudaStream_t stream
) = 0#

Commit the accepted recurrent/conv state per sequence after tree verification.

Parameters:

  • acceptedNodeIds – [activeBatchSize, maxAcceptLen] INT32 GPU verify-node id per accepted token.

  • acceptLengths – [activeBatchSize] INT32 GPU accepted-token count per sequence.

  • stateIndices – [activeBatchSize] INT32 GPU active-row to resident-slot map.

  • verifySize – The verify-tree node count; selects snapshot vs chunk-replay for gdn.

struct Config#

Geometry of the recurrent axis this policy commits. maxIntermediateSeqLen is the verify-pass token capacity; a zero value means no side buffers are allocated and every commit is a no-op.

Public Members

int32_t numRecurrentLayers = {0}#
int32_t maxBatchSize = {}#
int32_t recurrentStateNumHeads = {0}#
int32_t recurrentStateHeadDim = {0}#
int32_t recurrentStateSize = {0}#
int32_t convDim = {0}#
int32_t convKernel = {0}#
int32_t maxIntermediateSeqLen = {0}#
trt_edgellm::DataType recurrentStateType = {trt_edgellm::DataType::kHALF}#
trt_edgellm::DataType convStateType = {trt_edgellm::DataType::kHALF}#
int32_t recurrentStateNumGroups = {0}#

B/C group count. Consumed only by the mamba2 replay stash (the replay-B extent); gdn leaves it 0.

struct Config

Geometry of the recurrent axis this policy commits. maxIntermediateSeqLen is the verify-pass token capacity; a zero value means no side buffers are allocated and every commit is a no-op.

Public Members

int32_t numRecurrentLayers = {0}
int32_t maxBatchSize = {}
int32_t recurrentStateNumHeads = {0}
int32_t recurrentStateHeadDim = {0}
int32_t recurrentStateSize = {0}
int32_t convDim = {0}
int32_t convKernel = {0}
int32_t maxIntermediateSeqLen = {0}
trt_edgellm::DataType recurrentStateType = {trt_edgellm::DataType::kHALF}
trt_edgellm::DataType convStateType = {trt_edgellm::DataType::kHALF}
int32_t recurrentStateNumGroups = {0}

B/C group count. Consumed only by the mamba2 replay stash (the replay-B extent); gdn leaves it 0.