rwkv-tl

RWKV7 inference with TileLang fused CUDA kernels and CUDA Graph. The goal is to make decode and prefill faster than the pure PyTorch baseline and approach the performance of the Albatross reference implementation.

• Model: RWKV7-g1d (0.1B and 0.4B variants)
• Precision: float16 compute with float32 accumulation (DPLR state stays fp32), matching Albatross
• Paths:
  • Decode (T=1): fused TMIX/CMIX kernels plus GraphDecoder
  • Prefill (T>1): batched TMIX/CMIX kernels that turn token-wise GEMV into batched GEMM

Project layout

src/rwkv_tl/        # main implementation and fused kernels
script/             # benchmarking and profiling scripts
test/               # correctness and kernel tests
asset/              # tokenizer vocabulary

Install and test

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

Benchmark status

The numbers below were collected on an NVIDIA RTX 3060 (sm_86, 12GB), the target validation GPU, after the stateless refactor with the single-shot fused_dplr_T prefill kernel. rwkv_tl and pure_torch run the eager path; the benchmark harness routes through the raw methods so a sweep does not recompile a fresh graph per token count.

Key points:

• GraphDecoder is best for single-token decode latency.
• rwkv_tl is the only path that supports both batched prefill and decode.
• The single-shot fused_dplr_T prefill kernel made prefill latency flat across T (0.1B ~15-18 ms for all prefill cases); it now beats pure_torch by ~28x at 1x128.
• The Albatross reference (faster3a_2607) still leads prefill by ~2.2x (7.3 ms on 0.1B 1x128 vs 15.8 ms for rwkv_tl).
• Compiling prefill gives 1.11-1.43x on 0.1B, but recompiles a fresh graph per prompt length (minutes), so it stays eager. See script/benchmark_rwkv7.md and docs/benchmarks/rtx3060.md.

MX450 tuning (sm_75) now partially beats the sm75-adapted faster3a_2607

tl-mx450 (sm_75 tuning: fp32 prefill GEMMs + T<=16 tilelang fp16 rkv + CUDA-Graph decode) vs the sm75-adapted faster3a_2607 from yuyi2439/Albatross support/sm75 (0.1B / MX450, warmup=10, iters=20, single session):

tl-mx450 now leads (or ties) the sm75-adapted faster3a_2607 at every T. The wins stack three sm_75 findings: CUDA-Graph decode (stable 8.3ms T=1), CUDA-Graph prefill for T<=64 (small-T prefill was launch-bound: a constant ~2175 launches regardless of T; T=4 dropped 33 -> 11.7ms), and .contiguous() on transposed GEMM weights (non-contiguous cuBLAS operands are ~2.7x slower on Turing; T=128 prefill dropped 70.6 -> 43.4ms).

Why we win despite both sides using CUDA Graph: faster3a_2607 (its sm75 adaptation also captures per-stage torch.cuda.CUDAGraphs) still runs its prefill through fp16 tensor-core GEMMs (volta_fp16_s884gemm... ~39ms of 42.6ms at T=32), which are the pathological Turing fp16 cuBLAS kernels (~4-6x slower than fp32 for these shapes). tl-mx450 deliberately uses fp32 GEMMs for prefill, which is the correct sm_75 adaptation. This is an architecture-level difference, not a measurement artifact.

Run benchmark

.venv/bin/python script/benchmark_rwkv7.py \
  --project-checkpoint <checkpoint.pth> \
  --vocab asset/rwkv_vocab_v20230424.txt \
  --targets tl-fp16,pure-torch \
  --device cuda \
  --cases 1x1,1x8,1x32,2x1,8x1,8x8,16x16 \
  --warmup 10 --iters 20

On memory-constrained machines, split large sweeps into separate processes to avoid compiler-cache pressure and OOMs.