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 module

selective_scan

Selective scan kernel implementations for Mamba SSM prefill and decode.

Provides GPU and CPU forward-pass and update-step kernels for the selective scan (SSM) recurrence used by Mamba and Mamba-2, including the SSD (state-space duality) combined scan for variable-length batched prefill.

`comptime` values

`LOG2E`

comptime LOG2E = 1.4426950408889634

`MAX_DSTATE`

comptime MAX_DSTATE = 16

`Strides1D`

comptime Strides1D = IndexList[Int(1)]

`Strides2D`

comptime Strides2D = IndexList[Int(2)]

`Strides3D`

comptime Strides3D = IndexList[Int(3)]

`Strides4D`

comptime Strides4D = IndexList[Int(4)]

Functions

mamba_split_conv1d_scan_combined_cpu: CPU kernel for mamba_split_conv1d_scan_combined operation.
mamba_split_conv1d_scan_combined_gpu: GPU kernel for mamba_split_conv1d_scan_combined operation.
selective_scan_fwd_cpu: CPU kernel for selective scan forward pass.
selective_scan_fwd_cpu_minimal: Minimal CPU kernel for selective scan forward - no D, z, or delta_bias.
selective_scan_fwd_gpu: GPU kernel for selective scan forward pass.
selective_scan_fwd_gpu_minimal: Minimal GPU kernel for selective scan forward - no D, z, or delta_bias.
selective_scan_update_cpu: CPU kernel for selective scan update (single step).
selective_scan_update_gpu: GPU kernel for selective scan update (single step).
sigmoid: Optimized sigmoid using fast approximation for large negative values.
softplus: Softplus activation: log(1 + exp(x)) with numerical stability.
ssd_combined_cpu: CPU kernel for SSD combined operation.
ssd_combined_gpu: GPU kernel for SSD combined operation.

comptime values​

LOG2E​

MAX_DSTATE​

Strides1D​

Strides2D​

Strides3D​

Strides4D​

Functions​