# Mamba SSM Block

> Mamba State Space Model — selective SSM + causal conv gating, no attention (O(T) complexity)

Pick for very long sequences where attention's O(T²) cost is the bottleneck (DNA, audio, long-context LM). Tuning is trickier than transformers.

- Category: NLP/LLM
- Layers: 18
- Parameters: 168.27M
- Input shape (batchless): 1 × 1024
- Output shape: 1 × 1024 × 50280
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/mamba-block/model.json
- Open on the canvas: https://neurarch.com/?template=mamba-block

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | tokens | Input | shape=[1, 1024] | 1 × 1024 |
| 2 | embed | Embedding |  | 1 × 1024 × 1024 |
| 3 | norm_ssm | RMSNorm | normalizedShape=1024 | 1 × 1024 × 1024 |
| 4 | in_proj | Linear | outFeatures=4096 | 1 × 1024 × 4096 |
| 5 | to_channels | Permute |  | 1 × 4096 × 1024 |
| 6 | causal_conv | Conv1D | outChannels=4096, kernelSize=3, stride=1 | 1 × 4096 × 1024 |
| 7 | to_tokens | Permute |  | 1 × 1024 × 4096 |
| 8 | silu_x | Swish |  | 1 × 1024 × 4096 |
| 9 | ssm_scan | Mamba (SSM) |  | 1 × 1024 × 4096 |
| 10 | z_gate | Swish |  | 1 × 1024 × 4096 |
| 11 | gate_out | Multiply |  | 1 × 1024 × 4096 |
| 12 | out_proj | Linear | outFeatures=1024 | 1 × 1024 × 1024 |
| 13 | residual_1 | Add |  | 1 × 1024 × 1024 |
| 14 | norm_ffn | RMSNorm | normalizedShape=1024 | 1 × 1024 × 1024 |
| 15 | ffn | SwiGLU | embedDim=1024, intermediateSize=2048 | 1 × 1024 × 1024 |
| 16 | residual_2 | Add |  | 1 × 1024 × 1024 |
| 17 | lm_head | Linear | outFeatures=50280 | 1 × 1024 × 50280 |
| 18 | output | Output |  | 1 × 1024 × 50280 |

## Verifier findings

No finding. Shapes propagate end to end and no advisory rule fires.

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)
#
# WARNING: 2 layer(s) below are not yet supported by the PyTorch
# exporter and pass their input through UNCHANGED in forward():
#   - to_channels (permute)
#   - to_tokens (permute)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class MambaSSMBlock(nn.Module):
    def __init__(self):
        super().__init__()

        self.embedding_1 = nn.Embedding(50280, 1024)
        self.rmsNorm_1 = nn.RMSNorm(1024)
        self.linear_1 = nn.Linear(1048576, 4096)
        self.conv1d_1 = nn.Conv1d(4096, 4096, kernel_size=3, stride=1, padding=1)
        self.swish_1 = nn.SiLU()
        self.mamba_1 = nn.Identity()  # Mamba(d_model=4096, d_state=16, d_conv=4, expand=1), pip install mamba-ssm and swap in
        self.swish_2 = nn.SiLU()
        self.linear_2 = nn.Linear(4194304, 1024)
        self.rmsNorm_2 = nn.RMSNorm(1024)
        self.swiglu_1 = nn.ModuleDict({
            'gate_proj': nn.Linear(1024, 2048, bias=False),
            'up_proj':   nn.Linear(1024, 2048, bias=False),
            'down_proj': nn.Linear(2048, 1024, bias=False),
        })  # SwiGLU FFN (LLaMA-style)
        self.linear_3 = nn.Linear(1048576, 50280)

    def forward(self, x):
        # tokens shape: [1,1024]
        embedding_embed = self.embedding_1(x)
        rms_norm_norm1 = self.rmsNorm_1(embedding_embed)
        linear_n_proj = self.linear_1(rms_norm_norm1)
        # TODO: layer 'to_channels' (permute) is not yet supported by the exporter; passing through unchanged
        conv1d_l_conv = self.conv1d_1(linear_n_proj)
        # TODO: layer 'to_tokens' (permute) is not yet supported by the exporter; passing through unchanged
        swish_silu_x = self.swish_1(conv1d_l_conv)
        mamba_m_proj = self.mamba_1(swish_silu_x)
        swish_z_gate = self.swish_2(linear_n_proj)
        multiply_ltiply = mamba_m_proj * swish_z_gate
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md