maba-v1-architecture

<p align="center"> <img src="https://huggingface.co/AndrewThompson1233/maba-v1-architecture/resolve/main/assets/logo.svg" width="160" alt="Maba Logo" /> </p>

Maba Architecture: Sub-Quadratic Hybrid Linear-Recurrent Attention

Official specification, scaling topology, and reference implementation of the Maba neural network architecture. Maba combines Gated DeltaNet linear recurrence (GDN-2) with Grouped-Query Attention (GQA), 2-pass physical block recycling, and native multi-token prediction (MTP) speculative decoding.

Note: For the sparse latent-attention iteration with million-token context, see Maba v2 Architecture.

Checkpoint weights and benchmark evaluations (ARC, HellaSwag, Story Cloze) are published at AndrewThompson1233/maba-101m.

---

<p align="center"> <img src="https://huggingface.co/AndrewThompson1233/maba-v1-architecture/resolve/main/assets/architecture_comparison.svg" width="900" alt="Maba Architecture Feature Comparison" /> </p>

---

Architectural Principles

Standard transformers scale at O(N^2) memory and compute with sequence length. Pure linear RNNs and state-space models scale at O(1) state memory but struggle with associative recall across long token horizons.

Maba resolves this trade-off through a 3:1 macro-interleaved block structure:

• 75% Linear Recurrent Blocks (GDN-2): updates an input-dependent recurrent state matrix in O(1) memory per step with dynamic gating.
• 25% Grouped-Query Attention Blocks (GQA): provides exact associative retrieval and cross-context routing at low KV-cache overhead (4:1 query-to-KV compression).
• 2-Pass Physical Block Recycling: passes token representations through 20 physical blocks twice with layer-index positional conditioning, producing 40 effective layers from a 20-block parameter budget.
• Integrated Multi-Token Prediction (MTP): built-in speculative heads (k=2) verify subsequent tokens during generation without requiring external companion models.

---

Exact Parameter & Memory Breakdown (101M Reference Model)

1. Parameter Accounting

2. Weight Memory Footprint by Precision

3. KV-Cache and Recurrent State Scaling

Maba separates state memory into constant recurrent state (GDN-2) and compressed quadratic attention cache (GQA 4:1):

• Note: GDN-2 recurrent state is strictly O(1) constant: 15 blocks x 10 heads x (64 x 64 state) x 2 bytes = 1.17 MB. It never grows, regardless of sequence length.

---

Architectural Comparison

---

Scaling Specifications (100M to 30B)

The Maba architecture scales systematically from on-device 100M to large-scale 30B parameters, supporting context horizons up to 131k tokens.

For complete multi-scale topology configurations, closed-form parameter derivation equations, and architectural audits against 2026 foundation models (Qwen3.5, Muse-30B, Gemma4), see the dedicated scaling specification:

SCALING.md

---

Block Architecture

1. Factorized Token Embeddings

To avoid vocabulary parameters consuming core computation capacity, Maba factorizes the embedding matrix:

• W_emb: V x d_emb (32,768 x 128)
• W_proj_in: d_emb x dim (128 x 640)
• W_proj_out: dim x d_emb (640 x 128) This reduces embedding parameters to 4.36M (4.31% of budget), leaving 95.21% of weights dedicated to sequence modeling.

2. GDN-2 Recurrence Block (75% of Layers)

The Gated DeltaNet layer computes an input-dependent recurrent update over state matrix S of size (d_head x d_head):

• 1D depthwise causal convolution over projected inputs (kernel size 4).
• Data-dependent decay gate beta_t = sigmoid(W_beta x_t + b_beta).
• State update: S_t = S_{t-1} * beta_t + v_t (x) k_t^T.
• State readout: o_t = S_t q_t.

3. GQA Quadratic Attention Block (25% of Layers)

Every fourth block (blocks 3, 7, 11, 15, 19) is a Grouped-Query Attention block:

• 10 query heads and 2 key-value heads.
• Per-head RMSNorm applied to query and key vectors before dot-product attention.
• Rotary Position Embedding (RoPE) with theta = 500,000.

4. Feed-Forward Network (SwiGLU)

Each block contains an intermediate SwiGLU FFN:

• FFN(x) = (SiLU(x W_gate) * x W_up) W_down
• dim = 640, d_ffn = 1,728.

5. Gated Residual Connections

Instead of fixed addition, residual streams use a learned gating mechanism:

• y = norm(x) + sigmoid(w_gate) * block(norm(x))
• Initialized with bias = 2.0 (sigmoid approx 0.88), ensuring stable gradient flow at initialization while allowing layers to dynamically regulate residual contribution.

---

Native C++ Inference Engine

The repository includes a standalone C++ inference implementation in cpp/:

• Cache-aligned unit-stride row-major loop order for GDN-2 state updates.
• AVX2 / FMA vectorization with zero heap allocations during autoregressive generation.
• Step latency: 10.20 us per head update on x86_64.
• Numerical parity with PyTorch: maximum logit discrepancy strictly below 7.62e-5.

Build instructions:

cd cpp
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
cmake --build . -j$(nproc)
./test_numerical

---

Hardware Acceleration & Distributed Training

The architecture includes automated device detection and distributed execution in maba/hardware.py:

• NVIDIA CUDA: Multi-GPU training via PyTorch Distributed Data Parallel (DDP) with NCCL all-reduce.
• Apple Silicon: Metal Performance Shaders (MPS) auto-detection and acceleration.
• Google Cloud TPU: PyTorch/XLA auto-detection and execution.
• x86_64 AVX2: Fast CPU fallback with OpenMP multi-threading.
• Hybrid Optimizer: Built-in Muon (matrix parameters via Newton-Schulz orthogonalization) and AdamW (vectors and embeddings).

Check hardware status:

python3 -m maba.cli hardware

---

Tests

Run tests:

pytest tests/

---

Quickstart (Python)

Installation

git clone https://github.com/ivan-dev35/maba-v1-architecture.git
cd maba-v1-architecture
pip install -e .

Initializing the Architecture

import torch
from maba.config import Config
from maba.model import Model

# Initialize 101M reference architecture
cfg = Config.from_preset("100M")
model = Model(cfg)

# Forward pass
tokens = torch.randint(0, cfg.vocab_size, (1, 64))
outputs = model(tokens)
logits = outputs["logits"]
print(f"Logits shape: {logits.shape}")  # [1, 64, 32768]

High-Speed Speculative Generation

from maba.generate import spec_gen
from maba.tokenizer import Tokenizer

tok = Tokenizer()
output_text, acceptance_rate, steps = spec_gen(
    model,
    tok,
    prompt="Artificial intelligence architecture design",
    max_new_tokens=64
)