Skip to content

mamba-ssm: add layers #1051

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vasqu merged 6 commits into from
Aug 4, 2026
Merged

mamba-ssm: add layers #1051

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e62f113
layers for mamba2
vasqu Jul 30, 2026
9870de7
fix multi gpu issue
vasqu Jul 30, 2026
9531519
allow z for falcon h1
vasqu Jul 30, 2026
3ceaa3a
add mamba1 kernels
vasqu Jul 30, 2026
f84814a
add falcon mamba support
vasqu Aug 3, 2026
7b019b0
bump version
vasqu Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
207 changes: 206 additions & 1 deletion mamba-ssm/torch-ext/mamba_ssm/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,204 @@
import torch
import torch.nn as nn

from ._causal_conv1d import causal_conv1d_fn as cuda_causal_conv1d_fn
from ._causal_conv1d import causal_conv1d_update as cuda_causal_conv1d_update
from .ops import mamba_chunk_scan_combined as cuda_mamba_chunk_scan_combined
from .ops import mamba_split_conv1d_scan_combined as cuda_mamba_split_conv1d_scan_combined
from .ops import selective_state_update as cuda_selective_state_update
from .ops.selective_scan_interface import mamba_inner_fn as cuda_mamba_inner_fn
from .ops.selective_scan_interface import selective_scan_fn as cuda_selective_scan_fn


class mamba_inner_fn(nn.Module):
def forward(
self,
xz: torch.Tensor,
conv1d_weight: torch.Tensor,
conv1d_bias: torch.Tensor | None,
x_proj_weight: torch.Tensor,
delta_proj_weight: torch.Tensor,
out_proj_weight: torch.Tensor,
out_proj_bias: torch.Tensor | None,
A: torch.Tensor,
B: torch.Tensor | None = None,
C: torch.Tensor | None = None,
D: torch.Tensor | None = None,
delta_bias: torch.Tensor | None = None,
delta_softplus: bool = True,
b_rms_weight: torch.Tensor | None = None,
c_rms_weight: torch.Tensor | None = None,
dt_rms_weight: torch.Tensor | None = None,
b_c_dt_rms_eps: float = 1e-6,
**kwargs,
):
return cuda_mamba_inner_fn(
xz,
conv1d_weight,
conv1d_bias,
x_proj_weight,
delta_proj_weight,
out_proj_weight,
out_proj_bias,
A,
B,
C,
D,
delta_bias=delta_bias,
delta_softplus=delta_softplus,
b_rms_weight=b_rms_weight,
c_rms_weight=c_rms_weight,
dt_rms_weight=dt_rms_weight,
b_c_dt_rms_eps=b_c_dt_rms_eps,
)


class mamba_split_conv1d_scan_combined(nn.Module):
def forward(
self,
zxbcdt: torch.Tensor,
conv1d_weight: torch.Tensor,
conv1d_bias: torch.Tensor | None,
dt_bias: torch.Tensor,
A: torch.Tensor,
D: torch.Tensor,
chunk_size: int,
initial_states: torch.Tensor | None = None,
dt_limit: tuple[float, float] = (0.0, float("inf")),
return_final_states: bool = False,
activation: str = "silu",
rmsnorm_weight: torch.Tensor | None = None,
rmsnorm_eps: float = 1e-6,
outproj_weight: torch.Tensor | None = None,
outproj_bias: torch.Tensor | None = None,
headdim: int | None = None,
ngroups: int = 1,
norm_before_gate: bool = True,
**kwargs,
):
# For varlen
seq_idx = kwargs.pop("seq_idx", None)

return cuda_mamba_split_conv1d_scan_combined(
zxbcdt,
conv1d_weight,
conv1d_bias,
dt_bias,
A,
D=D,
chunk_size=chunk_size,
seq_idx=seq_idx,
activation=activation,
rmsnorm_weight=rmsnorm_weight,
rmsnorm_eps=rmsnorm_eps,
outproj_weight=outproj_weight,
outproj_bias=outproj_bias,
headdim=headdim,
ngroups=ngroups,
norm_before_gate=norm_before_gate,
return_final_states=return_final_states,
dt_limit=dt_limit,
initial_states=initial_states,
)


class mamba_chunk_scan_combined(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
dt: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
chunk_size: int,
D: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
initial_states: torch.Tensor | None = None,
dt_softplus: bool = False,
dt_limit: tuple[float, float] = (0.0, float("inf")),
return_final_states: bool = False,
**kwargs,
):
# For varlen
seq_idx = kwargs.pop("seq_idx", None)

return cuda_mamba_chunk_scan_combined(
hidden_states,
dt,
A,
B,
C,
D=D,
z=None,
dt_bias=dt_bias,
dt_softplus=dt_softplus,
chunk_size=chunk_size,
seq_idx=seq_idx,
return_final_states=return_final_states,
dt_limit=dt_limit,
initial_states=initial_states,
)


class selective_state_update(nn.Module):
def forward(
self,
state: torch.Tensor,
hidden_states: torch.Tensor,
dt: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
dt_softplus: bool = False,
z: torch.Tensor | None = None,
**kwargs,
):
return cuda_selective_state_update(
state,
hidden_states,
dt,
A,
B,
C,
D,
z=z,
dt_bias=dt_bias,
dt_softplus=dt_softplus,
)


class selective_scan_fn(nn.Module):
def forward(
self,
hidden_states: torch.Tensor,
dt: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
z: torch.Tensor | None = None,
delta_bias: torch.Tensor | None = None,
delta_softplus: bool = False,
return_last_state: bool = False,
# Unused here but fallbacks for torch only paths
use_mambapy: bool = False,
use_associative_scan: bool = False,
**kwargs,
):
return cuda_selective_scan_fn(
hidden_states,
dt,
A,
B,
C,
D=D,
z=z,
delta_bias=delta_bias,
delta_softplus=delta_softplus,
return_last_state=return_last_state,
)


class causal_conv1d_fn(nn.Module):
Expand Down Expand Up @@ -43,4 +240,12 @@ def forward(
)


__all__ = ["causal_conv1d_fn", "causal_conv1d_update"]
__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",
]
Comment on lines +243 to +251

@danieldk danieldk Aug 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

90 changes: 46 additions & 44 deletions mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_chunk_scan.py
View file
Open in desktop

@vasqu vasqu Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Original file line number Diff line number Diff line change
Expand Up @@ -1259,28 +1259,29 @@ def _chunk_scan_fwd(cb, x, dt, dA_cumsum, C, states, D=None, z=None, seq_idx=Non
batch * nchunks, nheads)
z_strides = ((z.stride(0), z.stride(1), z.stride(2), z.stride(3))
if z is not None else (0, 0, 0, 0))
_chunk_scan_fwd_kernel[grid](
cb, x, z, out, out_x, dt, dA_cumsum, seq_idx, C, states, D,
chunk_size, headdim, dstate,
batch, seqlen, nheads // ngroups,
cb.stride(0), cb.stride(1), cb.stride(2), cb.stride(3), cb.stride(4),
x.stride(0), x.stride(1), x.stride(2), x.stride(3),
z_strides[0], z_strides[1], z_strides[2], z_strides[3],
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
dt.stride(0), dt.stride(2), dt.stride(1), dt.stride(3),
dA_cumsum.stride(0), dA_cumsum.stride(2), dA_cumsum.stride(1), dA_cumsum.stride(3),
*((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
C.stride(0), C.stride(1), C.stride(2), C.stride(3),
states.stride(0), states.stride(1), states.stride(2), states.stride(3), states.stride(4),
D.stride(0) if D is not None else 0,
True,
D is not None,
D.dim() == 2 if D is not None else True,
BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
HAS_Z=z is not None,
HAS_SEQ_IDX=seq_idx is not None,
IS_TRITON_22=TRITON_22,
)
with torch.cuda.device(x.device.index):
_chunk_scan_fwd_kernel[grid](
cb, x, z, out, out_x, dt, dA_cumsum, seq_idx, C, states, D,
chunk_size, headdim, dstate,
batch, seqlen, nheads // ngroups,
cb.stride(0), cb.stride(1), cb.stride(2), cb.stride(3), cb.stride(4),
x.stride(0), x.stride(1), x.stride(2), x.stride(3),
z_strides[0], z_strides[1], z_strides[2], z_strides[3],
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
dt.stride(0), dt.stride(2), dt.stride(1), dt.stride(3),
dA_cumsum.stride(0), dA_cumsum.stride(2), dA_cumsum.stride(1), dA_cumsum.stride(3),
*((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
C.stride(0), C.stride(1), C.stride(2), C.stride(3),
states.stride(0), states.stride(1), states.stride(2), states.stride(3), states.stride(4),
D.stride(0) if D is not None else 0,
True,
D is not None,
D.dim() == 2 if D is not None else True,
BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
HAS_Z=z is not None,
HAS_SEQ_IDX=seq_idx is not None,
IS_TRITON_22=TRITON_22,
)
return out, out_x


Expand Down Expand Up @@ -1311,28 +1312,29 @@ def _chunk_scan_fwd_wip(cb, x, dt, dA_cumsum, C, B, states, D=None, z=None, seq_
grid = lambda META: (triton.cdiv(headdim, META['BLOCK_SIZE_N']), batch * nchunks, nheads)
z_strides = ((z.stride(0), z.stride(1), z.stride(2), z.stride(3))
if z is not None else (0, 0, 0, 0))
_chunk_scan_fwd_kernel_wip[grid](
cb, x, z, out, out_x, dt, dA_cumsum, seq_idx, C, B, states, D,
chunk_size, headdim, dstate,
batch, seqlen, nheads // ngroups,
cb.stride(0), cb.stride(1), cb.stride(2), cb.stride(3), cb.stride(4),
x.stride(0), x.stride(1), x.stride(2), x.stride(3),
z_strides[0], z_strides[1], z_strides[2], z_strides[3],
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
dt.stride(0), dt.stride(2), dt.stride(1), dt.stride(3),
dA_cumsum.stride(0), dA_cumsum.stride(2), dA_cumsum.stride(1), dA_cumsum.stride(3),
*((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
C.stride(0), C.stride(1), C.stride(2), C.stride(3),
B.stride(0), B.stride(1), B.stride(2), B.stride(3),
states.stride(0), states.stride(1), states.stride(2), states.stride(3), states.stride(4),
D.stride(0) if D is not None else 0,
D is not None,
D.dim() == 2 if D is not None else True,
BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
BLOCK_SIZE_M=128,
HAS_Z=z is not None,
HAS_SEQ_IDX=seq_idx is not None,
)
with torch.cuda.device(x.device.index):
_chunk_scan_fwd_kernel_wip[grid](
cb, x, z, out, out_x, dt, dA_cumsum, seq_idx, C, B, states, D,
chunk_size, headdim, dstate,
batch, seqlen, nheads // ngroups,
cb.stride(0), cb.stride(1), cb.stride(2), cb.stride(3), cb.stride(4),
x.stride(0), x.stride(1), x.stride(2), x.stride(3),
z_strides[0], z_strides[1], z_strides[2], z_strides[3],
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
dt.stride(0), dt.stride(2), dt.stride(1), dt.stride(3),
dA_cumsum.stride(0), dA_cumsum.stride(2), dA_cumsum.stride(1), dA_cumsum.stride(3),
*((seq_idx.stride(0), seq_idx.stride(1)) if seq_idx is not None else (0, 0)),
C.stride(0), C.stride(1), C.stride(2), C.stride(3),
B.stride(0), B.stride(1), B.stride(2), B.stride(3),
states.stride(0), states.stride(1), states.stride(2), states.stride(3), states.stride(4),
D.stride(0) if D is not None else 0,
D is not None,
D.dim() == 2 if D is not None else True,
BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
BLOCK_SIZE_M=128,
HAS_Z=z is not None,
HAS_SEQ_IDX=seq_idx is not None,
)
return out, out_x


Expand Down
Loading