rwkv7-kernel

RWKV-7 WKV recurrence for the Hugging Face kernels ecosystem: the single-token decode step runs a hand-written CUDA kernel (fp16) or a fused Triton kernel (bf16) that reads the recurrent state once per (batch, head); every other shape — prefill, packed rows, CPU, odd widths — takes the portable reference paths carried verbatim from the transformers RWKV-7 implementation, so this package is correct on any input.

Use

The package is importable directly from a hub snapshot:

from huggingface_hub import snapshot_download
import sys
sys.path.insert(0, snapshot_download("Hakureirm/rwkv7-kernel") + "/build/torch-universal")
import rwkv7_kernel as K

out, new_state = K.rwkv7_wkv(r, w_log, k, v, kk, a, state)

Loading through the kernels hub library (get_kernel(...)) is the intended end state and requires the kernel-type repo plus a version tag; this repo is currently the legacy model-type layout, and the Hakureirm/ namespace is not on the kernels-community allow-list, so that path is not live yet.

Same contract as the transformers RWKV-7 rwkv7_eager: vectors are [batch, seq_len, num_heads, head_dim], the state is [batch, num_heads, head_dim, head_dim] (key axis first), packed rows via cu_seq_lens_q, return (output, new_state).

Dispatch and honest scope

Three tiers, first matching one wins:

1. CUDA decode kernel (rwkv7_wkv.cu): fp16 vectors, head_dim 64, power-of-two heads, fp16/fp32 state. Updates the state in place (its pool semantics). Microbenchmark of the WKV step alone on an A800 (A100-class, sm80), 3 reps medianed: fp16 ~13× faster than the eager reference at batch 1–128 (13.0–13.5× across the sweep).
2. Triton decode kernel (fused_wkv_one): same decode domain, any vector dtype including bf16, fp32 state accumulation. bf16 (the official checkpoint dtype) measures ~7× faster than eager on the WKV step at batches 1–32 — real but below the fp16 CUDA path, which declines bf16.
3. Portable reference (rwkv7_recurrent / rwkv7_chunked / rwkv7_eager): everything else — prefill, packed rows, CPU, non-power-of-two widths.

Chunk-parallel prefill acceleration is future work.

Numerics

Checked against the reference paths in tests/, and against the real model end-to-end: on the official 1.5B checkpoint, a forward with this package wired in produces bit-identical logits to eager at both fp16 and bf16, and greedy generation is token-identical. The reference paths themselves are the ones validated against BlinkDL's own runtimes in huggingface/transformers#47787 (fp32 8.2e-6 worst-case over every position and per-layer state).

Attribution

See NOTICE: the reference paths come from the author's own transformers RWKV-7 submission; sparse_cmix.cu derives from BlinkDL / Albatross (Apache-2.0, chain documented); rwkv7_wkv.cu is original work carried from the author's RWKV-7 sglang overlay.