5.8 KiB

Raw Permalink Blame History

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

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

5.8 KiB Raw Permalink Blame History