For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

Mamba2ReplayFold

struct Mamba2ReplayFold[dt_softplus: Bool = True]

Commits a speculative verify's accepted rows to every Mamba-2 layer.

For each layer and request, sets the conv window to the last width - 1 accepted conv inputs and folds the accepted rows into the SSM state, in one launch with no output. The SSM state lands bit-identical to replaying the rows through mamba2_ssd_chunk_scan_varlen_fwd_inplace; see state_space.mamba2_replay_fold. Every request continues its live state. GPU only.

Tensor shapes: - conv_pool: (max_slots, conv_dim, width - 1) [MUT] - ssm_pool: (max_slots, nheads, head_dim, dstate) fp32 or bf16 [MUT] - conv_rows: (num_layers, batch) uint32 - ssm_rows: (num_layers, batch) uint32 - replay_bounds: (2 * batch,) int32, each request's accepted (start, end) rows - layers: 4 * num_layers tensors, in four runs of num_layers: each layer's in_proj output rows (rows, nheads * head_dim + conv_dim + nheads), its conv output rows (rows, conv_dim), its A as (1, nheads) and its dt_bias as (1, nheads). Every tensor's rows are contiguous.

Parameters

• ​dt_softplus (Bool): If True (default), apply softplus to dt + dt_bias.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[dtype: DType, conv_dtype: DType, state_dtype: DType, target: StaticString](conv_pool: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=conv_pool.static_spec], ssm_pool: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=ssm_pool.static_spec], conv_rows: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=conv_rows.static_spec], ssm_rows: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=ssm_rows.static_spec], replay_bounds: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=replay_bounds.static_spec], layers: VariadicTensors[IOSpec[_, _].Input, static_specs=layers.static_specs], ctx: DeviceContext)