rwkv-tl

An RWKV7 operator library (built on TileLang) plus ready-to-use stateless models. Import the package, point it at a checkpoint, and generate:

import rwkv_tl

model = rwkv_tl.rwkv7("model-0.4b.pth")          # tilelang on CUDA, torch elsewhere
out = model.generate("Once upon a time", max_tokens=128)
print(out)

Everything is stateless: models never own runtime state. Pass a State in and (optionally) get it back, or let generate create a fresh one.

User API

Build a model from a checkpoint path (or a pre-loaded RWKV7Weight):

model = rwkv_tl.rwkv7("model.pth")                       # backend auto-selected
model = rwkv_tl.rwkv7("model.pth", backend="torch")      # pure-PyTorch reference
model = rwkv_tl.RWKV7TL("model.pth")                     # explicit tilelang class

Text in, text out:

text = model.generate("The meaning of life is",
                      max_new_tokens=64, temperature=0.8, stop="\n\n")

Chat (messages through the packaged chat template):

answer = model.chat([
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is the capital of France?"},
], max_new_tokens=128)

Demo

examples/ holds runnable walk-throughs: basic_load_and_generate.py loads a checkpoint and generates text, chat.py chats through the template:

.venv/bin/python examples/basic_load_and_generate.py /path/to/rwkv7-0.1b.pth
.venv/bin/python examples/chat.py /path/to/rwkv7-0.1b.pth

Operator library

rwkv_tl.kernel is a TileLang operator library for building efficient RWKV implementations. Every factory is weight-bound: it takes the compile-time hyperparameters and the weights at construction and returns a callable that only needs activations/state at call time:

from rwkv_tl.kernel import ln_kernel, gemv_kernel

ln_pre = ln_kernel(C, DTYPE, ln_preW, ln_preB)   # weights captured here
x_ln = ln_pre(x0)                                 # call with activations only

Both granularities are supported:

• Fine-grained composable operators: ln_kernel, ln_per_row_kernel, gemv_kernel, gemv_batch_kernel.
• Coarse fused layer kernels: cmix_decode_kernel, cmix_prefill_kernel, tmix_decode_kernel, tmix_prefill_kernel.

The raw @tilelang.jit factories and shared macros (gemv_macro, gemv_main_macro, ...) remain available for custom fused chains. Weights are held by the wrapper, which is the hook for a future quantized-weight path (int8/any4 storage + dequant fused into the kernels).

Layout

src/rwkv_tl/        # published library: models, State, Tokenizer, kernel/
  core/             # low-level/inference modules (model/state/tokenizer/
                    # weight/cuda_graph); no references outside core
  kernel/           # weight-bound tilelang operator factories
  text_model.py     # RWKV7TextModel (exposed): composes a token model
                    # (self.model) + tokenizer; tokenize/generate/chat
  rwkv7_tl.py       # tilelang fused model
  rwkv7_torch.py    # pure-PyTorch reference model
  asset/            # packaged data: vocab + chat template
script/             # chat, benchmark, profiling scripts
examples/           # runnable usage examples
test/               # correctness and API tests
docs/               # benchmark reports and tuning notes (Chinese)

Install and test

cd rwkv-tl
uv sync
.venv/bin/python -m pytest test/ -v

Kernel correctness tests need CUDA and RWKV_CHECKPOINT_PATH:

RWKV_CHECKPOINT_PATH=/path/to/rwkv7-g1d-0.1b.pth .venv/bin/python -m pytest test/ -v

The user-facing text API and the pure-torch backend also run on CPU.

script/check_torch_vs_official.py additionally validates the pure-torch backend against the official RWKV-LM v7 demo (pure-torch path) on the same checkpoint — logits must agree on argmax and top-5 for batched and per-token decode:

.venv/bin/python script/check_torch_vs_official.py /path/to/rwkv7-0.1b.pth \
  --fast-path /path/to/RWKV-LM/RWKV-v7/rwkv_v7_demo.py

Performance

Decode and prefill use fused tilelang kernels with fp16 compute and fp32 accumulation (DPLR state stays fp32), CUDA-Graph accelerated on CUDA by default. Current numbers vs the Albatross reference implementation are in docs/runs/rtx3060.md (RTX 3060, the current target card). script/bench_tl_vs_torch.py measures tl vs pure-torch on CUDA (prefill sweep + decode), and script/bench_tl_vs_fast.py compares tl against the Albatross faster3a_2607 reference implementation.