ramnet-ops 0.1.0rc0__tar.gz

ramnet_ops-0.1.0rc0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025-2026 Kaicheng Xiao, Haotian Li, Liran Dong
+
+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.

ramnet_ops-0.1.0rc0/MANIFEST.in

@@ -0,0 +1,7 @@
+include LICENSE
+include README.md
+include RELEASING.md
+include csrc/*.cpp
+include csrc/*.cu
+include csrc/*.cuh
+include csrc/*.h

ramnet_ops-0.1.0rc0/PKG-INFO

@@ -0,0 +1,309 @@
+Metadata-Version: 2.4
+Name: ramnet-ops
+Version: 0.1.0rc0
+Summary: PyTorch and CUDA operators for RAM-Net attention
+Maintainer: Kaicheng Xiao
+License-Expression: MIT
+Project-URL: Source, https://github.com/MuonCat/RAM-Net
+Project-URL: Issues, https://github.com/MuonCat/RAM-Net/issues
+Project-URL: Documentation, https://github.com/MuonCat/RAM-Net/tree/main/packages/ramnet-ops
+Keywords: pytorch,cuda,attention,ram-net
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.5
+Requires-Dist: einops
+Requires-Dist: numpy
+Dynamic: license-file
+Dynamic: requires-dist
+
+<div align="center">
+
+<img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
+
+# ramnet-ops
+
+[![arXiv](https://img.shields.io/badge/arXiv-2602.11958-b31b1b.svg?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2602.11958)
+[![Project Page](https://img.shields.io/badge/Project-Page-1f6feb.svg)](https://muoncat.github.io/ramnet_web/)
+[![Weights](https://img.shields.io/badge/Weights-Hugging%20Face-FFD21E?logo=huggingface&logoColor=white)](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
+[![GitHub](https://img.shields.io/badge/GitHub-RAM--Net-181717?logo=github&logoColor=white)](https://github.com/MuonCat/RAM-Net)
+[![Models](https://img.shields.io/badge/Models-ramnet-82afe6?logo=pytorch&logoColor=white)](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet/README.md)
+[![Operators](https://img.shields.io/badge/Operators-ramnet--ops-76B900?logo=nvidia&logoColor=white)](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md)
+[![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
+
+</div>
+
+The attention computation core of RAM-Net: sparse address decoding, memory updates, readout, and autograd.
+Import as `ramnet_ops` for CUDA kernels and a PyTorch reference implementation.
+
+## Structure
+
+```text
+.
+|-- src/
+|   `-- ramnet_ops/
+|       |-- ram_attn.py           # Public API, input normalization, backend selection
+|       `-- impl/
+|           |-- torch_impl.py     # PyTorch reference
+|           |-- cuda_impl.py      # CUDA wrappers and segment layout
+|           |-- _cuda_autograd.py # FakeTensor and autograd registrations
+|           `-- _cuda_config.py   # Address and index-capacity checks
+`-- csrc/
+    |-- ram_attn_bind.cpp         # Torch schemas, CUDA registration, extension entry point
+    |-- ram_attn_ops.cpp          # Native wrappers, output allocation, kernel dispatch
+    |-- ram_attn_*.cu             # Address, segment, and autoregressive kernels
+    `-- torch_compat*             # Ordinary/stable Torch ABI interfaces
+```
+
+## Installation
+
+Requires:
+
+```text
+Python >= 3.10
+
+ramnet-ops
+|-- torch >= 2.5
+|-- einops
+`-- numpy
+```
+
+### 1. Install dependencies
+
+If PyTorch is not already installed, we **recommend PyTorch ≥ 2.10 with CUDA 13.0 (`cu130`)** for the default ops wheels below.
+See the [official installation guide](https://pytorch.org/get-started/locally/) for installation commands.
+Then install the remaining dependencies from TestPyPI / PyPI:
+
+```bash
+python -m pip install einops numpy
+```
+
+### 2. Install ramnet-ops
+
+The default `ramnet-ops` wheels use both **Python and PyTorch stable ABIs** and target **CUDA 13.0 (`cu130`)**.
+If your environment meets **all** of the following requirements, install directly from pip:
+
+- **Python:** CPython ≥ 3.10.
+- **PyTorch:** ≥ 2.10 with CUDA 13.0.
+- **GPU:** SM 8.9, 9.0, 10.0, or 12.0.
+- **Platform:** Windows x64 or Linux x86_64 with glibc ≥ 2.35.
+
+```bash
+python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
+```
+
+Prebuilt wheels need no local compiler or CUDA Toolkit.
+If your environment **does not meet all of the requirements above**, follow [Other Torch/CUDA builds (GitHub Releases)](#other-torchcuda-builds-github-releases) below to select another prebuilt wheel.
+
+#### Other Torch/CUDA builds (GitHub Releases)
+
+GitHub Releases provides additional wheels for older PyTorch versions, other CUDA builds, and GPUs outside the default `cu130` targets.
+Open [GitHub Releases](https://github.com/MuonCat/RAM-Net/releases), select the desired package version under `ramnet-ops/v<VERSION>`, and expand **Assets** to find its `.whl` files.
+Choose a wheel using [the selection guide below](#choose-a-matching-wheel), then copy its download link into `python -m pip install "<wheel URL>"`.
+You must select the Torch/CUDA variant yourself; pip does not infer it from your environment.
+
+With the matching PyTorch build already installed, the commands look like this for package version `0.1.0rc0`:
+
+Windows x64, **Torch 2.8.0 / CUDA 12.8** (ordinary ABI):
+
+```bash
+python -m pip install "https://github.com/MuonCat/RAM-Net/releases/download/ramnet-ops%2Fv0.1.0rc0/ramnet_ops-0.1.0rc0-1torch2_8_0cu128-cp310-abi3-win_amd64.whl"
+```
+
+Linux x86_64, **Torch ≥ 2.10 / CUDA 12.6** (stable ABI):
+
+```bash
+python -m pip install "https://github.com/MuonCat/RAM-Net/releases/download/ramnet-ops%2Fv0.1.0rc0/ramnet_ops-0.1.0rc0-1torchabi2_10cu126-cp310-abi3-manylinux_2_35_x86_64.whl"
+```
+
+##### Choose a matching wheel
+
+First, match the **Torch ABI** marker in the filename to your installed PyTorch version:
+
+- **Ordinary ABI (`torch2_6_0` or `torch2_8_0`):** requires exactly **PyTorch 2.6.0** or **2.8.0**, respectively.
+- **Stable ABI (`torchabi2_10`):** supports **PyTorch ≥ 2.10** with a matching CUDA/platform combination; the Torch 2.14 SDK used to build these wheels is not required at runtime.
+
+Next, find your **GPU row** and **PyTorch column** in the table below.
+Each cell lists the available CUDA builds: choose one matching your installed PyTorch's CUDA version (`torch.version.cuda`), such as `cu126` for CUDA 12.6.
+
+| [GPU CC](https://developer.nvidia.com/cuda/gpus) (representative GPUs) | PyTorch 2.6.0 | PyTorch 2.8.0 | PyTorch ≥ 2.10 |
+| --- | --- | --- | --- |
+| 7.5 (T4 / RTX 20) | 11.8 | ❌ | ❌ |
+| 8.0 (A100) | 11.8 / 12.6 | 12.6 | 12.6 |
+| 8.6 (RTX 30) | 11.8 / 12.6 | 12.6 | 12.6 |
+| 8.9 (RTX 40 / L40S) | 12.6 | 12.6 / 12.8 | 12.6 / **13.0 ★** |
+| 9.0 (H100 / H200) | 12.6 | 12.6 / 12.8 | 12.6 / **13.0 ★** |
+| 10.0 (B200 / GB200) | ❌ | 12.8 | **13.0 ★** |
+| 12.0 (RTX 50) | ❌ | 12.8 | **13.0 ★** |
+
+★ marks the default TestPyPI / PyPI build; ❌ means no prebuilt wheel is provided for that combination.
+
+This table describes the architectures included in our wheels, not every GPU supported by each CUDA Toolkit.
+
+Finally, match the **platform** suffix: `win_amd64` for Windows x64 or `manylinux_2_35_x86_64` for Linux x86_64 with glibc ≥ 2.35.
+All wheels use `cp310-abi3` for GIL-enabled **CPython ≥ 3.10**.
+Check the driver requirements below before installing; if no wheel matches your environment, see [From source](#from-source).
+
+##### NVIDIA drivers
+
+The table lists driver baselines from NVIDIA's CUDA GA release notes; use these or newer drivers that support your GPU.
+
+| CUDA | Windows driver | Linux driver |
+| --- | --- | --- |
+| [11.8](https://docs.nvidia.com/cuda/archive/11.8.0/cuda-toolkit-release-notes/index.html) | ≥ 522.06 | ≥ 520.61.05 |
+| [12.6](https://docs.nvidia.com/cuda/archive/12.6.0/cuda-toolkit-release-notes/index.html) | ≥ 560.76 | ≥ 560.28.03 |
+| [12.8](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-toolkit-release-notes/index.html) | ≥ 570.65 | ≥ 570.26 |
+| **[13.0 ★](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-toolkit-release-notes/index.html)** | **R580+** | **≥ 580.65.06** |
+
+CUDA 13.0 no longer bundles a Windows driver; NVIDIA requires the R580 branch or newer.
+CUDA 11.x/12.x may run on older drivers through [minor-version compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/minor-version-compatibility.html), subject to feature and PTX restrictions; newer GPUs may require newer drivers.
+
+#### From source
+
+Source installation supports both CUDA extension builds and PyTorch-only installations.
+
+##### With CUDA extension
+
+Use this option when no prebuilt wheel matches your environment or when modifying the CUDA operators.
+Prepare the build environment first:
+
+- **PyTorch:** install the intended CUDA build (Torch ≥ 2.5).
+- **CUDA Toolkit:** install a version compatible with that PyTorch build, including `nvcc`.
+- **C++ compiler:** use a compiler supported by the Toolkit; on Windows, run the commands in an **x64 Native Tools Command Prompt for VS 2022**.
+
+By default (`RAMNET_TORCH_STABLE_ABI=auto`), Torch 2.5–2.13 builds use the ordinary ABI and require that exact Torch version at runtime; Torch ≥ 2.14 builds use stable ABI targeting Torch ≥ 2.10.
+
+Install the build tools in the same Python environment as PyTorch:
+
+```bash
+python -m pip install "setuptools>=77.0.3" wheel ninja
+```
+
+Then choose one of the following sources.
+
+**Download and build with pip** — install the published source archive from TestPyPI / PyPI:
+
+```bash
+python -m pip install -i https://test.pypi.org/simple/ ramnet-ops --no-binary=ramnet-ops --no-build-isolation
+```
+
+**Clone from GitHub and build locally** — use the repository source, for example when modifying the operators:
+
+```bash
+git clone https://github.com/MuonCat/RAM-Net.git
+cd RAM-Net
+python -m pip install ./packages/ramnet-ops --no-build-isolation
+```
+
+The CUDA extension is compiled during installation.
+
+##### Environment variables
+
+Set these before source installation; ABI and architecture settings apply only when building the CUDA extension.
+
+| Variable | Default | Purpose |
+| --- | --- | --- |
+| `SKIP_CUDA` | `0` | Set to `1` to skip the CUDA extension and install only the Python package and PyTorch reference. |
+| `RAMNET_TORCH_STABLE_ABI` | `auto` | Select the ABI from the build SDK; `0` forces ordinary ABI; `1` forces stable ABI and requires Torch ≥ 2.14 to build. |
+| `TORCH_CUDA_ARCH_LIST` | GPUs visible during the build | Set the target GPU architectures explicitly, e.g. `8.9 9.0`. |
+
+Set `SKIP_CUDA=1` to install from source without building the native CUDA extension.
+PyTorch, `einops`, and `numpy` are still required; a CUDA Toolkit and C++ compiler are not needed.
+
+## Quick start
+
+`random_access_memory_attention` performs the sparse reads and writes of a RAM-Net layer, covering address decoding, CAPE, and the Gated Sparse Update.
+It takes address logits, values, and forgetting modulators, and returns the read outputs and the updated slot/mass state.
+
+The table maps paper notation to parameters and tensor dimensions in the example below.
+`B`, `T`, and `H` denote batch size, sequence length, and number of heads.
+
+| Paper quantity / symbol | Operator parameter / axis |
+| :---: | :---: |
+| Address order $U$ (digits per address) | `U` |
+| Address radix $d_p$ (values per digit) | `P` |
+| Value dimension $d_v$ | `C` |
+| Slots per head $d_p^U$ | `A = P ** U` |
+| Selected slots $K$ | `top_k` |
+
+This example computes attention outputs and the updated state, then backpropagates through both:
+
+```python
+import torch
+from ramnet_ops import random_access_memory_attention
+
+device = "cuda"  # Set to "cpu" for the PyTorch reference.
+backend = "cuda" if device == "cuda" else "torch"
+torch.manual_seed(0)
+
+B, T, H, U, P, C = 2, 32, 4, 3, 4, 32
+w_addr = torch.randn(B, T, H, U, P, device=device, requires_grad=True)
+r_addr = torch.randn(B, T, H, U, P, device=device, requires_grad=True)
+values = torch.randn(B, T, H, C, device=device, requires_grad=True)
+gamma = torch.randn(B, T, H, device=device, requires_grad=True)
+
+y, (slot, mass) = random_access_memory_attention(
+    w_addr, r_addr, values, gamma,
+    embed_num=2, top_k=8,
+    impl_backend=backend,
+)
+loss = y.square().mean() + slot.square().mean() + mass.square().mean()
+loss.backward()
+print(y.shape)     # (B, T, H, C)
+print(slot.shape)  # (B, H, P ** U, C)
+print(mass.shape)  # (B, H, P ** U)
+```
+
+`embed_num` selects how many heads, starting from head 0, use **Cyclic Address Positional Embedding (CAPE)** to cyclically shift read and write slot indices by token position.
+With `H=4` and `embed_num=2`, heads 0 and 1 use CAPE, while heads 2 and 3 address by content alone.
+Its range is `0` to `H`: `0` disables CAPE and `H` enables it for all heads.
+
+From the repository root, run [ops_smoke.py](../../examples/ops_smoke.py) with `python examples/ops_smoke.py --device cuda`, or use `--device cpu` for the PyTorch reference.
+This example demonstrates the API; it is not an experiment configuration.
+
+### Arguments
+
+- `w_addr`, `r_addr`: write and read address logits of shape `(…, T, …, H, U, P)`. The operator applies a softmax over each digit.
+- `values`: values to write, of shape `(…, T, …, H, C)`. The output `y` has the same shape and dtype.
+- `gamma`: forgetting modulator of shape `(…, T, …, H)`, computed by the model for each token. Larger values overwrite old slot content faster.
+- `init_ram`: the `(slot, mass)` state returned by a previous call, used to continue a sequence. Without it, every slot starts empty.
+- `inplace`: if `True`, updates the `init_ram` tensors in place instead of returning new ones. This requires contiguous FP32 state and does not support gradients.
+- `mask`: marks valid tokens. Masked tokens are not written to the state, and their outputs should be ignored.
+- `embed_num`: number of heads that use CAPE, counted from head 0. These heads shift each token's addresses by its position.
+- `top_k`: `K`, given as an int or a `(write_k, read_k)` pair. `None` keeps all slots and is supported only by the PyTorch backend.
+- `top_k_re_norm`: explicitly enables renormalization after Top-K. Accepts a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
+- `t_dim`: index of the time axis. Dimensions before `H` other than the time axis are broadcast as batch dimensions.
+- `impl_backend`: `'cuda'` or `'torch'`. By default, CUDA inputs use the CUDA backend when supported and otherwise fall back to PyTorch with a warning.
+
+## CUDA execution
+
+The CUDA backend automatically selects between **autoregressive** and **parallel** execution.
+
+**Autoregressive mode** handles single-token inputs (`T = 1`) when no input requires gradients.
+A dedicated kernel decodes the current token's addresses, updates the selected write slots, and reads the selected read slots directly.
+
+**Parallel mode** handles all other inputs, including training and prefill.
+It decodes all tokens' addresses in parallel, treating each selected read or write access as an **event**.
+Events targeting the same `(batch, head, slot)` are grouped into a **segment**.
+The kernels process independent segments and their value channels in parallel while preserving causal dependencies within each segment, then accumulate read contributions into the token outputs.
+
+## Autograd and compile
+
+Operators are registered under `torch.ops.ramnet_ops` with FakeTensor and first-order autograd support.
+Compile coverage targets fixed shapes/configuration with `dynamic=False`, tested through `aot_eager`;
+it does not establish Inductor or CUDA Graph support.
+
+Segment capacity is `B * H * A`; event capacity is `B * H * T_padded * (write_k + read_k)`.
+Masks do not reduce these allocations.
+Higher-order gradients are not supported.
+
+## Development
+
+See [RELEASING](RELEASING.md) for tests and publishing.
+
+MIT License. See [LICENSE](LICENSE).

ramnet_ops-0.1.0rc0/README.md

@@ -0,0 +1,285 @@
+<div align="center">
+
+<img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
+
+# ramnet-ops
+
+[![arXiv](https://img.shields.io/badge/arXiv-2602.11958-b31b1b.svg?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2602.11958)
+[![Project Page](https://img.shields.io/badge/Project-Page-1f6feb.svg)](https://muoncat.github.io/ramnet_web/)
+[![Weights](https://img.shields.io/badge/Weights-Hugging%20Face-FFD21E?logo=huggingface&logoColor=white)](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
+[![GitHub](https://img.shields.io/badge/GitHub-RAM--Net-181717?logo=github&logoColor=white)](https://github.com/MuonCat/RAM-Net)
+[![Models](https://img.shields.io/badge/Models-ramnet-82afe6?logo=pytorch&logoColor=white)](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet/README.md)
+[![Operators](https://img.shields.io/badge/Operators-ramnet--ops-76B900?logo=nvidia&logoColor=white)](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md)
+[![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
+
+</div>
+
+The attention computation core of RAM-Net: sparse address decoding, memory updates, readout, and autograd.
+Import as `ramnet_ops` for CUDA kernels and a PyTorch reference implementation.
+
+## Structure
+
+```text
+.
+|-- src/
+|   `-- ramnet_ops/
+|       |-- ram_attn.py           # Public API, input normalization, backend selection
+|       `-- impl/
+|           |-- torch_impl.py     # PyTorch reference
+|           |-- cuda_impl.py      # CUDA wrappers and segment layout
+|           |-- _cuda_autograd.py # FakeTensor and autograd registrations
+|           `-- _cuda_config.py   # Address and index-capacity checks
+`-- csrc/
+    |-- ram_attn_bind.cpp         # Torch schemas, CUDA registration, extension entry point
+    |-- ram_attn_ops.cpp          # Native wrappers, output allocation, kernel dispatch
+    |-- ram_attn_*.cu             # Address, segment, and autoregressive kernels
+    `-- torch_compat*             # Ordinary/stable Torch ABI interfaces
+```
+
+## Installation
+
+Requires:
+
+```text
+Python >= 3.10
+
+ramnet-ops
+|-- torch >= 2.5
+|-- einops
+`-- numpy
+```
+
+### 1. Install dependencies
+
+If PyTorch is not already installed, we **recommend PyTorch ≥ 2.10 with CUDA 13.0 (`cu130`)** for the default ops wheels below.
+See the [official installation guide](https://pytorch.org/get-started/locally/) for installation commands.
+Then install the remaining dependencies from TestPyPI / PyPI:
+
+```bash
+python -m pip install einops numpy
+```
+
+### 2. Install ramnet-ops
+
+The default `ramnet-ops` wheels use both **Python and PyTorch stable ABIs** and target **CUDA 13.0 (`cu130`)**.
+If your environment meets **all** of the following requirements, install directly from pip:
+
+- **Python:** CPython ≥ 3.10.
+- **PyTorch:** ≥ 2.10 with CUDA 13.0.
+- **GPU:** SM 8.9, 9.0, 10.0, or 12.0.
+- **Platform:** Windows x64 or Linux x86_64 with glibc ≥ 2.35.
+
+```bash
+python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
+```
+
+Prebuilt wheels need no local compiler or CUDA Toolkit.
+If your environment **does not meet all of the requirements above**, follow [Other Torch/CUDA builds (GitHub Releases)](#other-torchcuda-builds-github-releases) below to select another prebuilt wheel.
+
+#### Other Torch/CUDA builds (GitHub Releases)
+
+GitHub Releases provides additional wheels for older PyTorch versions, other CUDA builds, and GPUs outside the default `cu130` targets.
+Open [GitHub Releases](https://github.com/MuonCat/RAM-Net/releases), select the desired package version under `ramnet-ops/v<VERSION>`, and expand **Assets** to find its `.whl` files.
+Choose a wheel using [the selection guide below](#choose-a-matching-wheel), then copy its download link into `python -m pip install "<wheel URL>"`.
+You must select the Torch/CUDA variant yourself; pip does not infer it from your environment.
+
+With the matching PyTorch build already installed, the commands look like this for package version `0.1.0rc0`:
+
+Windows x64, **Torch 2.8.0 / CUDA 12.8** (ordinary ABI):
+
+```bash
+python -m pip install "https://github.com/MuonCat/RAM-Net/releases/download/ramnet-ops%2Fv0.1.0rc0/ramnet_ops-0.1.0rc0-1torch2_8_0cu128-cp310-abi3-win_amd64.whl"
+```
+
+Linux x86_64, **Torch ≥ 2.10 / CUDA 12.6** (stable ABI):
+
+```bash
+python -m pip install "https://github.com/MuonCat/RAM-Net/releases/download/ramnet-ops%2Fv0.1.0rc0/ramnet_ops-0.1.0rc0-1torchabi2_10cu126-cp310-abi3-manylinux_2_35_x86_64.whl"
+```
+
+##### Choose a matching wheel
+
+First, match the **Torch ABI** marker in the filename to your installed PyTorch version:
+
+- **Ordinary ABI (`torch2_6_0` or `torch2_8_0`):** requires exactly **PyTorch 2.6.0** or **2.8.0**, respectively.
+- **Stable ABI (`torchabi2_10`):** supports **PyTorch ≥ 2.10** with a matching CUDA/platform combination; the Torch 2.14 SDK used to build these wheels is not required at runtime.
+
+Next, find your **GPU row** and **PyTorch column** in the table below.
+Each cell lists the available CUDA builds: choose one matching your installed PyTorch's CUDA version (`torch.version.cuda`), such as `cu126` for CUDA 12.6.
+
+| [GPU CC](https://developer.nvidia.com/cuda/gpus) (representative GPUs) | PyTorch 2.6.0 | PyTorch 2.8.0 | PyTorch ≥ 2.10 |
+| --- | --- | --- | --- |
+| 7.5 (T4 / RTX 20) | 11.8 | ❌ | ❌ |
+| 8.0 (A100) | 11.8 / 12.6 | 12.6 | 12.6 |
+| 8.6 (RTX 30) | 11.8 / 12.6 | 12.6 | 12.6 |
+| 8.9 (RTX 40 / L40S) | 12.6 | 12.6 / 12.8 | 12.6 / **13.0 ★** |
+| 9.0 (H100 / H200) | 12.6 | 12.6 / 12.8 | 12.6 / **13.0 ★** |
+| 10.0 (B200 / GB200) | ❌ | 12.8 | **13.0 ★** |
+| 12.0 (RTX 50) | ❌ | 12.8 | **13.0 ★** |
+
+★ marks the default TestPyPI / PyPI build; ❌ means no prebuilt wheel is provided for that combination.
+
+This table describes the architectures included in our wheels, not every GPU supported by each CUDA Toolkit.
+
+Finally, match the **platform** suffix: `win_amd64` for Windows x64 or `manylinux_2_35_x86_64` for Linux x86_64 with glibc ≥ 2.35.
+All wheels use `cp310-abi3` for GIL-enabled **CPython ≥ 3.10**.
+Check the driver requirements below before installing; if no wheel matches your environment, see [From source](#from-source).
+
+##### NVIDIA drivers
+
+The table lists driver baselines from NVIDIA's CUDA GA release notes; use these or newer drivers that support your GPU.
+
+| CUDA | Windows driver | Linux driver |
+| --- | --- | --- |
+| [11.8](https://docs.nvidia.com/cuda/archive/11.8.0/cuda-toolkit-release-notes/index.html) | ≥ 522.06 | ≥ 520.61.05 |
+| [12.6](https://docs.nvidia.com/cuda/archive/12.6.0/cuda-toolkit-release-notes/index.html) | ≥ 560.76 | ≥ 560.28.03 |
+| [12.8](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-toolkit-release-notes/index.html) | ≥ 570.65 | ≥ 570.26 |
+| **[13.0 ★](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-toolkit-release-notes/index.html)** | **R580+** | **≥ 580.65.06** |
+
+CUDA 13.0 no longer bundles a Windows driver; NVIDIA requires the R580 branch or newer.
+CUDA 11.x/12.x may run on older drivers through [minor-version compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/minor-version-compatibility.html), subject to feature and PTX restrictions; newer GPUs may require newer drivers.
+
+#### From source
+
+Source installation supports both CUDA extension builds and PyTorch-only installations.
+
+##### With CUDA extension
+
+Use this option when no prebuilt wheel matches your environment or when modifying the CUDA operators.
+Prepare the build environment first:
+
+- **PyTorch:** install the intended CUDA build (Torch ≥ 2.5).
+- **CUDA Toolkit:** install a version compatible with that PyTorch build, including `nvcc`.
+- **C++ compiler:** use a compiler supported by the Toolkit; on Windows, run the commands in an **x64 Native Tools Command Prompt for VS 2022**.
+
+By default (`RAMNET_TORCH_STABLE_ABI=auto`), Torch 2.5–2.13 builds use the ordinary ABI and require that exact Torch version at runtime; Torch ≥ 2.14 builds use stable ABI targeting Torch ≥ 2.10.
+
+Install the build tools in the same Python environment as PyTorch:
+
+```bash
+python -m pip install "setuptools>=77.0.3" wheel ninja
+```
+
+Then choose one of the following sources.
+
+**Download and build with pip** — install the published source archive from TestPyPI / PyPI:
+
+```bash
+python -m pip install -i https://test.pypi.org/simple/ ramnet-ops --no-binary=ramnet-ops --no-build-isolation
+```
+
+**Clone from GitHub and build locally** — use the repository source, for example when modifying the operators:
+
+```bash
+git clone https://github.com/MuonCat/RAM-Net.git
+cd RAM-Net
+python -m pip install ./packages/ramnet-ops --no-build-isolation
+```
+
+The CUDA extension is compiled during installation.
+
+##### Environment variables
+
+Set these before source installation; ABI and architecture settings apply only when building the CUDA extension.
+
+| Variable | Default | Purpose |
+| --- | --- | --- |
+| `SKIP_CUDA` | `0` | Set to `1` to skip the CUDA extension and install only the Python package and PyTorch reference. |
+| `RAMNET_TORCH_STABLE_ABI` | `auto` | Select the ABI from the build SDK; `0` forces ordinary ABI; `1` forces stable ABI and requires Torch ≥ 2.14 to build. |
+| `TORCH_CUDA_ARCH_LIST` | GPUs visible during the build | Set the target GPU architectures explicitly, e.g. `8.9 9.0`. |
+
+Set `SKIP_CUDA=1` to install from source without building the native CUDA extension.
+PyTorch, `einops`, and `numpy` are still required; a CUDA Toolkit and C++ compiler are not needed.
+
+## Quick start
+
+`random_access_memory_attention` performs the sparse reads and writes of a RAM-Net layer, covering address decoding, CAPE, and the Gated Sparse Update.
+It takes address logits, values, and forgetting modulators, and returns the read outputs and the updated slot/mass state.
+
+The table maps paper notation to parameters and tensor dimensions in the example below.
+`B`, `T`, and `H` denote batch size, sequence length, and number of heads.
+
+| Paper quantity / symbol | Operator parameter / axis |
+| :---: | :---: |
+| Address order $U$ (digits per address) | `U` |
+| Address radix $d_p$ (values per digit) | `P` |
+| Value dimension $d_v$ | `C` |
+| Slots per head $d_p^U$ | `A = P ** U` |
+| Selected slots $K$ | `top_k` |
+
+This example computes attention outputs and the updated state, then backpropagates through both:
+
+```python
+import torch
+from ramnet_ops import random_access_memory_attention
+
+device = "cuda"  # Set to "cpu" for the PyTorch reference.
+backend = "cuda" if device == "cuda" else "torch"
+torch.manual_seed(0)
+
+B, T, H, U, P, C = 2, 32, 4, 3, 4, 32
+w_addr = torch.randn(B, T, H, U, P, device=device, requires_grad=True)
+r_addr = torch.randn(B, T, H, U, P, device=device, requires_grad=True)
+values = torch.randn(B, T, H, C, device=device, requires_grad=True)
+gamma = torch.randn(B, T, H, device=device, requires_grad=True)
+
+y, (slot, mass) = random_access_memory_attention(
+    w_addr, r_addr, values, gamma,
+    embed_num=2, top_k=8,
+    impl_backend=backend,
+)
+loss = y.square().mean() + slot.square().mean() + mass.square().mean()
+loss.backward()
+print(y.shape)     # (B, T, H, C)
+print(slot.shape)  # (B, H, P ** U, C)
+print(mass.shape)  # (B, H, P ** U)
+```
+
+`embed_num` selects how many heads, starting from head 0, use **Cyclic Address Positional Embedding (CAPE)** to cyclically shift read and write slot indices by token position.
+With `H=4` and `embed_num=2`, heads 0 and 1 use CAPE, while heads 2 and 3 address by content alone.
+Its range is `0` to `H`: `0` disables CAPE and `H` enables it for all heads.
+
+From the repository root, run [ops_smoke.py](../../examples/ops_smoke.py) with `python examples/ops_smoke.py --device cuda`, or use `--device cpu` for the PyTorch reference.
+This example demonstrates the API; it is not an experiment configuration.
+
+### Arguments
+
+- `w_addr`, `r_addr`: write and read address logits of shape `(…, T, …, H, U, P)`. The operator applies a softmax over each digit.
+- `values`: values to write, of shape `(…, T, …, H, C)`. The output `y` has the same shape and dtype.
+- `gamma`: forgetting modulator of shape `(…, T, …, H)`, computed by the model for each token. Larger values overwrite old slot content faster.
+- `init_ram`: the `(slot, mass)` state returned by a previous call, used to continue a sequence. Without it, every slot starts empty.
+- `inplace`: if `True`, updates the `init_ram` tensors in place instead of returning new ones. This requires contiguous FP32 state and does not support gradients.
+- `mask`: marks valid tokens. Masked tokens are not written to the state, and their outputs should be ignored.
+- `embed_num`: number of heads that use CAPE, counted from head 0. These heads shift each token's addresses by its position.
+- `top_k`: `K`, given as an int or a `(write_k, read_k)` pair. `None` keeps all slots and is supported only by the PyTorch backend.
+- `top_k_re_norm`: explicitly enables renormalization after Top-K. Accepts a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
+- `t_dim`: index of the time axis. Dimensions before `H` other than the time axis are broadcast as batch dimensions.
+- `impl_backend`: `'cuda'` or `'torch'`. By default, CUDA inputs use the CUDA backend when supported and otherwise fall back to PyTorch with a warning.
+
+## CUDA execution
+
+The CUDA backend automatically selects between **autoregressive** and **parallel** execution.
+
+**Autoregressive mode** handles single-token inputs (`T = 1`) when no input requires gradients.
+A dedicated kernel decodes the current token's addresses, updates the selected write slots, and reads the selected read slots directly.
+
+**Parallel mode** handles all other inputs, including training and prefill.
+It decodes all tokens' addresses in parallel, treating each selected read or write access as an **event**.
+Events targeting the same `(batch, head, slot)` are grouped into a **segment**.
+The kernels process independent segments and their value channels in parallel while preserving causal dependencies within each segment, then accumulate read contributions into the token outputs.
+
+## Autograd and compile
+
+Operators are registered under `torch.ops.ramnet_ops` with FakeTensor and first-order autograd support.
+Compile coverage targets fixed shapes/configuration with `dynamic=False`, tested through `aot_eager`;
+it does not establish Inductor or CUDA Graph support.
+
+Segment capacity is `B * H * A`; event capacity is `B * H * T_padded * (write_k + read_k)`.
+Masks do not reduce these allocations.
+Higher-order gradients are not supported.
+
+## Development
+
+See [RELEASING](RELEASING.md) for tests and publishing.
+
+MIT License. See [LICENSE](LICENSE).

ramnet_ops-0.1.0rc0/RELEASING.md

@@ -0,0 +1,47 @@
+# Releasing ramnet-ops
+
+`ramnet` and `ramnet-ops` are built and published independently, with versions read from their own `pyproject.toml`.
+Use `ramnet/v<VERSION>` for this package and `ramnet-ops/v<VERSION>` for the operators.
+
+## Tests
+
+Install the candidate wheel and its [dependencies](README.md#installation), then run from the repository root on a CUDA GPU:
+
+```bash
+python -m tests --part ops --device cuda
+```
+
+- Add `--fp16` and/or `--bf16` for optional low-precision checks; BF16 requires a supported GPU.
+- Use `--group contracts`, `--group behavior`, or `--group numerics` to select a group.
+- Seeds default to `0 1 2`; use `--seeds 0` for a focused rerun.
+- Use `--device cpu` for CPU-only checks; this does not validate the CUDA kernels.
+
+After public API changes, also check model integration with `ramnet` and its dependencies installed:
+
+```bash
+python -m tests --part models --device cuda
+```
+
+Tests run locally; the build and publish workflows do not run them.
+See the [test guide](../../tests/README.md) for selection options, parameter coverage, and acceptance criteria.
+
+## Release checklist
+
+- [ ] Set a new `version` in [pyproject.toml](pyproject.toml).
+- [ ] Update version references in [README → Other Torch/CUDA builds (GitHub Releases)](README.md#other-torchcuda-builds-github-releases)
+- [ ] Push the release commit to `main` and wait for [ramnet-ops build](../../.github/workflows/ramnet-ops-build.yml) to succeed.
+- [ ] Install a wheel from that run's artifacts and run the tests above.
+- [ ] Create and push `ramnet-ops/v<VERSION>` at the exact build commit; the tag version must match `pyproject.toml` at that commit.
+- [ ] Keep the successful build run ID and test results for publishing.
+
+## Publish
+
+Run [ramnet-ops publish](../../.github/workflows/ramnet-ops-publish.yml) manually with:
+
+- `build_run_id`: the numeric ID from the successful build's `/actions/runs/<ID>` URL.
+- `tag`: the full `ramnet-ops/v<VERSION>` tag.
+- `target`: `testpypi`, `pypi`, or `release` for GitHub Release only.
+
+Publish to TestPyPI first, verify installation and tests, then publish to PyPI using the same build ID and tag.
+The workflow checks build/tag/version agreement and reuses the original artifacts without rebuilding.
+Every target creates or verifies the GitHub Release first; an existing matching release is reused without uploading its assets again.