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

Mojo function

mamba2_ssd_chunk_scan_varlen_fwd_cpu

def mamba2_ssd_chunk_scan_varlen_fwd_cpu[kernel_dtype: DType, DSTATE: Int](nheads: Int, head_dim: Int, ngroups: Int, nheads_ngroups_ratio: Int, batch: Int, dt_softplus: Int8, x: TileTensor[kernel_dtype, Storage=x.Storage, address_space=x.address_space, linear_idx_type=x.linear_idx_type], dt: TileTensor[kernel_dtype, Storage=dt.Storage, address_space=dt.address_space, linear_idx_type=dt.linear_idx_type], A: TileTensor[kernel_dtype, Storage=A.Storage, address_space=A.address_space, linear_idx_type=A.linear_idx_type], B: TileTensor[kernel_dtype, Storage=B.Storage, address_space=B.address_space, linear_idx_type=B.linear_idx_type], C: TileTensor[kernel_dtype, Storage=C.Storage, address_space=C.address_space, linear_idx_type=C.linear_idx_type], D: TileTensor[kernel_dtype, Storage=D.Storage, address_space=D.address_space, linear_idx_type=D.linear_idx_type], dt_bias: TileTensor[kernel_dtype, Storage=dt_bias.Storage, address_space=dt_bias.address_space, linear_idx_type=dt_bias.linear_idx_type], initial_states: TileTensor[DType.float32, Storage=initial_states.Storage, address_space=initial_states.address_space, linear_idx_type=initial_states.linear_idx_type], y: TileTensor[kernel_dtype, Storage=y.Storage, address_space=y.address_space, linear_idx_type=y.linear_idx_type], final_states: TileTensor[DType.float32, Storage=final_states.Storage, address_space=final_states.address_space, linear_idx_type=final_states.linear_idx_type], query_start_loc: TileTensor[DType.int32, Storage=query_start_loc.Storage, address_space=query_start_loc.address_space, linear_idx_type=query_start_loc.linear_idx_type], has_initial_state: TileTensor[DType.bool, Storage=has_initial_state.Storage, address_space=has_initial_state.address_space, linear_idx_type=has_initial_state.linear_idx_type], x_strides: IndexList[Int(3)], dt_strides: IndexList[Int(2)], A_strides: IndexList[Int(1)], B_strides: IndexList[Int(3)], C_strides: IndexList[Int(3)], D_strides: IndexList[Int(1)], dt_bias_strides: IndexList[Int(1)], initial_states_strides: IndexList[Int(4)], y_strides: IndexList[Int(3)], final_states_strides: IndexList[Int(4)], ctx: Optional[DeviceContext] = None)

CPU reference for the Mamba-2 SSD varlen prefill scan.

This is the trusted reference for numerical-equivalence testing: it computes the same per-token recurrence in fp32. Parallelised over (b, h, p).