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

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
LiteRT CompiledModel (LITERT_CL)	GPU	1371 / 1371	~28 ms
TFLite benchmark_model (TfLiteGpuDelegateV2)	GPU (OpenCL)	1371 / 1371	283.6 ms
TFLite benchmark_model	CPU (XNNPACK, 4 threads)	—	402.5 ms

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator — the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

Snapdragon NPU (Hexagon)

The NPU is 1.08x faster than the GPU (62.06 ms against 66.92 ms) and loads 3.99x faster (282 ms against 1124 ms).

backend	inference (median / min)	load
NPU (Hexagon v81)	62.06 ms / 61.38 ms	282 ms
GPU (Adreno)	66.92 ms / 66.69 ms	1124 ms

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16), LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.67-0.69, where 1.0 is the throttling threshold.

The NPU rows here ran artifacts compiled ahead of time for SM8850 with QAIRT 2.47.0; the GPU rows ran the published files as they are. LiteRT can also compile for the NPU on the device at first load, which is what lets you ship the published file unchanged — that path and the ten runtime libraries it needs are in the NPU recipe, and we did not measure it here. GPU wiring is in the GPU recipe.

License

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