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_ops

Varlen selective scan operation registrations for Mamba SSM.

This module registers operations for variable-length selective scan:

• varlen_selective_scan_fwd: Forward pass for varlen sequences
• varlen_selective_state_update: State update for decode/autoregressive inference

comptime values

Strides1D

comptime Strides1D = IndexList[Int(1)]

Strides2D

comptime Strides2D = IndexList[Int(2)]

Strides3D

comptime Strides3D = IndexList[Int(3)]

Strides4D

comptime Strides4D = IndexList[Int(4)]

Structs

• ​VarlenSelectiveScanFwd: Variable-length selective scan forward pass.
• ​VarlenSelectiveStateUpdate: Varlen selective state update for autoregressive inference.

Functions

• ​varlen_selective_scan_fwd_shape: Returns the output shape for the varlen_selective_scan_fwd op.
• ​varlen_selective_state_update_shape: Returns the output shapes for the varlen_selective_state_update op.