Blame for README.md

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