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 function

mamba2_ssd_chunk_scan_varlen_fwd_inplace_shape

def mamba2_ssd_chunk_scan_varlen_fwd_inplace_shape[dtype: DType, state_dtype: DType](x: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec], dt: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec], A: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec], B: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec], C: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec], D: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec], dt_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec], ssm_pool: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=ssm_pool.static_spec], query_start_loc: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec], has_initial_state: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec], cache_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_indices.static_spec]) -> IndexList[Int(3)]

Computes the output shape for the mamba2_ssd_chunk_scan_varlen_fwd_inplace graph op.

Parameters:

• ​dtype (DType): Element type of the SSM input and parameter tensors (inferred).
• ​state_dtype (DType): Element type of the SSM state pool tensor (inferred).

Args:

• ​x (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec]): Packed input tensor of shape (total_len, nheads, head_dim).
• ​dt (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec]): Per-head time deltas of shape (total_len, nheads).
• ​A (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec]): Per-head scalar decay of shape (nheads,).
• ​B (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec]): Grouped input projection of shape (total_len, ngroups, dstate).
• ​C (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec]): Grouped output projection of shape (total_len, ngroups, dstate).
• ​D (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec]): Per-head skip connection of shape (nheads,); may be empty when unused.
• ​dt_bias (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec]): Per-head bias added to dt of shape (nheads,); may be empty when unused.
• ​ssm_pool (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=ssm_pool.static_spec]): Mutable SSM state pool of shape (max_slots, nheads, head_dim, dstate) in float32; final states are written in place at the slots indexed by cache_indices.
• ​query_start_loc (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec]): Cumulative sequence lengths of shape (batch + 1,) in int32.
• ​has_initial_state (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec]): Per-sequence flag of shape (batch,) in bool indicating whether to load the initial state from ssm_pool; may be empty when no initial states are used.
• ​cache_indices (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_indices.static_spec]): Per-sequence slot indices of shape (batch,) in uint32 selecting where in ssm_pool the final states are written.

Returns:

IndexList[Int(3)]