Metadata-Version: 2.1
Name: uattn
Version: 0.1.0
Summary: Unified Attention: single projection replacing Q/K/V with 67% fewer parameters
Author-email: Viraj Deshwal <viraj@reinforceai.com>
License: MIT
Project-URL: Homepage, https://github.com/ReinforceAI/unified-attention
Project-URL: Repository, https://github.com/ReinforceAI/unified-attention
Project-URL: Paper, https://github.com/ReinforceAI/yocto
Keywords: attention,transformer,unified-attention,efficient,llm
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.4.0
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Provides-Extra: fa3
Requires-Dist: flash-attn-3; extra == "fa3"
Provides-Extra: hopper
Requires-Dist: torch>=2.8.0; extra == "hopper"
Requires-Dist: flash-attn-3; extra == "hopper"

# Unified Attention

**One projection replaces Q, K, V. 67% fewer attention parameters.**

Standard self-attention uses three separate weight matrices (W_Q, W_K, W_V) to project the same input. Unified Attention uses one. The output splits into three functional bands that serve as query, key, and value. The bands form naturally during training through amplitude and phase differentiation in the weight matrix.

```python
from uattn import UnifiedAttention

attn = UnifiedAttention(dim=528, num_heads=4)
y = attn(x)  # [batch, seq, 528] -> [batch, seq, 528]
```

## Install

```bash
pip install uattn
```

From source:
```bash
pip install git+https://github.com/ReinforceAI/unified-attention.git
```

For FA3 (Hopper, 1.57x faster):
```bash
pip install uattn[fa3]
```

## Why

Attention routes tokens to each other. The MLP transforms them. In parameter-constrained models, every byte spent on routing is a byte not spent on computation. Unified Attention cuts the routing budget:

| | Standard Q/K/V | Unified |
|---|---|---|
| Attention projection params | 3d² | d² |
| Share of layer params | ~33% | ~18% |
| Compressed attention bytes (16MB model) | 5.10 MB | 2.82 MB |
| Compressed MLP bytes (16MB model) | 10.21 MB | 12.70 MB |

## Results

Validated on [OpenAI Parameter Golf](https://github.com/openai/parameter-golf) (16MB artifact, FineWeb):

| Setting | BPB | vs Standard SOTA |
|---------|-----|-------------------|
| 10-min record (8xH100) | 1.1412 | Rank 7, only novel architecture in top 10 |
| 1-hr unlimited (8xH100) | 1.1088 | Beats all submissions on both leaderboards |

Also tested at small scale (484K params, TinyStories): 9.58 perplexity, matching models 2-4x larger.

## Usage

### Standalone (owns its own weights)

```python
from uattn import UnifiedAttention

attn = UnifiedAttention(
    dim=528,              # must be divisible by 3
    num_heads=4,          # heads for multi-head attention
    rope_base=10000.0,    # RoPE base frequency
    seeking_gain_init=0.5,# per-head query scaling
    rope_fraction=1.0,    # fraction of head_dim for RoPE
)

x = torch.randn(batch, seq_len, 528)
y = attn(x)  # uses internal W_unified and W_out
```

### With external weight banks (for Muon optimizer)

```python
# Weights stored in contiguous 3D banks for batched optimization
unified_bank = nn.Parameter(torch.randn(K, d, d))
output_bank = nn.Parameter(torch.randn(K, d, d//3))

attn = UnifiedAttention(dim=d, num_heads=4, output_proj=False)
y = attn(x, unified_w=unified_bank[k], output_w=output_bank[k])
```

### Check backend

```python
from uattn import backend
print(backend())  # 'fa3', 'fa2', or 'sdpa'
```

## How It Works

```python
# Standard: 3 separate projections
q = W_q @ x;  k = W_k @ x;  v = W_v @ x

# Unified: 1 projection, 3 bands
unified = W_unified @ x
seeking, offering, content = unified.split(d//3, dim=-1)
```

RoPE is applied to seeking and offering bands only. Position affects routing, not content.

For FA3 (Hopper): head_dim is zero-padded to the nearest multiple of 8 before the kernel call and sliced back after. Padded zeros contribute nothing to dot products. Mathematically lossless.

## Citation

```bibtex
@misc{deshwal2026yocto,
  title={Attention Fields: Unified Projections for Efficient Language Models},
  author={Deshwal, Viraj},
  year={2026},
  url={https://github.com/ReinforceAI/yocto},
  howpublished={\url{https://github.com/ReinforceAI/unified-attention}}
}
```

## License

MIT
