Metadata-Version: 2.4
Name: opengdn
Version: 0.1.0
Summary: A correctness-first, MIT-licensed PyTorch implementation of Gated Delta Networks
Author: OpenGDN contributors
License: MIT License
        
        Copyright (c) 2026 OpenGDN contributors
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
License-File: NOTICE
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Requires-Dist: torch>=2.1
Provides-Extra: dev
Requires-Dist: numpy>=1.26; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Provides-Extra: triton
Requires-Dist: triton>=3.0; extra == 'triton'
Description-Content-Type: text/markdown

# OpenGDN

OpenGDN is a correctness-first, MIT-licensed PyTorch implementation of the
Gated Delta Network recurrence from [Yang et al. (ICLR 2025)](https://arxiv.org/abs/2412.06464).
Use the mixer as an ordinary PyTorch module. On supported CUDA/BF16 inputs it
selects the accelerated Triton path by default; everywhere else it safely uses
the exact PyTorch implementation.

## Install

For normal use:

```bash
pip install opengdn
```

For the CUDA/Triton acceleration, install the optional extra after installing
the PyTorch build appropriate for your CUDA system:

```bash
pip install 'opengdn[triton]'
```

For a source checkout and tests:

```bash
uv sync --extra dev --extra triton
```

## Quick start

```python
import torch
from opengdn import GatedDeltaMixer

mixer = GatedDeltaMixer(d_model=256, num_heads=4).cuda().bfloat16()
x = torch.randn(2, 32, 256, device="cuda", dtype=torch.bfloat16)
y, state = mixer(x, return_state=True)

# Continue an exact stream, or use state.detach() to start truncated BPTT.
next_y, next_state = mixer(x, state=state, return_state=True)
```

There is no CUDA backend switch to manage for ordinary use. The default uses
the fast path whenever its CUDA/BF16 contract is met and otherwise preserves
the same operator semantics with an exact fallback.

## Core operator reference

```python
import torch
from opengdn import gated_delta

B, T, H, K, V = 2, 32, 4, 16, 16
q = torch.randn(B, T, H, K)
k = torch.randn(B, T, H, K)
v = torch.randn(B, T, H, V)
alpha = torch.sigmoid(torch.randn(B, T, H))
beta = torch.sigmoid(torch.randn(B, T, H))
output, final_state = gated_delta(q, k, v, alpha, beta)
```

`q` and `k` have shape `[batch, time, heads, key_dim]`, `v` has shape
`[batch, time, heads, value_dim]`, gates have shape `[batch, time, heads]`,
and the recurrent state has shape `[batch, heads, value_dim, key_dim]`.
Passing `final_state` as `state=` in the next call implements exact streaming
continuation. The state participates in autograd; call `state.detach()` at
your chosen truncated-BPTT boundary.

`gated_delta_chunkwise(..., chunk_size=64)` provides a second PyTorch backend
using the paper's extended-WY chunk decomposition. It has the same output and
gradient contract as `gated_delta`; the token-by-token function remains the
correctness oracle. It is an explicit slower opt-out, intended for debugging
and verification: `GatedDeltaMixer(..., backend="chunkwise", chunk_size=64)`.

`gated_delta_triton(...)` is the default mixer backend for CUDA/BF16 training,
with the same operator and streaming-state contract. It targets contiguous inputs, a
64-token chunk, and key/value dimensions up to 512; `resolve_backend(...)`
reports whether that contract was selected or the exact chunkwise fallback ran.
Its native forward uses FP32 state arithmetic with BF16 public tensors. Its
first-order reverse is an exact 64-token extended-WY block VJP: it retains
FP32 chunk boundaries and recomputes compact 64x64 systems plus 64x(K+V)
solve factors, rather than retaining a token-state replay history. Higher-order
gradients replay the chunkwise PyTorch oracle.

For the d12 K=V=128 CUDA contract, 64-token reverse blocks compile with
Inductor by default and internal FP32 matrix products use TF32 tensor cores.
The fixed-shape Nanochat smoke also captures its training step in a CUDA graph
and uses Inductor's one-time CUDA kernel autotuner by default. These are
deliberately opt-out (`--no-*` smoke flags or `OPENGDN_DISABLE_*` environment
variables), and OpenGDN emits a performance warning when a slower path is
selected.

## Mixer

```python
from opengdn import GatedDeltaMixer

mixer = GatedDeltaMixer(d_model=256, num_heads=4, conv_kernel_size=4)
x = torch.randn(2, 32, 256)
y, state = mixer(x, return_state=True)
```

The mixer follows the paper's stated topology: projected causal short
convolutions and SiLU for Q/K/V, L2-normalized Q/K, per-head alpha/beta gates,
RMS-normalized recurrent output, and a SiLU output gate. The paper omits the
exact Mamba2 alpha parameterization, so this implementation deliberately uses
documented defaults: `alpha = sigmoid(alpha_logits)` and
`beta = sigmoid(beta_logits)`.

## Scope

For exact mixer streaming, the returned `GatedDeltaMixerState` includes both
the recurrent state and causal-convolution histories. Its `.detach()` method
sets a truncated-BPTT boundary for all state tensors together.

v0.1 supports dense, equal-length batch-first sequences on CPU and CUDA. It
does not yet support ragged batches or reset masks. The CUDA path is specialized
for the supported dense BF16 contract and falls back exactly when that contract
does not apply. See `NOTICE` for attribution and licensing boundaries.

## Nanochat integration smoke test

The optional smoke harness leaves Nanochat unchanged, builds a tiny real
Nanochat `GPT`, substitutes each attention mixer with `GatedDeltaMixer`, and
trains a synthetic next-token task. It verifies finite loss, gradients through
the OpenGDN path, and a decreasing loss; it is not a language-model benchmark.

```bash
uv run --extra dev --extra triton python examples/nanochat_smoke.py \
  --nanochat-root /path/to/nanochat
```

To demonstrate literal training-set memorization rather than a simple
transition rule, use a fixed corpus of random token documents. The harness
also reports the loss after temporarily disabling the trained GDN mixers.

```bash
uv run --extra dev --extra triton python examples/nanochat_smoke.py \
  --nanochat-root /path/to/nanochat --task fixed-corpus --documents 16 --steps 300
```

## Directional CUDA baseline

Profile either PyTorch backend at a Nanochat-relevant 768-wide, 6-head,
128-dimensional-head shape. This is deliberately a contention-sensitive
baseline, not a performance claim.

```bash
uv run --extra dev python benchmarks/profile_reference.py
uv run --extra dev python benchmarks/profile_reference.py --backend chunkwise
uv run --extra dev python benchmarks/profile_reference.py --backend chunkwise --compile
uv run --extra triton --extra dev python benchmarks/profile_reference.py --backend triton
```
