Mamba2 Commit Policy#
-
class Mamba2CommitPolicy : public trt_edgellm::rt::RecurrentCommitPolicy#
Mamba2 commit: reconstruct the accepted recurrent state by replaying the SSD scan from the committed state over the per-token dA/x/B/dt stash the verify pass wrote. Owns the replay stash and the B/C group count the stash extent needs.
Public Functions
- Mamba2CommitPolicy(
- Config config,
- std::vector<rt::Tensor> &liveRecurrent,
- std::vector<rt::Tensor> &liveConv,
- cudaStream_t stream
-
virtual bool hasIntermediateRecurrentStates() const noexcept override#
Whether verify-pass recurrent side buffers were allocated (spec-verify enabled).
-
virtual bool usesReplay() const noexcept override#
Whether the recurrent commit reconstructs by replay (mamba2) rather than scatter (gdn).
- virtual void commitLinear( ) override#
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
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.