flashinfer.mamba

Mamba / Mamba-2 state-space-model kernels. These wrap the selective scan and state-update primitives used in SSM blocks.

selective_state_update(state, x, dt, A, B, C, D)	Selective state update operation for Mamba layers (the generation phase).
checkpointing_ssu(state, old_x, old_B, ...)	Checkpointing SSU with MTP replay using matmul-based parallel token processing.