142 lines
5.8 KiB
Markdown
142 lines
5.8 KiB
Markdown
# RWKV-X Architecture
|
|
|
|
RWKV-X is a hybrid model architecture that interleaves RWKV-7 recurrent blocks (~75% of layers) with sparse attention blocks (~25%). It is fundamentally different from transformer models and uses safetensors format instead of GGUF.
|
|
|
|
## Source
|
|
|
|
The RWKV-X engine in Pantheon (`src/models/rwkvx.c`, ~2,500 LOC) is copied from the battle-tested Foundry implementation (`foundry/rwkvx/rwkvx.c`), with Metal, CUDA, HybridMem, training, batch, instrumentation, and OpenMP support stripped. A thin adapter layer wraps it into Pantheon's `pn_model`/`pn_state`/`pn_forward_fn` interface.
|
|
|
|
## Model Architecture
|
|
|
|
For the 0.2B model (12 RWKV + 4 attention layers), block execution order is:
|
|
|
|
```
|
|
R0, R1, R2, A0, R3, R4, R5, A1, R6, R7, R8, A2, R9, R10, R11, A3
|
|
```
|
|
|
|
Where `R` = RWKV-7 recurrent block, `A` = sparse attention block.
|
|
|
|
### Forward Pass
|
|
|
|
```
|
|
embed(token)
|
|
→ for each block in interleaved order:
|
|
if RWKV-7 block:
|
|
LayerNorm → TimeMix (WKV-7 kernel) → residual
|
|
→ LayerNorm → ChannelMix FFN → residual
|
|
if attention block:
|
|
LayerNorm → Multi-Head Attention (with KV cache) → residual
|
|
→ LayerNorm → ChannelMix FFN → residual
|
|
→ LayerNorm → LM head → logits
|
|
```
|
|
|
|
### WKV-7 Kernel (Generalized Delta Rule)
|
|
|
|
The core recurrent kernel, computed per-head per-timestep:
|
|
|
|
```
|
|
for each timestep t, for each head h:
|
|
state *= diag(w[t]) // decay
|
|
state += (a[t] @ state) * k^T // delta rule bonus
|
|
state += v[t] @ k[t]^T // key-value update
|
|
output[t] = r[t] @ state // readout
|
|
```
|
|
|
|
State is a `[head_size, head_size]` matrix per head per RWKV-7 layer — fixed size regardless of sequence length.
|
|
|
|
### Attention Blocks
|
|
|
|
Standard multi-head attention with KV cache. Uses cached keys and values up to a configurable max length (default 2048). No RoPE — position information comes from the surrounding RWKV-7 blocks.
|
|
|
|
## Weight Format
|
|
|
|
RWKV-X uses safetensors (not GGUF). Weights are loaded via Foundry's `fs_open()` (zero-copy mmap) and routed by prefix:
|
|
|
|
| Prefix | Block Type |
|
|
|--------|-----------|
|
|
| `rwkv.blocks.N.*` | RWKV-7 recurrent block N |
|
|
| `moba.blocks.N.*` | Sparse attention block N |
|
|
| `emb.weight` | Token embeddings |
|
|
| `head.weight` | LM head |
|
|
| `ln_out.*` | Final layer norm |
|
|
|
|
Weight transformations at load time:
|
|
- Linear weights are transposed (safetensors stores `[out, in]`, inference needs `[in, out]`)
|
|
- Singleton dimensions are squeezed
|
|
- `att.r_k` is flattened from `[n_head, head_size]` to `[n_embd]`
|
|
- Embeddings are pre-normalized (L2 norm per row)
|
|
|
|
## INT8 Quantization
|
|
|
|
When `--quantize` is passed, all linear projection weights are quantized to INT8 at load time using per-row symmetric quantization:
|
|
|
|
```
|
|
scale[row] = max(abs(row)) / 127
|
|
quantized[row][col] = round(float[row][col] / scale[row])
|
|
```
|
|
|
|
Inference uses a NEON-accelerated INT8 dot product kernel (`fast_matvec_q8`) that pre-quantizes the input vector once, then computes INT8 dot products for all output rows. This gives ~3x speedup on ARM.
|
|
|
|
## Tokenizer
|
|
|
|
RWKV-X uses a **trie-based World tokenizer** with 65,536 tokens. The tokenizer operates on raw UTF-8 bytes with greedy longest-match — fundamentally different from BPE or SentencePiece.
|
|
|
|
The vocab is loaded from a text file (`rwkv_vocab_v20230424.txt`) via the `--vocab` CLI flag. Each line contains a token ID, token length, and the token's byte representation as a Python byte literal.
|
|
|
|
The tokenizer is stored inside the RWKV-X model struct (not in `pn_model->tokenizer`). The engine routes to it via `pn_rwkvx_encode()`/`pn_rwkvx_decode()` when `config.arch == PN_ARCH_RWKVX`.
|
|
|
|
## Dual State
|
|
|
|
Each inference session maintains two state types in `pn_state->kv_cache`:
|
|
|
|
- **Recurrent state**: Fixed-size per-layer matrix `[n_heads, head_size, head_size]` for RWKV-7 blocks
|
|
- **KV cache**: Key-value store for sparse attention blocks (capped at `max_cache_len`)
|
|
- **Workspace**: Pre-allocated scratch buffers (16 buffers + ping-pong pair + INT8 buffer)
|
|
|
|
All state is encapsulated in a single opaque struct. `pn_state_reset()` zeros both the recurrent state and KV cache.
|
|
|
|
## Adapter API
|
|
|
|
The adapter layer in `rwkvx.c` maps Pantheon's interface to internal RWKV-X functions:
|
|
|
|
| Pantheon API | Internal Function |
|
|
|-------------|-------------------|
|
|
| `pn_rwkvx_load()` | `rx_model_load_internal()` |
|
|
| `pn_rwkvx_free()` | `rx_model_free_internal()` + tokenizer free |
|
|
| `pn_rwkvx_forward()` | `rx_forward_internal()` |
|
|
| `pn_rwkvx_state_new()` | `rx_state_new_internal()` |
|
|
| `pn_rwkvx_state_free()` | `rx_state_free_internal()` |
|
|
| `pn_rwkvx_state_reset()` | `rx_state_reset_internal()` |
|
|
| `pn_rwkvx_encode()` | `rx_tokenizer_encode()` |
|
|
| `pn_rwkvx_decode()` | `rx_tokenizer_decode()` |
|
|
|
|
## Usage
|
|
|
|
```bash
|
|
# FP32 inference
|
|
./pantheon --model ../models/base/RWKV-X-0.2B-64k-Base.safetensors \
|
|
--vocab foundry/rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
--prompt "Hello world" --max-tokens 32
|
|
|
|
# INT8 quantized inference (~3x faster decode on ARM)
|
|
./pantheon --model ../models/base/RWKV-X-0.2B-64k-Base.safetensors \
|
|
--vocab foundry/rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
--prompt "Hello world" --max-tokens 32 --quantize
|
|
```
|
|
|
|
## What Was Stripped
|
|
|
|
The Foundry RWKV-X engine has additional features not yet in Pantheon:
|
|
|
|
| Feature | Why Stripped |
|
|
|---------|-------------|
|
|
| Metal GPU (`fm_*`) | Pantheon doesn't have Metal model integration yet |
|
|
| CUDA GPU (`fc_*`) | Pantheon doesn't have CUDA model integration yet |
|
|
| HybridMem (`hm_*`) | Separate feature, not part of base inference |
|
|
| Training (`rx_forward_ex`, callbacks) | Training not in Pantheon scope yet |
|
|
| Batch forward (`rx_forward_batch`) | Prefill optimization, add later |
|
|
| Instrumentation (attn_log, embed_capture) | Debug tooling, add later |
|
|
| OpenMP (`#pragma omp`) | Pantheon uses threadpool, not OpenMP |
|
|
|
|
These can be added incrementally in future phases.
|