litert-community/Vision-RWKV-S-LiteRT

Vision-RWKV (VRWKV-S) — ImageNet classification on LiteRT GPU

The first RWKV-style vision backbone running its full forward pass on the LiteRT CompiledModel GPU delegate (no CPU fallback). Vision-RWKV (ICLR 2025, Apache-2.0) replaces softmax self-attention with a bidirectional WKV linear-attention scan — the vision counterpart of the RWKV language model. This is the VRWKV-S ImageNet-1K classifier (80.1% top-1); it is the vision companion to RWKV-7-World-0.1B-LiteRT.

• Architecture: VRWKV-S — 12 blocks, dim 384, patch 16, 14×14 = 196 tokens.
• Weights: OpenGVLab/Vision-RWKV · Apache-2.0.
• Size: 48 MB (fp16).

Top-5 ImageNet predictions on a Pixel 8a; the whole VRWKV-S backbone runs on the GPU.

I/O

• Inputs: image[1,3,224,224] NCHW (ImageNet-normalized, resize-256 → center-crop-224) and dist[1,1,196,196], the constant token-distance matrix dist[t,i] = |t−i|.
• Output: logits[1,1000] (ImageNet-1K).

GPU conversion

VRWKV's token mixer is a CUDA bi_wkv kernel. Because the token count is fixed (196), the bidirectional WKV is exactly a per-channel decay-biased attention:

L[c,t,i] = k[c,i] − (spatial_decay[c]/T)·|t−i| + (spatial_first[c]/T)·δ(t,i)
y[c,t]   = Σ_i softmax_i(L[c,t,·]) · v[c,i]

— C independent [T,T] attention matrices → plain 4D softmax + matmul, no sequential scan. Two things make it GPU-clean and small:

• The [C,T,T] decay bias w·dist (frozen w × constant dist) would be const-folded into a 59 MB-per-block flatbuffer constant — an unshippable 1.5 GB model that fp16 cannot shrink. Feeding the token-distance matrix as a runtime input (eye = relu(1 − dist)) keeps the bias a transient live tensor → 48 MB.
• VRWKV-S is post-norm (norm after the mixer); the LayerScale gamma is baked into the following norm's affine params; q-shift is pad+slice+concat (≤4D).

Pixel 8a: 1371/1371 nodes on the GPU delegate, 1 partition, ~28 ms/inference (fp16). Device fp16 top-1 matches desktop fp32 (bundled Samoyed sample: Samoyed 79%, top-5 identical; logits corr 0.9989). The Bi-WKV re-authoring is oracle-exact (matrix form vs the explicit bidirectional sum, corr 1.0000000).

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val model = CompiledModel.create(context.assets, "vrwkv_s_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()

inputs[0].writeFloat(imageNchw)   // [1,3,224,224] ImageNet-normalized
inputs[1].writeFloat(dist)        // [1,1,196,196], dist[t*196+i] = |t-i|
model.run(inputs, outputs)
val logits = outputs[0].readFloat()   // [1000] -> softmax + argmax

Python (LiteRT CompiledModel API)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

model = CompiledModel.from_file("vrwkv_s_fp16.tflite")
inputs = model.create_input_buffers(0)
outputs = model.create_output_buffers(0)

idx = np.arange(196, dtype=np.float32)
dist = np.abs(idx[:, None] - idx[None, :]).reshape(1, 1, 196, 196)
inputs[0].write(np.ascontiguousarray(image, np.float32))   # [1,3,224,224]
inputs[1].write(np.ascontiguousarray(dist, np.float32))
model.run_by_index(0, inputs, outputs)
logits = outputs[0].read(1000, np.float32)                 # argmax -> class

Files

File	Role
vrwkv_s_fp16.tflite	VRWKV-S step graph (fp16), 48 MB
imagenet_classes.txt	1000 ImageNet-1K labels

License

Apache-2.0 (Vision-RWKV / OpenGVLab). Converted with litert-torch.