Decouple from Triton about mamba HOT 4 CLOSED

Great to hear about the Squad result!

I don't have much experience with CPU or ONNX deployment. I mostly work with GPUs.
The model architecture is quite different from what's commonly used (Transformers), so there's probably quite a bit to do to make it compatible with the ecosystem.
Some of the operations are implemented in CUDA (selective scan), idk if that would work with ONNX.
The causal_conv1d has a CUDA implementation but you can also just call pytorch.
Triton is used for RMS norm I think, and you can replace that with a reference Pytorch implementation.

from mamba.

I was also interested in a CPU implementation; I was curious if the functions suffixed with _ref are reference implementations (like this one https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311) that can be used to guide new implementations?

In the past I had quite some success porting CUDA/OpenCL kernels to CPU using ISPC (which has very similar programming model) but I saw that C10 & CUB are used here and I don't know them very well.

from mamba.

Modified version to work on CPU is available on this fork

I chopped a lot of code, but it works (for inference only)

from mamba.

Congratulations! But I have a question that why the 130M model doesn't preform well in training, only 15% accuracy. Am I doing something wrong?

from mamba.