mamba-ssm: add layers - #1051
Conversation
There was a problem hiding this comment.
Triton corruption may happen otherwise; similar to liger a few days/weeks back
A lot of ppl did different fixes but never the true underlying fix
| __all__ = [ | ||
| "causal_conv1d_fn", | ||
| "causal_conv1d_update", | ||
| "mamba_inner_fn", | ||
| "mamba_split_conv1d_scan_combined", | ||
| "mamba_chunk_scan_combined", | ||
| "selective_state_update", | ||
| "selective_scan_fn", | ||
| ] |
There was a problem hiding this comment.
Adding functions on a compiled kernels (so not no-arch like torch-cuda) requires a version bump, since otherwise, there will be old variants (Torch 2.10, Torch 2.11) that do not have these functions.
|
/kernel-bot merge-and-upload mamba-ssm |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
Hub uploads: |
3 participants
As per title WIP, adding layer variations for mamba2 and mamba1 so we can deprecate lazy load kernels
Also fixes a multi gpu issue under triton (the same we had under liger), ref. state-spaces/mamba#1003