Why isn't selective_scan implemented in Triton? about mamba

[tridao]

Triton does have associative scan, but (1) it only supports scan with 1 float instead of a pair of floats, which is what we need (2) it doesn't have reverse scan yet (needed for the backward pass).
Does Pallas have those?

When Triton has more robust support for scan, it would be simpler to write it in Triton.

[albertfgu]

Triton did not support some of the operations. In particular their associative scan function is quite limited right now (and sometimes buggy / difficult to use); on the other hand CUDA has quite sophisticated support for scans (https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda).

I think it should be possible in principle to implement something similar in Pallas, but I don't know how efficient it would be.

[vmajor]

Ideally it would be nice to be truly hardware agnostic, in the spirit of the Mamba mission statement. Tying the architecture to CUDA may not be the best way.

[0xc1c4da]

Thanks for the response @tridao & @albertfgu, indeed Triton doesn't have the support, and Pallas wouldn't be much use here as it lowers to Triton. We'll work on a JAX interface to the CUDA kernels.