srush/annotated-mamba

Annotated version of the Mamba paper

Go to Blog Notebook Colab

Mamba: Linear-Time Sequence Modeling with Selective State Spaces https://arxiv.org/abs/2312.00752

Challenge

The triton version is still a lot slower than the mamba custom kernel. If anyone has ideas for speeding it up, I would love to hear them. The file final.py has a minimal version with benchmarking.

pip install mamba-ssm
pip install -U http://kermit.bounceme.net:8900/triton-3.0.0-cp310-cp310-linux_x86_64.whl

annotated-mamba's Issues

Minimal mamba in JAX

There is a minimal version in pytorch: https://github.com/johnma2006/mamba-minimal

Looking forward to a minimal version in JAX.

Then maybe we can try to accelerate it? (Not an expert in CUDA, is it possible to accelerate using mamba's original cuda code? )

Discretization is only approximative to first order in their code.

Note that in the original repo [^1], they only do the discretization up to first order Taylor approximation of the exponential [^2].
$$\bar{b} = \frac{ ( \exp ( \bar{a} ) - 1 ) \Delta b }{\bar{a}} \approx \Delta b ( 1 + \mathcal{O} ( (\Delta a) ) )$$

[^1] https://github.com/state-spaces/mamba/
[^2] https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/modules/mamba_simple.py#L239

associative ssm op order

i might be wrong, but in the colab/blog, i think the $\oplus$ op used to do the associative scan for the selective state space model should have, as the value of its first output, $a_2 a_1$ (rather than $a_1 a_2$), reflecting the fact that the leftmost $A$ transform gets applied first.

(Since the Mamba matrices are diagonal and therefore commutative it doesn't actually matter here I guess, I just found this initially confusing in the presentation).

It looks to be correct in the triton first_order_op but is I think reversed in the reference pytorch op and latex above it.

Thanks for this terrific writeup! It really clarified some things for me, thanks