Mamba-3 code release

[aakashlahoti]

Official code for "Mamba-3: Improved Sequence Modeling using State Space Principles".
Aakash Lahoti*, Kevin Y. Li*, Berlin Chen*, Caitlin Wang*, Aviv Bick, J. Zico Kolter, Tri Dao†, Albert Gu†

Paper: https://arxiv.org/abs/2603.15569

[skaae]

Hi,
Thank you for releasing the code! I wonder if you plan to implement varlen for the MIMO variant either with cu_seqlen or seq_idx?

Do you know if it would be difficult to implement varlen for MIMO if we want to give it a shot?

[tridao]

I feel it's not too difficult, but @berlinchen7 should know better

[berlinchen7]

Thanks for your interest! Indeed varlen MIMO shouldn’t be too difficult to implement. We just need to be careful not to introduce unnecessary overhead.

Support for varlen MIMO is currently in progress, and we plan to release it once the implementation is ready.

[skaae]

Sounds excellent. We'll start experimenting with the siso variant. Hopefully it works even better than the excellent mamba2 :) 🤞

Thank you!