SSM (Mamba) in Jax

[faresobeid]

I recently came across the Mamba paper https://arxiv.org/abs/2312.00752 and their hardware aware algorithm which allows for fast speeds for a state space model in recurrent mode. There are also many other SSM's following this principle like RWKV v6. The only implementations that actually have good speeds I've seen are in cuda. I was wondering if and how they can be implemented in Jax to utilise TPU's, one thing I've seen is Pallas but so far I've had no success which might be due to my inexperience with it. The main premise of these algorithms can be seen in appendix D of the paper:

I understand that this might be GPU specific but surely there is a way to replicate the effects on TPU in jax.

[EelcoHoogendoorn]

Here is an implementation of gateloop in JAX

I have not yet come to a full appreciation between these method but on a high level the ideas seem identical (make the coefficients of an SSM conditional on the input but not on the hidden state; thus retaining associative-scannability, while apparently getting sufficient modelling capacity to contest transformers, unlike fixed-coefficient SSMs).

Also curious to what extent this performs well in plain JAX. Or if it doesnt itd be a nice test case for pallas I guess.

[faresobeid]

Sorry for the late reply, in plain Jax all implementations I've seen or tried have been much slower than the cuda versions making using these models for training infeasible. I'm currently trying Pallas but keep getting errors about the vector layout #19058.

[EelcoHoogendoorn]

I guess 'training infeasible' means in the context of LLMs? Personally im interested in using it for very much small-data (but long sequence) applications, so a vanilla JAX implementation would already be very valuable to me. How has it been coming along in pallas though? Curious how it holds up for a use case like this.

[faresobeid]

Thank you for replying! Honestly since then I haven’t continued with Pallas as no one responded for my GitHub issue that I mentioned. Although I have gotten RWKV v6 which is a selective SSM like mamba working almost as fast as a GPT implementation in pure jax on TPU. However that doesn’t use the recurrent form and wouldn’t work for mamba. From my experience mamba should run at a decent speed for your small scale use cases. But still obviously much slower than a transformer. Although I’m planning to go back and try using Pallas again to hopefully make it work.

[davisyoshida]

We've been talking about Mamba in JAX over here: srush/annotated-mamba#1

I took a shot at implementing a kernel in Pallas, but I don't think I did it optimally.