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 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.