SSM / Mamba kernel

[superposition]

Description

Selective State Space Model kernel for efficient sequence modeling (linear time complexity).

Operation

# Mamba selective scan
# x: (B, T, D), A, B, C, D: SSM parameters
# dt: (B, T, D) input-dependent time step
output = selective_scan(x, dt, A, B, C, D)

Implementation

• Parallel associative scan
• Input-dependent discretization
• Fused with input projection
• Memory efficient (recompute in backward)

Key Algorithm

h[t] = A_bar * h[t-1] + B_bar * x[t]
y[t] = C * h[t] + D * x[t]
where A_bar, B_bar depend on dt[t]

Tests

• Matches sequential loop implementation
• Linear time complexity O(T*D)
• Gradient correct for all parameters
• gradcheck passes
• Works with long sequences (8K+)
• Handles variable sequence lengths

Used In

Mamba, Mamba-2, state space world models, efficient transformers

References

• Mamba paper: https://arxiv.org/abs/2312.00752
• Triton implementation: mamba-ssm repo