ramnet 0.1.0rc0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025-2026 Kaicheng Xiao, Haotian Li, Liran Dong
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,261 @@
1
+ Metadata-Version: 2.4
2
+ Name: ramnet
3
+ Version: 0.1.0rc0
4
+ Summary: RAM-Net layers and Hugging Face models
5
+ Maintainer: Kaicheng Xiao, Haotian Li, Liran Dong
6
+ License-Expression: MIT
7
+ Project-URL: Source, https://github.com/MuonCat/RAM-Net
8
+ Project-URL: Issues, https://github.com/MuonCat/RAM-Net/issues
9
+ Project-URL: Documentation, https://github.com/MuonCat/RAM-Net/tree/main/packages/ramnet
10
+ Requires-Python: >=3.10
11
+ Description-Content-Type: text/markdown
12
+ License-File: LICENSE
13
+ Requires-Dist: ramnet-ops>=0.1.0b0
14
+ Requires-Dist: torch>=2.5
15
+ Requires-Dist: einops
16
+ Requires-Dist: transformers>=4.57.3
17
+ Requires-Dist: flash-linear-attention>=0.3.2
18
+ Dynamic: license-file
19
+
20
+ <div align="center">
21
+
22
+ <img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
23
+
24
+ # ramnet
25
+
26
+ [![arXiv](https://img.shields.io/badge/arXiv-2602.11958-b31b1b.svg?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2602.11958)
27
+ [![Project Page](https://img.shields.io/badge/Project-Page-1f6feb.svg)](https://muoncat.github.io/ramnet_web/)
28
+ [![Weights](https://img.shields.io/badge/Weights-Hugging%20Face-FFD21E?logo=huggingface&logoColor=white)](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
29
+ [![GitHub](https://img.shields.io/badge/GitHub-RAM--Net-181717?logo=github&logoColor=white)](https://github.com/MuonCat/RAM-Net)
30
+ [![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)
31
+ [![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)
32
+ [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
33
+
34
+ </div>
35
+
36
+ RAM-Net attention layers and Hugging Face compatible language models.
37
+ This pure Python package uses [ramnet-ops](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md) for CUDA execution.
38
+ See the [repository overview](https://github.com/MuonCat/RAM-Net#overview) for the method and experimental results.
39
+
40
+ ## Structure
41
+
42
+ ```text
43
+ .
44
+ `-- src/
45
+ `-- ramnet/
46
+ |-- __init__.py # Public exports
47
+ |-- layers/
48
+ | |-- ram_attn_block.py # Attention layer and explicit state API
49
+ | `-- cache.py # Recurrent state and Hugging Face cache integration
50
+ `-- models/
51
+ |-- __init__.py # Hugging Face Auto class registration
52
+ |-- config_ramnet.py # RAMNetConfig
53
+ `-- modeling_ramnet.py # Blocks, backbone, causal LM, and generation
54
+ ```
55
+
56
+ Public exports: `RAMNetConfig`, `RAMNetModel`, `RAMNetForCausalLM`, `RAMNetAttention`, `RAMState`, and `RAMCache`.
57
+ Importing `ramnet` registers its Hugging Face Auto classes.
58
+
59
+ ## Installation
60
+
61
+ Requires:
62
+
63
+ ```text
64
+ Python >= 3.10
65
+
66
+ ramnet
67
+ |-- torch >= 2.5
68
+ |-- einops
69
+ |-- ramnet-ops >= 0.1.0b0
70
+ | `-- numpy
71
+ |-- transformers >= 4.57.3
72
+ `-- flash-linear-attention >= 0.3.2
73
+ ```
74
+
75
+ The attention layer uses the CUDA backend, so a compatible NVIDIA GPU and CUDA-enabled `ramnet-ops` build are **required**.
76
+
77
+ ### 1. Install dependencies
78
+
79
+ First install PyTorch following the [official installation guide](https://pytorch.org/get-started/locally/).
80
+ Choose PyTorch ≥ 2.10 with CUDA 13.0 for the default ops wheels, then install the remaining dependencies from TestPyPI / PyPI:
81
+
82
+ ```bash
83
+ python -m pip install einops numpy "transformers>=4.57.3" "flash-linear-attention>=0.3.2"
84
+ ```
85
+
86
+ ### 2. Install ramnet-ops
87
+
88
+ Select a CUDA build matching your PyTorch version, GPU, and platform using the [ops installation guide](../ramnet-ops/README.md#installation).
89
+ For environments matching the [default stable ABI wheel requirements](../ramnet-ops/README.md#2-install-ramnet-ops), install directly from TestPyPI / PyPI:
90
+
91
+ ```bash
92
+ python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
93
+ ```
94
+
95
+ ### 3. Install ramnet
96
+
97
+ The model package is pure Python and needs no compiler.
98
+ Install it from TestPyPI / PyPI:
99
+
100
+ ```bash
101
+ python -m pip install -i https://test.pypi.org/simple/ ramnet
102
+ ```
103
+
104
+ #### From source
105
+
106
+ With the dependencies above installed, choose one source.
107
+
108
+ **Download with pip** — install the published source archive from TestPyPI / PyPI:
109
+
110
+ ```bash
111
+ python -m pip install -i https://test.pypi.org/simple/ ramnet --no-binary=ramnet
112
+ ```
113
+
114
+ **Clone from GitHub and install locally**:
115
+
116
+ ```bash
117
+ git clone https://github.com/MuonCat/RAM-Net.git
118
+ cd RAM-Net
119
+ python -m pip install ./packages/ramnet
120
+ ```
121
+
122
+ If you also need to build the CUDA operators, follow [ops source installation](../ramnet-ops/README.md#from-source).
123
+
124
+ ## Quick start
125
+
126
+ `ramnet` provides the Hugging Face model interface, including Auto classes, language-modeling loss, and cached generation with `generate`.
127
+ The table maps paper notation to model configuration parameters.
128
+
129
+ | Paper quantity / symbol | Model parameter |
130
+ | :---: | :---: |
131
+ | Address order $U$ (digits per address) | `addr_order` |
132
+ | Address radix $d_p$ (values per digit) | `addr_radix` |
133
+ | Value dimension $d_v$ | `value_chs` |
134
+ | Slots per head $d_p^U$ | `addr_radix ** addr_order` (derived) |
135
+ | Selected slots $K$ | `addr_top_k` |
136
+
137
+ ### Training and generation
138
+
139
+ This example uses a small model with random weights, matching [model_smoke.py](../../examples/model_smoke.py):
140
+
141
+ ```python
142
+ import torch
143
+ from transformers import AutoModelForCausalLM
144
+ from ramnet import RAMNetConfig
145
+
146
+ config = RAMNetConfig(
147
+ hidden_size=128, num_attention_heads=4, num_hidden_layers=2,
148
+ addr_order=3, addr_radix=4,
149
+ embed_num=2, addr_top_k=8,
150
+ output_norm=True,
151
+ hidden_act="SwiGLU",
152
+ vocab_size=1024, pad_token_id=0,
153
+ dtype="float32",
154
+ )
155
+ model = AutoModelForCausalLM.from_config(config).cuda()
156
+ input_ids = torch.randint(3, config.vocab_size, (2, 32), device="cuda")
157
+ attention_mask = torch.ones_like(input_ids)
158
+
159
+ model.train()
160
+ output = model(
161
+ input_ids=input_ids, attention_mask=attention_mask,
162
+ labels=input_ids, use_cache=False,
163
+ )
164
+ output.loss.backward()
165
+ print(output.logits.shape) # (2, 32, 1024)
166
+
167
+ model.eval()
168
+ with torch.inference_mode():
169
+ tokens = model.generate(
170
+ input_ids=input_ids, attention_mask=attention_mask,
171
+ max_new_tokens=8, do_sample=False, use_cache=True,
172
+ )
173
+ print(tokens.shape)
174
+ ```
175
+
176
+ Run the same example from the repository root with `python examples/model_smoke.py`.
177
+
178
+ `embed_num=2` enables **Cyclic Address Positional Embedding (CAPE)** on the first two heads of each layer, shifting read and write slot indices by token position.
179
+ The other heads address by content alone; `0` disables CAPE and `num_attention_heads` enables it for all heads.
180
+ An integer applies to every layer; a list or tuple gives per-layer counts, with unspecified layers using `0`.
181
+ Set `embed_num` explicitly when changing the model depth or head count, since the default schedule targets 25 layers and 32 heads.
182
+
183
+ ### Save and load
184
+
185
+ Use the standard Hugging Face checkpoint methods after training:
186
+
187
+ ```python
188
+ model.save_pretrained("ramnet-checkpoint")
189
+ model = AutoModelForCausalLM.from_pretrained("ramnet-checkpoint").cuda().eval()
190
+ ```
191
+
192
+ Import `ramnet` before using Auto classes to load a checkpoint in a new process.
193
+
194
+ ## Standalone attention layer
195
+
196
+ ```python
197
+ import torch
198
+ from ramnet import RAMNetAttention
199
+
200
+ attn = RAMNetAttention(
201
+ hidden_size=128, n_heads=4, addr_order=3, addr_radix=4,
202
+ embed_num=2, top_k=8, output_norm=True,
203
+ )
204
+ attn._init_weights()
205
+ attn = attn.cuda()
206
+ x = torch.randn(2, 32, 128, device="cuda")
207
+ slot = torch.zeros(2, 4, 4 ** 3, 32, device="cuda")
208
+ mass = torch.full((2, 4, 4 ** 3), 1 / (4 ** 3), device="cuda")
209
+ y, state = attn.forward_with_state(x, init_ram=(slot, mass))
210
+ slot, mass = state.slot, state.mass
211
+ print(y.shape) # (2, 32, 128)
212
+ print(slot.shape) # (2, 4, 64, 32)
213
+ print(mass.shape) # (2, 4, 64)
214
+ ```
215
+
216
+ Standalone layers need the explicit initialization above; full models apply it automatically.
217
+ The standard attention `forward` accepts `(B, T, D)` inputs and returns `(output, None, past_key_values)`, following the FLA layer convention.
218
+ Set `layer_idx` when constructing attention that uses `RAMCache`.
219
+
220
+ Both attention and complete model blocks provide `forward_with_state`:
221
+
222
+ - `attn.forward_with_state(..., init_ram=state)` runs the attention layer.
223
+ - `model.model.layers[i].forward_with_state(..., init_state=state)` runs the full block, including normalization, MLP, and residual connections.
224
+
225
+ Both accept a `(slot, mass)` pair or `RAMState(slot, mass)` and return `(output, RAMState)`; these APIs also expose `t_dim` and `inplace`.
226
+ The default `inplace=False` preserves input states and supports differentiation through consecutive calls.
227
+ Explicit-state callers supply continuation `position_ids` for CAPE.
228
+
229
+ ## Configuration
230
+
231
+ Blocks use pre-norm RMSNorm; `fuse_norm` selects its fused implementation.
232
+ `hidden_act="fuse_SwiGLU"` (default) or `"SwiGLU"` selects FLA's `GatedMLP`, with or without fusion.
233
+ `"GELU"`, `"SiLU"`, and `"Swish"` select a standard two-layer MLP; `"SiLU"` and `"Swish"` use the same activation.
234
+ Models support `input_ids` or `inputs_embeds`, `return_dict`, hidden-state outputs, and `logits_to_keep`.
235
+
236
+ - `addr_top_k` sets the number of selected slots for both writing and reading; a `(write_k, read_k)` pair specifies separate counts.
237
+ - `addr_top_k_re_norm` explicitly enables renormalization after Top-K: a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
238
+ - `output_norm` enables head-wise RMSNorm; `mass_aware` enables mass-weighted reads.
239
+
240
+ See [RAMNetConfig](src/ramnet/models/config_ramnet.py) for the full parameter list and defaults.
241
+
242
+ ## Cache and sequence handling
243
+
244
+ `config.use_cache` controls the evaluation default; training disables caching unless explicitly requested.
245
+
246
+ - `use_cache=True` carries recurrent state during generation.
247
+ The attention layer automatically uses inplace updates for initialized, gradient-free inference caches.
248
+ - `RAMCache()` allocates states on first use and stores per-layer `recurrent_state` entries.
249
+ It supports beam reordering, batch selection/repetition, and reset.
250
+ Logical CAPE positions are tracked per sample, so a full accumulated mask or a query-only mask can be used after padded prefill.
251
+ Explicit `position_ids` override automatic positions.
252
+ - `use_cache=False` reads a supplied cache without modifying or returning it.
253
+ Gradient checkpointing requires an uncached sequence and disables cache updates during training.
254
+ - Cache rollback, speculative decoding, static/offloaded KV-cache implementations, and full-graph compiled generation are not supported.
255
+ - Packed `cu_seqlens` inputs and attention-matrix outputs (`output_attentions=True`) are not supported.
256
+
257
+ ## Development
258
+
259
+ See [RELEASING](RELEASING.md) for tests and publishing.
260
+
261
+ MIT License. See [LICENSE](LICENSE).
@@ -0,0 +1,242 @@
1
+ <div align="center">
2
+
3
+ <img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
4
+
5
+ # ramnet
6
+
7
+ [![arXiv](https://img.shields.io/badge/arXiv-2602.11958-b31b1b.svg?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2602.11958)
8
+ [![Project Page](https://img.shields.io/badge/Project-Page-1f6feb.svg)](https://muoncat.github.io/ramnet_web/)
9
+ [![Weights](https://img.shields.io/badge/Weights-Hugging%20Face-FFD21E?logo=huggingface&logoColor=white)](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
10
+ [![GitHub](https://img.shields.io/badge/GitHub-RAM--Net-181717?logo=github&logoColor=white)](https://github.com/MuonCat/RAM-Net)
11
+ [![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)
12
+ [![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)
13
+ [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
14
+
15
+ </div>
16
+
17
+ RAM-Net attention layers and Hugging Face compatible language models.
18
+ This pure Python package uses [ramnet-ops](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md) for CUDA execution.
19
+ See the [repository overview](https://github.com/MuonCat/RAM-Net#overview) for the method and experimental results.
20
+
21
+ ## Structure
22
+
23
+ ```text
24
+ .
25
+ `-- src/
26
+ `-- ramnet/
27
+ |-- __init__.py # Public exports
28
+ |-- layers/
29
+ | |-- ram_attn_block.py # Attention layer and explicit state API
30
+ | `-- cache.py # Recurrent state and Hugging Face cache integration
31
+ `-- models/
32
+ |-- __init__.py # Hugging Face Auto class registration
33
+ |-- config_ramnet.py # RAMNetConfig
34
+ `-- modeling_ramnet.py # Blocks, backbone, causal LM, and generation
35
+ ```
36
+
37
+ Public exports: `RAMNetConfig`, `RAMNetModel`, `RAMNetForCausalLM`, `RAMNetAttention`, `RAMState`, and `RAMCache`.
38
+ Importing `ramnet` registers its Hugging Face Auto classes.
39
+
40
+ ## Installation
41
+
42
+ Requires:
43
+
44
+ ```text
45
+ Python >= 3.10
46
+
47
+ ramnet
48
+ |-- torch >= 2.5
49
+ |-- einops
50
+ |-- ramnet-ops >= 0.1.0b0
51
+ | `-- numpy
52
+ |-- transformers >= 4.57.3
53
+ `-- flash-linear-attention >= 0.3.2
54
+ ```
55
+
56
+ The attention layer uses the CUDA backend, so a compatible NVIDIA GPU and CUDA-enabled `ramnet-ops` build are **required**.
57
+
58
+ ### 1. Install dependencies
59
+
60
+ First install PyTorch following the [official installation guide](https://pytorch.org/get-started/locally/).
61
+ Choose PyTorch ≥ 2.10 with CUDA 13.0 for the default ops wheels, then install the remaining dependencies from TestPyPI / PyPI:
62
+
63
+ ```bash
64
+ python -m pip install einops numpy "transformers>=4.57.3" "flash-linear-attention>=0.3.2"
65
+ ```
66
+
67
+ ### 2. Install ramnet-ops
68
+
69
+ Select a CUDA build matching your PyTorch version, GPU, and platform using the [ops installation guide](../ramnet-ops/README.md#installation).
70
+ For environments matching the [default stable ABI wheel requirements](../ramnet-ops/README.md#2-install-ramnet-ops), install directly from TestPyPI / PyPI:
71
+
72
+ ```bash
73
+ python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
74
+ ```
75
+
76
+ ### 3. Install ramnet
77
+
78
+ The model package is pure Python and needs no compiler.
79
+ Install it from TestPyPI / PyPI:
80
+
81
+ ```bash
82
+ python -m pip install -i https://test.pypi.org/simple/ ramnet
83
+ ```
84
+
85
+ #### From source
86
+
87
+ With the dependencies above installed, choose one source.
88
+
89
+ **Download with pip** — install the published source archive from TestPyPI / PyPI:
90
+
91
+ ```bash
92
+ python -m pip install -i https://test.pypi.org/simple/ ramnet --no-binary=ramnet
93
+ ```
94
+
95
+ **Clone from GitHub and install locally**:
96
+
97
+ ```bash
98
+ git clone https://github.com/MuonCat/RAM-Net.git
99
+ cd RAM-Net
100
+ python -m pip install ./packages/ramnet
101
+ ```
102
+
103
+ If you also need to build the CUDA operators, follow [ops source installation](../ramnet-ops/README.md#from-source).
104
+
105
+ ## Quick start
106
+
107
+ `ramnet` provides the Hugging Face model interface, including Auto classes, language-modeling loss, and cached generation with `generate`.
108
+ The table maps paper notation to model configuration parameters.
109
+
110
+ | Paper quantity / symbol | Model parameter |
111
+ | :---: | :---: |
112
+ | Address order $U$ (digits per address) | `addr_order` |
113
+ | Address radix $d_p$ (values per digit) | `addr_radix` |
114
+ | Value dimension $d_v$ | `value_chs` |
115
+ | Slots per head $d_p^U$ | `addr_radix ** addr_order` (derived) |
116
+ | Selected slots $K$ | `addr_top_k` |
117
+
118
+ ### Training and generation
119
+
120
+ This example uses a small model with random weights, matching [model_smoke.py](../../examples/model_smoke.py):
121
+
122
+ ```python
123
+ import torch
124
+ from transformers import AutoModelForCausalLM
125
+ from ramnet import RAMNetConfig
126
+
127
+ config = RAMNetConfig(
128
+ hidden_size=128, num_attention_heads=4, num_hidden_layers=2,
129
+ addr_order=3, addr_radix=4,
130
+ embed_num=2, addr_top_k=8,
131
+ output_norm=True,
132
+ hidden_act="SwiGLU",
133
+ vocab_size=1024, pad_token_id=0,
134
+ dtype="float32",
135
+ )
136
+ model = AutoModelForCausalLM.from_config(config).cuda()
137
+ input_ids = torch.randint(3, config.vocab_size, (2, 32), device="cuda")
138
+ attention_mask = torch.ones_like(input_ids)
139
+
140
+ model.train()
141
+ output = model(
142
+ input_ids=input_ids, attention_mask=attention_mask,
143
+ labels=input_ids, use_cache=False,
144
+ )
145
+ output.loss.backward()
146
+ print(output.logits.shape) # (2, 32, 1024)
147
+
148
+ model.eval()
149
+ with torch.inference_mode():
150
+ tokens = model.generate(
151
+ input_ids=input_ids, attention_mask=attention_mask,
152
+ max_new_tokens=8, do_sample=False, use_cache=True,
153
+ )
154
+ print(tokens.shape)
155
+ ```
156
+
157
+ Run the same example from the repository root with `python examples/model_smoke.py`.
158
+
159
+ `embed_num=2` enables **Cyclic Address Positional Embedding (CAPE)** on the first two heads of each layer, shifting read and write slot indices by token position.
160
+ The other heads address by content alone; `0` disables CAPE and `num_attention_heads` enables it for all heads.
161
+ An integer applies to every layer; a list or tuple gives per-layer counts, with unspecified layers using `0`.
162
+ Set `embed_num` explicitly when changing the model depth or head count, since the default schedule targets 25 layers and 32 heads.
163
+
164
+ ### Save and load
165
+
166
+ Use the standard Hugging Face checkpoint methods after training:
167
+
168
+ ```python
169
+ model.save_pretrained("ramnet-checkpoint")
170
+ model = AutoModelForCausalLM.from_pretrained("ramnet-checkpoint").cuda().eval()
171
+ ```
172
+
173
+ Import `ramnet` before using Auto classes to load a checkpoint in a new process.
174
+
175
+ ## Standalone attention layer
176
+
177
+ ```python
178
+ import torch
179
+ from ramnet import RAMNetAttention
180
+
181
+ attn = RAMNetAttention(
182
+ hidden_size=128, n_heads=4, addr_order=3, addr_radix=4,
183
+ embed_num=2, top_k=8, output_norm=True,
184
+ )
185
+ attn._init_weights()
186
+ attn = attn.cuda()
187
+ x = torch.randn(2, 32, 128, device="cuda")
188
+ slot = torch.zeros(2, 4, 4 ** 3, 32, device="cuda")
189
+ mass = torch.full((2, 4, 4 ** 3), 1 / (4 ** 3), device="cuda")
190
+ y, state = attn.forward_with_state(x, init_ram=(slot, mass))
191
+ slot, mass = state.slot, state.mass
192
+ print(y.shape) # (2, 32, 128)
193
+ print(slot.shape) # (2, 4, 64, 32)
194
+ print(mass.shape) # (2, 4, 64)
195
+ ```
196
+
197
+ Standalone layers need the explicit initialization above; full models apply it automatically.
198
+ The standard attention `forward` accepts `(B, T, D)` inputs and returns `(output, None, past_key_values)`, following the FLA layer convention.
199
+ Set `layer_idx` when constructing attention that uses `RAMCache`.
200
+
201
+ Both attention and complete model blocks provide `forward_with_state`:
202
+
203
+ - `attn.forward_with_state(..., init_ram=state)` runs the attention layer.
204
+ - `model.model.layers[i].forward_with_state(..., init_state=state)` runs the full block, including normalization, MLP, and residual connections.
205
+
206
+ Both accept a `(slot, mass)` pair or `RAMState(slot, mass)` and return `(output, RAMState)`; these APIs also expose `t_dim` and `inplace`.
207
+ The default `inplace=False` preserves input states and supports differentiation through consecutive calls.
208
+ Explicit-state callers supply continuation `position_ids` for CAPE.
209
+
210
+ ## Configuration
211
+
212
+ Blocks use pre-norm RMSNorm; `fuse_norm` selects its fused implementation.
213
+ `hidden_act="fuse_SwiGLU"` (default) or `"SwiGLU"` selects FLA's `GatedMLP`, with or without fusion.
214
+ `"GELU"`, `"SiLU"`, and `"Swish"` select a standard two-layer MLP; `"SiLU"` and `"Swish"` use the same activation.
215
+ Models support `input_ids` or `inputs_embeds`, `return_dict`, hidden-state outputs, and `logits_to_keep`.
216
+
217
+ - `addr_top_k` sets the number of selected slots for both writing and reading; a `(write_k, read_k)` pair specifies separate counts.
218
+ - `addr_top_k_re_norm` explicitly enables renormalization after Top-K: a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
219
+ - `output_norm` enables head-wise RMSNorm; `mass_aware` enables mass-weighted reads.
220
+
221
+ See [RAMNetConfig](src/ramnet/models/config_ramnet.py) for the full parameter list and defaults.
222
+
223
+ ## Cache and sequence handling
224
+
225
+ `config.use_cache` controls the evaluation default; training disables caching unless explicitly requested.
226
+
227
+ - `use_cache=True` carries recurrent state during generation.
228
+ The attention layer automatically uses inplace updates for initialized, gradient-free inference caches.
229
+ - `RAMCache()` allocates states on first use and stores per-layer `recurrent_state` entries.
230
+ It supports beam reordering, batch selection/repetition, and reset.
231
+ Logical CAPE positions are tracked per sample, so a full accumulated mask or a query-only mask can be used after padded prefill.
232
+ Explicit `position_ids` override automatic positions.
233
+ - `use_cache=False` reads a supplied cache without modifying or returning it.
234
+ Gradient checkpointing requires an uncached sequence and disables cache updates during training.
235
+ - Cache rollback, speculative decoding, static/offloaded KV-cache implementations, and full-graph compiled generation are not supported.
236
+ - Packed `cu_seqlens` inputs and attention-matrix outputs (`output_attentions=True`) are not supported.
237
+
238
+ ## Development
239
+
240
+ See [RELEASING](RELEASING.md) for tests and publishing.
241
+
242
+ MIT License. See [LICENSE](LICENSE).
@@ -0,0 +1,30 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0.3", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "ramnet"
7
+ version = "0.1.0rc0"
8
+ description = "RAM-Net layers and Hugging Face models"
9
+ maintainers = [{name = "Kaicheng Xiao"}, {name = "Haotian Li"}, {name = "Liran Dong"}]
10
+ readme = "README.md"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ requires-python = ">=3.10"
14
+ dependencies = [
15
+ "ramnet-ops>=0.1.0b0",
16
+ "torch>=2.5",
17
+ "einops",
18
+ "transformers>=4.57.3",
19
+ "flash-linear-attention>=0.3.2",
20
+ ]
21
+
22
+ [project.urls]
23
+ Source = "https://github.com/MuonCat/RAM-Net"
24
+ Issues = "https://github.com/MuonCat/RAM-Net/issues"
25
+ Documentation = "https://github.com/MuonCat/RAM-Net/tree/main/packages/ramnet"
26
+
27
+ [tool.setuptools.packages.find]
28
+ where = ["src"]
29
+ include = ["ramnet", "ramnet.*"]
30
+ namespaces = false
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,11 @@
1
+ from .layers import RAMNetAttention, RAMCache, RAMState
2
+ from .models import RAMNetConfig, RAMNetModel, RAMNetForCausalLM
3
+
4
+ __all__ = [
5
+ 'RAMCache',
6
+ 'RAMState',
7
+ 'RAMNetAttention',
8
+ 'RAMNetConfig',
9
+ 'RAMNetModel',
10
+ 'RAMNetForCausalLM',
11
+ ]
@@ -0,0 +1,4 @@
1
+ from .ram_attn_block import RAMNetAttention
2
+ from .cache import RAMCacheLayer, RAMCache, RAMState
3
+
4
+ __all__ = ['RAMNetAttention', 'RAMCacheLayer', 'RAMCache', 'RAMState']