Hi @gabrieldemarmiesse!

Thanks for taking a look at the performance of this kernel. Just for your reference and some general context about how this was built and some general constraints to keep in mind. When I built this kernel I did use https://github.com/Dao-AILab/causal-conv1d as a direct reference point and parity comparison for the output and performance of the kernel. It was developed on a single node H100, but it also was tuned on the MI350X series of GPUs. That being said I'm sure there's room for improvement and convergence on higher performance seeing there's been some updates to causal-conv1d as recently as a few weeks ago.

One feature I will note that is important for this kernel is support for both Channel First and Channel Last derivations (in the original CUDA they are similar but slightly different implementations).

The layout also took inspiration from Modular's conv1d when it was originally implemented, and if you look at what the modular team has done for their insanely fast conv kernels, you might find some additional improvements that I was not able to. Best of luck and I'm happy to leave a review when you move this out of draft.

Also please make sure you test the Mamba1 models before moving this out of draft to prevent any regressions 🙏

Thanks @ulmentflam for the context, related to that I have a question, there are a few functions that were added and never called anywhere:

  - causal_conv1d_channel_last_fwd_cpu
  - causal_conv1d_channel_last_fwd_cpu_no_bias
  - causal_conv1d_channel_last_fwd_cpu_with_seq_idx
  - causal_conv1d_channel_last_fwd_cpu_no_bias_with_seq_idx
  - causal_conv1d_channel_last_fwd_gpu
  - causal_conv1d_channel_last_fwd_gpu_no_bias
  - causal_conv1d_channel_last_fwd_gpu_with_seq_idx
  - causal_conv1d_channel_last_fwd_gpu_no_bias_with_seq_idx

  - causal_conv1d_channel_first_fwd_cpu_no_bias
  - causal_conv1d_channel_first_fwd_gpu_no_bias
  - causal_conv1d_channel_first_fwd_gpu_with_seq_idx
  - causal_conv1d_channel_first_fwd_gpu_no_bias_with_seq_idx

  - causal_conv1d_update_cpu_no_bias (called only in test/state_space + test/gpu/state_space)
  - causal_conv1d_update_gpu_no_bias (called only in test/gpu/state_space)

do you have more information about what they're used for? I don't see any call site. Maybe it was something in draft that was never used?

Yes, so they are not needed for Mamba1 and it seems they landed without their tests. They are actually for full parity with the forward pass kernels with causal-conv1d. I have an internal repo I was using that has a full kernel implementation of the causal-conv1d library that these functions were taken from for Mamba support. Let me grab what I can surface from those tests. It is correct that the update kernels are all that Mamba needs, which is the channel_first_fwd suite.

Thanks, I'll let you do the cleanup and wiring, since I don't know what functions I should touch. many of those functions have overlapping capabilities, and I guess maintainers don't want to maintain functions that are used only in unit tests. Ping me when they're wired to the rest of the codebase, and I can start from there! There is also massive code duplication, so if we could address that too, that would be great!

Will do. Most of the duplication was a forcing function of what needed to be included at compile time and issues with passing pointers which I believe have been resolved.

I just did a cleanup pass on selective scan kernel with some new (or new to me) syntax features and will take one here as well.

I've created #6632 that can serve as a reference. One point to note is that on the GB10 (Blackwell) this implementation is currently faster then Tri's implementation. However, that's an unfair comparison as the causal_conv1d kernel is only optimized for Datacenter Hopper. I have not tested it on any of the datacenter cards nor AMD for a fair comparison.