ramnet-ops

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

.
|-- 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:

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 for installation commands. Then install the remaining dependencies from TestPyPI / PyPI:

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.

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) 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, 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, 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):

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):

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.

★ 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.

NVIDIA drivers

The table lists driver baselines from NVIDIA's CUDA GA release notes; use these or newer drivers that support your GPU.

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, 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:

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:

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:

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.

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.

This example computes attention outputs and the updated state, then backpropagates through both:

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 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 for tests and publishing.

MIT License. See LICENSE.