A few questions about this Pallas mamba kernel

[davisyoshida]

(@sharadmv)

I made an attempt at implementing the Mamba block forward pass in Pallas here, but as you can see I had to drop down into Triton a bunch of times. A few questions:

1. I wrote some custom lowering rules for bitcast_convert_type and convert_element_type (to handle uint values), would you guys like a PR with support for these?
2. Not necessary for this version of the kernel, but is there any chance of getting strided loads supported? When I was trying to implement slicing I had to write custom lowering rules to do the pointer math for tl.load.
3. Any general advice on ways this could be made less hacky?

[sharadmv]

Thanks a lot for pushing on this Davis! I'm super excited about Mamba and your progress.

I wrote some custom lowering rules for bitcast_convert_type and convert_element_type (to handle uint values), would you guys like a PR with support for these?

Of course, that would be amazing! Just ask that you write a couple tests in pallas_test.py.

Not necessary for this version of the kernel, but is there any chance of getting strided loads supported? When I was trying to implement slicing I had to write custom lowering rules to do the pointer math for tl.load.

Yea we should support this. Let's perhaps open a new thread about the best way to surface these types of loads.

Any general advice on ways this could be made less hacky?

Re: the mamba_scan primitive. Is is possible to repurpose jax.lax.scan for this? @mattjj. Not sure if we need a proper associative_scan primitive for this (even if it's just used in Pallas). Also, re AD. I think you're best off writing a backwards pass pallas kernel for this. I think ADing the pallas kernel is going to be tricky (currently only fwd mode works and only pretty simple kernels will AD efficiently once transposes are supported).

[davisyoshida]

I just meant adding the Triton lowering for bitcast_convert_type and the uint types. I don't know if it's worth you guys' time to unblock me on running the pallas tests on my setup.

[sharadmv]

I'm assuming if you can't run the tests, then you can't run other things though. 😔

[davisyoshida]

Well, the version set you posted the other day lets me use pallas, it's just not up to date.

[superbobry]

Can you share the error you are getting from XLA?

[davisyoshida]

I think that particular error might have been fixed when I upgraded my CUDA pip packages. I'm still having trouble getting a consistent version set though.

At least part of the problem seems to be that jax-triton needs an older version of triton that JAX is expecting at HEAD. For example, jax-triton tries to import get_kernel_name from triton.compiler.backends.cuda, which was removed triton-lang/triton@8f64d4c, but it also tries to use triton._C.libtriton.ir.load_dialects, which wasn't added until triton-lang/triton@7d3f045.