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_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, Engine=x.Engine, address_space=x.address_space, linear_idx_type=x.linear_idx_type], dt: TileTensor[kernel_dtype, Engine=dt.Engine, address_space=dt.address_space, linear_idx_type=dt.linear_idx_type], A: TileTensor[kernel_dtype, Engine=A.Engine, address_space=A.address_space, linear_idx_type=A.linear_idx_type], B: TileTensor[kernel_dtype, Engine=B.Engine, address_space=B.address_space, linear_idx_type=B.linear_idx_type], C: TileTensor[kernel_dtype, Engine=C.Engine, address_space=C.address_space, linear_idx_type=C.linear_idx_type], D: TileTensor[kernel_dtype, Engine=D.Engine, address_space=D.address_space, linear_idx_type=D.linear_idx_type], dt_bias: TileTensor[kernel_dtype, Engine=dt_bias.Engine, address_space=dt_bias.address_space, linear_idx_type=dt_bias.linear_idx_type], initial_states: TileTensor[.float32, Engine=initial_states.Engine, address_space=initial_states.address_space, linear_idx_type=initial_states.linear_idx_type], y: TileTensor[kernel_dtype, Engine=y.Engine, address_space=y.address_space, linear_idx_type=y.linear_idx_type], final_states: TileTensor[.float32, Engine=final_states.Engine, address_space=final_states.address_space, linear_idx_type=final_states.linear_idx_type], query_start_loc: TileTensor[.int32, Engine=query_start_loc.Engine, address_space=query_start_loc.address_space, linear_idx_type=query_start_loc.linear_idx_type], has_initial_state: TileTensor[.bool, Engine=has_initial_state.Engine, address_space=has_initial_state.address_space, linear_idx_type=has_initial_state.linear_idx_type], 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).