Ilikemechuri/quail-test

RWKV-7 "g1d" 0.1B — HuggingFace port (OLMo tokenizer)

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

If it keeps falling back, check the cache folder first!!!!!!!!!!!!!!!!!!!!Basically, download the 'cuda' folder and place it inside the 'cache' folder.

A trust_remote_code=True HuggingFace wrapper around the BlinkDL RWKV-7 "Goose" g1d 0.1B checkpoint (rwkv7-g1d-0.1b-20260129-ctx8192.pth), re-headed for the OLMo tokenizer.

• The 12 transformer-style RWKV-7 blocks keep their pretrained weights.
• The embedding and lm-head are re-initialized (from RWKV's 65536-token vocab to OLMo's 100278-token vocab), so they are untrained and need fine-tuning.
• Time-mixing runs on a fused CUDA kernel (forward and backward) when a GPU
  • CUDA toolchain are available, and transparently falls back to pure PyTorch otherwise.
• generate() runs in RNN mode with a recurrent state cache — after the prompt is absorbed into the state, each new token is a single-token forward instead of a full-sequence recompute. Stateful inference runs on a second, forward-only "wkv7s" CUDA kernel (state in/out), so prefill and decode are kernel-speed too.

Because emb/head are freshly initialized, generate() produces gibberish until you fine-tune on OLMo-tokenized data. The model body is pretrained; only the vocabulary projection is new.

---

Files

File	Purpose
config.json	Serialized RWKV7Config (dims + auto_map to the remote-code classes).
configuration_rwkv7.py	RWKV7Config — all architecture hyperparameters.
modeling_rwkv7.py	Model code: CUDA-kernel dispatch + PyTorch fallback, recurrent state cache for generation, RWKV7Model, RWKV7ForCausalLM.
cuda/wkv7_cuda.cu, cuda/wkv7_op.cpp	The fused bf16 "wind_backstepping" RWKV-7 kernel (forward + backward), copied from RWKV-v7/train_temp/cuda/. Used on the stateless (training) path.
cuda/wkv7s.cu, cuda/wkv7s_op.cpp	The stateful, forward-only "wkv7s" inference kernel (fp32 state in/out, arbitrary T, no chunk padding), adapted from RWKV-LM's inference kernel — patched here for bf16 (upstream typedef was fp16) and B > 1 (batch-aware state indexing). Used on the stateful (generation) path.
model.safetensors	Converted weights (bf16, ~244M params).
tokenizer.json, tokenizer_config.json	OLMo tokenizer (vocab 100278, GPT2-style BPE).
generation_config.json	Default generation settings (eos/pad ids).
convert.py	Reproduces model.safetensors from the original .pth.
verify.py	End-to-end smoke test (load / forward / parity / backward / generate).

Architecture (from config.json)

field	value	meaning
num_hidden_layers	12	RWKV-7 blocks
hidden_size	768	embedding dim C
head_size	64	→ 12 heads (H = C / head_size)
intermediate_size	3072	channel-mix (FFN) hidden
decay_lora / aaa_lora / mv_lora / gate_lora	64 / 64 / 32 / 128	LoRA ranks for w / a / v / g
vocab_size	100278	OLMo tokenizer size (re-initialized emb/head)
chunk_len	16	CUDA kernel chunk length; sequence is padded to a multiple of this
use_cuda_kernel	true	prefer the fused kernel when possible

---

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PATH = "/workspace/rwkv7-g1d-olmo"

tok = AutoTokenizer.from_pretrained(PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    PATH, trust_remote_code=True, dtype=torch.bfloat16
).cuda().eval()

ids = tok("The Eiffel tower is in the city of", return_tensors="pt").input_ids.cuda()
with torch.no_grad():
    logits = model(ids).logits          # (1, T, 100278)

# training / backward
model.train
#🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨labels[labels == tokenizer.pad_token_id] = -100🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
out = model(ids, labels=ids)            # shifted cross-entropy loss
out.loss.backward()                     # gradients flow through the CUDA kernel

Forcing the PyTorch fallback

Set the flag on the config (useful on CPU, non-bf16, or to debug the kernel):

model.config.use_cuda_kernel = False    # every RWKV-7 op now runs in pure PyTorch

The kernel is also skipped automatically when: no CUDA device, input dtype is not bfloat16, or the kernel fails to compile — in all cases the fallback is used and a one-line notice is printed.

Stateful generation (RNN mode)

generate() uses the state cache automatically — nothing to configure:

out_ids = model.generate(ids, max_new_tokens=64, do_sample=False)

You can also drive the state manually (streaming, chunked prefill, state save/reuse, "infinite context" style loops):

with torch.no_grad():
    out = model(ids, use_cache=True)          # prefill: prompt -> state
    next_tok = out.logits[:, -1].argmax(-1, keepdim=True)
    out = model(next_tok, state=out.state)    # passing `state` implies use_cache

outputs.state is a list of per-layer tuples (att_x_prev, wkv_state, ffn_x_prev):

tensor	shape	dtype	role
att_x_prev	(B, C)	input dtype	last post-ln1 token, time-shift for the time-mix
wkv_state	(B, H, N, N)	fp32	the WKV matrix state
ffn_x_prev	(B, C)	input dtype	last post-ln2 token, time-shift for the channel-mix

States are ordinary tensors — you can clone() them to snapshot/branch a context, or persist them to disk and resume later.

Which implementation runs when? Stateful calls under torch.no_grad() (i.e. normal generation) use the forward-only wkv7s kernel — one kernel launch per forward, any T, no chunk padding — so both long-prompt prefill and per-token decode are fast. wkv7s has no backward, so a stateful call with autograd enabled (e.g. TBPTT-style training that carries state) silently drops to the sequential PyTorch path instead. The kernel mutates its state buffer in place; the wrapper always clones the incoming state first, so a state you're holding on to is never corrupted (snapshot/branching stays safe).

Zero token-shift state + zero WKV state is exactly equivalent to the zero-padded non-cached forward, so prefill→decode with the cache reproduces the full-recompute logits.

---

How the CUDA kernels + fallback work (modeling_rwkv7.py)

The time-mixing recurrence is dispatched by run_rwkv7(r, w, k, v, a, b, config, initial_state=None, output_final_state=False), which returns (y, final_state_or_None). Dispatch is three-way:

call	implementation
stateless (training / plain forward), CUDA + bf16	wind_backstepping (fused fwd + bwd)
stateful (cache), CUDA + bf16, autograd off	wkv7s (fwd-only, state in/out)
everything else (CPU, non-bf16, grad-enabled stateful, compile failure)	pure PyTorch

1. Training kernel (_rwkv7_cuda, "wind_backstepping") — used when config.use_cuda_kernel is set, the tensors are CUDA + bf16, the call is stateless (initial_state is None and no final state requested), and _try_load_cuda_kernel() succeeds. On first call it JIT-compiles cuda/wkv7_op.cpp + cuda/wkv7_cuda.cu via torch.utils.cpp_extension.load (flags -D_C_=head_size, -D_CHUNK_LEN_=chunk_len) and registers the wind_backstepping op. Compilation happens once per process and is cached; any failure is caught and flips the model to the fallback.

   • _WindBackstepping is a torch.autograd.Function:
     • forward → wind_backstepping.forward (produces y plus the saved state s and sa needed for backprop).
     • backward → wind_backstepping.backward (returns gradients for all six inputs w, q, k, v, z, b), so training works end-to-end on the kernel.
   • The sequence length is padded to a multiple of chunk_len (16) and sliced back.

2. Inference kernel (_rwkv7_cuda_stateful, "wkv7s") — used for stateful calls (an incoming state and/or output_final_state=True) when the tensors are CUDA + bf16 and autograd is disabled (torch.no_grad() — the kernel has no backward). JIT-compiled from cuda/wkv7s_op.cpp + cuda/wkv7s.cu (flag -D_N_=head_size), registered as the wkv7s op.

   • The time loop lives inside the kernel: a single launch processes any T (no chunk_len padding), fully parallel over B·H blocks × N threads.
   • The fp32 state (B, H, N, N) is read at the start and written back at the end (in-place); the Python wrapper clones the incoming state so the caller's tensor is never mutated.
   • Patched relative to the upstream RWKV-LM inference kernel: bf16 typedef was actually at::Half (fp16) upstream → switched to at::BFloat16, and the state indexing was B = 1-only → made batch-aware.

3. Fallback path (_rwkv7_pytorch) — a plain sequential-over-time implementation of the same recurrence state = state*exp(-exp(w)) + state·aᵀ·b + vᵀ·k, y = state·r, in fp32. It is fully differentiable through ordinary autograd (no custom backward needed), accepts an initial_state, and can return the final state — it covers every case the kernels can't (CPU, non-bf16, grad-enabled stateful calls).

All three paths take the raw (pre-exp) decay w and apply exp(-exp(w)) internally, and implement the identical per-head recurrence (s[i,j] = s[i,j]·w_j + v_i·k_j + (Σⱼa_j·s[i,j])·b_j), so they are numerically interchangeable. verify.py confirms kernel-vs-fallback parity (identical bf16 logits and top-1 prediction).

Model structure

RWKV7ForCausalLM → .rwkv (RWKV7Model) + .head (lm-head). RWKV7Model holds emb, blocks[0..11] (RWKV7Block = ln1 + att time-mix, ln2 + ffn channel-mix; block 0 also has ln0), and ln_out. State-dict keys mirror the original RWKV layout under the rwkv. prefix (e.g. rwkv.blocks.0.att.receptance.weight).

Generation runs in RNN mode: the model carries a per-layer recurrent state (see the state table above) through the state= kwarg / outputs.state field — the same convention as transformers' Rwkv and Mamba models, so GenerationMixin propagates it between steps automatically (state is in ALL_CACHE_NAMES; a _update_model_kwargs_for_generation override covers older versions). prepare_inputs_for_generation feeds only the last token once a state exists. Passing state= implies use_cache=True; gradient checkpointing forces it off during training. Stateful forwards run on the wkv7s kernel under no_grad, PyTorch otherwise.

---

Reproducing the conversion (convert.py)

# 1) download the original checkpoint (already done in /workspace)
wget https://huggingface.co/BlinkDL/rwkv7-g1/resolve/main/rwkv7-g1d-0.1b-20260129-ctx8192.pth \
     -O /workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth

# 2) convert -> /workspace/rwkv7-g1d-olmo
python convert.py

What convert.py does:

1. Loads the OLMo tokenizer to read the target vocab size (100278).
2. Loads the .pth and infers all dims from the tensor shapes → builds RWKV7Config.
3. Remaps the original RWKV keys to the HF module layout (rwkv. prefix), dropping emb.weight / head.weight, and loads them with strict=False (asserts there are no unexpected or unmatched keys besides emb/head).
4. Re-initializes the embedding (uniform(±1e-4)) and head (orthogonal, gain 0.5·√(vocab/hidden)) for the new vocabulary — RWKV's own init scheme.
5. Saves weights (safetensors), config, tokenizer, and copies the remote-code files.

Verifying (verify.py)

python verify.py

Checks: load via AutoModelForCausalLM(trust_remote_code=True), CUDA-kernel forward, kernel-vs-fallback parity, cached-vs-uncached parity (prefill + stateful decode vs full recompute), a backward pass (gradients on emb / attention / decay-LoRA), and a short greedy generate() through the state cache.

Note: 399/402 parameters receive gradients — the 3 without are blocks.0.att.v0/v1/v2 (the value-residual params are unused in layer 0 by design).

---

Requirements

• PyTorch with CUDA (bf16-capable GPU; tested on RTX 3060 / CUDA 13.0) for the kernel; CPU/other works via the fallback.
• transformers >= 5, safetensors.
• A working CUDA toolchain (nvcc) for first-call kernel JIT compilation; if absent, the model still runs on the PyTorch fallback.

Credits

RWKV-7 architecture and the original checkpoint/kernels by BlinkDL — https://github.com/BlinkDL/RWKV-LM. Training kernel copied from RWKV-v7/train_temp/cuda/; the stateful wkv7s inference kernel adapted from RWKV-LM's inference code (patched here for bf16 + batched state).