mamba3-ssm 0.1.0__tar.gz

mamba3_ssm-0.1.0/LICENSE

@@ -0,0 +1,20 @@
+MIT License
+
+Copyright (c) 2026 Aiyoniganmaaiya
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

mamba3_ssm-0.1.0/PKG-INFO

@@ -0,0 +1,164 @@
+Metadata-Version: 2.4
+Name: mamba3-ssm
+Version: 0.1.0
+Summary: Mamba-3: Improved Sequence Modeling using State Space Principles
+Home-page: https://github.com/Aiyoniganmaaiya/mamba3-ssm
+Author: Aiyoniganmaaiya
+Author-email: Aiyoniganmaaiya <gerintacc004@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/Aiyoniganmaaiya/mamba3-ssm
+Project-URL: Repository, https://github.com/Aiyoniganmaaiya/mamba3-ssm
+Keywords: mamba,ssm,state-space-model,transformer,pytorch
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Requires-Dist: einops>=0.7
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: author
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+# Mamba-3: Improved Sequence Modeling using State Space Principles
+
+A clean, readable, from-scratch PyTorch implementation of **Mamba-3** — a selective state space model that addresses three core limitations of Mamba-2. No Triton/CUDA kernels; designed for understanding and reproducing the algorithm.
+
+**Paper:** [Mamba-3: Improved Sequence Modeling using State Space Principles](https://arxiv.org/abs/2603.15569)
+**Authors:** Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu
+
+## Installation
+
+```bash
+pip install mamba3-ssm
+```
+
+Or install from source:
+
+```bash
+pip install git+https://github.com/Aiyoniganmaaiya/mamba3-ssm.git
+```
+
+## Quick Start
+
+```python
+import torch
+from mamba3_ssm import Mamba3, MambaLMHeadModel, MambaConfig
+
+# ── SISO mode (standard) ──────────────────────────────
+model = Mamba3(
+    d_model=256,
+    d_state=64,
+    expand=2,
+    headdim=32,
+    is_mimo=False,
+)
+x = torch.randn(2, 128, 256)
+y = model(x)  # (2, 128, 256)
+
+# ── MIMO mode (better decode efficiency) ──────────────
+model_mimo = Mamba3(
+    d_model=256,
+    d_state=64,
+    expand=2,
+    headdim=32,
+    is_mimo=True,
+    mimo_rank=4,
+)
+y = model_mimo(x)  # same I/O shape
+
+# ── Autoregressive decode (one token at a time) ───────
+angle_state, ssm_state, bx_prev = model.allocate_inference_cache(batch_size=2)
+u = torch.randn(2, 256)
+out, angle_state, ssm_state, bx_prev = model.step(
+    u, angle_state, ssm_state, bx_prev
+)
+
+# ── Full language model ───────────────────────────────
+cfg = MambaConfig(
+    d_model=2048,
+    n_layer=24,
+    vocab_size=50277,
+    ssm_cfg={"is_mimo": True, "mimo_rank": 4},
+)
+lm = MambaLMHeadModel(cfg)
+logits = lm(torch.randint(0, 50277, (1, 512)))  # (1, 512, vocab_size)
+```
+
+## Core Ideas
+
+### 1. Exponential-Trapezoidal Discretization
+
+Mamba-2 used Zero-Order Hold (exponential-Euler), a first-order approximation. Mamba-3 adopts the **trapezoidal rule**, averaging the `B*x` contribution at times `t-1` and `t`:
+
+```
+h_t = exp(A·dt_t) · h_{t-1} + dt_t · σ(trap_t) · (B_t·x_t + B_{t-1}·x_{t-1}) / 2
+```
+
+`trap` is a learned sigmoid gate blending between Euler (`trap≈0`) and full trapezoidal (`trap≈1`).
+
+### 2. Complex-Valued (Rotary) State Space
+
+Real-valued SSM hidden states cannot easily represent oscillatory patterns. Mamba-3 applies **RoPE** to B and C projections, giving the state an effective complex-valued structure that tracks phase-dependent dependencies.
+
+### 3. Multi-Input Multi-Output (MIMO) Formulation
+
+Mamba-2 is SISO with state `(H, P, D)` — during decode the GPU is memory-bandwidth bound. **MIMO** reuses a shared `(H, D)` state for `R` rank streams, multiplying FLOPs/byte by `R`:
+
+| | SISO | MIMO |
+|---|---|---|
+| State shape | `(H, P, D)` | `(H, D)` |
+| Update | outer product `x ⊗ B` | sum of R rank-1 terms |
+
+## Project Structure
+
+```
+mamba3_ssm/
+├── __init__.py   # Public API
+├── config.py     # MambaConfig / SSMConfig dataclasses
+├── ops.py        # RMSNorm, RoPE, SSM scan (SISO + MIMO)
+├── layer.py      # Mamba3 module (forward + step + inference cache)
+├── block.py      # MambaBlock, MLPBlock, MambaLMHeadModel
+├── tests.py      # 10 sanity checks
+└── utils.py      # Parameter counting
+```
+
+## Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `d_model` | — | Token embedding dimension |
+| `d_state` | 128 | SSM state size per head (D) |
+| `expand` | 2 | Inner dim multiplier; `d_inner = expand * d_model` |
+| `headdim` | 64 | Features per SSM head (P) |
+| `is_mimo` | False | Enable MIMO formulation |
+| `mimo_rank` | 4 | Number of parallel MIMO streams (R) |
+| `rope_fraction` | 0.5 | Fraction of state dims that rotate |
+
+## Testing
+
+```bash
+python -m mamba3_ssm.tests
+```
+
+10/10 sanity checks pass, including shape tests, numerical consistency (step-by-step decode matches forward), gradient flow, and edge cases.
+
+## Dependencies
+
+- `torch>=2.0`
+- `einops>=0.7`
+
+## License
+
+MIT
+
+## References
+
+- Lahoti et al., *Mamba-3: Improved Sequence Modeling using State Space Principles*, 2026. [arXiv:2603.15569](https://arxiv.org/abs/2603.15569)
+- Official implementation: [state-spaces/mamba](https://github.com/state-spaces/mamba)

mamba3_ssm-0.1.0/README.md

@@ -0,0 +1,136 @@
+# Mamba-3: Improved Sequence Modeling using State Space Principles
+
+A clean, readable, from-scratch PyTorch implementation of **Mamba-3** — a selective state space model that addresses three core limitations of Mamba-2. No Triton/CUDA kernels; designed for understanding and reproducing the algorithm.
+
+**Paper:** [Mamba-3: Improved Sequence Modeling using State Space Principles](https://arxiv.org/abs/2603.15569)
+**Authors:** Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu
+
+## Installation
+
+```bash
+pip install mamba3-ssm
+```
+
+Or install from source:
+
+```bash
+pip install git+https://github.com/Aiyoniganmaaiya/mamba3-ssm.git
+```
+
+## Quick Start
+
+```python
+import torch
+from mamba3_ssm import Mamba3, MambaLMHeadModel, MambaConfig
+
+# ── SISO mode (standard) ──────────────────────────────
+model = Mamba3(
+    d_model=256,
+    d_state=64,
+    expand=2,
+    headdim=32,
+    is_mimo=False,
+)
+x = torch.randn(2, 128, 256)
+y = model(x)  # (2, 128, 256)
+
+# ── MIMO mode (better decode efficiency) ──────────────
+model_mimo = Mamba3(
+    d_model=256,
+    d_state=64,
+    expand=2,
+    headdim=32,
+    is_mimo=True,
+    mimo_rank=4,
+)
+y = model_mimo(x)  # same I/O shape
+
+# ── Autoregressive decode (one token at a time) ───────
+angle_state, ssm_state, bx_prev = model.allocate_inference_cache(batch_size=2)
+u = torch.randn(2, 256)
+out, angle_state, ssm_state, bx_prev = model.step(
+    u, angle_state, ssm_state, bx_prev
+)
+
+# ── Full language model ───────────────────────────────
+cfg = MambaConfig(
+    d_model=2048,
+    n_layer=24,
+    vocab_size=50277,
+    ssm_cfg={"is_mimo": True, "mimo_rank": 4},
+)
+lm = MambaLMHeadModel(cfg)
+logits = lm(torch.randint(0, 50277, (1, 512)))  # (1, 512, vocab_size)
+```
+
+## Core Ideas
+
+### 1. Exponential-Trapezoidal Discretization
+
+Mamba-2 used Zero-Order Hold (exponential-Euler), a first-order approximation. Mamba-3 adopts the **trapezoidal rule**, averaging the `B*x` contribution at times `t-1` and `t`:
+
+```
+h_t = exp(A·dt_t) · h_{t-1} + dt_t · σ(trap_t) · (B_t·x_t + B_{t-1}·x_{t-1}) / 2
+```
+
+`trap` is a learned sigmoid gate blending between Euler (`trap≈0`) and full trapezoidal (`trap≈1`).
+
+### 2. Complex-Valued (Rotary) State Space
+
+Real-valued SSM hidden states cannot easily represent oscillatory patterns. Mamba-3 applies **RoPE** to B and C projections, giving the state an effective complex-valued structure that tracks phase-dependent dependencies.
+
+### 3. Multi-Input Multi-Output (MIMO) Formulation
+
+Mamba-2 is SISO with state `(H, P, D)` — during decode the GPU is memory-bandwidth bound. **MIMO** reuses a shared `(H, D)` state for `R` rank streams, multiplying FLOPs/byte by `R`:
+
+| | SISO | MIMO |
+|---|---|---|
+| State shape | `(H, P, D)` | `(H, D)` |
+| Update | outer product `x ⊗ B` | sum of R rank-1 terms |
+
+## Project Structure
+
+```
+mamba3_ssm/
+├── __init__.py   # Public API
+├── config.py     # MambaConfig / SSMConfig dataclasses
+├── ops.py        # RMSNorm, RoPE, SSM scan (SISO + MIMO)
+├── layer.py      # Mamba3 module (forward + step + inference cache)
+├── block.py      # MambaBlock, MLPBlock, MambaLMHeadModel
+├── tests.py      # 10 sanity checks
+└── utils.py      # Parameter counting
+```
+
+## Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `d_model` | — | Token embedding dimension |
+| `d_state` | 128 | SSM state size per head (D) |
+| `expand` | 2 | Inner dim multiplier; `d_inner = expand * d_model` |
+| `headdim` | 64 | Features per SSM head (P) |
+| `is_mimo` | False | Enable MIMO formulation |
+| `mimo_rank` | 4 | Number of parallel MIMO streams (R) |
+| `rope_fraction` | 0.5 | Fraction of state dims that rotate |
+
+## Testing
+
+```bash
+python -m mamba3_ssm.tests
+```
+
+10/10 sanity checks pass, including shape tests, numerical consistency (step-by-step decode matches forward), gradient flow, and edge cases.
+
+## Dependencies
+
+- `torch>=2.0`
+- `einops>=0.7`
+
+## License
+
+MIT
+
+## References
+
+- Lahoti et al., *Mamba-3: Improved Sequence Modeling using State Space Principles*, 2026. [arXiv:2603.15569](https://arxiv.org/abs/2603.15569)
+- Official implementation: [state-spaces/mamba](https://github.com/state-spaces/mamba)

mamba3_ssm-0.1.0/mamba3_ssm/__init__.py

@@ -0,0 +1,23 @@
+"""
+Mambo-3: Improved Sequence Modeling using State Space Principles
+
+Clean PyTorch implementation of the Mamba-3 SSM architecture.
+Paper: https://arxiv.org/abs/2603.15569
+"""
+
+from .config import MambaConfig, SSMConfig
+from .layer import Mamba3
+from .block import MambaLMHeadModel, MambaBlock
+from .ops import RMSNorm, apply_rope, ssm_scan_siso, ssm_scan_mimo
+
+__all__ = [
+    "MambaConfig",
+    "SSMConfig",
+    "Mamba3",
+    "MambaLMHeadModel",
+    "MambaBlock",
+    "RMSNorm",
+    "apply_rope",
+    "ssm_scan_siso",
+    "ssm_scan_mimo",
+]

mamba3_ssm-0.1.0/mamba3_ssm/block.py

@@ -0,0 +1,101 @@
+"""
+Building blocks for stacking Mamba-3 layers into a full model:
+  - MambaBlock:    RMSNorm → Mamba3 → residual add
+  - MLPBlock:      SwiGLU feed-forward (optional)
+  - Embed / LMHead: Token embedding and language-model head
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass, field
+
+from .ops import RMSNorm
+from .layer import Mamba3
+from .config import MambaConfig
+
+
+class MambaBlock(nn.Module):
+    """Single Mamba-3 residual block: Norm → Mixer → residual."""
+
+    def __init__(self, d_model: int, ssm_cfg: dict, device=None, dtype=None):
+        super().__init__()
+        self.norm = RMSNorm(d_model)
+        self.mixer = Mamba3(d_model=d_model, **ssm_cfg, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.mixer(self.norm(x))
+
+
+class MLPBlock(nn.Module):
+    """SwiGLU-style feed-forward layer. Used when d_intermediate > 0."""
+
+    def __init__(self, d_model: int, d_intermediate: int, device=None, dtype=None):
+        super().__init__()
+        factory = {"device": device, "dtype": dtype}
+        self.fc1 = nn.Linear(d_model, 2 * d_intermediate, bias=False, **factory)
+        self.fc2 = nn.Linear(d_intermediate, d_model, bias=False, **factory)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, val = self.fc1(x).chunk(2, dim=-1)
+        return self.fc2(F.silu(gate) * val)
+
+
+class MambaLMHeadModel(nn.Module):
+    """Stacked Mamba-3 language model.
+
+    Architecture:
+        Embedding → n_layer × (MambaBlock [+ MLPBlock]) → RMSNorm → LMHead
+
+    The LM head weight can be tied to the embedding weight.
+    Vocab size is padded up to the nearest multiple of pad_vocab_size_multiple.
+    """
+
+    def __init__(self, config: MambaConfig, device=None, dtype=None):
+        factory = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.config = config
+
+        # Pad vocab size
+        vocab_size = config.vocab_size
+        r = vocab_size % config.pad_vocab_size_multiple
+        if r != 0:
+            vocab_size += config.pad_vocab_size_multiple - r
+        self.vocab_size = vocab_size
+
+        self.embedding = nn.Embedding(vocab_size, config.d_model, **factory)
+
+        self.layers = nn.ModuleList([
+            MambaBlock(config.d_model, config.ssm_cfg, **factory)
+            for _ in range(config.n_layer)
+        ])
+
+        if config.d_intermediate > 0:
+            self.mlp_norms = nn.ModuleList(
+                [RMSNorm(config.d_model) for _ in range(config.n_layer)]
+            )
+            self.mlp_layers = nn.ModuleList([
+                MLPBlock(config.d_model, config.d_intermediate, **factory)
+                for _ in range(config.n_layer)
+            ])
+        else:
+            self.mlp_norms = None
+            self.mlp_layers = None
+
+        self.norm_f = RMSNorm(config.d_model)
+        # LM head: projects from vocab_size to d_model (tied with embedding)
+        self.lm_head = nn.Linear(vocab_size, config.d_model, bias=False, **factory)
+        if config.tie_embeddings:
+            self.lm_head.weight = self.embedding.weight
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Args: input_ids (batch, seq_len) → logits (batch, seq_len, vocab_size)"""
+        x = self.embedding(input_ids)
+        for i, block in enumerate(self.layers):
+            x = block(x)
+            if self.mlp_layers is not None:
+                x = x + self.mlp_layers[i](self.mlp_norms[i](x))
+        x = self.norm_f(x)
+        return self.lm_head(x)

mamba3_ssm-0.1.0/mamba3_ssm/config.py

@@ -0,0 +1,61 @@
+"""
+Configuration dataclasses for Mamba-3 models.
+
+Paper reference: arXiv.2603.15569
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any
+
+
+@dataclass
+class SSMConfig:
+    """Configuration for the SSM core inside Mamba-3."""
+
+    # State dimensions
+    d_state: int = 128        # D: SSM state size per head
+    expand: int = 2            # Inner dimension multiplier; d_inner = expand * d_model
+    headdim: int = 64          # P: per-head feature dimension; nheads = d_inner / headdim
+    ngroups: int = 1           # G: number of groups for B/C projection sharing
+
+    # RoPE
+    rope_fraction: float = 0.5  # Fraction of state dims that rotate (0.5 or 1.0)
+
+    # Time step bounds
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dt_init_floor: float = 1e-4
+    A_floor: float = 1e-4
+
+    # MIMO
+    is_mimo: bool = False
+    mimo_rank: int = 4         # R: number of MIMO streams
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            k: v for k, v in self.__dict__.items()
+        }
+
+
+@dataclass
+class MambaConfig:
+    """Full model configuration for Mamba-3 language models."""
+
+    d_model: int = 2560
+    d_intermediate: int = 0    # >0 adds an MLP sub-layer after each SSM block
+    n_layer: int = 64
+    vocab_size: int = 50277
+
+    # SSM config (passed to each Mamba3 layer)
+    ssm_cfg: Dict[str, Any] = field(default_factory=dict)
+
+    # Norm / residual options
+    rms_norm: bool = True
+    residual_in_fp32: bool = True
+    fused_add_norm: bool = True
+    pad_vocab_size_multiple: int = 8
+    tie_embeddings: bool = True
+
+    # Layer configuration
+    attn_layer_idx: List[int] = field(default_factory=list)
+    attn_cfg: Dict[str, Any] = field(default_factory=dict)