Module mamba3 

Mamba-3: trapezoidal SSD with data-dependent RoPE and MIMO.

Mamba-3

Mamba-3 extends Mamba-2 with three independent additions (each works alone or combined): trapezoidal discretisation, data-dependent RoPE on the B/C projections, and MIMO (multiple-input multiple-output) rank expansion. See mamba3 for the full combined math.

Two SSD pathways

The trapezoidal recurrence is realised by two interchangeable algorithms, selected at runtime by which cache variant is supplied:

• [double_ssd] — splits the trapezoid into two standard SSD calls (simple, easy to verify; ~2× the intra-chunk memory).
• [single_ssd] — one SSD call in the official-kernel form (≈ half the training memory; the cache’s SSM accumulator has different mid-sequence semantics).

[cache] holds the enum that dispatches between them; [ssd_path] selects the pathway-agnostic algorithm (Minimal / Serial / SerialRecalculated).

Modules

backwards

Blanket Mamba3BackendExt implementation for autodiff backends.

cache

Mamba-3 Cache and Pathway Selection

double_ssd

Double-SSD pathway (VikramLex-style)

helpers 🔒

Shared helpers used by both Mamba3::forward and Mamba3::step. They isolate three blocks that previously appeared in both methods at different ranks:

mamba3

Mamba-3 SSM Block — Exponential-Trapezoidal SSD with Data-Dependent RoPE

moments

Physical-frame state moments for Mamba-3 (serial chunkwise, exact)

prelude

Public re-exports for Mamba-3.

quat_scan

Memory-efficient quaternion cumulative-product scan (custom backward)

rotation

Quaternion (k=4) rotational state — the non-abelian generalisation of RoPE

single_ssd

Single-SSD pathway (official-kernel form)

ssd_path

Pathway-agnostic SSD algorithm selection (Mamba-3)

step_constant 🔒

Constant-input stepping shortcuts — step_n_approx / step_infinite

Traits

Mamba3AutodiffBackendExt

Marker for an autodiff-capable backend that also implements the SSD extension trait (so the custom memory-efficient backward is available).

Mamba3BackendExt

Backend capability required to run Mamba-3.