---
license: apache-2.0
base_model: OpenGVLab/Vision-RWKV
base_model_relation: quantized
library_name: litert
pipeline_tag: image-classification
tags:
  - litert
  - tflite
  - android
  - on-device
  - gpu
  - rwkv
  - vision-rwkv
  - image-classification
---

Measured on device (edge-compat): Raspberry Pi 5 · LiteRT 2.2.0.dev20260804 · CPU/XNNPACK, 4 threads · 1083 ms p50 (2026-08-31); browser · Chromium 151 on M4 Max · LiteRT.js 2.5.3 · WebGPU · 27.5 ms p50 · output differs from CPU (max rel diff 25) (2026-08-11). Record: https://github.com/john-rocky/edge-compat/blob/main/cards/vision-rwkv-s/CARD.md

# 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](https://github.com/OpenGVLab/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](https://huggingface.co/litert-community/RWKV-7-World-0.1B-LiteRT).

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

![Vision-RWKV on-device classification](hero.png)

*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)

```kotlin
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)

```python
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`](https://ai.google.dev/edge/litert/models/measurement) 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](https://github.com/john-rocky/hf-to-litertlm/blob/main/docs/android-npu.md), and we did not measure it here. GPU wiring is in the [GPU recipe](https://github.com/john-rocky/hf-to-litertlm/blob/main/docs/android-gpu.md).

## Raspberry Pi 5 (CPU)

Measured on a Raspberry Pi 5 Model B Rev 1.1 (8 GB, Raspberry Pi OS 64-bit) with the LiteRT [`benchmark_model`](https://ai.google.dev/edge/litert/models/measurement) tool from litert-cli-nightly 0.2.0.dev20260805: CPU inference (XNNPACK, 4 threads), 3 invocations per file of 10 warm-up plus 50 timed runs (the tool caps a phase at 150 s, so very slow graphs run fewer — the Runs column is the actual timed total). The latency is the median across invocations; the spread is the min–max over all timed runs. No thermal throttling occurred during these runs (`vcgencmd get_throttled` stayed `0x0`).

| File | Inference (median) | Spread (min–max) | Runs | Peak memory |
|---|---:|---:|---:|---:|
| `vrwkv_s_fp16.tflite` | 1,082.6 ms | 1,075.0–1,128.4 ms | 150 | 277 MB |

## License

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