lite-mamba 0.1.4__tar.gz

lite_mamba-0.1.4/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.4
+Name: lite-mamba
+Version: 0.1.4
+Summary: Pure-PyTorch lightweight Mamba with multi-dilated causal conv front-end
+Home-page: https://github.com/Mrrobi/lite_mamba
+Author: Md Robiuddin
+Author-email: mrrobi040@gmail.com
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/Mrrobi/lite_mamba
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: einops>=0.6
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# lite-mamba
+
+A minimal, pure-PyTorch version of Mamba with a multi-dilated causal depthwise conv front-end. No CUDA/Triton build needed; works on CPU or GPU with standard PyTorch ops.
+
+## Install
+```bash
+pip install torch einops
+pip install lite-mamba
+```
+
+## Usage
+```python
+from lite_mamba import Mamba
+import torch
+
+x = torch.randn(2, 128, 512)  # (batch, seq, d_model)
+m = Mamba(d_model=512, d_conv=3, conv_dilations=(1,2,4,8))
+y = m(x)
+print(y.shape)  # (2, 128, 512)
+```
+
+## API quick reference
+`Mamba(d_model, d_state=16, d_conv=4, conv_dilations=(1,), expand=2, dt_rank="auto", dt_min=0.001, dt_max=0.1, dt_init="random", dt_scale=1.0, dt_init_floor=1e-4, conv_bias=True, bias=False, use_fast_path=False, layer_idx=None, device=None, dtype=None)`
+
+- `d_model` (int, required): input/output embedding size.
+- `d_state` (int, default 16): SSM state dimension per channel. Larger gives longer memory; increases compute.
+- `d_conv` (int, default 4): depthwise conv kernel size for each branch.
+- `conv_dilations` (tuple[int], default `(1,)`): dilation per branch. Multiple values create parallel dilated convs; effective receptive field is `(d_conv-1)*dilation`.
+- `expand` (float, default 2): inner width multiplier; sets `d_inner = expand * d_model`.
+- `dt_rank` (int or "auto", default "auto"): rank of delta projection. "auto" sets `ceil(d_model/16)`.
+- `dt_min`, `dt_max` (float, defaults 1e-3 / 1e-1): log-uniform range for delta initialization.
+- `dt_init` ("random" | "constant", default "random") and `dt_scale`, `dt_init_floor`: control delta init magnitude/stability.
+- `conv_bias` (bool, default True): include bias in depthwise convs.
+- `bias` (bool, default False): include bias in input/output linear projections.
+- `use_fast_path` (bool): ignored in this pure-PyTorch build; kept for API compatibility.
+- `layer_idx` (int | None): identifier for streaming cache registration; required when using `allocate_inference_cache` + `inference_params`.
+- `device`, `dtype`: standard module factory kwargs.
+
+### Inference / streaming helpers
+- `allocate_inference_cache(batch_size, max_seqlen, dtype=None)`: preallocates conv and SSM state buffers for step-wise decoding.
+- `step(hidden_states, conv_state, ssm_state)`: single-token forward (expects `hidden_states` with shape `(B, 1, d_model)`).
+- `forward(..., inference_params)`: if `inference_params` has cached states (with `key_value_memory_dict` and `seqlen_offset`), uses them for streaming.
+
+## Highlights
+- Multi-branch causal dilated convs (weighted sum via learned gates).
+- Pure Python: no custom C++/CUDA or Triton kernels.
+- Streaming support via per-branch conv states and SSM state caching.
+
+## Practical setups
+- **Local modeling / small context**: `d_conv=3`, `conv_dilations=(1,2,4)`, `d_state=8–16`, `expand=2`.
+- **Longer context**: widen `conv_dilations` (e.g., `(1,2,4,8,16)`) or increase `d_state` to 32; expect higher memory/compute.
+- **Streaming/AR decoding**: call `allocate_inference_cache` once per layer, pass `inference_params` during forward; use `step` inside your generation loop.
+- **Stability first**: keep `dt_min` >= 1e-4 and `dt_init_floor` small; leave defaults unless you observe drift or exploding activations.
+
+## Notes
+- Set different `conv_dilations` to adjust receptive field; keep kernels small (e.g., 3–5) to avoid excessive padding.
+- `use_fast_path` flag is ignored here (kept for API compatibility).
+- Reference selective scan is implemented in PyTorch for portability; faster fused kernels are omitted intentionally.
+
+## License
+Apache-2.0

lite_mamba-0.1.4/README.md

@@ -0,0 +1,61 @@
+# lite-mamba
+
+A minimal, pure-PyTorch version of Mamba with a multi-dilated causal depthwise conv front-end. No CUDA/Triton build needed; works on CPU or GPU with standard PyTorch ops.
+
+## Install
+```bash
+pip install torch einops
+pip install lite-mamba
+```
+
+## Usage
+```python
+from lite_mamba import Mamba
+import torch
+
+x = torch.randn(2, 128, 512)  # (batch, seq, d_model)
+m = Mamba(d_model=512, d_conv=3, conv_dilations=(1,2,4,8))
+y = m(x)
+print(y.shape)  # (2, 128, 512)
+```
+
+## API quick reference
+`Mamba(d_model, d_state=16, d_conv=4, conv_dilations=(1,), expand=2, dt_rank="auto", dt_min=0.001, dt_max=0.1, dt_init="random", dt_scale=1.0, dt_init_floor=1e-4, conv_bias=True, bias=False, use_fast_path=False, layer_idx=None, device=None, dtype=None)`
+
+- `d_model` (int, required): input/output embedding size.
+- `d_state` (int, default 16): SSM state dimension per channel. Larger gives longer memory; increases compute.
+- `d_conv` (int, default 4): depthwise conv kernel size for each branch.
+- `conv_dilations` (tuple[int], default `(1,)`): dilation per branch. Multiple values create parallel dilated convs; effective receptive field is `(d_conv-1)*dilation`.
+- `expand` (float, default 2): inner width multiplier; sets `d_inner = expand * d_model`.
+- `dt_rank` (int or "auto", default "auto"): rank of delta projection. "auto" sets `ceil(d_model/16)`.
+- `dt_min`, `dt_max` (float, defaults 1e-3 / 1e-1): log-uniform range for delta initialization.
+- `dt_init` ("random" | "constant", default "random") and `dt_scale`, `dt_init_floor`: control delta init magnitude/stability.
+- `conv_bias` (bool, default True): include bias in depthwise convs.
+- `bias` (bool, default False): include bias in input/output linear projections.
+- `use_fast_path` (bool): ignored in this pure-PyTorch build; kept for API compatibility.
+- `layer_idx` (int | None): identifier for streaming cache registration; required when using `allocate_inference_cache` + `inference_params`.
+- `device`, `dtype`: standard module factory kwargs.
+
+### Inference / streaming helpers
+- `allocate_inference_cache(batch_size, max_seqlen, dtype=None)`: preallocates conv and SSM state buffers for step-wise decoding.
+- `step(hidden_states, conv_state, ssm_state)`: single-token forward (expects `hidden_states` with shape `(B, 1, d_model)`).
+- `forward(..., inference_params)`: if `inference_params` has cached states (with `key_value_memory_dict` and `seqlen_offset`), uses them for streaming.
+
+## Highlights
+- Multi-branch causal dilated convs (weighted sum via learned gates).
+- Pure Python: no custom C++/CUDA or Triton kernels.
+- Streaming support via per-branch conv states and SSM state caching.
+
+## Practical setups
+- **Local modeling / small context**: `d_conv=3`, `conv_dilations=(1,2,4)`, `d_state=8–16`, `expand=2`.
+- **Longer context**: widen `conv_dilations` (e.g., `(1,2,4,8,16)`) or increase `d_state` to 32; expect higher memory/compute.
+- **Streaming/AR decoding**: call `allocate_inference_cache` once per layer, pass `inference_params` during forward; use `step` inside your generation loop.
+- **Stability first**: keep `dt_min` >= 1e-4 and `dt_init_floor` small; leave defaults unless you observe drift or exploding activations.
+
+## Notes
+- Set different `conv_dilations` to adjust receptive field; keep kernels small (e.g., 3–5) to avoid excessive padding.
+- `use_fast_path` flag is ignored here (kept for API compatibility).
+- Reference selective scan is implemented in PyTorch for portability; faster fused kernels are omitted intentionally.
+
+## License
+Apache-2.0

lite_mamba-0.1.4/lite_mamba/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.0"
+
+from .mamba_simple import Mamba

lite_mamba-0.1.4/lite_mamba/causal_dilated_conv1d.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn.functional as F
+
+
+def _apply_activation(x, activation):
+    if activation is None or activation == "identity":
+        return x
+    if activation in ("silu", "swish"):
+        return F.silu(x)
+    if activation == "relu":
+        return F.relu(x)
+    raise ValueError(f"Unsupported activation: {activation}")
+
+
+def causal_dilated_conv1d_fn(x, weight, bias=None, activation=None, dilation=1):
+    """
+    Depthwise causal 1D convolution with dilation.
+    x: (B, D, L)
+    weight: (D, W)
+    bias: (D,) or None
+    """
+    if dilation < 1:
+        raise ValueError(f"dilation must be >= 1, got {dilation}")
+    pad = dilation * (weight.shape[-1] - 1)
+    x = F.pad(x, (pad, 0))
+    y = F.conv1d(
+        x,
+        weight.unsqueeze(1),
+        bias=bias,
+        stride=1,
+        padding=0,
+        dilation=dilation,
+        groups=weight.shape[0],
+    )
+    return _apply_activation(y, activation)
+
+
+def causal_dilated_conv1d_update(x, conv_state, weight, bias=None, activation=None, dilation=1):
+    """
+    Single-step causal dilated conv update.
+    x: (B, D)
+    conv_state: (B, D, S) where S = dilation * (W - 1) + 1
+    weight: (D, W)
+    bias: (D,) or None
+    """
+    if dilation < 1:
+        raise ValueError(f"dilation must be >= 1, got {dilation}")
+    conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))
+    conv_state[:, :, -1] = x
+
+    idx = torch.arange(0, weight.shape[-1], device=conv_state.device) * dilation
+    pos = conv_state.shape[-1] - 1 - idx
+    values = conv_state.index_select(-1, pos)
+    y = torch.sum(values * weight.unsqueeze(0), dim=-1)
+    if bias is not None:
+        y = y + bias
+    return _apply_activation(y, activation)

lite_mamba-0.1.4/lite_mamba/mamba_simple.py

@@ -0,0 +1,282 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from einops import rearrange, repeat
+
+from .selective_scan import selective_scan_fn
+from .causal_dilated_conv1d import causal_dilated_conv1d_fn, causal_dilated_conv1d_update
+
+
+class Mamba(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        d_state=16,
+        d_conv=4,
+        conv_dilations=(1,),
+        expand=2,
+        dt_rank="auto",
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init="random",
+        dt_scale=1.0,
+        dt_init_floor=1e-4,
+        conv_bias=True,
+        bias=False,
+        use_fast_path=False,  # kept for API; ignored in this pure-Python build
+        layer_idx=None,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.conv_dilations = tuple(conv_dilations)
+        self.num_conv_branches = len(self.conv_dilations)
+        self.conv_state_lens = [(self.d_conv - 1) * d + 1 for d in self.conv_dilations]
+        self.expand = expand
+        self.d_inner = int(self.expand * self.d_model)
+        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
+        self.use_fast_path = False
+        self.layer_idx = layer_idx
+
+        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)
+
+        self.conv1d_layers = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    in_channels=self.d_inner,
+                    out_channels=self.d_inner,
+                    bias=conv_bias,
+                    kernel_size=d_conv,
+                    groups=self.d_inner,
+                    padding=d * (d_conv - 1),
+                    dilation=d,
+                    **factory_kwargs,
+                )
+                for d in self.conv_dilations
+            ]
+        )
+        self.conv1d = self.conv1d_layers[0]
+        self.conv_gates = nn.Parameter(torch.ones(self.num_conv_branches))
+
+        self.activation = "silu"
+        self.act = nn.SiLU()
+
+        self.x_proj = nn.Linear(
+            self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
+        )
+        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True, **factory_kwargs)
+
+        dt_init_std = self.dt_rank**-0.5 * dt_scale
+        if dt_init == "constant":
+            nn.init.constant_(self.dt_proj.weight, dt_init_std)
+        elif dt_init == "random":
+            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+
+        dt = torch.exp(
+            torch.rand(self.d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(inv_dt)
+        self.dt_proj.bias._no_reinit = True
+
+        A = repeat(
+            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
+            "n -> d n",
+            d=self.d_inner,
+        ).contiguous()
+        A_log = torch.log(A)
+        self.A_log = nn.Parameter(A_log)
+        self.A_log._no_weight_decay = True
+
+        self.D = nn.Parameter(torch.ones(self.d_inner, device=device))
+        self.D._no_weight_decay = True
+
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+
+    def forward(self, hidden_states, inference_params=None):
+        batch, seqlen, dim = hidden_states.shape
+
+        conv_state, ssm_state = None, None
+        if inference_params is not None:
+            conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
+            if inference_params.seqlen_offset > 0:
+                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
+                return out
+
+        xz = rearrange(
+            self.in_proj.weight @ rearrange(hidden_states, "b l d -> d (b l)"),
+            "d (b l) -> b d l",
+            l=seqlen,
+        )
+        if self.in_proj.bias is not None:
+            xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), "d -> d 1")
+
+        A = -torch.exp(self.A_log.float())
+        use_multi_branch = self.num_conv_branches > 1
+
+        x, z = xz.chunk(2, dim=1)
+        # Convs
+        conv_outputs = []
+        if conv_state is not None:
+            for state, conv_layer, dilation in zip(conv_state, self.conv1d_layers, self.conv_dilations):
+                pad_len = state.shape[-1] - x.shape[-1]
+                if pad_len >= 0:
+                    state.copy_(F.pad(x, (pad_len, 0)))
+                else:
+                    state.copy_(x[..., -state.shape[-1] :])
+                if dilation > 1:
+                    xi = causal_dilated_conv1d_fn(
+                        x=x,
+                        weight=rearrange(conv_layer.weight, "d 1 w -> d w"),
+                        bias=conv_layer.bias,
+                        activation=self.activation,
+                        dilation=dilation,
+                    )
+                else:
+                    xi = self.act(conv_layer(x)[..., :seqlen])
+                conv_outputs.append(xi)
+        else:
+            for conv_layer, dilation in zip(self.conv1d_layers, self.conv_dilations):
+                if dilation > 1:
+                    xi = causal_dilated_conv1d_fn(
+                        x=x,
+                        weight=rearrange(conv_layer.weight, "d 1 w -> d w"),
+                        bias=conv_layer.bias,
+                        activation=self.activation,
+                        dilation=dilation,
+                    )
+                else:
+                    xi = self.act(conv_layer(x)[..., :seqlen])
+                conv_outputs.append(xi)
+        gate = torch.softmax(self.conv_gates, dim=0)
+        x = sum(g * xi for g, xi in zip(gate, conv_outputs))
+
+        x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))
+        dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
+        dt = self.dt_proj.weight @ dt.t()
+        dt = rearrange(dt, "d (b l) -> b d l", l=seqlen)
+        B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
+        C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
+        y = selective_scan_fn(
+            x,
+            dt,
+            A,
+            B,
+            C,
+            self.D.float(),
+            z=z,
+            delta_bias=self.dt_proj.bias.float(),
+            delta_softplus=True,
+            return_last_state=ssm_state is not None,
+        )
+        if ssm_state is not None:
+            y, last_state = y
+            ssm_state.copy_(last_state)
+        y = rearrange(y, "b d l -> b l d")
+        out = self.out_proj(y)
+        return out
+
+    def step(self, hidden_states, conv_state, ssm_state):
+        dtype = hidden_states.dtype
+        assert hidden_states.shape[1] == 1
+        xz = self.in_proj(hidden_states.squeeze(1))
+        x, z = xz.chunk(2, dim=-1)
+
+        conv_outputs = []
+        for state, conv_layer, dilation in zip(conv_state, self.conv1d_layers, self.conv_dilations):
+            if dilation > 1:
+                xi = causal_dilated_conv1d_update(
+                    x,
+                    state,
+                    rearrange(conv_layer.weight, "d 1 w -> d w"),
+                    conv_layer.bias,
+                    self.activation,
+                    dilation=dilation,
+                )
+            else:
+                state.copy_(torch.roll(state, shifts=-1, dims=-1))
+                state[:, :, -1] = x
+                xi = torch.sum(state * rearrange(conv_layer.weight, "d 1 w -> d w"), dim=-1)
+                if conv_layer.bias is not None:
+                    xi = xi + conv_layer.bias
+                xi = self.act(xi).to(dtype=dtype)
+            conv_outputs.append(xi)
+        gate = torch.softmax(self.conv_gates, dim=0)
+        x = sum(g * xi for g, xi in zip(gate, conv_outputs))
+
+        x_db = self.x_proj(x)
+        dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1)
+        dt = F.linear(dt, self.dt_proj.weight)
+        A = -torch.exp(self.A_log.float())
+
+        dt = F.softplus(dt + self.dt_proj.bias.to(dtype=dt.dtype))
+        dA = torch.exp(torch.einsum("bd,dn->bdn", dt, A))
+        dB = torch.einsum("bd,bn->bdn", dt, B)
+        ssm_state.copy_(ssm_state * dA + rearrange(x, "b d -> b d 1") * dB)
+        y = torch.einsum("bdn,bn->bd", ssm_state.to(dtype), C)
+        y = y + self.D.to(dtype) * x
+        y = y * self.act(z)
+
+        out = self.out_proj(y)
+        return out.unsqueeze(1), conv_state, ssm_state
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        device = self.out_proj.weight.device
+        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+        conv_state = [
+            torch.zeros(
+                batch_size,
+                self.d_model * self.expand,
+                state_len,
+                device=device,
+                dtype=conv_dtype,
+            )
+            for state_len in self.conv_state_lens
+        ]
+        ssm_dtype = self.dt_proj.weight.dtype if dtype is None else dtype
+        ssm_state = torch.zeros(
+            batch_size, self.d_model * self.expand, self.d_state, device=device, dtype=ssm_dtype
+        )
+        return conv_state, ssm_state
+
+    def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+        assert self.layer_idx is not None
+        if self.layer_idx not in inference_params.key_value_memory_dict:
+            conv_state = [
+                torch.zeros(
+                    batch_size,
+                    self.d_model * self.expand,
+                    state_len,
+                    device=self.conv1d.weight.device,
+                    dtype=self.conv1d.weight.dtype,
+                )
+                for state_len in self.conv_state_lens
+            ]
+            ssm_state = torch.zeros(
+                batch_size,
+                self.d_model * self.expand,
+                self.d_state,
+                device=self.dt_proj.weight.device,
+                dtype=self.dt_proj.weight.dtype,
+            )
+            inference_params.key_value_memory_dict[self.layer_idx] = (conv_state, ssm_state)
+        else:
+            conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_idx]
+            if initialize_states:
+                for s in conv_state:
+                    s.zero_()
+                ssm_state.zero_()
+        return conv_state, ssm_state

lite_mamba-0.1.4/lite_mamba/selective_scan.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+# Pure-PyTorch reference selective scan (no CUDA/Triton deps)
+def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
+                     return_last_state=False):
+    return selective_scan_ref(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
+
+
+def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
+                      return_last_state=False):
+    dtype_in = u.dtype
+    u = u.float()
+    delta = delta.float()
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None].float()
+    if delta_softplus:
+        delta = F.softplus(delta)
+    batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
+    is_variable_B = B.dim() >= 3
+    is_variable_C = C.dim() >= 3
+    if A.is_complex():
+        if is_variable_B:
+            B = torch.view_as_complex(rearrange(B.float(), "... (L two) -> ... L two", two=2))
+        if is_variable_C:
+            C = torch.view_as_complex(rearrange(C.float(), "... (L two) -> ... L two", two=2))
+    else:
+        B = B.float()
+        C = C.float()
+    x = A.new_zeros((batch, dim, dstate))
+    ys = []
+    deltaA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))
+    if not is_variable_B:
+        deltaB_u = torch.einsum('bdl,dn,bdl->bdln', delta, B, u)
+    else:
+        if B.dim() == 3:
+            deltaB_u = torch.einsum('bdl,bnl,bdl->bdln', delta, B, u)
+        else:
+            B = repeat(B, "B G N L -> B (G H) N L", H=dim // B.shape[1])
+            deltaB_u = torch.einsum('bdl,bdnl,bdl->bdln', delta, B, u)
+    if is_variable_C and C.dim() == 4:
+        C = repeat(C, "B G N L -> B (G H) N L", H=dim // C.shape[1])
+    last_state = None
+    for i in range(u.shape[2]):
+        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
+        if not is_variable_C:
+            y = torch.einsum('bdn,dn->bd', x, C)
+        else:
+            if C.dim() == 3:
+                y = torch.einsum('bdn,bn->bd', x, C[:, :, i])
+            else:
+                y = torch.einsum('bdn,bdn->bd', x, C[:, :, :, i])
+        if i == u.shape[2] - 1:
+            last_state = x
+        if y.is_complex():
+            y = y.real * 2
+        ys.append(y)
+    y = torch.stack(ys, dim=2) # (batch dim L)
+    out = y if D is None else y + u * rearrange(D, "d -> d 1")
+    if z is not None:
+        out = out * F.silu(z)
+    out = out.to(dtype=dtype_in)
+    if return_last_state:
+        return out, last_state
+    return out

lite_mamba-0.1.4/lite_mamba.egg-info/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.4
+Name: lite-mamba
+Version: 0.1.4
+Summary: Pure-PyTorch lightweight Mamba with multi-dilated causal conv front-end
+Home-page: https://github.com/Mrrobi/lite_mamba
+Author: Md Robiuddin
+Author-email: mrrobi040@gmail.com
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/Mrrobi/lite_mamba
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: einops>=0.6
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# lite-mamba
+
+A minimal, pure-PyTorch version of Mamba with a multi-dilated causal depthwise conv front-end. No CUDA/Triton build needed; works on CPU or GPU with standard PyTorch ops.
+
+## Install
+```bash
+pip install torch einops
+pip install lite-mamba
+```
+
+## Usage
+```python
+from lite_mamba import Mamba
+import torch
+
+x = torch.randn(2, 128, 512)  # (batch, seq, d_model)
+m = Mamba(d_model=512, d_conv=3, conv_dilations=(1,2,4,8))
+y = m(x)
+print(y.shape)  # (2, 128, 512)
+```
+
+## API quick reference
+`Mamba(d_model, d_state=16, d_conv=4, conv_dilations=(1,), expand=2, dt_rank="auto", dt_min=0.001, dt_max=0.1, dt_init="random", dt_scale=1.0, dt_init_floor=1e-4, conv_bias=True, bias=False, use_fast_path=False, layer_idx=None, device=None, dtype=None)`
+
+- `d_model` (int, required): input/output embedding size.
+- `d_state` (int, default 16): SSM state dimension per channel. Larger gives longer memory; increases compute.
+- `d_conv` (int, default 4): depthwise conv kernel size for each branch.
+- `conv_dilations` (tuple[int], default `(1,)`): dilation per branch. Multiple values create parallel dilated convs; effective receptive field is `(d_conv-1)*dilation`.
+- `expand` (float, default 2): inner width multiplier; sets `d_inner = expand * d_model`.
+- `dt_rank` (int or "auto", default "auto"): rank of delta projection. "auto" sets `ceil(d_model/16)`.
+- `dt_min`, `dt_max` (float, defaults 1e-3 / 1e-1): log-uniform range for delta initialization.
+- `dt_init` ("random" | "constant", default "random") and `dt_scale`, `dt_init_floor`: control delta init magnitude/stability.
+- `conv_bias` (bool, default True): include bias in depthwise convs.
+- `bias` (bool, default False): include bias in input/output linear projections.
+- `use_fast_path` (bool): ignored in this pure-PyTorch build; kept for API compatibility.
+- `layer_idx` (int | None): identifier for streaming cache registration; required when using `allocate_inference_cache` + `inference_params`.
+- `device`, `dtype`: standard module factory kwargs.
+
+### Inference / streaming helpers
+- `allocate_inference_cache(batch_size, max_seqlen, dtype=None)`: preallocates conv and SSM state buffers for step-wise decoding.
+- `step(hidden_states, conv_state, ssm_state)`: single-token forward (expects `hidden_states` with shape `(B, 1, d_model)`).
+- `forward(..., inference_params)`: if `inference_params` has cached states (with `key_value_memory_dict` and `seqlen_offset`), uses them for streaming.
+
+## Highlights
+- Multi-branch causal dilated convs (weighted sum via learned gates).
+- Pure Python: no custom C++/CUDA or Triton kernels.
+- Streaming support via per-branch conv states and SSM state caching.
+
+## Practical setups
+- **Local modeling / small context**: `d_conv=3`, `conv_dilations=(1,2,4)`, `d_state=8–16`, `expand=2`.
+- **Longer context**: widen `conv_dilations` (e.g., `(1,2,4,8,16)`) or increase `d_state` to 32; expect higher memory/compute.
+- **Streaming/AR decoding**: call `allocate_inference_cache` once per layer, pass `inference_params` during forward; use `step` inside your generation loop.
+- **Stability first**: keep `dt_min` >= 1e-4 and `dt_init_floor` small; leave defaults unless you observe drift or exploding activations.
+
+## Notes
+- Set different `conv_dilations` to adjust receptive field; keep kernels small (e.g., 3–5) to avoid excessive padding.
+- `use_fast_path` flag is ignored here (kept for API compatibility).
+- Reference selective scan is implemented in PyTorch for portability; faster fused kernels are omitted intentionally.
+
+## License
+Apache-2.0

lite_mamba-0.1.4/lite_mamba.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+setup.py
+lite_mamba/__init__.py
+lite_mamba/causal_dilated_conv1d.py
+lite_mamba/mamba_simple.py
+lite_mamba/selective_scan.py
+lite_mamba.egg-info/PKG-INFO
+lite_mamba.egg-info/SOURCES.txt
+lite_mamba.egg-info/dependency_links.txt
+lite_mamba.egg-info/requires.txt
+lite_mamba.egg-info/top_level.txt

lite_mamba-0.1.4/lite_mamba.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

lite_mamba-0.1.4/lite_mamba.egg-info/requires.txt

@@ -0,0 +1,2 @@
+torch>=2.0
+einops>=0.6

lite_mamba-0.1.4/lite_mamba.egg-info/top_level.txt

@@ -0,0 +1 @@
+lite_mamba

lite_mamba-0.1.4/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"

lite_mamba-0.1.4/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

lite_mamba-0.1.4/setup.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+ROOT = Path(__file__).parent
+README = (ROOT / "README.md").read_text(encoding="utf-8")
+
+setup(
+    name="lite-mamba",
+    version="0.1.4",
+    description="Pure-PyTorch lightweight Mamba with multi-dilated causal conv front-end",
+    long_description=README,
+    long_description_content_type="text/markdown",
+    author="Md Robiuddin",
+    author_email="mrrobi040@gmail.com",
+    url="https://github.com/Mrrobi/lite_mamba",
+    project_urls={
+        "Homepage": "https://github.com/Mrrobi/lite_mamba",
+    },
+    packages=find_packages(exclude=("tests", "docs")),
+    python_requires=">=3.9",
+    install_requires=[
+        "torch>=2.0",
+        "einops>=0.6",
+    ],
+    license="Apache-2.0",
+    classifiers=[
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3 :: Only",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Intended Audience :: Science/Research",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)