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

varlen_selective_scan

Variable-length selective scan kernels for Mamba SSM architecture.

comptime values

LOG2E

comptime LOG2E = 1.4426950408889634

MAX_DSTATE

comptime MAX_DSTATE = 256

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

• ​varlen_selective_scan_fwd_cpu: CPU kernel for variable-length selective scan.
• ​varlen_selective_scan_fwd_gpu: GPU kernel for variable-length selective scan.
• ​varlen_selective_state_update_cpu: CPU kernel for varlen selective state update.
• ​varlen_selective_state_update_gpu: GPU kernel for selective state update with multi-head support.