srush / annotated-mamba

Annotated version of the Mamba paper

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Go to Blog Notebook Colab

image

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

About

Annotated version of the Mamba paper

License:MIT License


Languages

Language:Jupyter Notebook 98.6%Language:Python 1.4%Language:Makefile 0.1%