aabbdev commited on
Commit
4a4cfc1
·
verified ·
1 Parent(s): 5e31b0e

Publish RWKV7-1.5B-20260805

Browse files
.gitattributes CHANGED
@@ -1,35 +1,3 @@
1
- *.7z filter=lfs diff=lfs merge=lfs -text
2
- *.arrow filter=lfs diff=lfs merge=lfs -text
3
- *.bin filter=lfs diff=lfs merge=lfs -text
4
- *.bz2 filter=lfs diff=lfs merge=lfs -text
5
- *.ckpt filter=lfs diff=lfs merge=lfs -text
6
- *.ftz filter=lfs diff=lfs merge=lfs -text
7
- *.gz filter=lfs diff=lfs merge=lfs -text
8
- *.h5 filter=lfs diff=lfs merge=lfs -text
9
- *.joblib filter=lfs diff=lfs merge=lfs -text
10
- *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
- *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
- *.model filter=lfs diff=lfs merge=lfs -text
13
- *.msgpack filter=lfs diff=lfs merge=lfs -text
14
- *.npy filter=lfs diff=lfs merge=lfs -text
15
- *.npz filter=lfs diff=lfs merge=lfs -text
16
- *.onnx filter=lfs diff=lfs merge=lfs -text
17
- *.ot filter=lfs diff=lfs merge=lfs -text
18
- *.parquet filter=lfs diff=lfs merge=lfs -text
19
- *.pb filter=lfs diff=lfs merge=lfs -text
20
- *.pickle filter=lfs diff=lfs merge=lfs -text
21
- *.pkl filter=lfs diff=lfs merge=lfs -text
22
- *.pt filter=lfs diff=lfs merge=lfs -text
23
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
25
  *.safetensors filter=lfs diff=lfs merge=lfs -text
26
- saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
- *.tar.* filter=lfs diff=lfs merge=lfs -text
28
- *.tar filter=lfs diff=lfs merge=lfs -text
29
- *.tflite filter=lfs diff=lfs merge=lfs -text
30
- *.tgz filter=lfs diff=lfs merge=lfs -text
31
- *.wasm filter=lfs diff=lfs merge=lfs -text
32
- *.xz filter=lfs diff=lfs merge=lfs -text
33
- *.zip filter=lfs diff=lfs merge=lfs -text
34
- *.zst filter=lfs diff=lfs merge=lfs -text
35
- *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  *.safetensors filter=lfs diff=lfs merge=lfs -text
2
+ *.pth filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
LICENSE ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and do not
117
+ modify the License. You may add Your own attribution notices
118
+ within Derivative Works that You distribute, alongside or as an
119
+ addendum to the NOTICE text from the Work, provided that such
120
+ additional attribution notices cannot be construed as modifying
121
+ the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer, and
167
+ charge a fee for, acceptance of support, warranty, indemnity, or
168
+ other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
NOTICE ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ RWKV-7 model release
2
+ Source: BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth
3
+ Bundled Transformers RWKV-7 code: huggingface/transformers@4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2 (Apache-2.0).
4
+ Exported inference bundle licensed under Apache-2.0.
README.md ADDED
@@ -0,0 +1,279 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ pipeline_tag: text-generation
4
+ license: apache-2.0
5
+ language:
6
+ - "en"
7
+ - "zh"
8
+ - "fr"
9
+ - "es"
10
+ - "de"
11
+ - "pt"
12
+ - "ru"
13
+ - "it"
14
+ - "ja"
15
+ - "ko"
16
+ - "vi"
17
+ - "ar"
18
+ datasets:
19
+ - "HuggingFaceFW/fineweb-edu"
20
+ - "mlfoundations/dclm-baseline-1.0"
21
+ - "cerebras/SlimPajama-627B"
22
+ - "EleutherAI/pile"
23
+ - "bigcode/starcoderdata"
24
+ - "oscar-corpus/OSCAR-2301"
25
+ tags:
26
+ - rwkv
27
+ - rwkv7
28
+ - recurrent
29
+ - causal-lm
30
+ - conversational
31
+ ---
32
+
33
+ <!-- markdownlint-disable first-line-h1 -->
34
+ <!-- markdownlint-disable html -->
35
+
36
+ <div align="center">
37
+ <a href="https://www.rwkv.com/">
38
+ <img src="https://www.rwkv.com/images/avatar.png" width="140" alt="RWKV logo" />
39
+ </a>
40
+ <h1>RWKV7-1.5B-20260805</h1>
41
+ <p><strong>RWKV-7 “Goose” · constant-state recurrent language modeling</strong></p>
42
+ </div>
43
+
44
+ <div align="center">
45
+ <a href="https://www.rwkv.com/"><img alt="Website" src="https://img.shields.io/badge/Website-RWKV-16a7c9" /></a>
46
+ <a href="https://huggingface.co/BlinkDL"><img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-BlinkDL-ffc107" /></a>
47
+ <a href="https://github.com/BlinkDL/RWKV-LM"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-RWKV--LM-181717?logo=github" /></a>
48
+ <a href="https://arxiv.org/abs/2503.14456v2"><img alt="RWKV-7 paper" src="https://img.shields.io/badge/Paper-arXiv%3A2503.14456-b31b1b" /></a>
49
+ <a href="LICENSE"><img alt="License" src="https://img.shields.io/badge/License-apache-2.0-4c8bf5" /></a>
50
+ </div>
51
+
52
+ ---
53
+
54
+ ## Model introduction
55
+
56
+ This is an official BlinkDL release of **RWKV-7 Goose** in Hugging Face
57
+ Transformers format. RWKV-7 is an attention-free recurrent architecture with a
58
+ constant-size recurrent state and constant inference work per generated token.
59
+ Training remains parallelizable.
60
+
61
+ This checkpoint is a **base model** pretrained with web, code, synthetic, instruction, chat, and reasoning data. It is suitable for evaluation, post-training, and fine-tuning; the included chat template is a prompt interface, not a claim that the checkpoint is a safety-aligned assistant.
62
+
63
+ The Transformers integration, conversion, release packaging, Fast Tokenizer, and
64
+ optional TileLang inference implementation are distributed with this release.
65
+
66
+ ## Highlights
67
+
68
+ - **Constant recurrent state:** memory does not grow like an attention KV cache.
69
+ - **Bundled Transformers integration:** auditable remote configuration and modeling
70
+ modules provide generation, recurrent cache continuation, training, and LoRA
71
+ workflows on Transformers 5.15+.
72
+ - **Exact Fast Tokenizer:** self-contained Rust-backed `tokenizer.json`, generated
73
+ from the canonical RWKV World byte vocabulary during conversion.
74
+ - **Chat-ready:** `chat_template.jinja` supports system, multi-turn, thinking, and
75
+ strict model-generated tool-call prompts.
76
+ - **Optional optimized runtime:** the isolated [`inference/`](inference/) bundle
77
+ provides PyTorch fallback and TileLang acceleration without changing the
78
+ standard model root.
79
+
80
+ ## Model overview
81
+
82
+ | Field | Value |
83
+ | --- | --- |
84
+ | Repository | `aabbdev/RWKV7-1.5B-20260805` |
85
+ | Architecture class | `Rwkv7ForCausalLM` |
86
+ | Public size label | `1.5`B |
87
+ | Source parameters | `1,527,668,736` |
88
+ | Serialized parameters | `1,527,668,736` |
89
+ | Synthesized compatibility tensors | `0` |
90
+ | Layers | `24` |
91
+ | Hidden / FFN size | `2048` / `8192` |
92
+ | Heads / head size | `32` / `64` |
93
+ | Vocabulary | `65536` |
94
+ | Training context | `16384 tokens` |
95
+ | Weight dtype | `bfloat16` |
96
+ | Numerical conversion | `source dtype preserved` |
97
+ | Metadata profile | `g1i` |
98
+ | Metadata provenance | `locked-profile` |
99
+ | Source checkpoint | [`BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth`](https://huggingface.co/BlinkDL/rwkv7-g1/blob/ede85bf8ab2e59aff7d7ca909fbbc73317866d89/rwkv7-g1i-1.5b-20260805-ctx16384.pth) |
100
+ | Source SHA-256 | `32ef7b5bf4dc8bde843cf26dfad809a1f527e2e76a9e790e7d406e71bcd785da` |
101
+
102
+ ## Transformers quickstart
103
+
104
+ The repository includes `configuration_rwkv7.py` and `modeling_rwkv7.py`, adapted
105
+ from the Transformers RWKV-7 integration at commit
106
+ [`4ad9ed0`](https://github.com/huggingface/transformers/commit/4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2).
107
+ Review those files and pin a model-repository revision in production. Passing
108
+ `trust_remote_code=True` selects this bundled implementation even when the local
109
+ Transformers installation also provides native RWKV-7 support.
110
+
111
+ ```python
112
+ import torch
113
+ from transformers import (
114
+ AutoModelForCausalLM,
115
+ AutoTokenizer,
116
+ PreTrainedConfig,
117
+ )
118
+
119
+ model_id = "aabbdev/RWKV7-1.5B-20260805"
120
+ tokenizer = AutoTokenizer.from_pretrained(
121
+ model_id,
122
+ config=PreTrainedConfig(),
123
+ )
124
+ model = AutoModelForCausalLM.from_pretrained(
125
+ model_id,
126
+ trust_remote_code=True,
127
+ dtype=torch.bfloat16,
128
+ )
129
+ ```
130
+
131
+ The recurrent cache returned by the model can be passed back for incremental
132
+ decoding. Use an `attention_mask` for padded batches.
133
+
134
+ ## Chat quickstart
135
+
136
+ ```python
137
+ import re
138
+
139
+ import torch
140
+ from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedConfig
141
+
142
+
143
+ THINK_RE = re.compile(r"\A<think>?\s*(.*?)\s*</think>?", re.DOTALL)
144
+
145
+
146
+ def assistant_content(completion, thinking, *, close_incomplete=False):
147
+ prefix = "<think" if thinking else "<think></think"
148
+ reply = prefix + completion
149
+ thinking_block = THINK_RE.match(reply)
150
+ if thinking:
151
+ if thinking_block is not None or not close_incomplete:
152
+ return reply.strip()
153
+ return f"{reply.rstrip()}\n</think>".strip()
154
+ return "" if thinking_block is None else reply[thinking_block.end():].strip()
155
+
156
+ model_id = "aabbdev/RWKV7-1.5B-20260805"
157
+ tokenizer = AutoTokenizer.from_pretrained(
158
+ model_id,
159
+ config=PreTrainedConfig(),
160
+ )
161
+ model = AutoModelForCausalLM.from_pretrained(
162
+ model_id,
163
+ trust_remote_code=True,
164
+ dtype=torch.bfloat16,
165
+ ).to("cuda")
166
+
167
+ messages = [{"role": "user", "content": "Explain why RWKV uses constant state."}]
168
+ thinking = False
169
+ max_new_tokens = 256
170
+ inputs = tokenizer.apply_chat_template(
171
+ messages,
172
+ tokenize=True,
173
+ add_generation_prompt=True,
174
+ thinking=thinking,
175
+ return_dict=True,
176
+ return_tensors="pt",
177
+ ).to(model.device)
178
+
179
+ output = model.generate(
180
+ **inputs,
181
+ max_new_tokens=max_new_tokens,
182
+ do_sample=True,
183
+ temperature=1.0,
184
+ top_p=0.5,
185
+ eos_token_id=0,
186
+ pad_token_id=0,
187
+ stop_strings=["\n\nUser:"],
188
+ tokenizer=tokenizer,
189
+ )
190
+ completion = tokenizer.decode(
191
+ output[0, inputs["input_ids"].shape[1]:],
192
+ skip_special_tokens=True,
193
+ )
194
+ completion = completion.split("\n\nUser:", 1)[0]
195
+ reached_token_limit = output.shape[1] - inputs["input_ids"].shape[1] >= max_new_tokens
196
+ print(
197
+ assistant_content(
198
+ completion,
199
+ thinking,
200
+ close_incomplete=reached_token_limit,
201
+ )
202
+ )
203
+ ```
204
+
205
+ Set `thinking=True` for the RWKV thinking prefix. The intentional generation
206
+ prefixes are `Assistant: <think></think` and `Assistant: <think`; do not append a
207
+ closing `>` to them. The post-processing above reconstructs that prefix before
208
+ removing an empty thinking block or preserving an enabled one. If generation hits
209
+ the token limit inside thinking, it closes the displayed block before returning it.
210
+ Reference stops are token ID `0` and `\n\nUser:`.
211
+
212
+ Strip trailing spaces from user input. The official RWKV prompt guide is available
213
+ in [`RWKV7-G1x-templates.txt`](https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/RWKV7-G1x-templates.txt).
214
+
215
+ ## Optimized local inference
216
+
217
+ Install the versions listed in `inference/requirements.txt`, then run the bundled
218
+ interactive chat:
219
+
220
+ ```bash
221
+ python inference/generate.py --model aabbdev/RWKV7-1.5B-20260805 --backend auto --interactive
222
+ ```
223
+
224
+ Or independent prompts separated by blank lines:
225
+
226
+ ```bash
227
+ python inference/generate.py \
228
+ --model aabbdev/RWKV7-1.5B-20260805 \
229
+ --backend auto \
230
+ --input-file prompts.txt
231
+ ```
232
+
233
+ `--backend auto` uses validated exact optimized boundaries and otherwise falls
234
+ back to PyTorch. Full explicit TileLang execution can change floating-point
235
+ operation order and requires checkpoint-, dtype-, shape-, and device-specific
236
+ parity validation.
237
+
238
+ ## Tokenizer
239
+
240
+ The model root contains one self-contained tokenizer artifact: `tokenizer.json`.
241
+ Textual `vocab.json` and `rwkv_vocab_v20230424.txt` files are intentionally omitted
242
+ because they would duplicate the tokenizer used by Transformers. The tokenizer is
243
+ loaded natively as `PreTrainedTokenizerFast` and never executes remote Python code.
244
+ The explicit generic config prevents `AutoTokenizer` from probing the remote model
245
+ configuration and emitting a harmless model-type fallback warning.
246
+
247
+ ## Intended use and limitations
248
+
249
+ - This is a base causal language model. Quality, instruction following, and
250
+ language behavior depend on the checkpoint and downstream prompting or
251
+ post-training.
252
+ - Assisted or speculative decoding that requires recurrent-cache rollback is not
253
+ supported without retaining prior state snapshots.
254
+ - Optimized support depends on GPU architecture, dtype, batch, and shape.
255
+ Unsupported `auto` configurations fall back to pure PyTorch.
256
+ - Explicit full TileLang execution can change floating-point operation order and
257
+ requires checkpoint-, dtype-, shape-, and device-specific parity validation.
258
+ - No safety, bias, toxicity, factuality, or high-stakes-use evaluation is claimed
259
+ by this model card.
260
+
261
+ ## License and provenance
262
+
263
+ The model weights use the locked profile license `apache-2.0`. The exported inference bundle is licensed separately under [Apache-2.0](LICENSE). The bundled Transformers configuration and modeling modules
264
+ retain their Apache-2.0 headers. See [`NOTICE`](NOTICE) and the source checkpoint
265
+ link above for provenance.
266
+
267
+ ## Citation
268
+
269
+ ```bibtex
270
+ @misc{peng2025250314456,
271
+ title = {RWKV-7 "Goose" with Expressive Dynamic State Evolution},
272
+ author = {Bo Peng and Ruichong Zhang and Daniel Goldstein and Eric Alcaide and Xingjian Du and Haowen Hou and Jiaju Lin and Jiaxing Liu and Janna Lu and William Merrill and Guangyu Song and Kaifeng Tan and Saiteja Utpala and Nathan Wilce and Johan S. Wind and Tianyi Wu and Daniel Wuttke and Christian Zhou-Zheng},
273
+ year = {2025},
274
+ eprint = {2503.14456v2},
275
+ archivePrefix = {arXiv},
276
+ primaryClass = {cs.CL},
277
+ url = {https://arxiv.org/abs/2503.14456v2},
278
+ }
279
+ ```
chat_template.jinja ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- set add_generation_prompt = add_generation_prompt | default(true) -%}
2
+ {%- set thinking = thinking | default(false) -%}
3
+ {%- set bos_token = bos_token | default('', true) -%}
4
+ {%- set tools = tools | default([], true) -%}
5
+ {%- set ns = namespace(system_prompt='') -%}
6
+ {%- for message in messages -%}
7
+ {%- if message.role == 'system' -%}
8
+ {%- set ns.system_prompt = message.content | trim -%}
9
+ {%- endif -%}
10
+ {%- endfor -%}
11
+ {{- bos_token -}}
12
+ {%- if ns.system_prompt or tools | length > 0 -%}
13
+ {{ 'System: ' }}{{ ns.system_prompt }}
14
+ {%- if tools | length > 0 -%}
15
+ {%- if ns.system_prompt %}{{ '\n' }}{%- endif -%}
16
+ {{ 'Tools:\n' -}}
17
+ {{ tools | tojson }}
18
+ {{ '\nWhen using a tool, return only a compact JSON function call in a ```json block, like {"name":"calculator","arguments":{"expression":"2+2"}}. The `name` field must be top-level, never inside `arguments`. Do not copy the tool schema into arguments. Otherwise answer normally.' }}
19
+ {%- endif -%}
20
+ {{ '\n\n' }}
21
+ {%- endif -%}
22
+ {%- for message in messages -%}
23
+ {%- if message.role == 'user' -%}
24
+ {{ 'User: ' ~ (message.content | trim) ~ '\n\n' }}
25
+ {%- elif message.role == 'assistant' -%}
26
+ {%- set content = message.content | default('', true) | trim -%}
27
+ {{ 'Assistant:' }}
28
+ {%- if message.tool_calls is defined and message.tool_calls | length > 0 -%}
29
+ {%- if content %}{{ ' ' ~ content ~ '\n' }}{%- endif -%}
30
+ {%- for tool_call in message.tool_calls -%}
31
+ {%- if tool_call.function is defined -%}
32
+ {%- set name = tool_call.function.name | default('') -%}
33
+ {%- set args = tool_call.function.arguments | default({}, true) -%}
34
+ {%- else -%}
35
+ {%- set name = tool_call.name | default('') -%}
36
+ {%- set args = tool_call.arguments | default({}, true) -%}
37
+ {%- endif -%}
38
+ {{ ' ```json\n' -}}
39
+ {{ '{"name": ' }}{{ name | tojson }}{{ ', "arguments": ' }}{% if args is string %}{{ args }}{% else %}{{ args | tojson }}{% endif %}{{ '}\n' -}}
40
+ {{ '```' }}{{ '\n' if not loop.last else '' }}
41
+ {%- endfor -%}
42
+ {%- elif content -%}
43
+ {{ ' ' ~ content }}
44
+ {%- endif -%}
45
+ {{ '\n\n' }}
46
+ {%- elif message.role == 'tool' -%}
47
+ {{ 'User: Function output:\n' ~ (message.content | trim) ~ '\n\n' }}
48
+ {%- endif -%}
49
+ {%- endfor -%}
50
+ {%- if add_generation_prompt -%}
51
+ {%- if thinking -%}
52
+ {{ 'Assistant: <think' }}
53
+ {%- else -%}
54
+ {{ 'Assistant: <think></think' }}
55
+ {%- endif -%}
56
+ {%- endif -%}
config.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "a_low_rank_dim": 96,
3
+ "architectures": [
4
+ "Rwkv7ForCausalLM"
5
+ ],
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_rwkv7.Rwkv7Config",
8
+ "AutoModel": "modeling_rwkv7.Rwkv7Model",
9
+ "AutoModelForCausalLM": "modeling_rwkv7.Rwkv7ForCausalLM"
10
+ },
11
+ "bos_token_id": 0,
12
+ "decay_low_rank_dim": 96,
13
+ "dtype": "bfloat16",
14
+ "eos_token_id": 0,
15
+ "gate_low_rank_dim": 256,
16
+ "head_dim": 64,
17
+ "hidden_size": 2048,
18
+ "intermediate_size": 8192,
19
+ "model_type": "rwkv7",
20
+ "norm_bias": true,
21
+ "norm_eps": 1e-05,
22
+ "num_heads": 32,
23
+ "num_hidden_layers": 24,
24
+ "pad_token_id": 0,
25
+ "tie_word_embeddings": false,
26
+ "use_cache": true,
27
+ "v_low_rank_dim": 64,
28
+ "vocab_size": 65536,
29
+ "wkv_implementation": "eager",
30
+ "wkv_state_dtype": "float32"
31
+ }
configuration_rwkv7.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """RWKV-7 (Goose) model configuration."""
15
+
16
+ from huggingface_hub.dataclasses import strict
17
+
18
+ from transformers.configuration_utils import PreTrainedConfig
19
+ from transformers.utils import auto_docstring
20
+
21
+
22
+ @auto_docstring(
23
+ checkpoint="RWKV/RWKV7-1.5B-20260805",
24
+ custom_intro="""
25
+ Configuration for [`Rwkv7Model`], an all-recurrent (attention-free) RWKV-7 "Goose"
26
+ model. Instantiating with the defaults yields the ~0.1B RWKV-7 configuration.
27
+
28
+ Parameter names follow the upstream RWKV reference implementation
29
+ (`BlinkDL/RWKV-LM`) rather than a renamed variant.
30
+ """,
31
+ )
32
+ @strict
33
+ class Rwkv7Config(PreTrainedConfig):
34
+ r"""
35
+ vocab_size (`int`, *optional*, defaults to 65536):
36
+ Vocabulary size (RWKV "world" tokenizer).
37
+ hidden_size (`int`, *optional*, defaults to 768):
38
+ Model width `C`.
39
+ num_hidden_layers (`int`, *optional*, defaults to 12):
40
+ Number of blocks.
41
+ head_dim (`int`, *optional*, defaults to 64):
42
+ Width of one WKV head. `hidden_size` must be divisible by it.
43
+ num_heads (`int`, *optional*, defaults to 12):
44
+ Number of WKV heads; must equal `hidden_size // head_dim`.
45
+ decay_low_rank_dim (`int`, *optional*):
46
+ Rank of the decay (`w`) LoRA. When omitted, follows the RWKV-LM recipe:
47
+ `max(32, round(2.5 * sqrt(hidden_size) / 32) * 32)`.
48
+ a_low_rank_dim (`int`, *optional*):
49
+ Rank of the in-context-learning-rate (`a`) LoRA. Uses the same default
50
+ formula as `decay_low_rank_dim`.
51
+ v_low_rank_dim (`int`, *optional*):
52
+ Rank of the value-residual (`v`) LoRA. Defaults to
53
+ `max(32, round(1.7 * sqrt(hidden_size) / 32) * 32)`. Unused on layer 0,
54
+ which *produces* `v_first` instead of mixing towards it.
55
+ gate_low_rank_dim (`int`, *optional*):
56
+ Rank of the output-gate (`g`) LoRA. Defaults to
57
+ `max(32, round(5 * sqrt(hidden_size) / 32) * 32)`.
58
+ intermediate_size (`int`, *optional*):
59
+ Channel-mix inner width. Defaults to `4 * hidden_size`.
60
+ norm_eps (`float`, *optional*, defaults to 1e-05):
61
+ Epsilon of every LayerNorm/GroupNorm in the model.
62
+ norm_bias (`bool`, *optional*, defaults to `True`):
63
+ Whether the norms carry a bias.
64
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
65
+ Whether to tie the input embedding and the LM head.
66
+ use_cache (`bool`, *optional*, defaults to `True`):
67
+ Whether to return the recurrent state.
68
+ wkv_state_dtype (`str`, *optional*, defaults to `"float32"`):
69
+ Precision the recurrent WKV state is carried and accumulated in,
70
+ independently of the activation dtype. The recurrence is unrolled over
71
+ the whole sequence, so a narrow state drifts; `"float32"` with fp16
72
+ activations is the combination the reference implementation uses.
73
+ `"float16"`/`"bfloat16"` trade that for a smaller state.
74
+ wkv_implementation (`str`, *optional*, defaults to `"eager"`):
75
+ Which WKV recurrence to use, by name, from
76
+ `models.rwkv7.modeling_rwkv7.RWKV7_WKV_FUNCTIONS`. `"eager"` is the
77
+ exact portable PyTorch path and preserves the reference token order for
78
+ both prefill and decode. `"chunked"` is an explicit faster
79
+ chunk-parallel rearrangement whose floating-point order differs. Register
80
+ an entry in that mapping to plug in a fused or varlen kernel without
81
+ forking the model.
82
+ bos_token_id (`int`, *optional*, defaults to 0):
83
+ Beginning-of-sequence id. The RWKV world tokenizer has no dedicated BOS
84
+ token and the reference implementation prepends nothing, so this exists to
85
+ satisfy `GenerationMixin` rather than to be emitted.
86
+ eos_token_id (`int`, *optional*, defaults to 0):
87
+ End-of-sequence id, id 0 in the RWKV world vocabulary.
88
+ pad_token_id (`int`, *optional*, defaults to 0):
89
+ Padding id, the same id 0. Set deliberately rather than left `None`:
90
+ `generate` needs one to pad a batch, and without it a batched call either
91
+ raised or fell back to the eos id with a warning on every step.
92
+
93
+ ```python
94
+ >>> from transformers import Rwkv7Config, Rwkv7Model
95
+
96
+ >>> configuration = Rwkv7Config()
97
+ >>> model = Rwkv7Model(configuration)
98
+ >>> configuration = model.config
99
+ ```"""
100
+
101
+ model_type = "rwkv7"
102
+ keys_to_ignore_at_inference = ["state"]
103
+
104
+ vocab_size: int = 65536
105
+ hidden_size: int = 768
106
+ num_hidden_layers: int = 12
107
+ head_dim: int = 64
108
+ num_heads: int = 12
109
+ decay_low_rank_dim: int | None = None
110
+ a_low_rank_dim: int | None = None
111
+ v_low_rank_dim: int | None = None
112
+ gate_low_rank_dim: int | None = None
113
+ # `None` rather than a number: a literal default is correct for the default
114
+ # `hidden_size` and silently wrong for every other one, so a config built as
115
+ # `Rwkv7Config(hidden_size=4096, num_heads=64)` would come back with a channel-mix
116
+ # four times narrower than the architecture it names. `__post_init__` resolves it,
117
+ # and the resolved value is written to `config.json` either way.
118
+ intermediate_size: int | None = None
119
+ norm_eps: float = 1e-5
120
+ norm_bias: bool = True
121
+ tie_word_embeddings: bool = False
122
+ use_cache: bool = True
123
+ wkv_state_dtype: str = "float32"
124
+ wkv_implementation: str = "eager"
125
+ bos_token_id: int | None = 0
126
+ eos_token_id: int | None = 0
127
+ pad_token_id: int | None = 0
128
+
129
+ def __post_init__(self, **kwargs):
130
+ rank_unit = self.hidden_size**0.5 / 32
131
+ if self.decay_low_rank_dim is None:
132
+ self.decay_low_rank_dim = max(32, round(2.5 * rank_unit) * 32)
133
+ if self.a_low_rank_dim is None:
134
+ self.a_low_rank_dim = max(32, round(2.5 * rank_unit) * 32)
135
+ if self.v_low_rank_dim is None:
136
+ self.v_low_rank_dim = max(32, round(1.7 * rank_unit) * 32)
137
+ if self.gate_low_rank_dim is None:
138
+ self.gate_low_rank_dim = max(32, round(5 * rank_unit) * 32)
139
+ if self.intermediate_size is None:
140
+ self.intermediate_size = 4 * self.hidden_size
141
+ if self.wkv_state_dtype not in ("float32", "float16", "bfloat16"):
142
+ raise ValueError(f"wkv_state_dtype must be float32/float16/bfloat16, got {self.wkv_state_dtype}")
143
+ if self.hidden_size % self.head_dim != 0:
144
+ raise ValueError(f"hidden_size {self.hidden_size} must be divisible by head_dim {self.head_dim}")
145
+ if self.num_heads != self.hidden_size // self.head_dim:
146
+ raise ValueError(
147
+ f"num_heads must be hidden_size // head_dim = {self.hidden_size // self.head_dim}, "
148
+ f"got {self.num_heads}"
149
+ )
150
+ super().__post_init__(**kwargs)
151
+
152
+
153
+ __all__ = ["Rwkv7Config"]
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 0,
3
+ "eos_token_id": 0,
4
+ "pad_token_id": 0,
5
+ "use_cache": true
6
+ }
inference/generate.py ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import importlib
5
+ import os
6
+ import re
7
+ import sys
8
+ from types import ModuleType
9
+ from pathlib import Path
10
+ from typing import Any
11
+
12
+ import torch
13
+ from transformers import StoppingCriteria, StoppingCriteriaList
14
+
15
+ if __package__ in {None, ""}:
16
+ package_name = "_rwkv7_release_inference"
17
+ package = ModuleType(package_name)
18
+ package.__package__ = package_name
19
+ package.__path__ = [str(Path(__file__).resolve().parent)]
20
+ sys.modules[package_name] = package
21
+ load_model_and_tokenizer = importlib.import_module(
22
+ f"{package_name}.model_loader"
23
+ ).load_model_and_tokenizer
24
+ else:
25
+ from .model_loader import load_model_and_tokenizer
26
+
27
+
28
+ DTYPES = {
29
+ "bfloat16": torch.bfloat16,
30
+ "float16": torch.float16,
31
+ "float32": torch.float32,
32
+ }
33
+ STOP_TEXT = "\n\nUser:"
34
+ THINK_RE = re.compile(r"\A<think>?\s*(.*?)\s*</think>?", re.DOTALL)
35
+
36
+
37
+ class StopOnText(StoppingCriteria):
38
+ def __init__(self, tokenizer: Any, prompt_length: int, stop_text: str) -> None:
39
+ self.tokenizer = tokenizer
40
+ self.prompt_length = prompt_length
41
+ self.stop_text = stop_text
42
+ self.matched = False
43
+
44
+ def __call__(
45
+ self,
46
+ input_ids: torch.LongTensor,
47
+ scores: torch.FloatTensor,
48
+ **kwargs: Any,
49
+ ) -> bool:
50
+ del scores, kwargs
51
+ completion = self.tokenizer.decode(
52
+ input_ids[0, self.prompt_length :],
53
+ skip_special_tokens=False,
54
+ )
55
+ self.matched = self.stop_text in completion
56
+ return self.matched
57
+
58
+
59
+ def _assistant_content(
60
+ completion: str, thinking: bool, *, close_incomplete: bool = False
61
+ ) -> str:
62
+ prefix = "<think" if thinking else "<think></think"
63
+ reply = prefix + completion
64
+ thinking_block = THINK_RE.match(reply)
65
+ if thinking:
66
+ if thinking_block is not None or not close_incomplete:
67
+ return reply.strip()
68
+ return f"{reply.rstrip()}\n</think>".strip()
69
+ return "" if thinking_block is None else reply[thinking_block.end() :].strip()
70
+
71
+
72
+ def _prompt_ids(tokenizer: Any, messages: list[dict[str, str]], thinking: bool):
73
+ tokens = tokenizer.apply_chat_template(
74
+ messages,
75
+ tokenize=True,
76
+ add_generation_prompt=True,
77
+ thinking=thinking,
78
+ return_tensors="pt",
79
+ )
80
+ if hasattr(tokens, "input_ids"):
81
+ tokens = tokens.input_ids
82
+ elif isinstance(tokens, dict):
83
+ tokens = tokens["input_ids"]
84
+ if tokens.ndim == 1:
85
+ tokens = tokens.unsqueeze(0)
86
+ return tokens
87
+
88
+
89
+ @torch.inference_mode()
90
+ def generate_completion(
91
+ model: Any,
92
+ tokenizer: Any,
93
+ messages: list[dict[str, str]],
94
+ *,
95
+ device: str,
96
+ max_new_tokens: int,
97
+ temperature: float,
98
+ top_p: float,
99
+ thinking: bool,
100
+ ) -> str:
101
+ input_ids = _prompt_ids(tokenizer, messages, thinking).to(device)
102
+ prompt_length = input_ids.shape[1]
103
+ stop_on_text = StopOnText(tokenizer, prompt_length, STOP_TEXT)
104
+ generation: dict[str, Any] = {
105
+ "input_ids": input_ids,
106
+ "attention_mask": torch.ones_like(input_ids),
107
+ "max_new_tokens": max_new_tokens,
108
+ "do_sample": temperature > 0,
109
+ "eos_token_id": 0,
110
+ "pad_token_id": 0,
111
+ "stopping_criteria": StoppingCriteriaList([stop_on_text]),
112
+ }
113
+ if temperature > 0:
114
+ generation["temperature"] = temperature
115
+ generation["top_p"] = top_p
116
+ output = model.generate(**generation)
117
+ completion_ids = output[0, prompt_length:]
118
+ completion = tokenizer.decode(completion_ids, skip_special_tokens=True)
119
+ if STOP_TEXT in completion:
120
+ completion = completion.split(STOP_TEXT, 1)[0]
121
+ reached_token_limit = (
122
+ completion_ids.shape[0] >= max_new_tokens and not stop_on_text.matched
123
+ )
124
+ return _assistant_content(
125
+ completion,
126
+ thinking,
127
+ close_incomplete=reached_token_limit,
128
+ )
129
+
130
+
131
+ def _interactive(
132
+ model: Any,
133
+ tokenizer: Any,
134
+ args: argparse.Namespace,
135
+ ) -> None:
136
+ messages: list[dict[str, str]] = []
137
+ print("RWKV-7 Goose — /clear resets the conversation, /exit quits.")
138
+ while True:
139
+ try:
140
+ prompt = input(">>> ")
141
+ except EOFError:
142
+ break
143
+ if prompt == "/exit":
144
+ break
145
+ if prompt == "/clear":
146
+ messages.clear()
147
+ continue
148
+ prompt = prompt.strip()
149
+ if not prompt:
150
+ continue
151
+ messages.append({"role": "user", "content": prompt})
152
+ completion = generate_completion(
153
+ model,
154
+ tokenizer,
155
+ messages,
156
+ device=args.device,
157
+ max_new_tokens=args.max_new_tokens,
158
+ temperature=args.temperature,
159
+ top_p=args.top_p,
160
+ thinking=args.thinking,
161
+ )
162
+ print(completion)
163
+ messages.append({"role": "assistant", "content": completion})
164
+
165
+
166
+ def _file_prompts(
167
+ model: Any,
168
+ tokenizer: Any,
169
+ args: argparse.Namespace,
170
+ ) -> None:
171
+ text = Path(args.input_file).read_text(encoding="utf-8")
172
+ prompts = [prompt.strip() for prompt in re.split(r"\n\s*\n", text) if prompt.strip()]
173
+ if not prompts:
174
+ raise ValueError("input file contains no prompts")
175
+ for prompt in prompts:
176
+ completion = generate_completion(
177
+ model,
178
+ tokenizer,
179
+ [{"role": "user", "content": prompt}],
180
+ device=args.device,
181
+ max_new_tokens=args.max_new_tokens,
182
+ temperature=args.temperature,
183
+ top_p=args.top_p,
184
+ thinking=args.thinking,
185
+ )
186
+ print(f"Prompt: {prompt}")
187
+ print(f"Completion: {completion}")
188
+ print()
189
+
190
+
191
+ def parse_args() -> argparse.Namespace:
192
+ parser = argparse.ArgumentParser(description="Generate with RWKV-7 Goose")
193
+ parser.add_argument("--model", required=True, help="Hub repo ID or local model directory")
194
+ mode = parser.add_mutually_exclusive_group(required=True)
195
+ mode.add_argument("--interactive", action="store_true")
196
+ mode.add_argument("--input-file")
197
+ parser.add_argument("--device", default="cuda")
198
+ parser.add_argument("--dtype", choices=("auto", *DTYPES), default="auto")
199
+ parser.add_argument("--state-dtype", choices=DTYPES, default="float32")
200
+ parser.add_argument("--backend", choices=("auto", "torch", "tilelang"), default="auto")
201
+ parser.add_argument("--max-new-tokens", type=int, default=300)
202
+ parser.add_argument("--temperature", type=float, default=1.0)
203
+ parser.add_argument("--top-p", type=float, default=0.5)
204
+ parser.add_argument("--seed", type=int, default=33377335)
205
+ parser.add_argument("--thinking", action="store_true")
206
+ return parser.parse_args()
207
+
208
+
209
+ def main() -> None:
210
+ args = parse_args()
211
+ if int(os.getenv("WORLD_SIZE", "1")) != 1:
212
+ raise RuntimeError("the bundled runtime supports one process and one GPU")
213
+ if int(os.getenv("RANK", "0")) != 0 or int(os.getenv("LOCAL_RANK", "0")) != 0:
214
+ raise RuntimeError("RANK and LOCAL_RANK must be zero")
215
+ if args.max_new_tokens <= 0:
216
+ raise ValueError("max-new-tokens must be positive")
217
+ if args.temperature < 0:
218
+ raise ValueError("temperature must be non-negative")
219
+ if not 0 < args.top_p <= 1:
220
+ raise ValueError("top-p must be in (0, 1]")
221
+ torch.manual_seed(args.seed)
222
+ model, tokenizer = load_model_and_tokenizer(
223
+ args.model,
224
+ device=args.device,
225
+ dtype=None if args.dtype == "auto" else DTYPES[args.dtype],
226
+ backend=args.backend,
227
+ state_dtype=args.state_dtype,
228
+ )
229
+ model.set_kernel_backend(args.backend)
230
+ if args.backend == "tilelang":
231
+ model.prepare_inference_weights()
232
+ if args.interactive:
233
+ _interactive(model, tokenizer, args)
234
+ else:
235
+ _file_prompts(model, tokenizer, args)
236
+
237
+
238
+ if __name__ == "__main__":
239
+ main()
inference/kernel.py ADDED
The diff for this file is too large to render. See raw diff
 
inference/model_loader.py ADDED
@@ -0,0 +1,178 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ from pathlib import Path
5
+
6
+ import torch
7
+ from huggingface_hub import snapshot_download
8
+ from safetensors.torch import load_file
9
+ from transformers import AutoTokenizer, PreTrainedConfig
10
+
11
+ from .runtime import RWKV7Config, RWKV7ForCausalLM
12
+
13
+
14
+ SQUEEZE_MARKERS = (
15
+ ".x_",
16
+ ".k_",
17
+ "att.r",
18
+ "att.w",
19
+ "att.v0",
20
+ "att.v1",
21
+ "att.v2",
22
+ "att.a",
23
+ "att.g",
24
+ )
25
+
26
+
27
+ def _resolve_model(model: str) -> tuple[Path, bool]:
28
+ local = Path(model).expanduser()
29
+ if local.is_dir():
30
+ return local.resolve(), False
31
+ return Path(snapshot_download(model)), True
32
+
33
+
34
+ def _optimized_config(native: dict, backend: str, state_dtype: str) -> RWKV7Config:
35
+ return RWKV7Config(
36
+ vocab_size=native["vocab_size"],
37
+ hidden_size=native["hidden_size"],
38
+ num_hidden_layers=native["num_hidden_layers"],
39
+ head_size=native["head_dim"],
40
+ intermediate_size=native["intermediate_size"],
41
+ decay_lora_rank=native["decay_low_rank_dim"],
42
+ a_lora_rank=native["a_low_rank_dim"],
43
+ gate_lora_rank=native["gate_low_rank_dim"],
44
+ value_lora_rank=native["v_low_rank_dim"],
45
+ layer_norm_epsilon=native.get("norm_eps", 1e-5),
46
+ use_cache=native.get("use_cache", True),
47
+ kernel_backend=backend,
48
+ recurrent_state_dtype=state_dtype,
49
+ tie_word_embeddings=native.get("tie_word_embeddings", False),
50
+ bos_token_id=native.get("bos_token_id"),
51
+ eos_token_id=native.get("eos_token_id", 0),
52
+ pad_token_id=native.get("pad_token_id", 0),
53
+ )
54
+
55
+
56
+ def _native_key_to_optimized(key: str) -> str:
57
+ return key.removeprefix("rwkv7.")
58
+
59
+
60
+ def _native_tensor_to_optimized(key: str, tensor: torch.Tensor) -> torch.Tensor:
61
+ return tensor.squeeze() if any(marker in key for marker in SQUEEZE_MARKERS) else tensor
62
+
63
+
64
+ def _checked_weight(path: Path, *, hub_snapshot: bool) -> Path:
65
+ if path.is_symlink():
66
+ if not hub_snapshot:
67
+ raise RuntimeError(f"local safetensor must not be a symlink: {path.name}")
68
+ path = path.resolve(strict=True)
69
+ if not path.is_file():
70
+ raise RuntimeError(f"safetensor must be a regular file: {path.name}")
71
+ return path
72
+
73
+
74
+ def _weight_plan(
75
+ model_dir: Path, *, hub_snapshot: bool
76
+ ) -> list[tuple[Path, set[str] | None]]:
77
+ present = sorted(model_dir.glob("model*.safetensors"))
78
+ if not present:
79
+ raise FileNotFoundError(f"no safetensors found in {model_dir}")
80
+ checked = {
81
+ path.name: _checked_weight(path, hub_snapshot=hub_snapshot) for path in present
82
+ }
83
+ index_path = model_dir / "model.safetensors.index.json"
84
+ if not index_path.is_file():
85
+ if [path.name for path in present] != ["model.safetensors"]:
86
+ raise RuntimeError("multiple safetensors require model.safetensors.index.json")
87
+ return [(checked["model.safetensors"], None)]
88
+
89
+ index = json.loads(index_path.read_text(encoding="utf-8"))
90
+ weight_map = index.get("weight_map")
91
+ if not isinstance(weight_map, dict) or not weight_map:
92
+ raise RuntimeError("model.safetensors.index.json has no weight_map")
93
+ keys_by_file: dict[str, set[str]] = {}
94
+ for key, filename in weight_map.items():
95
+ if not isinstance(key, str) or not isinstance(filename, str):
96
+ raise RuntimeError("invalid weight_map entry")
97
+ keys_by_file.setdefault(filename, set()).add(key)
98
+ if set(keys_by_file) != {path.name for path in present}:
99
+ raise RuntimeError("indexed and present safetensors files differ")
100
+ return [(checked[filename], keys_by_file[filename]) for filename in sorted(keys_by_file)]
101
+
102
+
103
+ def _load_native_model(
104
+ model_dir: Path,
105
+ native_config: dict,
106
+ backend: str,
107
+ state_dtype: str,
108
+ *,
109
+ hub_snapshot: bool,
110
+ ) -> RWKV7ForCausalLM:
111
+ config = _optimized_config(native_config, backend, state_dtype)
112
+ with torch.device("meta"):
113
+ model = RWKV7ForCausalLM(config)
114
+ expected = set(model.state_dict())
115
+ seen: set[str] = set()
116
+ for weight_file, indexed_keys in _weight_plan(
117
+ model_dir, hub_snapshot=hub_snapshot
118
+ ):
119
+ native_shard = load_file(weight_file, device="cpu")
120
+ if indexed_keys is not None and set(native_shard) != indexed_keys:
121
+ raise RuntimeError(f"tensor keys in {weight_file.name} do not match the index")
122
+ shard = {}
123
+ for native_key, tensor in native_shard.items():
124
+ optimized_key = _native_key_to_optimized(native_key)
125
+ if optimized_key in seen or optimized_key in shard:
126
+ raise RuntimeError(f"duplicate optimized tensor key: {optimized_key}")
127
+ shard[optimized_key] = _native_tensor_to_optimized(optimized_key, tensor)
128
+ unexpected = sorted(set(shard) - expected)
129
+ if unexpected:
130
+ raise RuntimeError(f"native checkpoint has unexpected keys: {unexpected}")
131
+ model.load_state_dict(shard, strict=False, assign=True)
132
+ seen.update(shard)
133
+ missing = sorted(expected - seen)
134
+ if missing:
135
+ raise RuntimeError(f"native checkpoint is missing optimized runtime keys: {missing}")
136
+ return model
137
+
138
+
139
+ def _load_tokenizer(model_dir: Path):
140
+ return AutoTokenizer.from_pretrained(
141
+ model_dir,
142
+ config=PreTrainedConfig(),
143
+ local_files_only=True,
144
+ )
145
+
146
+
147
+ def load_model_and_tokenizer(
148
+ model: str,
149
+ *,
150
+ device: str,
151
+ dtype: torch.dtype | None,
152
+ backend: str,
153
+ state_dtype: str,
154
+ ):
155
+ model_dir, hub_snapshot = _resolve_model(model)
156
+ native_config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
157
+ architectures = set(native_config.get("architectures", []))
158
+ if architectures != {"Rwkv7ForCausalLM"}:
159
+ raise ValueError(f"unsupported RWKV-7 architecture: {sorted(architectures)}")
160
+ loaded = _load_native_model(
161
+ model_dir,
162
+ native_config,
163
+ backend,
164
+ state_dtype,
165
+ hub_snapshot=hub_snapshot,
166
+ )
167
+ if dtype is None:
168
+ dtype_name = str(native_config.get("dtype", "bfloat16")).removeprefix("torch.")
169
+ try:
170
+ dtype = {
171
+ "bfloat16": torch.bfloat16,
172
+ "float16": torch.float16,
173
+ "float32": torch.float32,
174
+ }[dtype_name]
175
+ except KeyError as error:
176
+ raise ValueError(f"unsupported model dtype: {dtype_name}") from error
177
+ loaded = loaded.to(device=device, dtype=dtype).eval()
178
+ return loaded, _load_tokenizer(model_dir)
inference/requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ transformers>=5.15,<6
2
+ huggingface-hub>=0.34
3
+ safetensors>=0.5
4
+ jinja2>=3.1,<4
5
+ tilelang==0.1.12
inference/runtime.py ADDED
The diff for this file is too large to render. See raw diff
 
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:84ccbb857c84e00cefc48b233937ada79c411e491df25fb21aed23237f39a14f
3
+ size 3055418240
modeling_rwkv7.py ADDED
@@ -0,0 +1,1186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The RWKV team and The HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """PyTorch RWKV-7 ("Goose") model."""
15
+
16
+ import math
17
+ from dataclasses import dataclass
18
+
19
+ import torch
20
+ from torch import nn
21
+
22
+ from transformers import initialization as init
23
+ from transformers.cache_utils import Cache, LinearAttentionLayer
24
+ from transformers.generation import GenerationMixin
25
+ from transformers.modeling_layers import GradientCheckpointingLayer
26
+ from transformers.modeling_utils import PreTrainedModel
27
+ from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, logging
28
+ from .configuration_rwkv7 import Rwkv7Config
29
+
30
+
31
+ logger = logging.get_logger(__name__)
32
+
33
+
34
+ def _rwkv7_lora_orthogonal_(tensor: torch.Tensor, scale: float = 0.1) -> None:
35
+ """RWKV-LM compensates the gain when a low-rank output factor is tall."""
36
+ rows, columns = tensor.shape
37
+ gain = math.sqrt(rows / columns) if rows > columns else 1.0
38
+ init.orthogonal_(tensor, gain=gain * scale)
39
+
40
+
41
+ class Rwkv7TokenShift(nn.Module):
42
+ """`prev_token(x)`: the previous token's hidden state, zero at sequence start.
43
+
44
+ Carries one vector per layer per stream across forward calls, which is what
45
+ makes incremental decoding exact rather than approximate.
46
+ """
47
+
48
+ def forward(
49
+ self,
50
+ x: torch.Tensor,
51
+ shift_state: torch.Tensor | None,
52
+ cu_seq_lens: torch.Tensor | None = None,
53
+ keep: torch.Tensor | None = None,
54
+ ):
55
+ # x: [batch, seq_len, hidden]; shift_state: [batch, hidden] or None
56
+ if shift_state is None:
57
+ prev = torch.zeros_like(x[:, :1])
58
+ else:
59
+ prev = shift_state[:, None]
60
+ shifted = torch.cat([prev, x[:, :-1]], dim=1)
61
+ if cu_seq_lens is not None:
62
+ # In a packed row the token before a segment's first one belongs to the
63
+ # PREVIOUS sequence. Resetting the recurrent state per segment does not
64
+ # cover this (the shift reaches back through it), so the first token of
65
+ # each segment gets the zero shift a sequence start is supposed to see.
66
+ positions = torch.arange(x.shape[1], device=x.device)
67
+ starts = (positions[:, None] == cu_seq_lens[None, :-1]).any(dim=1)
68
+ shifted = torch.where(starts[None, :, None], torch.zeros_like(shifted), shifted)
69
+ if keep is None:
70
+ return shifted, x[:, -1]
71
+ # The state handed back must be the last REAL token, not the last position.
72
+ # They coincide under left padding, which is why this went unnoticed; under
73
+ # right padding `x[:, -1]` is a blanked pad, so a continuation resumes from
74
+ # zero instead of from where the sequence actually got to.
75
+ # Integer arithmetic, not the activation dtype. `keep` arrives cast to the
76
+ # model's dtype, and bf16 carries eight mantissa bits: past position 256 the
77
+ # products are no longer distinct, `argmax` returns the first of a tie, and
78
+ # the state comes back from a token several places short of the last real one
79
+ # -- 1022 for 1023 at length 1024, 4088 for 4095 at 4096. An all-ones mask is
80
+ # enough to trigger it, which is what `generate` sends, so this was not
81
+ # confined to padded batches.
82
+ positions = torch.arange(x.shape[1], device=x.device)
83
+ last_real = ((keep.squeeze(-1) > 0) * positions).argmax(dim=-1)
84
+ return shifted, x[torch.arange(x.shape[0], device=x.device), last_real]
85
+
86
+
87
+ # e^-0.5. The decay LoRA emits w_log = -INV_SQRT_E * sigmoid(...), so w_log lies in
88
+ # (-e^-0.5, 0) and the per-step decay exp(w_log) lies in (exp(-e^-0.5), 1), i.e.
89
+ # (0.5452, 1). See RWKV-7 reference. Note the floor is exp(-e^-0.5), not e^-0.5: this
90
+ # comment said the latter, and `rwkv7_chunked` built its chunk_size bound on it.
91
+ _INV_SQRT_E = 0.6065306597126334
92
+
93
+
94
+ def _rwkv7_matmul(x: torch.Tensor, weight: torch.Tensor, tokenwise: bool) -> torch.Tensor:
95
+ """Preserve decode GEMM shapes across a sequence when exact execution is requested."""
96
+ if not tokenwise or x.ndim < 3 or x.shape[1] == 1:
97
+ return x @ weight
98
+ return torch.stack([x[:, token_index] @ weight for token_index in range(x.shape[1])], dim=1)
99
+
100
+
101
+ def _rwkv7_linear(module: nn.Linear, x: torch.Tensor, tokenwise: bool) -> torch.Tensor:
102
+ if not tokenwise or x.ndim < 3 or x.shape[1] == 1:
103
+ return module(x)
104
+ return torch.stack([module(x[:, token_index]) for token_index in range(x.shape[1])], dim=1)
105
+
106
+
107
+ def rwkv7_recurrent(
108
+ r: torch.Tensor,
109
+ w_log: torch.Tensor,
110
+ k: torch.Tensor,
111
+ v: torch.Tensor,
112
+ kk: torch.Tensor,
113
+ a: torch.Tensor,
114
+ state: torch.Tensor,
115
+ compute_dtype: torch.dtype = torch.float32,
116
+ ) -> tuple[torch.Tensor, torch.Tensor]:
117
+ """Reference RWKV-7 WKV recurrence (generalised delta rule).
118
+
119
+ All inputs are `[batch, seq_len, num_heads, head_dim]`; `state` is
120
+ `[batch, num_heads, head_dim, head_dim]` and is carried across calls.
121
+
122
+ Per token, with `S` the per-head state::
123
+
124
+ decay = exp(w_log)
125
+ vk = v @ k
126
+ ab = (-kk) @ (kk * a)
127
+ S = S * decay + S @ ab + vk
128
+ out = S @ r
129
+
130
+ The state axes are (value, key), as in the reference implementation. Products
131
+ made from activations keep the activation dtype; the carried-state arithmetic
132
+ uses `compute_dtype`, and a narrower configured state is stored at every token
133
+ boundary before the output projection. Those boundaries are observable in a
134
+ recurrent BF16 model, so an algebraically transposed FP32 implementation is not
135
+ equivalent.
136
+ """
137
+ batch, seq_len, num_heads, head_dim = r.shape
138
+ dtype = r.dtype
139
+ state_dtype = state.dtype
140
+ state = state.to(compute_dtype)
141
+ out = torch.empty(batch, seq_len, num_heads, head_dim, device=r.device, dtype=compute_dtype)
142
+
143
+ for t in range(seq_len):
144
+ decay = torch.exp(w_log[:, t])[:, :, None, :] # [B, H, 1, K]
145
+ kk_t = kk[:, t]
146
+ value_key = v[:, t][..., None] @ k[:, t][:, :, None, :]
147
+ anti_value_key = (-kk_t)[..., None] @ (kk_t * a[:, t])[:, :, None, :]
148
+ matrix = state.to(compute_dtype)
149
+ state = matrix * decay.to(compute_dtype)
150
+ state = state + matrix @ anti_value_key.to(compute_dtype)
151
+ state = state + value_key.to(compute_dtype)
152
+ if state_dtype != compute_dtype:
153
+ state = state.to(state_dtype)
154
+ out[:, t] = (state.to(dtype) @ r[:, t][..., None]).squeeze(-1).to(compute_dtype)
155
+
156
+ return out.to(dtype), state
157
+
158
+
159
+ def rwkv7_chunked(
160
+ r: torch.Tensor,
161
+ w_log: torch.Tensor,
162
+ k: torch.Tensor,
163
+ v: torch.Tensor,
164
+ kk: torch.Tensor,
165
+ a: torch.Tensor,
166
+ state: torch.Tensor,
167
+ chunk_size: int = 64,
168
+ compute_dtype: torch.dtype = torch.float32,
169
+ ) -> tuple[torch.Tensor, torch.Tensor]:
170
+ """Chunk-parallel form of [`rwkv7_recurrent`], for sequences.
171
+
172
+ The step is `S_t = A_t S_{t-1} + k_t v_t^T` with `A_t = diag(w_t) - b_t kk_t^T`
173
+ and `b_t = kk_t * a_t`, i.e. diagonal-plus-rank-one. Substituting
174
+ `S_t = diag(c_t) P_t` with the running decay `c_t = prod_{s<=t} w_s` removes the
175
+ diagonal part and leaves a plain delta rule::
176
+
177
+ P_t = (I - b~_t q~_t^T) P_{t-1} + k~_t v_t^T
178
+ b~ = b / c, k~ = k / c, q~ = kk * c_{t-1}, r~ = r * c
179
+
180
+ Writing `u_t = q~_t^T P_{t-1}` turns a whole chunk into one unit-lower-triangular
181
+ system, so the chunk is a handful of matmuls instead of `chunk_size` sequential
182
+ steps; only the chunk-to-chunk carry stays serial::
183
+
184
+ (I + tril(Q~ B~^T, -1)) U = Q~ P_0 + tril(Q~ K~^T, -1) V
185
+ O = R~ P_0 + tril(R~ K~^T, 0) V - tril(R~ B~^T, 0) U
186
+ P_C = P_0 + K~^T V - B~^T U
187
+
188
+ `chunk_size` is bounded by that division by `c`, but by OVERFLOW rather than by
189
+ precision, and the difference is worth a factor of four. The per-step decay is at
190
+ least `exp(-e^-0.5)` = 0.5452, so at the worst case, every channel pinned at that
191
+ floor, `1/c` grows like `e^(e^-0.5 * chunk_size)`. That reaches 7.2e16 at 64 and
192
+ fp32 tops out near 3.4e38, so there are twenty-odd decades of headroom left at the
193
+ default, and the ceiling is `ln(finfo.max) / e^-0.5` = 146 in fp32.
194
+
195
+ That derivation is worth stating carefully because the first version of it was
196
+ wrong in a way the numbers hid. It said the decay floor was `e^-0.5` rather than
197
+ `exp(-e^-0.5)`, hence growth like `e^(0.5 * chunk_size)`, which puts the ceiling at
198
+ 177, and measured at the decay floor this function returns all-NaN from 147 up. The
199
+ quoted 7.2e16 came from the correct law all along (`e^(0.5*64)` is 7.9e13), so the
200
+ arithmetic had been done right and written up wrong, and following the prose rather
201
+ than the number led into the overflow band. `chunk_size` is checked below rather
202
+ than left to that reasoning.
203
+
204
+ Precision does not degrade with it, because the substitution is a similarity
205
+ transform: whatever `1/c` inflates, `c` deflates again on the way out, and a
206
+ common scale factor does not move a floating-point relative error. Measured at
207
+ the decay floor, `T=256`, against the sequential form: chunk 16 gives 2.0e-07 and
208
+ chunk 64 gives 2.7e-07, both fp32 noise, while the recurrence runs 3.9x faster
209
+ (3.771 ms -> 0.913 ms). 16 was the conservative reading of the same bound.
210
+ """
211
+ batch, seq_len, num_heads, head_dim = r.shape
212
+ dtype = r.dtype
213
+ r, w_log, k, v, kk, a = (t.to(compute_dtype) for t in (r, w_log, k, v, kk, a))
214
+ state_dtype = state.dtype
215
+ # The chunk derivation below uses the transposed (key, value) form internally.
216
+ state = state.transpose(-2, -1).to(compute_dtype)
217
+ # Everything that does not touch the carried state is computed for ALL chunks at
218
+ # once. Only the state recurrence is serial, and it was pulling the rest of the
219
+ # arithmetic into the Python loop with it: at T=256 and chunk 64 that was four
220
+ # iterations of ten-odd launches each, plus three constant matrices rebuilt every
221
+ # iteration. The loop below now does the state-dependent terms and nothing else.
222
+ # Never pad past the sequence: a 16-token prefill grouped into 64-wide chunks does
223
+ # four times the arithmetic for the same answer, which cost 1x16 nine points before
224
+ # this line existed. The old loop avoided it by shortening its last chunk.
225
+ chunk_size = min(chunk_size, seq_len)
226
+ # Refuse rather than return NaN. `1/c` reaches `e^(e^-0.5 * chunk_size)` at the
227
+ # decay floor, so the widest chunk this dtype can carry is derived from the dtype
228
+ # rather than written down: 146 in fp32, 12 in fp16. A caller who raises
229
+ # `chunk_size` for a longer prefill gets an error naming the limit instead of
230
+ # all-NaN logits several layers later, which is what the previous version did.
231
+ widest = int(math.log(torch.finfo(compute_dtype).max) / _INV_SQRT_E)
232
+ if chunk_size > widest:
233
+ raise ValueError(
234
+ f"chunk_size={chunk_size} overflows {compute_dtype}: the running decay's "
235
+ f"reciprocal grows like e^(e^-0.5 * chunk_size), so the widest chunk that "
236
+ f"stays finite is {widest}. Lower chunk_size, or raise compute_dtype."
237
+ )
238
+ chunks = (seq_len + chunk_size - 1) // chunk_size
239
+ padded = chunks * chunk_size
240
+ if padded != seq_len:
241
+ pad = (0, 0, 0, 0, 0, padded - seq_len)
242
+ r, k, v, kk, a = (torch.nn.functional.pad(t, pad) for t in (r, k, v, kk, a))
243
+ # A pad step must be the identity for the recurrence: decay 1, nothing added.
244
+ w_log = torch.nn.functional.pad(w_log, pad)
245
+ grouped = lambda t: t.reshape(batch, chunks, chunk_size, num_heads, head_dim) # noqa: E731
246
+ rg, kg, vg, kkg, ag, wg = (grouped(t) for t in (r, k, v, kk, a, w_log))
247
+
248
+ w_c = torch.exp(wg)
249
+ c = torch.cumprod(w_c, dim=2)
250
+ c_prev = c / w_c
251
+ bg = kkg * ag
252
+ b_t, k_t = bg / c, kg / c
253
+ q_t, r_t = kkg * c_prev, rg * c
254
+
255
+ span = chunk_size
256
+ tri = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(-1)
257
+ causal = torch.ones(span, span, device=r.device, dtype=r.dtype).tril(0)
258
+ eye = torch.eye(span, device=r.device, dtype=r.dtype)
259
+
260
+ # [batch, chunks, heads, span, span] -- one launch each instead of one per chunk.
261
+ qb = torch.einsum("bcthn,bcshn->bchts", q_t, b_t) * tri
262
+ qk = torch.einsum("bcthn,bcshn->bchts", q_t, k_t) * tri
263
+ rk = torch.einsum("bcthn,bcshn->bchts", r_t, k_t) * causal
264
+ rb = torch.einsum("bcthn,bcshn->bchts", r_t, b_t) * causal
265
+ lhs = eye + qb
266
+ qkv = torch.einsum("bchts,bcshv->bchtv", qk, vg)
267
+ rkv = torch.einsum("bchts,bcshv->bchtv", rk, vg)
268
+ c_last = c[:, :, -1].unsqueeze(-1)
269
+
270
+ # `k~^T v` carries no state either, so it joins the batched half above; only the
271
+ # two products that read the carried state and the solve that depends on them are
272
+ # genuinely serial.
273
+ kv = torch.einsum("bcthn,bcthv->bchnv", k_t, vg)
274
+ # Both state products read the SAME state, so they are one matmul over a stacked
275
+ # token axis rather than two -- the loop is short and launch-bound, and this is one
276
+ # launch per chunk instead of two.
277
+ qr = torch.cat([q_t, r_t], dim=2)
278
+
279
+ outputs = []
280
+ for i in range(chunks):
281
+ qr_s = torch.einsum("bthn,bhnv->bhtv", qr[:, i], state)
282
+ rhs = qr_s[:, :, :chunk_size] + qkv[:, i]
283
+ u = torch.linalg.solve_triangular(lhs[:, i], rhs, upper=False, unitriangular=True)
284
+ out_c = qr_s[:, :, chunk_size:] + rkv[:, i] - torch.einsum("bhts,bhsv->bhtv", rb[:, i], u)
285
+ outputs.append(out_c.permute(0, 2, 1, 3))
286
+ state = c_last[:, i] * (state + kv[:, i] - torch.einsum("bthn,bhtv->bhnv", b_t[:, i], u))
287
+
288
+ state = state.transpose(-2, -1).contiguous().to(state_dtype)
289
+ return torch.cat(outputs, dim=1)[:, :seq_len].to(dtype), state
290
+
291
+
292
+ def rwkv7_eager(
293
+ r: torch.Tensor,
294
+ w_log: torch.Tensor,
295
+ k: torch.Tensor,
296
+ v: torch.Tensor,
297
+ kk: torch.Tensor,
298
+ a: torch.Tensor,
299
+ state: torch.Tensor,
300
+ cu_seq_lens: torch.Tensor | None = None,
301
+ **kwargs,
302
+ ) -> tuple[torch.Tensor, torch.Tensor]:
303
+ """Portable exact WKV, evaluated in reference token order.
304
+
305
+ With `cu_seq_lens` the row is a *packed* batch: several independent sequences
306
+ laid end to end, the layout a varlen kernel consumes. Each one has to start
307
+ from a fresh state, so they are run in turn and their outputs concatenated;
308
+ the state returned is the last segment's, which is the one a continuation
309
+ would resume from. This is the reference behaviour, not the fast path: a
310
+ fused varlen kernel does the same segments in one launch, which is exactly
311
+ the kind of implementation [`RWKV7_WKV_FUNCTIONS`] exists to let you drop in.
312
+
313
+ This default deliberately stays sequential for a multi-token prefill. The
314
+ chunk-parallel rearrangement is much faster, but changes floating-point operation
315
+ order and is therefore available only through the explicit `"chunked"` registry
316
+ entry. Automatic execution must remain the reference users can compare it to.
317
+ """
318
+ if cu_seq_lens is not None:
319
+ bounds = cu_seq_lens.tolist()
320
+ outputs = []
321
+ for start, stop in zip(bounds[:-1], bounds[1:]):
322
+ if stop <= start:
323
+ continue
324
+ segment_state = torch.zeros_like(state)
325
+ out, state = rwkv7_eager(
326
+ r[:, start:stop],
327
+ w_log[:, start:stop],
328
+ k[:, start:stop],
329
+ v[:, start:stop],
330
+ kk[:, start:stop],
331
+ a[:, start:stop],
332
+ segment_state,
333
+ )
334
+ outputs.append(out)
335
+ return torch.cat(outputs, dim=1), state
336
+
337
+ return rwkv7_recurrent(r, w_log, k, v, kk, a, state)
338
+
339
+
340
+ def rwkv7_parallel(
341
+ r: torch.Tensor,
342
+ w_log: torch.Tensor,
343
+ k: torch.Tensor,
344
+ v: torch.Tensor,
345
+ kk: torch.Tensor,
346
+ a: torch.Tensor,
347
+ state: torch.Tensor,
348
+ cu_seq_lens: torch.Tensor | None = None,
349
+ **kwargs,
350
+ ) -> tuple[torch.Tensor, torch.Tensor]:
351
+ """Explicit chunk-parallel WKV with the same packed-sequence contract as eager."""
352
+ if cu_seq_lens is not None:
353
+ bounds = cu_seq_lens.tolist()
354
+ outputs = []
355
+ for start, stop in zip(bounds[:-1], bounds[1:]):
356
+ if stop <= start:
357
+ continue
358
+ out, state = rwkv7_chunked(
359
+ r[:, start:stop],
360
+ w_log[:, start:stop],
361
+ k[:, start:stop],
362
+ v[:, start:stop],
363
+ kk[:, start:stop],
364
+ a[:, start:stop],
365
+ torch.zeros_like(state),
366
+ )
367
+ outputs.append(out)
368
+ return torch.cat(outputs, dim=1), state
369
+ return rwkv7_chunked(r, w_log, k, v, kk, a, state)
370
+
371
+
372
+ RWKV7_WKV_FUNCTIONS = {"eager": rwkv7_eager, "chunked": rwkv7_parallel}
373
+
374
+
375
+ class Rwkv7Attention(nn.Module):
376
+ """RWKV-7 time-mixing block (the recurrent replacement for self-attention)."""
377
+
378
+ def __init__(self, config: Rwkv7Config, layer_id: int):
379
+ super().__init__()
380
+ self.config = config
381
+ self.layer_id = layer_id
382
+ C = config.hidden_size
383
+ self.hidden_size = C
384
+ self.num_heads = config.num_heads
385
+ self.head_dim = config.head_dim
386
+ self.state_dtype = getattr(torch, config.wkv_state_dtype)
387
+
388
+ self.time_shift = Rwkv7TokenShift()
389
+
390
+ # per-channel token-shift mixes (kept at Bo's (1, 1, C) shape)
391
+ self.x_r = nn.Parameter(torch.zeros(1, 1, C))
392
+ self.x_w = nn.Parameter(torch.zeros(1, 1, C))
393
+ self.x_k = nn.Parameter(torch.zeros(1, 1, C))
394
+ self.x_v = nn.Parameter(torch.zeros(1, 1, C))
395
+ self.x_a = nn.Parameter(torch.zeros(1, 1, C))
396
+ self.x_g = nn.Parameter(torch.zeros(1, 1, C))
397
+
398
+ # LoRA factors as raw tensors, exactly as the reference stores them
399
+ self.w1 = nn.Parameter(torch.zeros(C, config.decay_low_rank_dim))
400
+ self.w2 = nn.Parameter(torch.zeros(config.decay_low_rank_dim, C))
401
+ self.w0 = nn.Parameter(torch.zeros(1, 1, C))
402
+ self.a1 = nn.Parameter(torch.zeros(C, config.a_low_rank_dim))
403
+ self.a2 = nn.Parameter(torch.zeros(config.a_low_rank_dim, C))
404
+ self.a0 = nn.Parameter(torch.zeros(1, 1, C))
405
+ self.g1 = nn.Parameter(torch.zeros(C, config.gate_low_rank_dim))
406
+ self.g2 = nn.Parameter(torch.zeros(config.gate_low_rank_dim, C))
407
+ # The value-residual LoRA exists on every layer in a reference checkpoint,
408
+ # but layer 0 PRODUCES v_first instead of mixing towards it, so its copy is
409
+ # never read. It is registered anyway so that loading is lossless.
410
+ self.v1 = nn.Parameter(torch.zeros(C, config.v_low_rank_dim))
411
+ self.v2 = nn.Parameter(torch.zeros(config.v_low_rank_dim, C))
412
+ self.v0 = nn.Parameter(torch.zeros(1, 1, C))
413
+
414
+ self.k_k = nn.Parameter(torch.zeros(1, 1, C))
415
+ self.k_a = nn.Parameter(torch.zeros(1, 1, C))
416
+ self.r_k = nn.Parameter(torch.zeros(config.num_heads, config.head_dim))
417
+
418
+ self.receptance = nn.Linear(C, C, bias=False)
419
+ self.key = nn.Linear(C, C, bias=False)
420
+ self.value = nn.Linear(C, C, bias=False)
421
+ self.output = nn.Linear(C, C, bias=False)
422
+ # GroupNorm over heads, matching the reference's per-head normalisation.
423
+ #
424
+ # The reference hardcodes `eps=64e-5`, which is `head_dim * 1e-5` at its
425
+ # head_dim of 64 -- NOT `num_heads * 1e-5`. The two coincide only when a
426
+ # model happens to have as many heads as channels per head, which for
427
+ # head_dim 64 means hidden_size 4096 exactly. Scaling by `num_heads` is
428
+ # therefore right on the 7.2B and wrong everywhere else: 2x low on the 1.5B,
429
+ # 5.33x low on the 0.1B. It is written against `head_dim` because that is
430
+ # the axis GroupNorm actually reduces over, so the constant tracks the
431
+ # reference at any width.
432
+ self.ln_x = nn.GroupNorm(config.num_heads, C, eps=config.norm_eps * config.head_dim)
433
+
434
+ def forward(
435
+ self,
436
+ hidden_states: torch.Tensor,
437
+ v_first: torch.Tensor | None,
438
+ shift_state: torch.Tensor | None,
439
+ wkv_state: torch.Tensor | None,
440
+ keep: torch.Tensor | None = None,
441
+ cu_seq_lens: torch.Tensor | None = None,
442
+ ):
443
+ batch, seq_len, C = hidden_states.shape
444
+ H, N = self.num_heads, self.head_dim
445
+ tokenwise = self.config.wkv_implementation != "chunked"
446
+
447
+ shifted, new_shift_state = self.time_shift(hidden_states, shift_state, cu_seq_lens, keep)
448
+ delta = shifted - hidden_states
449
+ xr = hidden_states + self.x_r * delta
450
+ xw = hidden_states + self.x_w * delta
451
+ xk = hidden_states + self.x_k * delta
452
+ xv = hidden_states + self.x_v * delta
453
+ xa = hidden_states + self.x_a * delta
454
+ xg = hidden_states + self.x_g * delta
455
+
456
+ r = _rwkv7_linear(self.receptance, xr, tokenwise)
457
+ k = _rwkv7_linear(self.key, xk, tokenwise)
458
+ v = _rwkv7_linear(self.value, xv, tokenwise)
459
+
460
+ w_log, a, g, v_gate = self.lora_gates(xw, xa, xg, None if self.layer_id == 0 else xv, tokenwise)
461
+
462
+ if keep is not None:
463
+ # A padding position has to leave the recurrent state exactly as it
464
+ # found it, which takes three things and not one.
465
+ #
466
+ # The decay is held at w = exp(0) = 1, so the transition is the identity.
467
+ #
468
+ # `k` and `v` are zeroed EXPLICITLY rather than relying on the blanked
469
+ # hidden state to make them zero: the projections are bias-free, so that
470
+ # only holds for a pad with nothing before it.
471
+ # A pad that FOLLOWS a real token still receives that token's hidden
472
+ # state through the token shift, so `delta = shifted - 0` is non-zero,
473
+ # and so are `k` and `v`. The update term `k v^T` then entered the state
474
+ # on every right-padded batch and on any left-padded batch continued
475
+ # from a carried shift state.
476
+ w_log = w_log * keep
477
+ k = k * keep
478
+ v = v * keep
479
+
480
+ if self.layer_id == 0:
481
+ v_first = v
482
+ else:
483
+ v = v + (v_first - v) * v_gate
484
+
485
+ kk = k * self.k_k
486
+ kk = torch.nn.functional.normalize(kk.view(batch, seq_len, H, N), dim=-1, p=2.0).view(batch, seq_len, C)
487
+ if keep is not None:
488
+ # A blanked padding position makes `k` exactly zero, so this normalises a
489
+ # zero vector. `F.normalize` divides by `max(norm, 1e-12)`, and 1e-12 is
490
+ # below the smallest fp16 subnormal -- so in fp16 the divisor really is
491
+ # zero and every padded row comes out NaN. (In fp32 it is representable
492
+ # and the whole thing looks fine, which is why a test on a small fp32
493
+ # model missed it entirely.) `where` is used rather than a multiply
494
+ # because NaN * 0 is still NaN.
495
+ kk = torch.where(keep.bool().expand_as(kk), kk, torch.zeros_like(kk))
496
+ k = k * (1 + (a - 1) * self.k_a)
497
+
498
+ if wkv_state is None:
499
+ wkv_state = torch.zeros(batch, H, N, N, device=r.device, dtype=self.state_dtype)
500
+
501
+ def _heads(t):
502
+ return t.view(batch, seq_len, H, N)
503
+
504
+ # Named rather than indexed: an unregistered key would otherwise surface as a
505
+ # bare `KeyError: 'chunked'` from inside the forward, several frames from the
506
+ # config field that caused it, and a caller looping over shapes would record it
507
+ # as "this shape did not run" instead of "this model was never built". The
508
+ # registry is open by design, so this cannot be validated in the config.
509
+ try:
510
+ wkv = RWKV7_WKV_FUNCTIONS[self.config.wkv_implementation]
511
+ except KeyError:
512
+ raise ValueError(
513
+ f"wkv_implementation={self.config.wkv_implementation!r} is not registered. "
514
+ f"Known: {sorted(RWKV7_WKV_FUNCTIONS)}. Register your own with "
515
+ "`RWKV7_WKV_FUNCTIONS['name'] = fn` before building the model."
516
+ ) from None
517
+ y, wkv_state = wkv(
518
+ _heads(r),
519
+ _heads(w_log),
520
+ _heads(k),
521
+ _heads(v),
522
+ _heads(kk),
523
+ _heads(a),
524
+ wkv_state,
525
+ cu_seq_lens=cu_seq_lens,
526
+ )
527
+
528
+ # `reshape`, not `view`: what comes back from the WKV is whatever layout that
529
+ # implementation produced, and a registered kernel -- or inductor, which is
530
+ # free to pick its own -- can hand back a `[batch, seq_len, heads, head_dim]`
531
+ # that is strided as if it were `[batch, heads, seq_len, head_dim]`. `view`
532
+ # then raises, and it raises only when batch and seq_len are BOTH greater
533
+ # than one, which is why a compiled 16x16 forward failed while 1xT and Bx1
534
+ # both passed.
535
+ y = self.ln_x(y.reshape(batch * seq_len, C)).view(batch, seq_len, C)
536
+ # r·k·r_k summed per head, broadcast back over the head's value channels
537
+ bonus = ((_heads(r) * _heads(k) * self.r_k).sum(dim=-1, keepdim=True) * _heads(v)).reshape(batch, seq_len, C)
538
+ y = _rwkv7_linear(self.output, (y + bonus) * g, tokenwise)
539
+ return y, v_first, new_shift_state, wkv_state
540
+
541
+ def lora_gates(self, xw, xa, xg, xv, tokenwise):
542
+ """The four low-rank chains, as one unit.
543
+
544
+ Each is a rank-r down projection, an activation, and an up projection, and
545
+ they differ only in where the activation sits: `w` has `tanh` between the
546
+ two, `g` has `sigmoid` between them and nothing after, `a` and `v` have
547
+ nothing between and `sigmoid` after. Together they are a dozen tiny
548
+ matrix-vector products per layer whose weights are a rounding error of the
549
+ model's bytes -- so at batch 1 they cost latency, not bandwidth, which is
550
+ the shape a fused kernel improves and a portable implementation cannot.
551
+
552
+ Kept as one method for exactly that reason: it is the unit worth replacing,
553
+ and replacing it should not need a fork. `xv` is None on layer 0, which
554
+ produces `v_first` rather than mixing towards it and never reads that chain.
555
+ """
556
+ w_log = -_INV_SQRT_E * torch.sigmoid(
557
+ _rwkv7_matmul(torch.tanh(_rwkv7_matmul(xw, self.w1, tokenwise)), self.w2, tokenwise) + self.w0
558
+ )
559
+ a = torch.sigmoid(_rwkv7_matmul(_rwkv7_matmul(xa, self.a1, tokenwise), self.a2, tokenwise) + self.a0)
560
+ g = _rwkv7_matmul(torch.sigmoid(_rwkv7_matmul(xg, self.g1, tokenwise)), self.g2, tokenwise)
561
+ v_gate = (
562
+ None
563
+ if xv is None
564
+ else torch.sigmoid(_rwkv7_matmul(_rwkv7_matmul(xv, self.v1, tokenwise), self.v2, tokenwise) + self.v0)
565
+ )
566
+ return w_log, a, g, v_gate
567
+
568
+
569
+ class Rwkv7FeedForward(nn.Module):
570
+ """RWKV-7 channel-mixing block: squared-ReLU over a single token shift."""
571
+
572
+ def __init__(self, config: Rwkv7Config, layer_id: int):
573
+ super().__init__()
574
+ self.config = config
575
+ self.layer_id = layer_id
576
+ C = config.hidden_size
577
+ self.time_shift = Rwkv7TokenShift()
578
+ self.x_k = nn.Parameter(torch.zeros(1, 1, C))
579
+ self.key = nn.Linear(C, config.intermediate_size, bias=False)
580
+ self.value = nn.Linear(config.intermediate_size, C, bias=False)
581
+
582
+ def forward(
583
+ self,
584
+ hidden_states: torch.Tensor,
585
+ shift_state: torch.Tensor | None,
586
+ cu_seq_lens: torch.Tensor | None = None,
587
+ keep: torch.Tensor | None = None,
588
+ ):
589
+ tokenwise = self.config.wkv_implementation != "chunked"
590
+ shifted, new_shift_state = self.time_shift(hidden_states, shift_state, cu_seq_lens, keep)
591
+ xk = hidden_states + self.x_k * (shifted - hidden_states)
592
+ inner = torch.relu(_rwkv7_linear(self.key, xk, tokenwise)) ** 2
593
+ out = self._project(inner, tokenwise)
594
+ return out, new_shift_state
595
+
596
+ def _project(self, inner: torch.Tensor, tokenwise: bool) -> torch.Tensor:
597
+ return _rwkv7_linear(self.value, inner, tokenwise)
598
+
599
+
600
+ class Rwkv7CacheLayer(LinearAttentionLayer):
601
+ """One block's slice of the recurrent state: the WKV matrix and two token shifts.
602
+
603
+ Everything a beam search or a batched generate needs to do to a cache is a
604
+ permutation of its batch axis, and for RWKV-7 that is the whole job -- the state
605
+ is O(1) in sequence length, so there is no time axis to gather along and no
606
+ length bookkeeping to keep consistent.
607
+
608
+ The slot layout is [`Rwkv7Cache`]'s; this class only moves whatever is in them.
609
+ """
610
+
611
+ def lazy_initialization(self, conv_states=None, recurrent_states=None, state_idx: int = 0, **kwargs) -> None:
612
+ super().lazy_initialization(conv_states, recurrent_states, state_idx, **kwargs)
613
+ # Upstream records device/dtype only on the conv branch, since a linear
614
+ # attention layer normally has a convolution in front of it. This model has
615
+ # none, so without this every recurrent-only layer would keep `device=None`
616
+ # and `reorder_cache`'s `beam_idx.to(self.device)` would fail.
617
+ if recurrent_states is not None and self.device is None:
618
+ self.dtype, self.device = recurrent_states.dtype, recurrent_states.device
619
+
620
+ def allocate(
621
+ self, batch: int, shapes: dict[int, tuple], device, dtypes: dict[int, torch.dtype], state_dtype
622
+ ) -> None:
623
+ """Create every slot up front, zeroed, at pinned addresses.
624
+
625
+ The lazy path allocates a slot the first time it is written, which is inside
626
+ the compiled region -- and a buffer that first appears there cannot be given
627
+ a static address, so inductor declines CUDA graphs for a recurrent decode,
628
+ the one workload that most needs them.
629
+ """
630
+ for slot, shape in shapes.items():
631
+ buffer = torch.zeros((batch, *shape), device=device, dtype=dtypes[slot])
632
+ if not torch.compiler.is_compiling():
633
+ torch._dynamo.mark_static_address(buffer)
634
+ self.recurrent_states[slot] = buffer
635
+ self.is_recurrent_states_initialized[slot] = True
636
+ self.has_previous_state[slot] = True
637
+ self.dtype, self.device = state_dtype, device
638
+
639
+ def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
640
+ """Permute the batch onto the beams that survived, without moving house.
641
+
642
+ Upstream rebinds each slot to the `index_select` result. That is correct but
643
+ it hands back a freshly allocated tensor, so the address pinned at allocation
644
+ is gone for the rest of the generation and the compiled decode quietly loses
645
+ its CUDA graphs. Copying back into the same buffer costs one temporary and
646
+ keeps the pinning.
647
+ """
648
+ for slot in range(self.number_of_states):
649
+ buffer = self.recurrent_states[slot]
650
+ if self.is_recurrent_states_initialized[slot] and buffer is not None:
651
+ buffer.copy_(buffer.index_select(0, beam_idx.to(buffer.device)))
652
+
653
+ def batch_repeat_interleave(self, repeats: int) -> None:
654
+ """Fan each sequence's state out to `repeats` copies, for one-prompt-many-samples.
655
+
656
+ This changes the batch size, so unlike `reorder_cache` it cannot preserve the
657
+ pinned addresses -- the buffers are necessarily new ones. Callers that then
658
+ compile should re-pin, which `Rwkv7Model.allocate_state` does.
659
+ """
660
+ for slot in range(self.number_of_states):
661
+ buffer = self.recurrent_states[slot]
662
+ if self.is_recurrent_states_initialized[slot] and buffer is not None:
663
+ self.recurrent_states[slot] = buffer.repeat_interleave(repeats, dim=0)
664
+
665
+ def batch_select_indices(self, indices: torch.Tensor) -> None:
666
+ """Keep only the given batch rows. Same address caveat as `batch_repeat_interleave`."""
667
+ for slot in range(self.number_of_states):
668
+ buffer = self.recurrent_states[slot]
669
+ if self.is_recurrent_states_initialized[slot] and buffer is not None:
670
+ self.recurrent_states[slot] = buffer[indices, ...]
671
+
672
+
673
+ class Rwkv7Cache(Cache):
674
+ """The recurrent state of every block, as a `Cache`.
675
+
676
+ It replaces a KV cache and is a constant size: one `[num_heads, head_dim,
677
+ head_dim]` matrix and two `[hidden]` token shifts per layer per sequence,
678
+ whatever the context length. That is the property the architecture is for, so
679
+ `get_max_length()` is -1 (no limit) and nothing here grows as tokens arrive.
680
+ """
681
+
682
+ # Slot numbering inside one block's layer. All three are recurrent states:
683
+ # RWKV-7's token shift is a one-token history, not a convolution window, so it
684
+ # lives in a recurrent slot, where the update is a plain copy rather than the
685
+ # rolling concatenate a conv slot would do.
686
+ WKV, ATT_SHIFT, FFN_SHIFT = 0, 1, 2
687
+
688
+ def __init__(
689
+ self,
690
+ config: Rwkv7Config,
691
+ batch_size: int | None = None,
692
+ device: torch.device | str | None = None,
693
+ dtype: torch.dtype | None = None,
694
+ ):
695
+ super().__init__(layers=[Rwkv7CacheLayer(number_of_states=3) for _ in range(config.num_hidden_layers)])
696
+ self.config = config
697
+ if batch_size is not None:
698
+ self.allocate(batch_size, device, dtype)
699
+
700
+ def allocate(self, batch_size: int, device=None, dtype=None) -> "Rwkv7Cache":
701
+ config = self.config
702
+ shapes = {
703
+ self.WKV: (config.num_heads, config.head_dim, config.head_dim),
704
+ self.ATT_SHIFT: (config.hidden_size,),
705
+ self.FFN_SHIFT: (config.hidden_size,),
706
+ }
707
+ # The WKV state carries the whole history and is the one place where a
708
+ # narrower dtype actually costs accuracy, so it is configured separately
709
+ # from the activation dtype the shifts follow.
710
+ dtypes = {
711
+ self.WKV: getattr(torch, config.wkv_state_dtype),
712
+ self.ATT_SHIFT: dtype,
713
+ self.FFN_SHIFT: dtype,
714
+ }
715
+ for layer in self.layers:
716
+ assert isinstance(layer, Rwkv7CacheLayer)
717
+ layer.allocate(batch_size, shapes, device, dtypes, dtypes[self.WKV])
718
+ return self
719
+
720
+ def read(self, layer_idx: int):
721
+ """This block's `(att_shift, ffn_shift, wkv)`, each None before allocation."""
722
+ layer = self.layers[layer_idx]
723
+ assert isinstance(layer, Rwkv7CacheLayer)
724
+ states = layer.recurrent_states
725
+ return states[self.ATT_SHIFT], states[self.FFN_SHIFT], states[self.WKV]
726
+
727
+ def write(self, layer_idx: int, att_shift: torch.Tensor, ffn_shift: torch.Tensor, wkv: torch.Tensor) -> None:
728
+ states = (
729
+ (self.ATT_SHIFT, att_shift),
730
+ (self.FFN_SHIFT, ffn_shift),
731
+ (self.WKV, wkv),
732
+ )
733
+ # Copying into the pre-allocated slot is what keeps its address fixed, which
734
+ # is what lets a captured CUDA graph replay the decode loop. It is also an
735
+ # in-place write on a tensor autograd may be holding, and backward through a
736
+ # cached forward then dies on "a variable needed for gradient computation has
737
+ # been modified by an inplace operation". Rebinding instead costs nothing
738
+ # here, because a training step is not the workload that wants graphs.
739
+ if torch.is_grad_enabled() and any(state.requires_grad for _, state in states):
740
+ layer = self.layers[layer_idx]
741
+ assert isinstance(layer, Rwkv7CacheLayer)
742
+ for slot, state in states:
743
+ layer.recurrent_states[slot] = state
744
+ layer.is_recurrent_states_initialized[slot] = True
745
+ layer.has_previous_state[slot] = True
746
+ return
747
+ for slot, state in states:
748
+ self.update_recurrent_state(state, layer_idx, slot)
749
+
750
+
751
+ class Rwkv7Block(GradientCheckpointingLayer):
752
+ def __init__(self, config: Rwkv7Config, layer_id: int):
753
+ super().__init__()
754
+ self.config = config
755
+ self.layer_id = layer_id
756
+ eps, bias = config.norm_eps, config.norm_bias
757
+ # Layer 0 carries the extra input norm of the reference implementation.
758
+ if layer_id == 0:
759
+ self.ln0 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
760
+ self.ln1 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
761
+ self.ln2 = nn.LayerNorm(config.hidden_size, eps=eps, bias=bias)
762
+ self.att = Rwkv7Attention(config, layer_id)
763
+ self.ffn = Rwkv7FeedForward(config, layer_id)
764
+
765
+ def forward(
766
+ self,
767
+ hidden_states: torch.Tensor,
768
+ v_first: torch.Tensor | None,
769
+ state: Rwkv7Cache | None,
770
+ keep: torch.Tensor | None = None,
771
+ cu_seq_lens: torch.Tensor | None = None,
772
+ ):
773
+ if self.layer_id == 0:
774
+ hidden_states = self.ln0(hidden_states)
775
+
776
+ att_shift, ffn_shift, wkv = state.read(self.layer_id) if state is not None else (None, None, None)
777
+
778
+ # Padding is blanked AFTER each norm, not before: a LayerNorm maps the zero
779
+ # vector to its own bias, so masking the residual stream instead would let
780
+ # every pad position come back to life on the way into the next mixer -- and
781
+ # a live pad both moves the state and leaks into the next token's shift.
782
+ attn_in = self.ln1(hidden_states)
783
+ if keep is not None:
784
+ attn_in = attn_in * keep
785
+ attn_out, v_first, att_shift, wkv = self.att(attn_in, v_first, att_shift, wkv, keep, cu_seq_lens)
786
+ hidden_states = hidden_states + attn_out
787
+
788
+ ffn_in = self.ln2(hidden_states)
789
+ if keep is not None:
790
+ ffn_in = ffn_in * keep
791
+ ffn_out, ffn_shift = self.ffn(ffn_in, ffn_shift, cu_seq_lens, keep)
792
+ hidden_states = hidden_states + ffn_out
793
+
794
+ if state is not None:
795
+ # `update_recurrent_state` copies into the pre-allocated slot rather than
796
+ # rebinding it, so the buffers keep fixed addresses across steps -- which
797
+ # is what lets a captured CUDA graph replay the decode loop.
798
+ state.write(self.layer_id, att_shift, ffn_shift, wkv)
799
+ return hidden_states, v_first, state
800
+
801
+
802
+ @dataclass
803
+ class Rwkv7Output(ModelOutput):
804
+ r"""
805
+ state (`Rwkv7Cache`, *optional*):
806
+ The recurrent state of every block. Feed it back to continue a sequence; it
807
+ replaces the KV cache and is a constant size, whatever the context length.
808
+ attentions (always `None`):
809
+ Present so that code written against the common output shape does not have to
810
+ special-case this model. There is no attention here to report -- the mixing is
811
+ a recurrence, not a score matrix -- so it is `None` whatever `output_attentions`
812
+ is set to, rather than an empty tuple pretending to be a per-layer list.
813
+ """
814
+
815
+ last_hidden_state: torch.FloatTensor | None = None
816
+ state: Rwkv7Cache | None = None
817
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
818
+ attentions: None = None
819
+
820
+
821
+ @dataclass
822
+ class Rwkv7CausalLMOutput(ModelOutput):
823
+ r"""
824
+ state (`Rwkv7Cache`, *optional*):
825
+ The recurrent state, as in [`Rwkv7Output`].
826
+ """
827
+
828
+ loss: torch.FloatTensor | None = None
829
+ logits: torch.FloatTensor | None = None
830
+ state: Rwkv7Cache | None = None
831
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
832
+ attentions: None = None
833
+
834
+
835
+ @auto_docstring
836
+ class Rwkv7PreTrainedModel(PreTrainedModel):
837
+ config: Rwkv7Config
838
+ base_model_prefix = "rwkv7"
839
+ _no_split_modules = ["Rwkv7Block"]
840
+ supports_gradient_checkpointing = True
841
+ _is_stateful = True
842
+ # Beam search reorders through `Rwkv7Cache.reorder_cache`. Defining
843
+ # `_reorder_cache` here instead would take precedence over it in `generate` and
844
+ # bypass the cache's own bookkeeping, so it is deliberately absent.
845
+
846
+ @torch.no_grad()
847
+ def _init_weights(self, module):
848
+ """Reproduce RWKV-LM's x070 constructor and final `generate_init_weight` pass."""
849
+ super()._init_weights(module)
850
+ if isinstance(module, Rwkv7Attention):
851
+ layer_id = module.layer_id
852
+ num_layers = module.config.num_hidden_layers
853
+ hidden_size = module.config.hidden_size
854
+ head_dim = module.config.head_dim
855
+
856
+ ratio_0_to_1 = layer_id / (num_layers - 1) if num_layers > 1 else 0.0
857
+ ratio_1_to_almost0 = 1.0 - layer_id / num_layers
858
+ channel = torch.arange(hidden_size, dtype=torch.float32, device=module.x_r.device)
859
+ ddd = (channel / hidden_size).reshape(1, 1, hidden_size)
860
+ linear = channel / max(hidden_size - 1, 1) - 0.5
861
+ if head_dim > 1:
862
+ zigzag = ((channel % head_dim) - (head_dim - 1) / 2) / ((head_dim - 1) / 2)
863
+ zigzag = zigzag * zigzag.abs()
864
+ else:
865
+ zigzag = torch.zeros_like(channel)
866
+ time_decay = -6 + 6 * (channel / max(hidden_size - 1, 1)) ** (1 + ratio_0_to_1**0.3)
867
+
868
+ init.copy_(module.x_r, 1 - ddd ** (0.2 * ratio_1_to_almost0))
869
+ init.copy_(module.x_w, 1 - ddd ** (0.9 * ratio_1_to_almost0))
870
+ init.copy_(module.x_k, 1 - ddd ** (0.7 * ratio_1_to_almost0))
871
+ init.copy_(module.x_v, 1 - ddd ** (0.7 * ratio_1_to_almost0))
872
+ init.copy_(module.x_a, 1 - ddd ** (0.9 * ratio_1_to_almost0))
873
+ init.copy_(module.x_g, 1 - ddd ** (0.2 * ratio_1_to_almost0))
874
+
875
+ init.zeros_(module.w1)
876
+ _rwkv7_lora_orthogonal_(module.w2)
877
+ init.copy_(module.w0, (time_decay + 0.5 + zigzag * 2.5).reshape(1, 1, hidden_size))
878
+ init.zeros_(module.a1)
879
+ _rwkv7_lora_orthogonal_(module.a2)
880
+ init.copy_(module.a0, (-0.19 + zigzag * 0.3 + linear * 0.4).reshape(1, 1, hidden_size))
881
+ init.zeros_(module.v1)
882
+ _rwkv7_lora_orthogonal_(module.v2)
883
+ init.copy_(module.v0, (0.73 - linear * 0.4).reshape(1, 1, hidden_size))
884
+ init.zeros_(module.g1)
885
+ _rwkv7_lora_orthogonal_(module.g2)
886
+ init.copy_(module.k_k, (0.71 - linear * 0.1).reshape(1, 1, hidden_size))
887
+ init.constant_(module.k_a, 1.02)
888
+ init.constant_(module.r_k, -0.04)
889
+
890
+ init.orthogonal_(module.key.weight, gain=0.1)
891
+ init.zeros_(module.output.weight)
892
+ init.constant_(module.ln_x.weight, ((layer_id + 1) / num_layers) ** 0.7)
893
+ if module.ln_x.bias is not None:
894
+ init.zeros_(module.ln_x.bias)
895
+ elif isinstance(module, Rwkv7FeedForward):
896
+ ratio_1_to_almost0 = 1.0 - module.layer_id / module.config.num_hidden_layers
897
+ channel = torch.arange(module.config.hidden_size, dtype=torch.float32, device=module.x_k.device)
898
+ ddd = (channel / module.config.hidden_size).reshape(1, 1, module.config.hidden_size)
899
+ init.copy_(module.x_k, 1 - ddd ** (ratio_1_to_almost0**4))
900
+ init.zeros_(module.value.weight)
901
+ elif isinstance(module, nn.Linear):
902
+ init.orthogonal_(module.weight)
903
+ if module.bias is not None:
904
+ init.zeros_(module.bias)
905
+ elif isinstance(module, nn.Embedding):
906
+ init.uniform_(module.weight, a=-1e-4, b=1e-4)
907
+
908
+
909
+ @auto_docstring
910
+ class Rwkv7Model(Rwkv7PreTrainedModel):
911
+ def __init__(self, config: Rwkv7Config):
912
+ super().__init__(config)
913
+ self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
914
+ self.blocks = nn.ModuleList([Rwkv7Block(config, i) for i in range(config.num_hidden_layers)])
915
+ self.ln_out = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
916
+ self.gradient_checkpointing = False
917
+ self.post_init()
918
+
919
+ def get_input_embeddings(self):
920
+ return self.emb
921
+
922
+ def set_input_embeddings(self, new_embeddings):
923
+ self.emb = new_embeddings
924
+
925
+ def allocate_state(self, batch: int, device=None, dtype=None) -> Rwkv7Cache:
926
+ """A zeroed cache, plus everything else the decode path would build lazily.
927
+
928
+ Call this before compiling, and pass the result in as `state=`.
929
+ `mark_static_address` cannot run during tracing, so anything first allocated
930
+ *inside* the compiled region stays unpinned, and inductor declines CUDA graphs
931
+ for a region that mutates its inputs. Starting from `state=None` loses the
932
+ state buffers that way, and the decode then runs several times slower while
933
+ saying so only in a line of warning.
934
+ """
935
+ state = self._empty_state(
936
+ batch,
937
+ device if device is not None else self.emb.weight.device,
938
+ dtype if dtype is not None else self.emb.weight.dtype,
939
+ )
940
+ return state
941
+
942
+ def _empty_state(self, batch: int, device, dtype) -> Rwkv7Cache:
943
+ return Rwkv7Cache(self.config, batch_size=batch, device=device, dtype=dtype)
944
+
945
+ @can_return_tuple
946
+ @auto_docstring
947
+ def forward(
948
+ self,
949
+ input_ids: torch.LongTensor | None = None,
950
+ attention_mask: torch.LongTensor | None = None,
951
+ inputs_embeds: torch.FloatTensor | None = None,
952
+ state: Rwkv7Cache | None = None,
953
+ cu_seq_lens: torch.LongTensor | None = None,
954
+ use_cache: bool | None = None,
955
+ output_attentions: bool | None = None,
956
+ output_hidden_states: bool | None = None,
957
+ **kwargs,
958
+ ) -> Rwkv7Output:
959
+ r"""
960
+ attention_mask (`torch.LongTensor`, *optional*):
961
+ 1 on real tokens, 0 on padding. Read as the **tail** of whatever is
962
+ given, so a decode step may hand over the whole conversation's mask and
963
+ only the last position is used. A prefix chunk must therefore slice its
964
+ own mask to match the `input_ids` it passes, or it silently masks the
965
+ wrong positions. There is nothing here for a mask to hide behind: an
966
+ all-recurrent model feeds pads through the recurrence like any other
967
+ token unless they are neutralised.
968
+ state (`Rwkv7Cache`, *optional*):
969
+ Recurrent state returned by a previous call; pass it back to continue
970
+ the sequence. Allocated on the first forward if omitted. Ignored as a
971
+ *history* when `cu_seq_lens` is given -- see there.
972
+ cu_seq_lens (`torch.LongTensor`, *optional*):
973
+ Cumulative sequence lengths for a *packed* batch: several sequences
974
+ concatenated into one row instead of padded to a rectangle, starting at
975
+ 0 and ending at `seq_len`, and non-decreasing. Each segment then decodes
976
+ from a fresh recurrent state, as if it had been run on its own. This is
977
+ the varlen layout; use it instead of padding when the lengths vary a
978
+ lot, since a recurrent model pays for pad tokens in time as well as
979
+ memory.
980
+
981
+ A packed batch is a set of *new* sequences, not a continuation, so a
982
+ `state` passed alongside contributes its shape and dtype and nothing
983
+ else -- its contents are not read. What comes back is the last segment's
984
+ state, which is the one a continuation of this row would resume from.
985
+ """
986
+ output_hidden_states = (
987
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
988
+ )
989
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
990
+ # Gradient checkpointing replays the forward during backward. If the cache is
991
+ # live, the replay re-reads a state the first pass already advanced, and the
992
+ # gradients that come back are wrong rather than absent. `GradientCheckpointingLayer`
993
+ # neutralises this itself, but only for a cache arriving as a keyword named
994
+ # `use_cache` / `past_key_values` / `layer_past`; this model hands its state to
995
+ # the block positionally, so none of those guards see it and the whole thing
996
+ # has to be caught here.
997
+ if self.gradient_checkpointing and self.training and use_cache:
998
+ logger.warning_once(
999
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
1000
+ )
1001
+ use_cache = False
1002
+
1003
+ if (input_ids is None) == (inputs_embeds is None):
1004
+ raise ValueError("Specify exactly one of input_ids or inputs_embeds")
1005
+ if inputs_embeds is None:
1006
+ inputs_embeds = self.emb(input_ids)
1007
+
1008
+ if use_cache and state is None:
1009
+ state = self._empty_state(inputs_embeds.shape[0], inputs_embeds.device, inputs_embeds.dtype)
1010
+
1011
+ # An attention-free recurrence has to be told where the padding is: pads are
1012
+ # fed through the recurrence like any other token, so a left-padded batch
1013
+ # (what `generate` produces) would otherwise start every short row from a
1014
+ # state the pads had already moved. A single decoded token is never padding,
1015
+ # so the decode step takes none of this.
1016
+ keep = None
1017
+ if attention_mask is not None:
1018
+ # The mask is taken as the TAIL of whatever is passed, so a decode step
1019
+ # can hand over the whole conversation's mask. A prefix chunk must slice
1020
+ # its own mask to match, or it silently reads the wrong positions.
1021
+ #
1022
+ # Applied at seq_len == 1 as well. Skipping it there assumes a single
1023
+ # decoded token is never padding -- true of `generate`, but an
1024
+ # assumption about the caller rather than a property
1025
+ # of the model, and a fully-masked 1-token row was moving the state.
1026
+ keep = attention_mask[:, -inputs_embeds.shape[1] :, None].to(inputs_embeds.dtype)
1027
+
1028
+ if cu_seq_lens is not None:
1029
+ # Checked rather than trusted: a malformed boundary list does not fail,
1030
+ # it silently splits the recurrence in the wrong places and returns
1031
+ # fluent output computed from the wrong states.
1032
+ if inputs_embeds.shape[0] != 1:
1033
+ raise ValueError(
1034
+ f"cu_seq_lens describes one packed row, but got batch size {inputs_embeds.shape[0]}. "
1035
+ "Pack the sequences into a single row, or use attention_mask with a padded batch."
1036
+ )
1037
+ if cu_seq_lens.ndim != 1 or cu_seq_lens[0] != 0 or cu_seq_lens[-1] != inputs_embeds.shape[1]:
1038
+ raise ValueError(
1039
+ f"cu_seq_lens must be 1-D, start at 0 and end at seq_len ({inputs_embeds.shape[1]}); "
1040
+ f"got {cu_seq_lens.tolist()}"
1041
+ )
1042
+ # Endpoints alone do not pin the list down. A pair that goes backwards
1043
+ # is skipped rather than rejected further in, so the segments emit fewer
1044
+ # tokens than came in and the row silently changes length -- and one that
1045
+ # merely repeats a boundary contributes an empty segment, which is
1046
+ # harmless but is never what the caller meant.
1047
+ if bool((cu_seq_lens[1:] <= cu_seq_lens[:-1]).any()):
1048
+ raise ValueError(
1049
+ "cu_seq_lens must be strictly increasing (each segment needs at least one token); "
1050
+ f"got {cu_seq_lens.tolist()}"
1051
+ )
1052
+
1053
+ hidden_states = inputs_embeds
1054
+ v_first = None
1055
+ all_hidden_states = () if output_hidden_states else None
1056
+
1057
+ for block in self.blocks:
1058
+ if output_hidden_states:
1059
+ all_hidden_states = all_hidden_states + (hidden_states,)
1060
+ hidden_states, v_first, state = block(hidden_states, v_first, state, keep, cu_seq_lens)
1061
+
1062
+ hidden_states = self.ln_out(hidden_states)
1063
+ if output_hidden_states:
1064
+ all_hidden_states = all_hidden_states + (hidden_states,)
1065
+
1066
+ return Rwkv7Output(last_hidden_state=hidden_states, state=state, hidden_states=all_hidden_states)
1067
+
1068
+
1069
+ @auto_docstring
1070
+ class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
1071
+ _tied_weights_keys = {"head.weight": "rwkv7.emb.weight"}
1072
+
1073
+ def __init__(self, config: Rwkv7Config):
1074
+ super().__init__(config)
1075
+ self.rwkv7 = Rwkv7Model(config)
1076
+ self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1077
+ self.post_init()
1078
+
1079
+ @torch.no_grad()
1080
+ def _init_weights(self, module):
1081
+ super()._init_weights(module)
1082
+ # Composite-model dispatch sends the base RWKV7 modules through their own
1083
+ # initializer, so every Linear reaching this override is an LM head. Using
1084
+ # its actual shape also covers a newly resized, not-yet-installed head.
1085
+ if isinstance(module, nn.Linear):
1086
+ gain = 0.5 * math.sqrt(module.out_features / module.in_features)
1087
+ if module.out_features <= module.in_features:
1088
+ gain = 0.5
1089
+ init.orthogonal_(module.weight, gain=gain)
1090
+
1091
+ def get_output_embeddings(self):
1092
+ return self.head
1093
+
1094
+ def set_output_embeddings(self, new_embeddings):
1095
+ self.head = new_embeddings
1096
+
1097
+ def prepare_inputs_for_generation(
1098
+ self,
1099
+ input_ids,
1100
+ state=None,
1101
+ inputs_embeds=None,
1102
+ next_sequence_length=None,
1103
+ is_first_iteration=False,
1104
+ **kwargs,
1105
+ ):
1106
+ # `state is not None` does not by itself mean decode: callers can provide
1107
+ # an empty preallocated state for the initial prompt, or a carried state
1108
+ # followed by a multi-token continuation. GenerationMixin tells us how
1109
+ # much of the input is new through `next_sequence_length`.
1110
+ if next_sequence_length is not None:
1111
+ input_ids = input_ids[:, -next_sequence_length:]
1112
+ if inputs_embeds is not None:
1113
+ inputs_embeds = inputs_embeds[:, -next_sequence_length:]
1114
+ elif state is not None and not is_first_iteration:
1115
+ input_ids = input_ids[:, -1:]
1116
+ model_inputs = (
1117
+ {"input_ids": input_ids}
1118
+ if inputs_embeds is None or not is_first_iteration
1119
+ else {"inputs_embeds": inputs_embeds}
1120
+ )
1121
+ model_inputs["state"] = state
1122
+ # Everything else the caller passed goes through, minus what this model does
1123
+ # not take. An allowlist of specific names instead would make
1124
+ # `generate(output_hidden_states=True)` return a tuple of `None`: the flag
1125
+ # would be dropped here, the forward would never see it, and generate collects the
1126
+ # nothing it got back. Any user kwarg met the same fate, silently. The two
1127
+ # excluded here are `labels`, which would make generate compute a loss it
1128
+ # never reads, and the KV-cache bookkeeping that belongs to models with a KV
1129
+ # cache; this one carries its history in `state`.
1130
+ skip = ("labels", "past_key_values", "cache_position")
1131
+ model_inputs.update({k: v for k, v in kwargs.items() if k not in skip and k not in model_inputs})
1132
+ return model_inputs
1133
+
1134
+ @can_return_tuple
1135
+ @auto_docstring
1136
+ def forward(
1137
+ self,
1138
+ input_ids: torch.LongTensor | None = None,
1139
+ attention_mask: torch.LongTensor | None = None,
1140
+ inputs_embeds: torch.FloatTensor | None = None,
1141
+ state: Rwkv7Cache | None = None,
1142
+ labels: torch.LongTensor | None = None,
1143
+ use_cache: bool | None = None,
1144
+ output_attentions: bool | None = None,
1145
+ output_hidden_states: bool | None = None,
1146
+ logits_to_keep: int | torch.Tensor = 0,
1147
+ **kwargs,
1148
+ ) -> Rwkv7CausalLMOutput:
1149
+ r"""
1150
+ state (`Rwkv7Cache`, *optional*):
1151
+ Recurrent state returned by a previous call.
1152
+ logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0):
1153
+ Compute the head over only the last `logits_to_keep` positions, or over the
1154
+ positions this tensor indexes; 0 means all of them. Worth having on a
1155
+ recurrent model for the same reason as on any other: a prefill needs one
1156
+ row of logits and the vocabulary is the widest matrix in the model, so
1157
+ running the head over the whole prompt is the single largest avoidable cost
1158
+ in a prefill. Until this argument existed it was swallowed by `**kwargs`,
1159
+ so `generate` declined to pass it and a caller who passed it was quietly
1160
+ ignored.
1161
+ """
1162
+ outputs = self.rwkv7(
1163
+ input_ids=input_ids,
1164
+ attention_mask=attention_mask,
1165
+ inputs_embeds=inputs_embeds,
1166
+ state=state,
1167
+ use_cache=use_cache,
1168
+ output_hidden_states=output_hidden_states,
1169
+ **kwargs,
1170
+ )
1171
+ keep = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
1172
+ hidden_to_project = outputs.last_hidden_state[:, keep, :]
1173
+ logits = _rwkv7_linear(
1174
+ self.head,
1175
+ hidden_to_project,
1176
+ self.config.wkv_implementation != "chunked",
1177
+ )
1178
+
1179
+ loss = None
1180
+ if labels is not None:
1181
+ loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
1182
+
1183
+ return Rwkv7CausalLMOutput(loss=loss, logits=logits, state=outputs.state, hidden_states=outputs.hidden_states)
1184
+
1185
+
1186
+ __all__ = ["Rwkv7Cache", "Rwkv7PreTrainedModel", "Rwkv7Model", "Rwkv7ForCausalLM"]
release-manifest.json ADDED
@@ -0,0 +1,332 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "builder": {
3
+ "asset_set": "2026.08.06-r1",
4
+ "assets_sha256": "3c3f6237fcd9dc2ec8ba0b61b626ef72924ebe42bdc6dc49e21a5bfeb19bb96b",
5
+ "version": "0.2.0"
6
+ },
7
+ "conversion": {
8
+ "explicit_cast": false,
9
+ "serialized_parameter_count": 1527668736,
10
+ "source_float_dtypes": [
11
+ "bfloat16"
12
+ ],
13
+ "source_parameter_count": 1527668736,
14
+ "synthesized_tensors": [],
15
+ "target_float_dtypes": [
16
+ "bfloat16"
17
+ ],
18
+ "tensor_count": 798,
19
+ "tensor_map_sha256": "03131ced241e7b0f86869363b3969ead6ef58462f23d9f04e56e26d00565aefc"
20
+ },
21
+ "files": {
22
+ ".gitattributes": {
23
+ "role": "metadata",
24
+ "sha256": "edfb97afcbabd3c89eb1d19070521d403f4644ecacfec826e14c2dc68cf419d8",
25
+ "size_bytes": 134
26
+ },
27
+ "LICENSE": {
28
+ "role": "metadata",
29
+ "sha256": "475641f33cfd755f19245d2eb6a6d8fb7fdca60835cfa369b7e3774460409cc4",
30
+ "size_bytes": 11357
31
+ },
32
+ "NOTICE": {
33
+ "role": "metadata",
34
+ "sha256": "7f8e7f0f3bf0a8e7caf13132770cfe841e5eec10203c18798f8051e96d234108",
35
+ "size_bytes": 250
36
+ },
37
+ "README.md": {
38
+ "role": "model_card",
39
+ "sha256": "b7795213e1fcf1d9e7de41fa027a0b2d311b2dacfa5a5ba9f62c4c2484bd729f",
40
+ "size_bytes": 10260
41
+ },
42
+ "chat_template.jinja": {
43
+ "role": "tokenizer",
44
+ "sha256": "f0d8f8161f03d38d0bf68b1ea9a1293064c0fe2b60dfc2874008154c33861c07",
45
+ "size_bytes": 2333
46
+ },
47
+ "config.json": {
48
+ "role": "model_config",
49
+ "sha256": "442d1155144fd8331bac86250131cc1a874d8d41c1504f930e0b522a99b615aa",
50
+ "size_bytes": 748
51
+ },
52
+ "configuration_rwkv7.py": {
53
+ "role": "model_code",
54
+ "sha256": "a4a2adcd02a101cc2a17f2452f8e56f4a8d8348be0d19ab989e908a83921c123",
55
+ "size_bytes": 7318
56
+ },
57
+ "generation_config.json": {
58
+ "role": "model_config",
59
+ "sha256": "358114ae91c5eee6fe10b41a22a9ee95a9b1c20ed86d0666617aa049386d63a3",
60
+ "size_bytes": 87
61
+ },
62
+ "inference/generate.py": {
63
+ "role": "inference",
64
+ "sha256": "9e4f3461ade971a96a276ec4e18683d6d8df6420b09bf3e63cc499b6e119aa05",
65
+ "size_bytes": 7804
66
+ },
67
+ "inference/kernel.py": {
68
+ "role": "inference",
69
+ "sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6",
70
+ "size_bytes": 165253
71
+ },
72
+ "inference/model_loader.py": {
73
+ "role": "inference",
74
+ "sha256": "c48abe28abde91ede570498372cd6b9aa6a05080525df7f5045613849ed9bb61",
75
+ "size_bytes": 6452
76
+ },
77
+ "inference/requirements.txt": {
78
+ "role": "inference",
79
+ "sha256": "9cc4f89e2edc45da7a04cdf8e4c1da83b682a6bf969e0eed7b3dec91128d925a",
80
+ "size_bytes": 93
81
+ },
82
+ "inference/runtime.py": {
83
+ "role": "inference",
84
+ "sha256": "064c53ce5dfe811ea873ad9976f8f0839b55e515e604adf3c3e28fd4c02de784",
85
+ "size_bytes": 158191
86
+ },
87
+ "model.safetensors": {
88
+ "role": "weights",
89
+ "sha256": "84ccbb857c84e00cefc48b233937ada79c411e491df25fb21aed23237f39a14f",
90
+ "size_bytes": 3055418240
91
+ },
92
+ "modeling_rwkv7.py": {
93
+ "role": "model_code",
94
+ "sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
95
+ "size_bytes": 57446
96
+ },
97
+ "tokenizer.json": {
98
+ "role": "tokenizer",
99
+ "sha256": "1c0bc8e5f355de60a4587a1e9abb9aa4ef65508ff59641344e7f9f10f76feacf",
100
+ "size_bytes": 1580973
101
+ },
102
+ "tokenizer_config.json": {
103
+ "role": "tokenizer",
104
+ "sha256": "58a994975a1a64327cd7ee71f230d4a5f290e7260957b18840e9ecf0d5d57d0e",
105
+ "size_bytes": 407
106
+ }
107
+ },
108
+ "identity": {
109
+ "parameter_label": "1.5",
110
+ "release_date": "20260805",
111
+ "training_context_length": 16384
112
+ },
113
+ "metadata": {
114
+ "context_length": 16384,
115
+ "datasets": [
116
+ "HuggingFaceFW/fineweb-edu",
117
+ "mlfoundations/dclm-baseline-1.0",
118
+ "cerebras/SlimPajama-627B",
119
+ "EleutherAI/pile",
120
+ "bigcode/starcoderdata",
121
+ "oscar-corpus/OSCAR-2301"
122
+ ],
123
+ "languages": [
124
+ "en",
125
+ "zh",
126
+ "fr",
127
+ "es",
128
+ "de",
129
+ "pt",
130
+ "ru",
131
+ "it",
132
+ "ja",
133
+ "ko",
134
+ "vi",
135
+ "ar"
136
+ ],
137
+ "license": "apache-2.0",
138
+ "profile": "g1i",
139
+ "provenance": "locked-profile"
140
+ },
141
+ "model_code": {
142
+ "format_version": 1,
143
+ "source_repository": "https://github.com/huggingface/transformers.git",
144
+ "source_revision": "4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2",
145
+ "sources": {
146
+ "configuration_rwkv7.py": {
147
+ "asset_path": "model_code/configuration_rwkv7.py",
148
+ "output_sha256": "a4a2adcd02a101cc2a17f2452f8e56f4a8d8348be0d19ab989e908a83921c123",
149
+ "repository_path": "src/transformers/models/rwkv7/configuration_rwkv7.py",
150
+ "source_sha256": "6f5b92c5fe7498ad22b0054a2f735a7ca82e7577436f4ad32f0fc27d1e900fdd"
151
+ },
152
+ "modeling_rwkv7.py": {
153
+ "asset_path": "model_code/modeling_rwkv7.py",
154
+ "output_sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
155
+ "repository_path": "src/transformers/models/rwkv7/modeling_rwkv7.py",
156
+ "source_sha256": "3e8e5af7c4eba0b5de1496aef44773d7ac1bb4d96756e6f55efaf29453d67952"
157
+ }
158
+ },
159
+ "transformers_min_version": "5.15"
160
+ },
161
+ "profile": {
162
+ "checkpoint": "g1i-1.5b-20260805",
163
+ "family": "g1i"
164
+ },
165
+ "runtime": {
166
+ "files": {
167
+ "kernel.py": {
168
+ "sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6"
169
+ },
170
+ "runtime.py": {
171
+ "sha256": "064c53ce5dfe811ea873ad9976f8f0839b55e515e604adf3c3e28fd4c02de784"
172
+ }
173
+ },
174
+ "format_version": 4,
175
+ "kernel": {
176
+ "format_version": 3,
177
+ "output": "inference/kernel.py",
178
+ "output_sha256": "88379d8979b908a01af1f06b89d8fc624d332796db6044ec4a027d7cecab1cb6",
179
+ "sources": {
180
+ "kernel_tilelang_decode.py": {
181
+ "exports": [
182
+ "Any",
183
+ "_HEAD_SIZE",
184
+ "_NEXP_HALF_LOG2_E",
185
+ "_NLOG2_E",
186
+ "_ROTATOR1_SIGNED",
187
+ "_TWO_NEG_41",
188
+ "_build_cmix_add_layernorm_mix_program",
189
+ "_build_cmix_binned_finalize_program",
190
+ "_build_cmix_finalize_program",
191
+ "_build_cmix_layernorm_mix_program",
192
+ "_build_cmix_sparse_atomic_program",
193
+ "_build_cmix_sparse_binned_program",
194
+ "_build_cmix_sparse_split_program",
195
+ "_build_cmix_value_program",
196
+ "_build_ffn_program",
197
+ "_build_gemv_program",
198
+ "_build_key_gate_program",
199
+ "_build_post_state_program",
200
+ "_build_rankout_program",
201
+ "_build_rankout_reduced_program",
202
+ "_build_rkv_program",
203
+ "_build_tmix_layernorm_mix6_program",
204
+ "_build_wkv_program",
205
+ "_build_wkv_w0_t1_program",
206
+ "_cmix_value_out",
207
+ "_compiled_cmix_add_layernorm_mix",
208
+ "_compiled_cmix_binned_finalize",
209
+ "_compiled_cmix_finalize",
210
+ "_compiled_cmix_layernorm_mix",
211
+ "_compiled_cmix_sparse_atomic",
212
+ "_compiled_cmix_sparse_binned",
213
+ "_compiled_cmix_sparse_split",
214
+ "_compiled_cmix_value",
215
+ "_compiled_ffn",
216
+ "_compiled_gemv",
217
+ "_compiled_key_gate",
218
+ "_compiled_post_state",
219
+ "_compiled_rankout",
220
+ "_compiled_rankout_reduced",
221
+ "_compiled_rkv",
222
+ "_compiled_tmix_layernorm_mix6",
223
+ "_compiled_wkv",
224
+ "_compiled_wkv_w0_t1",
225
+ "_dtype_name",
226
+ "_ffn_out",
227
+ "_post_state_out",
228
+ "_rankout_out",
229
+ "_require_contiguous_cuda",
230
+ "_rkv_out",
231
+ "_wkv_kernel_out",
232
+ "_wkv_out",
233
+ "_wkv_precise_out",
234
+ "_wkv_w0_t1_out",
235
+ "clear_tilelang_kernel_caches",
236
+ "cuda_arch_key",
237
+ "lru_cache",
238
+ "torch"
239
+ ],
240
+ "namespace": "decode",
241
+ "sha256": "b582151a6d2bd345cc5f87415543370e027c8f943a4dfe57a18300f08560bd73"
242
+ },
243
+ "kernel_tilelang_state.py": {
244
+ "exports": [
245
+ "Any",
246
+ "BFLOAT16",
247
+ "CUDA_GET_DEVICE_CAPABILITY",
248
+ "EXACT_FUSED_STATE_PROJECTION_CAPABILITIES",
249
+ "FLOAT16",
250
+ "FLOAT32",
251
+ "IS_GRAD_ENABLED",
252
+ "TORCH_STACK",
253
+ "TORCH_WHERE",
254
+ "TORCH_ZEROS_LIKE",
255
+ "_compiled_backward_kernel",
256
+ "_compiled_fast_state_scan",
257
+ "_compiled_fused_state_update",
258
+ "_compiled_kernel",
259
+ "_compiled_low_precision_state_kernel",
260
+ "_compiled_post_state",
261
+ "_compiled_state_projection",
262
+ "_compiled_x_mix",
263
+ "_dtype_name",
264
+ "_require_exact_fused_state_projection",
265
+ "build_fast_state_scan_program",
266
+ "build_fused_state_update_program",
267
+ "build_low_precision_state_program",
268
+ "build_post_state_program",
269
+ "build_state_backward_program",
270
+ "build_state_program",
271
+ "build_state_projection_program",
272
+ "build_x_mix_program",
273
+ "clear_tilelang_state_kernel_caches",
274
+ "cuda_arch_key",
275
+ "exact_fused_state_projection_supported",
276
+ "lru_cache",
277
+ "tilelang_fast_state_scan",
278
+ "tilelang_fused_state_update",
279
+ "tilelang_low_precision_state_finalize",
280
+ "tilelang_post_state",
281
+ "tilelang_state_finalize",
282
+ "tilelang_state_finalize_backward",
283
+ "tilelang_state_projection",
284
+ "tilelang_state_scan",
285
+ "tilelang_state_update",
286
+ "tilelang_x_mix",
287
+ "torch",
288
+ "torch_state_scan",
289
+ "torch_state_update"
290
+ ],
291
+ "namespace": "state",
292
+ "sha256": "a3597160cb8d7317151200d05b241af7edcad2928c3731687764d914c28bf21a"
293
+ }
294
+ }
295
+ },
296
+ "transforms": {
297
+ "configuration_rwkv7.py": {
298
+ "output_sha256": "60a060199526f5a6c4d5a2436520eed97e9f21f2ea3cb6e6584d9b5fcb6f21f4",
299
+ "source_sha256": "60a060199526f5a6c4d5a2436520eed97e9f21f2ea3cb6e6584d9b5fcb6f21f4"
300
+ },
301
+ "custom_ops.py": {
302
+ "output_sha256": "8240d4eebc9decbd5ebf8fbcc3abc0ee6441b0eb4854b8d7bc18172cd2babf17",
303
+ "source_sha256": "fb86e5c9dfc9deb5a1be821ce02c3a390ce5c95ff55feded6add2422f0dc3623"
304
+ },
305
+ "kernel_dispatch.py": {
306
+ "output_sha256": "0290559437016ba4e2c2a26f7c7a0985b6fbab991535b7ae65ba18eb421fa9e1",
307
+ "source_sha256": "0290559437016ba4e2c2a26f7c7a0985b6fbab991535b7ae65ba18eb421fa9e1"
308
+ },
309
+ "modeling_rwkv7.py": {
310
+ "output_sha256": "ad7f8c981d21a3b0361cd26a2c68234a1bfbacc31fe52b49c3fe9d79e5b32844",
311
+ "source_sha256": "7e1e57a4516efe78a74f971ada0fe58cb71c87102479ed7516b6d3585964ca21"
312
+ },
313
+ "state.py": {
314
+ "output_sha256": "d13d81af6a301a70024494f03b95006bf5f2471ba0b57fd7f1f94bd54f70f2ef",
315
+ "source_sha256": "d13d81af6a301a70024494f03b95006bf5f2471ba0b57fd7f1f94bd54f70f2ef"
316
+ },
317
+ "tilelang_decode.py": {
318
+ "output_sha256": "313120b38178642c17ef16531c4510de87dc0f0d79d80f70b6ade8818b4dfbe7",
319
+ "source_sha256": "fe63e6332037cacbd0400d82f65fd6f1c282e31d92906c5f34d39eb9e83eb642"
320
+ }
321
+ }
322
+ },
323
+ "schema_version": 6,
324
+ "source": {
325
+ "filename": "rwkv7-g1i-1.5b-20260805-ctx16384.pth",
326
+ "kind": "huggingface",
327
+ "reference": "BlinkDL/rwkv7-g1/rwkv7-g1i-1.5b-20260805-ctx16384.pth",
328
+ "revision": "ede85bf8ab2e59aff7d7ca909fbbc73317866d89",
329
+ "sha256": "32ef7b5bf4dc8bde843cf26dfad809a1f527e2e76a9e790e7d406e71bcd785da",
330
+ "size_bytes": 3055444605
331
+ }
332
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "<|endoftext|>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ }
11
+ },
12
+ "backend": "tokenizers",
13
+ "eos_token": "<|endoftext|>",
14
+ "pad_token": "<|endoftext|>",
15
+ "padding_side": "left",
16
+ "tokenizer_class": "PreTrainedTokenizerFast",
17
+ "unk_token": "<|endoftext|>"
18
+ }