--- license: apache-2.0 library_name: pytorch tags: - state-space-models - s4 - mamba - sequence-modeling - pytorch - benchmark --- # torch-dimensions **An N-dimensional model is a 1-D mixer plus a plan for sweeping it over an N-D lattice.** State-space models, RNNs, transformers and convolutions over arbitrary lattices, behind one API. - **Code:** https://github.com/triadastra/torch-dimensions - **PyPI:** `pip install torch-dimensions` - **License:** Apache-2.0, © 2026 Celsia Juilyn Fan ```python import torch_dimensions as td model = td.S4ND(d_model=64, n_layers=12, dim=2, shape=(32, 32)) loss = model(x).pow(2).mean() # x: (B, T, 32, 32, 64) loss.backward() # ordinary autograd; nothing custom to call ``` `td.S4`, `td.S4D`, `td.Mamba`, `td.Mamba2` and `td.Mamba3` run the **original authors' code**, redistributed verbatim under `torch_dimensions/_vendor/` with every patched line tagged and a manifest pinning upstream commits and sha256s. `portable=True` selects a pure-torch build with a smaller dependency footprint. --- ## What is in this repository Sixteen models trained identically on three devices, plus the artifacts needed to check that the comparison means anything. | directory | contents | |---|---| | `CUDA bench/` | 16 checkpoints trained on an **NVIDIA RTX 5090** (Blackwell, sm_120; torch 2.12.1+cu130), plus `cuda_check.txt` | | `MPS bench/` | the same 16 on an **Apple Mac Studio, M1 Ultra** (Metal / MPS; torch 2.13.0) | | `CPU bench/` | the same 16 on that Mac's **CPU** (arm64, torch 2.13.0) — the same-machine control for MPS | | `CUDA agree/`, `MPS agree/`, `CPU agree/` | one forward and one backward from fixed weights, **no optimiser** — outputs and every gradient, saved as tensors | | `init weights/` | the single set of starting weights every device loads | | `docs/device-comparison.png` | the figure below, and its `.svg` | | `AGREEMENT.md` | do the two devices compute the same thing? | | `COMPARISON.md` | do they train to the same place, and how fast? | | `*/SCORECARD.md` | per-machine model ranking, one column per question, no combined score | | `BENCH-README.md`, `BENCHMARK-DESIGN.md` | the design, and what the comparison cannot show | Each checkpoint is `/weights.pt` (a `{"model": ..., "head": ...}` state dict) beside `/metrics.json`. ### Loading a checkpoint ```python import json, torch, torch.nn as nn, torch_dimensions as td lat = td.Lattice(shape=(6, 8), names=("h", "w"), valid=..., time=True) model = td.S4D(32, 4, lat, d_input=1, d_state=16) head = nn.Linear(32, 1) blob = torch.load("CUDA bench/s4d_upstream_2d/weights.pt", weights_only=True) model.load_state_dict(blob["model"]) # strict; tested in CI head.load_state_dict(blob["head"]) ``` The exact lattice and constructor for every entry is in [`benchmarks/pretrain.py`](https://github.com/triadastra/torch-dimensions/blob/main/benchmarks/pretrain.py)'s `ZOO`, and `tests/test_published_checkpoints.py` loads all 32 with `strict=True` on every commit — these weights are a compatibility contract, not a snapshot. **These are not pretrained models to fine-tune.** They are small models (12k–141k parameters) trained 300 steps on one synthetic task — a cumulative sum along a lattice axis — for the purpose of comparing devices and mechanisms. Use them to reproduce the comparison, not as a starting point. --- ## Evaluation and Device Comparison ![Evaluation and device comparison across CUDA, CPU and MPS](docs/device-comparison.png) Sixteen models, three devices, two benchmarks, from one shared set of starting weights on identical data. Regenerate with `python benchmarks/figure.py` — every panel reads artifact directories that already exist, so the figure cannot drift from `AGREEMENT.md` and `COMPARISON.md`. | device | hardware | torch | artifacts | |---|---|---|---| | **CUDA** | NVIDIA RTX 5090 (Blackwell, sm_120) | 2.12.1+cu130 | `CUDA bench/`, `CUDA agree/` | | **MPS** | Apple Mac Studio, M1 Ultra (Metal) | 2.13.0 | `MPS bench/`, `MPS agree/` | | **CPU** | Apple M1 Ultra (arm64) | 2.13.0 | `CPU bench/`, `CPU agree/` | **CPU and MPS are the same machine and the same torch build**, so a difference between them is the device and nothing else. CPU-vs-CUDA crosses machines and torch versions. Reading the pairs together separates what the hardware did from what the software version did — neither pair alone can. ### What the throughput panel shows Three findings, none of which is "the GPU is faster": | model | CPU | MPS | CUDA | CUDA vs CPU | |---|---|---|---|---| | `mamba_upstream_3d` | 0.1 | 26.8 | 23.7 | **455×** | | `mamba2_2d` | 0.1 | 5.6 | 24.3 | **285×** | | `mamba_upstream_2d` | 0.2 | 40.1 | 23.8 | 95× | | `tcn_2d_sparse` | 5.2 | 90.5 | 87.3 | 17× | | `cnn_2d_sparse` | 12.7 | 106.5 | 104.4 | 8× | | `mamba3_2d` | 5.3 | 14.6 | 21.0 | 3.9× | | `s4_upstream_2d` | 39.4 | 39.0 | 28.3 | **0.7×** | | `s4d_upstream_2d` | 42.7 | 46.9 | 32.7 | **0.8×** | *training steps per second, 300 steps, batch 4* **The upstream reference scans are the bottleneck, not the hardware.** Mamba-1 and Mamba-2 fall to 0.1–0.2 steps/s on CPU because `selective_scan_ref` and `ssd_minimal_discrete` are sequential Python loops over sequence length. That is what the fused CUDA kernels exist to avoid, and it is why a CPU pass over this zoo takes hours while a GPU pass takes minutes. **Our Mamba-3 transcription is 30–50× faster on CPU than the authors' Mamba-1 and Mamba-2 references** (5.3 steps/s against 0.1–0.2), because `mixers/mamba3_compat.py` computes the recurrence as a chunked, vectorised scan rather than a Python loop. Mamba-3 ships Triton-only upstream, so there was no reference loop to copy — the transcription had to be written, and it was written vectorised. **The RTX 5090 loses to an M1 Ultra CPU on the small S4 models** (0.7–0.8×). These are 13k–17k-parameter models on a 6×8 lattice: the axial sweep issues many small sequential kernels, launch overhead dominates every one of them, and the GPU never fills. That is a property of this benchmark's size rather than of the hardware — a defensible throughput number needs the warmup-and-repeats protocol designed in `BENCHMARK-DESIGN.md` and not yet implemented. ## Results ### Do the two devices compute the same thing? **Worst float32 output difference across every device pair: 3.11e-06.** Thirteen of sixteen models are at or below **1e-06**. In float64 the devices agree to **2.2e-16** — the last bit. Full per-model table in `AGREEMENT.md`. Two corrections had to happen before that number meant anything, and both generalise to anyone benchmarking across devices: **CUDA does not run float32 by default.** `torch.backends.cudnn.allow_tf32` ships as `True`, so cuDNN convolutions and RNNs execute in TF32 — 10 mantissa bits against float32's 23. That accounted for the *entire* original 1.96e-04 gap, and explains why only LSTM, GRU, CNN and TCN were affected while attention, Mamba and S4 sat at ~2e-07 throughout (`matmul.allow_tf32` ships as `False`). | model | TF32 on (torch's default) | TF32 off | + cuDNN off | |---|---|---|---| | `tcn_2d_sparse` | 1.16e-04 | **1.13e-07** | — | | `cnn_2d_sparse` | 1.96e-04 | **2.34e-07** | — | | `gru_2d_sparse` | 1.82e-04 | 3.13e-06 | **1.50e-07** | | `lstm_2d_sparse` | 1.21e-04 | 2.66e-06 | **1.73e-07** | The residual for the three RNNs is not precision but *algorithm* — cuDNN's fused LSTM/GRU is a different implementation. That is left visible rather than configured away, because the fused kernel is what runs. **A fixed seed does not give identical S4 weights across platforms.** `hippo.nplr` diagonalises with `torch.linalg.eigh`. Eigenvalues are unique and matched across macOS and Linux to twelve decimals — which is exactly why this hid, since `A_imag` looked perfect. Eigenvectors are fixed only up to a phase, so `B` and `P` differed by a relative **1.5** and **0.53**: two valid S4 initialisations, two different models. The comparison was reporting 2.6e-01 for the vendored S4D, unchanged in float64. With shared weights the same pair agrees at 4e-07, and after 300 optimiser steps their losses agree to 8.9e-07. **One gradient cannot meet a float32 bound, by construction.** The derivative with respect to an SSM frequency sums oscillating terms that nearly cancel: `A_imag` differs by **1.35e-04 in float32 and 4.87e-15 in float64**. Eleven orders between the two is what cancellation looks like; a genuinely different computation does not shrink when mantissa bits are added. **MPS has no float64 at all**, so on Apple silicon float32 is the only column there is. ### CUDA verification `scripts/cuda_check.py` runs every CUDA claim the library makes as one file. On the RTX 5090: **13 passed, 0 failed, 1 skipped**, and the full test suite **1211 passed, 0 failed**. Highlights: - `prefer_upstream` returns **True** on real hardware — the per-call dispatch between the authors' fused kernels and the portable path, verified for the first time. - The vendored **S4 (DPLR)** agrees with CPU at **1.9e-07 including L=64**, where MPS lands exactly on the Nyquist pole — so the guard that makes S4 work on Metal is provably inert on CUDA. - The rank-1 LSTM is still **bitwise identical** to its pre-norm residual under cuDNN. **Still not established:** Mamba-3's PyTorch transcription has never been compared against the Triton kernel it was transcribed from. `mamba-ssm` has no sm_120 wheel and does not build against CUDA 13, so even on the 5090 the fused Mamba entry points were never importable. The transcription is validated three other ways — against an independently written recurrent form (3.6e-15 in float64), a direct sum in the `trap → 1` limit (2.2e-16), and gradcheck — but not against its own kernel. An Ampere or Ada card would close it. ### Reproducing ```bash pip install "torch-dimensions[dev,upstream]" # on each machine, once python benchmarks/agreement.py --out "MPS agree" --init "init weights" python benchmarks/pretrain.py --out "MPS bench" --init "init weights" # then, anywhere python benchmarks/compare_agreement.py "MPS agree" "CUDA agree" --out AGREEMENT.md python benchmarks/compare.py "MPS bench" "CUDA bench" --out COMPARISON.md python benchmarks/figure.py --out docs/device-comparison.png python scripts/cuda_check.py ``` Pass the **same** `--init` directory on every machine: the seed alone is not enough, for the reason above. --- ## Attribution Full third-party attribution is in [`NOTICE`](https://github.com/triadastra/torch-dimensions/blob/main/NOTICE). In short: the files under `_vendor/` are redistributed verbatim from [state-spaces/s4](https://github.com/state-spaces/s4) and [state-spaces/mamba](https://github.com/state-spaces/mamba) (both Apache-2.0) with their licenses alongside and every patched line tagged; the portable mixer mathematics derives from the same two repositories; [CaFA](https://github.com/BaratiLab/CaFA) (MIT) contributed two kernel ideas, independently implemented. `ScanPlan.paired()` produces the same schedule as [Mamba-ND](https://github.com/jacklishufan/Mamba-ND)'s reference and was verified against it — that repository states **no license**, so no code from it is used or redistributed here.