| --- |
| license: apache-2.0 |
| pipeline_tag: text-generation |
| language: |
| - ja |
| - en |
| - zh |
| base_model: |
| - RWKV/RWKV7-G1l-2.9B-20261031 |
| - RWKV/RWKV7-G1l-7.2B-20261031 |
| - RWKV/RWKV7-G1l-13.3B-20261031 |
| tags: |
| - rwkv |
| - rwkv7 |
| - speculative-decoding |
| - multi-token-prediction |
| - mtp |
| - draft-model |
| --- |
| |
| # RWKV7-MTP-Exp: MTP draft sidecars for RWKV-7 G1l |
|
|
| **Experimental** multi-token-prediction (MTP) drafters for speculative decoding with RWKV-7 G1l models. |
|
|
| - **What a sidecar is:** a small add-on network that reads the frozen base model's hidden states and proposes the next few tokens. The base model then verifies all of them in a single forward pass. |
| - **Quality:** the output distribution is exactly the base model's. The sidecar changes only speed. |
| - **Speed:** with [RWKV-Infer-Next](#usage-with-rwkv-infer-next), single-stream decoding becomes **1.8β2.3Γ faster**. That includes the 13.3B model on a single 16 GB GPU. |
|
|
| | Folder | Base model | Base layers / width | Tapped layers | Training tokens | Size (bf16) | |
| |---|---|---|---|---|---| |
| | `RWKV7-G1l-2.9B-20261031/` | [RWKV/RWKV7-G1l-2.9B-20261031](https://huggingface.co/RWKV/RWKV7-G1l-2.9B-20261031) | 32 / 2560 | 9, 19 | 1.0B | 208 MB (103.6M params) | |
| | `RWKV7-G1l-7.2B-20261031/` | [RWKV/RWKV7-G1l-7.2B-20261031](https://huggingface.co/RWKV/RWKV7-G1l-7.2B-20261031) | 32 / 4096 | 9, 19 | 1.0B | 496 MB (247.6M params) | |
| | `RWKV7-G1l-13.3B-20261031/` | [RWKV/RWKV7-G1l-13.3B-20261031](https://huggingface.co/RWKV/RWKV7-G1l-13.3B-20261031) | 61 / 4096 | 17, 36 | 1.0B | 496 MB (247.6M params) | |
|
|
| > A sidecar works **only with the base model it was trained on**. Same width is not enough: a sidecar trained for the 7.2B |
| > model loads on the 13.3B model but gives no speed-up (1.03 tokens per forward). `mtp.json` records the base model, |
| > and RWKV-Infer-Next refuses a mismatch. |
|
|
| ## Measured speed |
|
|
| Setup: |
| - Engine: RWKV-Infer-Next, int6 weights. |
| - Workload: single stream, 6 chat prompts (Japanese chat, proofreading, translation, English chat, code, Chinese), greedy and sampled. |
|
|
| Speeds are the mean over the 6 prompts, greedy / sampled. |
|
|
| | Base model | GPUs | GPU memory used | Without MTP | With this sidecar | Speed-up | |
| |---|---|---|---|---|---| |
| | G1l 2.9B | **one 16 GB GPU** | 5.1 GB | 240 / 230 tok/s | **436 / 443 tok/s** | **1.8β1.9Γ** | |
| | G1l 2.9B | two 16 GB GPUs (tensor parallel) | β 4 GB each | 205 / 204 tok/s | **403 / 382 tok/s** | **1.9β2.0Γ** | |
| | G1l 7.2B | **one 16 GB GPU** | 9.4 GB | 119 / 116 tok/s | **243 / 249 tok/s** | **2.0β2.1Γ** | |
| | G1l 7.2B | two 16 GB GPUs (tensor parallel) | β 6 GB each | 141 / 140 tok/s | **269 / 275 tok/s** | **1.9β2.0Γ** | |
| | G1l 13.3B | **one 16 GB GPU** (`--max-batch 16`) | 14.6 GB | 64 / 63 tok/s | **143 / 143 tok/s** | **2.2β2.3Γ** | |
| | G1l 13.3B | two 16 GB GPUs (tensor parallel) | β 9.5 GB each | 78 / 76 tok/s | **162 / 162 tok/s** | **2.1Γ** | |
|
|
| - **The 13.3B model with its sidecar fits one 16 GB GPU** in int6. Cap the concurrent sequences at 16 (`--max-batch 16`); the default of 32 reserves recurrent state and CUDA graphs that do not fit. |
| - **The 2.9B model is faster on one GPU** than split over two, with or without MTP. |
| - The two-GPU rows are PCIe GPUs without P2P. |
|
|
| - **By task** (greedy, one GPU): acceptance is highest on proofreading and code, and lowest on English chat and Chinese. |
|
|
| | Base model | Proofreading | Code | Japanese chat | Translation | Chinese | English chat | |
| |---|---|---|---|---|---|---| |
| | 2.9B | 559 | 498 | 426 | 407 | 402 | 323 tok/s | |
| | 7.2B | 268 | 239 | 256 | 253 | 231 | 211 tok/s | |
| | 13.3B | 171 | 156 | 142 | 154 | 117 | 115 tok/s | |
|
|
| - **Training length:** speed-ups grow with training. The 1.0B-token sidecars replace the earlier ones. Measured in the same session on one GPU (greedy / sampled): |
|
|
| | Base model | Earlier sidecar | This sidecar (1.0B tokens) | Gain | |
| |---|---|---|---| |
| | 7.2B | 300M tokens: 226 / 223 tok/s | 243 / 249 tok/s | +7β12% | |
| | 13.3B | 270M tokens: 136 / 133 tok/s | 143 / 143 tok/s | +5β7% | |
|
|
| Earlier checkpoints: 7.2B gave 1.4Γ after 25M tokens; 13.3B gave 1.5Γ after 30M and 1.7Γ after 90M. |
|
|
| ## Architecture |
|
|
| Each sidecar is **one transformer-style block**, applied K times per decoding step. The block is gated sliding-window attention followed by an RWKV-7 channel-mix. At every step it reads: |
| - the base model's last hidden state; |
| - two of the base model's intermediate residual streams (EAGLE-3-style fusion); |
| - the embedding of the token it is conditioned on. |
|
|
| ``` |
| base model (RWKV-7, frozen) |
| token ββΊ emb ββΊ ln0 ββ¬ββΊ layer 0 β¦ layer L-1 ββΊ ln_out ββΊ h_j ββββββββββββΊ head ββΊ p(token j+1) |
| β β outputs of layers f1, f2 |
| β βΌ |
| β m1_j, m2_j |
| βββΊ e_t = ln0(emb(t)) (shared with the base model) |
| |
| drafter (this repository): one block, depth k = 1..K drafts the token k positions ahead |
| depth 1 : u = fproj Β· [ LN(h_j) ; LN(m1_j) ; LN(m2_j) ] mid-layer fusion |
| depth k : u = LN_d(o_{k-1}) its own previous output (depth adapter) |
| x = LN( proj Β· [ u ; LN(e_{j+k}) ] + dep[k-1] ) |
| x = x + GatedSWA( LN(x) ) 16 query / 4 KV heads Γ 128, q/k RMSNorm, partial RoPE, window 512 |
| x = x + ChannelMix( LN(x) ) RWKV-7 channel-mix: token shift, reluΒ², 4Γ hidden |
| o_k = LN(x) ββΊ base model's head ββΊ draft token k |
| ``` |
|
|
| | Component | Definition | |
| |---|---| |
| | Token mixer | **Gated sliding-window attention**: grouped-query attention with 16 query heads and 4 KV heads of dim 128 (2048-wide), per-head RMSNorm on q and k, rotate-half RoPE on the first 32 dims (ΞΈ = 1e6), sigmoid output gate (`o = W_o (attn β Ο(W_g x))`), window of 512 keys | |
| | Channel mixer | RWKV-7 channel-mix: `x + (x_prev β x) β x_k` token shift, `W_v Β· relu(W_k Β·)Β²`, hidden size 4C | |
| | Depth-1 input | the base model's final hidden state `h` (after `ln_out`), fused with the **output** residual streams of two base layers through `fproj` (3C β C) | |
| | Deeper inputs | the drafter's own previous output through a separate LayerNorm, plus a learned per-depth vector (`dep`) | |
| | Embedding and head | the base model's `ln0(emb)` and output head, shared and frozen (not stored here) | |
| | Draft vocabulary | `draft_vocab`: token ids by training frequency. Drafting over the top 16384 rows covers about 97% of the token stream and makes the draft head 4Γ cheaper | |
| | Normalisation | LayerNorm (eps 1e-5) everywhere except q/k RMSNorm (eps 1e-6); no biases | |
|
|
| **Why gated SWA attention rather than an RWKV block?** A preliminary screen trained drafters with the same recipe and compared depth-1 greedy acceptance: |
|
|
| | Drafter layer | chat | text | |
| |---|---|---| |
| | RWKV time-mix | 54.6% | 55.8% | |
| | attention | 64.1% | 74.9% | |
|
|
| The attention layer drafts much better, and the sliding window keeps its state bounded. The state is a KV ring of 512 + K entries per sequence, so the system keeps RWKV's constant-memory property. Rolling back a rejected draft costs nothing. |
|
|
| **Inference.** One speculative step: |
| 1. Draft d1β¦dK by chaining the block. |
| 2. The base model processes `[t, d1 β¦ dK]` in one forward. |
| 3. The sampler picks a token at every position from the **base model's** distribution. |
| 4. The drafts are accepted up to the first mismatch. |
| 5. The base model's recurrent state is rolled back to the accepted prefix. |
|
|
| Every emitted token is a sample from the base model given the accepted prefix, so the output distribution is the same as plain decoding, greedy or sampled. |
|
|
| ## Training |
|
|
| | Item | Setting | |
| |---|---| |
| | Base model | frozen (bf16, or weight-only int8 where it did not fit) | |
| | Target | the base model's own next-token distribution (KL distillation), not the data tokens | |
| | Unrolling | 3 drafting depths per position, trained on the drafter's own outputs (training-time test); equal weights | |
| | Context | 512-token windows | |
| | Optimiser | AdamW, Ξ² = (0.9, 0.99), **no weight decay**, lr 1e-3 with 200 warm-up steps, gradient clip 1.0, fp32 master weights | |
| | Initialisation | the block starts as an identity (zero output projections); fusion starts as "final hidden only" | |
| | Data | a Japanese-weighted multilingual mix: Japanese / English / Chinese web and encyclopedic text, code, math, and chat and translation conversations. The 13.3B mix was β 49% Japanese (incl. translation), 30% English (incl. math), 12% Chinese, 9% code | |
|
|
| ## Usage with RWKV-Infer-Next |
|
|
| [RWKV-Infer-Next](https://github.com/OpenMOSE/RWKV-Infer-Next) (an inference server for RWKV-7): |
| - tensor parallel, continuous batching on CUDA graphs, int6 / int8 / bf16 weights; |
| - OpenAI-compatible API, chat templates, thinking and tool calls, WebUI. |
|
|
| It downloads the base model and the matching sidecar from this repository automatically. |
|
|
| ```bash |
| pip install -r requirements.txt |
| python -m rwkv_next --model RWKV/RWKV7-G1l-2.9B-20261031 --tp 1 # one GPU (β 5 GB) |
| python -m rwkv_next --model RWKV/RWKV7-G1l-7.2B-20261031 --tp 1 # one 16 GB GPU |
| python -m rwkv_next --model RWKV/RWKV7-G1l-13.3B-20261031 --tp 1 --max-batch 16 # one 16 GB GPU |
| python -m rwkv_next --model RWKV/RWKV7-G1l-13.3B-20261031 --tp 2 # two GPUs |
| # WebUI: http://localhost:8000/ui API: http://localhost:8000/v1/chat/completions |
| ``` |
|
|
| What happens at start: |
| 1. The base model is downloaded once into the Hugging Face cache. |
| 2. `RWKV_MTP_PATH=auto` (the default) looks for `mtp/` in the model folder. When there is none, it reads `index.json` from `RWKV_MTP_REPO` (default `OpenMOSE/RWKV7-MTP-Exp`) and downloads only the matching folder. |
| 3. The status page shows the sidecar in use. |
|
|
| **Local models.** A local model folder named like the base model (e.g. `.../RWKV7-G1l-7.2B-20261031`) gets its sidecar from this repository the same way. Alternatively, place a folder from this repository in the model folder as `mtp/`: |
|
|
| ```bash |
| huggingface-cli download OpenMOSE/RWKV7-MTP-Exp --include "RWKV7-G1l-7.2B-20261031/*" --local-dir /tmp/mtp |
| mv /tmp/mtp/RWKV7-G1l-7.2B-20261031 /path/to/RWKV7-G1l-7.2B-20261031/mtp |
| ``` |
|
|
| **Newer RWKV-7 7B / 13B releases** that have no sidecar of their own can borrow one with `--mtp-force` (`RWKV_MTP_FORCE=1`). It picks the sidecar of the same width with the closest layer count. |
| - The output is unchanged; only the speed depends on how close the new model is to the base model. |
| - A poor match can be slower than no MTP. Measured: the 7.2B sidecar forced onto the 13.3B model ran at 66 tok/s, against 78 without MTP. |
|
|
| **Settings:** |
| - `RWKV_MTP_PATH=off` disables MTP. |
| - `RWKV_MTP_REPO=` (empty) prevents any download. |
| - `RWKV_MTP_K` sets the drafts per step by batch size. The default `4:3,16:2,0:0` drafts 3 tokens for β€ 4 concurrent streams, 2 for β€ 16, and none beyond, where the GPU is already saturated. |
|
|
| ## Files |
|
|
| ``` |
| index.json which folder serves which base model |
| RWKV7-G1l-2.9B-20261031/mtp.safetensors drafter weights (bf16) + draft_vocab (int64) |
| RWKV7-G1l-2.9B-20261031/mtp.json {"meta": {...}} |
| RWKV7-G1l-7.2B-20261031/... |
| RWKV7-G1l-13.3B-20261031/... |
| ``` |
|
|
| `mtp.json` meta: |
|
|
| | Key | Meaning | |
| |---|---| |
| | `format` / `version` | `rwkv-mtp` / `1` | |
| | `base_model`, `n_layer`, `n_embd` | the base model this sidecar belongs to | |
| | `fuse_layers` | the two tapped base layers (0-based; the layer's output stream) | |
| | `kind`, `depth`, `dd`, `ctx` | `attn`, training depth 3, depth adapter present, window 512 | |
| | `target`, `tokens` | `kl`, training tokens | |
|
|
| Tensor names follow the definition above (`lh`, `lf.{0,1}`, `fproj`, `lhd`, `dep`, `le`, `proj`, `blk.ln{0,1,2}`, `attn.{q,k,v,g,o}`, |
| `attn.{q,k}_norm`, `blk.ffn.{x_k,key,value}`, `lo`). |
|
|
| ## Limitations |
|
|
| - **Experimental.** The architecture is fixed (format v1), but these checkpoints may be replaced by longer-trained ones. |
| - **Base model.** Each sidecar is tied to its base model revision (`RWKV7-G1l-*-20261031`). |
| - **Where the speed-up is large:** for few concurrent streams, where decoding is memory-bound. Large batches gain little and are left to plain decoding by default. |
| - **Where acceptance is lower:** English chat and Chinese. |
| - **Exact output for a given seed:** a K+1-token verify rounds differently from a 1-token step in bf16 and int6, so greedy outputs can differ from plain decoding at near-ties. The distribution is unchanged. |
|
|
| ## Acknowledgements |
|
|
| This model was made possible through the computational resources and technical support provided by |
| **Recursal AI** and **Featherless AI**. We would like to express our deepest gratitude. |
|
|
| ## License |
|
|
| Apache-2.0, the same as the base models. |
|
|
