From e62f113948147274262e1e3d84f774d19eed0288 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 30 Jul 2026 15:19:18 +0200 Subject: [PATCH 1/6] layers for mamba2 --- mamba-ssm/torch-ext/mamba_ssm/layers.py | 127 +++++++++++++++++++++++- 1 file changed, 126 insertions(+), 1 deletion(-) diff --git a/mamba-ssm/torch-ext/mamba_ssm/layers.py b/mamba-ssm/torch-ext/mamba_ssm/layers.py index fc410c276..3f65beee3 100644 --- a/mamba-ssm/torch-ext/mamba_ssm/layers.py +++ b/mamba-ssm/torch-ext/mamba_ssm/layers.py @@ -1,7 +1,126 @@ +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 + + +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, + **kwargs, + ): + return cuda_selective_state_update( + state, + hidden_states, + dt, + A, + B, + C, + D, + z=None, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + ) class causal_conv1d_fn(nn.Module): @@ -43,4 +162,10 @@ def forward( ) -__all__ = ["causal_conv1d_fn", "causal_conv1d_update"] +__all__ = [ + "causal_conv1d_fn", + "causal_conv1d_update", + "mamba_split_conv1d_scan_combined", + "mamba_chunk_scan_combined", + "selective_state_update", +] From 9870de7b5e16de08eead3cc527cea2af12b8d556 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 30 Jul 2026 16:12:51 +0200 Subject: [PATCH 2/6] fix multi gpu issue --- .../mamba_ssm/ops/triton/ssd_chunk_scan.py | 90 ++++++++++--------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_chunk_scan.py b/mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_chunk_scan.py index f4e65ad10..d155e9e95 100644 --- a/mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_chunk_scan.py +++ b/mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_chunk_scan.py @@ -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 @@ -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 From 95315195f7cc451b48466c23877ea3bf4485b7ac Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 30 Jul 2026 16:25:52 +0200 Subject: [PATCH 3/6] allow z for falcon h1 --- mamba-ssm/torch-ext/mamba_ssm/layers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mamba-ssm/torch-ext/mamba_ssm/layers.py b/mamba-ssm/torch-ext/mamba_ssm/layers.py index 3f65beee3..bfb96485a 100644 --- a/mamba-ssm/torch-ext/mamba_ssm/layers.py +++ b/mamba-ssm/torch-ext/mamba_ssm/layers.py @@ -107,6 +107,7 @@ def forward( 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( @@ -117,7 +118,7 @@ def forward( B, C, D, - z=None, + z=z, dt_bias=dt_bias, dt_softplus=dt_softplus, ) From 3ceaa3a3f7427f717ec38ce8a3326aeaefc5c2f5 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 30 Jul 2026 20:09:41 +0200 Subject: [PATCH 4/6] add mamba1 kernels --- mamba-ssm/torch-ext/mamba_ssm/layers.py | 71 +++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/mamba-ssm/torch-ext/mamba_ssm/layers.py b/mamba-ssm/torch-ext/mamba_ssm/layers.py index bfb96485a..ad72bf25f 100644 --- a/mamba-ssm/torch-ext/mamba_ssm/layers.py +++ b/mamba-ssm/torch-ext/mamba_ssm/layers.py @@ -6,6 +6,43 @@ 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, + **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, + ) class mamba_split_conv1d_scan_combined(nn.Module): @@ -124,6 +161,38 @@ def forward( ) +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): def forward( self, @@ -166,7 +235,9 @@ def forward( __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", ] From f84814a1d4a1c4f71d005d1f48fbd44c503fcb58 Mon Sep 17 00:00:00 2001 From: vasqu Date: Mon, 3 Aug 2026 22:52:25 +0200 Subject: [PATCH 5/6] add falcon mamba support --- mamba-ssm/torch-ext/mamba_ssm/layers.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mamba-ssm/torch-ext/mamba_ssm/layers.py b/mamba-ssm/torch-ext/mamba_ssm/layers.py index ad72bf25f..a53d73c91 100644 --- a/mamba-ssm/torch-ext/mamba_ssm/layers.py +++ b/mamba-ssm/torch-ext/mamba_ssm/layers.py @@ -26,6 +26,10 @@ def forward( 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( @@ -42,6 +46,10 @@ def forward( 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, ) From 7b019b052037ecbbd1d37794f62036cc19e636a8 Mon Sep 17 00:00:00 2001 From: vasqu Date: Tue, 4 Aug 2026 13:14:55 +0200 Subject: [PATCH 6/6] bump version --- mamba-ssm/build.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mamba-ssm/build.toml b/mamba-ssm/build.toml index b52829306..a7edd5a55 100644 --- a/mamba-ssm/build.toml +++ b/mamba-ssm/build.toml @@ -1,6 +1,6 @@ [general] name = "mamba-ssm" -version = 1 +version = 2 license = "Apache-2.0" backends = ["cuda"] python-depends = ["einops"]