Metadata-Version: 2.4
Name: mcr-attention
Version: 3.0.0
Summary: MCR-Attention V3.0: Multi-Scale Compressed Recurrent Attention — efficient autoregressive language modeling with linear-cost multi-scale recurrent memories and adaptive attention routing.
Author-email: "Gerson Fabian Buenahora Ormaza (BUEORM)" <dalusx64@gmail.com>
License-Expression: AGPL-3.0-or-later
Project-URL: Homepage, https://github.com/bueormnew/MCRA
Project-URL: Repository, https://github.com/bueormnew/MCRA
Project-URL: Documentation, https://github.com/bueormnew/MCRA#readme
Project-URL: Changelog, https://github.com/bueormnew/MCRA/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/bueormnew/MCRA/issues
Keywords: attention,recurrent,language-model,linear-attention,state-space,transformer,pytorch,deep-learning,mcr-attention,bueorm,mamba,rwkv
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: hypothesis>=6.0; extra == "dev"
Provides-Extra: benchmark
Requires-Dist: matplotlib>=3.5; extra == "benchmark"
Provides-Extra: hf
Requires-Dist: transformers>=4.35.0; extra == "hf"
Requires-Dist: huggingface-hub>=0.20.0; extra == "hf"
Provides-Extra: notebook
Requires-Dist: transformers>=4.35.0; extra == "notebook"
Requires-Dist: datasets>=2.0; extra == "notebook"
Requires-Dist: tiktoken>=0.5; extra == "notebook"
Requires-Dist: matplotlib>=3.5; extra == "notebook"
Provides-Extra: all
Requires-Dist: mcr-attention[benchmark,dev,hf]; extra == "all"
Dynamic: license-file

# MCR-Attention V3.0

**Multi-Scale Compressed Recurrent Attention** — an autoregressive language
model architecture that replaces quadratic self-attention with **K multi-scale
linear recurrences** and a tiny **1×K attention router**, stacked into
configurable depth.

> **Mission:** keep using attention — but simple and *linear everywhere*.
> **O(N) training**, **O(1) per-token streaming inference**, GPT-style
> autoregressive generation, optimized.

[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![License: AGPL-3.0-or-later](https://img.shields.io/badge/license-AGPL--3.0--or--later-red.svg)](./LICENSE)
[![Tests](https://img.shields.io/badge/tests-passing-brightgreen.svg)](./tests)

---

## Why MCR-Attention?

Transformers scale quadratically with sequence length. MCR-Attention keeps the
*idea* of attention (a learned query attending over context) but operates over
**K compressed recurrent memory vectors** instead of all N tokens. Each memory
scale is an independent linear recurrence with a learned decay, so:

| Mode | Time | Memory |
|------|------|--------|
| **Training** (parallel scan) | **O(N)** | O(N · layers · K · d_state) |
| **Streaming inference** | **O(1) per token** | O(layers · K · d_state) |

Measured on the bundled benchmark suite (CPU, d_model=128, 2 layers):

```
[1] Parallel training cost: time ratio 27x over 32x length   -> ~LINEAR  PASS
[2] Streaming latency: 1.49 -> 1.49 -> 1.49 -> 1.70 ms/token -> FLAT     PASS (O(1))
[3] Streaming memory: 1.0 KB constant                        -> FLAT     PASS (O(1))
[4] Train/inference identity: cosine 0.99999988              -> IDENTICAL PASS
```

## What's new in V3

V3 keeps the philosophy and fixes the real issues of V2:

- **Stacked blocks (depth)** — repeat the `{memory + router} + {FFN}` block
  `num_layers` times. Streaming stays O(1); capacity goes up.
- **Train == inference** — the parallel training forward and the streaming
  inference step compute *exactly the same math* (V2 diverged for long-memory
  scales).
- **Geometric decay init** — each scale starts with its target memory horizon.
- **O(N) scan** — a single fused `cumsum` (with a stable fallback).
- **Real HuggingFace integration** — `AutoModelForCausalLM` compatible.
- **V2 compatibility** — old models still load and run, bit-exact.

See [CHANGELOG.md](./CHANGELOG.md) for the full diff.

---

## Installation

```bash
pip install mcr-attention
```

Optional extras:

```bash
pip install "mcr-attention[dev]"        # pytest, hypothesis
pip install "mcr-attention[hf]"         # transformers + huggingface-hub
pip install "mcr-attention[benchmark]"  # matplotlib
pip install "mcr-attention[all]"
```

From source:

```bash
git clone https://github.com/bueormnew/MCRA.git
cd MCRA
pip install -e ".[dev]"
```

---

## Quick start

### Create a model

```python
from mcr_attention import MCRConfig, create_model

config = MCRConfig(vocab_size=50257, d_model=512, num_layers=6)
model = create_model(config)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")
```

### Train (causal multi-position LM)

```python
import torch
from mcr_attention import Trainer
from mcr_attention.api import TrainingConfig

model = create_model(vocab_size=50257, d_model=512, num_layers=6)

def data_fn():
    # Return a batch of token sequences [B, N]; the trainer auto-shifts targets
    # (predicts token t+1 from position t) so every position is a training signal.
    return torch.randint(0, 50257, (16, 512))

trainer = Trainer(model, TrainingConfig(
    num_steps=10000, learning_rate=3e-4, warmup_steps=500,
    lr_schedule="cosine", use_amp=True, save_dir="ckpt/", save_interval=1000,
))
stats = trainer.fit(data_fn)
print(f"Final loss: {stats['final_loss']:.4f}")
```

> You may also return `(input_ids, targets)` — `targets` can be `[B, N]`
> (multi-position) or `[B]` (single last-token, V2-compatible).

### Save / load

```python
from mcr_attention import save_model, load_model

save_model(model, "my_model/")            # config.json + model.pt + metadata.json
model = load_model("my_model/", device="auto")  # auto-detects V2 vs V3
```

### Generate (O(1) per token streaming)

```python
from mcr_attention import generate, GenerationConfig

tokens = generate(model, prompt=[1, 2, 3], config=GenerationConfig(
    max_tokens=200, temperature=0.8, top_k=50, top_p=0.9,
    repetition_penalty=1.1, seed=42,
))
# Greedy: temperature=0.0
```

---

## Architecture

```
Token IDs
   │
   ▼
[Embedding + RoPE]                 ── position-aware token representations
   │
   ▼
┌─────────────── repeated num_layers times (pre-norm residuals) ───────────┐
│                                                                          │
│   m = RMSNorm(x)                                                         │
│   C = MultiScaleMemory(m)        ── K linear recurrences:                │
│   │        h_t = a_k · h_{t-1} + B_k · m_t   (parallel scan / 1 step)    │
│   │        C ∈ R^{K × d}  (one compressed memory per scale)              │
│   ▼                                                                      │
│   z = AttentionRouter(C, m)      ── 1 query × K keys softmax attention   │
│   x = x + Dropout(z)              (residual)                             │
│   │                                                                      │
│   ▼                                                                      │
│   x = x + Dropout(FFN(RMSNorm(x)))  ── SwiGLU / GELU position-wise MLP   │
│                                                                          │
└──────────────────────────────────────────────────────────────────────────┘
   │
   ▼
RMSNorm → LanguageHead → logits [B, N, vocab]
```

**Training** runs the parallel scan over the whole sequence and produces logits
at every position. **Streaming** runs one recurrence step per token. The two are
mathematically identical at each position — there is no train/inference gap.

### Multi-scale memory

Each scale `k` has a learned decay `a_k = exp(-softplus(λ_k))` and is initialized
so its effective horizon is roughly `L_k = scale_min · scale_factor^(k-1)` tokens
(short scales forget fast, long scales retain). With the defaults this gives
horizons of 256, 512, …, 32768 tokens out of the box.

### Attention router

A **1-query × K-key** attention: the query is formed from the current token and
the longest-scale memory; keys/values are the K memory vectors. It is O(K)
(cheap) regardless of sequence length.

---

## Configuration

```python
MCRConfig(
    vocab_size=50257,     # token vocabulary
    d_model=512,          # residual-stream width (must be even for RoPE)
    d_state=128,          # recurrent state width per scale
    num_scales=8,         # number of memory scales (K)
    num_layers=6,         # depth (stacked blocks)
    scale_min=256,        # nominal shortest horizon
    scale_factor=2,       # geometric ratio between horizons
    rope_base=10000.0,    # RoPE base frequency
    dropout=0.1,
    use_input_gate=True,  # learned sigmoid input gate per scale
    use_rope=True,
    init_scales_geometric=True,
    ffn_type="swiglu",    # "swiglu" or "gelu"
    ffn_hidden=0,         # 0 = auto (8/3 · d_model)
    tie_embeddings=False,
    mlp_hidden=0,         # head MLP width (0 = auto: 4 · d_model)
    max_seq_len=32768,
)
```

---

## Fine-tuning

```python
from mcr_attention import load_model
from mcr_attention.finetune import finetune, FinetuneConfig

model = load_model("my_model/")
stats = finetune(model, data_fn, FinetuneConfig(
    learning_rate=1e-4, num_steps=500, freeze_embeddings=True, freeze_memory=True,
))
```

Selective freezing: `freeze_embeddings`, `freeze_memory`, `freeze_router`,
`freeze_ffn`, or `freeze_all_except(model, ["head"])`.

---

## HuggingFace integration

Real `transformers` integration (requires `mcr-attention[hf]`):

```python
from mcr_attention.hf import MCRForCausalLM, MCRPretrainedConfig, push_to_hub, from_pretrained

# AutoModel-compatible
model = MCRForCausalLM(MCRPretrainedConfig(vocab_size=50257, num_layers=6, d_model=512))

# Hub
push_to_hub(model, "username/my-mcr", tokenizer=tokenizer)
model = from_pretrained("username/my-mcr")
print(model.generate_text("Once upon a time", max_new_tokens=200))
```

`MCRForCausalLM` supports the standard HF `model.generate(...)` and works with
`AutoModelForCausalLM.from_pretrained`.

---

## Compatibility: keep using V2 models

V3 loads V2 checkpoints **bit-exact** via the legacy architecture, and can
migrate them to V3:

```python
from mcr_attention import load_model, convert_v2_to_v3

old = load_model("v2_model/")        # returns the legacy V2 model, identical behavior
convert_v2_to_v3("v2_model/", "v3_model/", num_layers=1)  # warm-start migration
```

Version is auto-detected from `config.json` (`arch_version`). Old checkpoints
without a version load as V2.

---

## CLI

```bash
mcr-attention info path/to/model/                       # show model info
mcr-attention benchmark --lengths 256 1024 4096 --layers 2   # speed/memory
mcr-attention generate path/to/model/ --prompt 1 2 3 --max-tokens 50 --temperature 0.8
```

---

## Benchmarks

Run the mission-verification suite:

```bash
python benchmarks/run_benchmark.py --d-model 128 --layers 2
```

It checks the four pillars and reports PASS/FAIL:
1. **Linear training cost** — parallel forward time grows ~linearly with N.
2. **O(1) streaming** — per-token latency is flat as N grows.
3. **O(1) streaming memory** — peak memory is flat.
4. **Train/inference identity** — parallel logits == streaming logits.

Synthetic-quality tasks (copy / needle / associative recall) live in
`mcr_attention.benchmarks`.

---

## Development

```bash
pip install -e ".[dev]"
pytest tests/ -v                 # 100+ tests, all green
python benchmarks/run_benchmark.py
```

## License

**GNU AGPL-3.0-or-later**. See [LICENSE](./LICENSE).

- **Author:** Gerson Fabian Buenahora Ormaza (BUEORM)
- **GitHub:** [bueormnew/MCRA](https://github.com/bueormnew/MCRA)
- **Contact:** dalusx64@gmail.com

MCR-Attention is free software: you can redistribute it and/or modify it under
the terms of the GNU Affero General Public License as published by the Free
Software Foundation, either version 3 of the License, or (at your option) any
later version. Network use (e.g. offering the model as a service) triggers the
source-disclosure obligation — see section 13 of the AGPL.
