Qwen3-Reranker-0.6B — LiteRT on-device RAG reranker (fully GPU)

Qwen3-Reranker-0.6B (Apache-2.0), the 2025 SOTA small reranker, re-authored to run entirely on the LiteRT CompiledModel GPU (ML Drift). Given a query and candidate documents, it scores each by relevance (P("yes")) and reorders them — the reranking half of an on-device RAG pipeline.

Pairs with litert-community/Qwen3-Embedding-0.6B-LiteRT: embed → retrieve top-k → rerank, all on-device, no server.

Like the embedder it is a single forward pass (no generation, no KV cache) → a plain .tflite, not a .litertlm. Verified on a Pixel 8a / Tensor G3: all nodes on the GPU delegate, P(yes) parity ref 0.9995 / dev 0.9994 vs the HF fp32 reference.

Files

file	purpose	runs on
`qwen3rerank_gpu_fp16.tflite`	28-layer Qwen3 decoder + baked 2-logit head, `inputs_embeds[1,256,1024] → logits[1,256,2]`	GPU
`embeddings_fp16.bin`	tied token-embedding table `[151669,1024]` fp16, for the host-side lookup	host
`vocab.json`, `merges.txt`	Qwen byte-level BPE tokenizer	host

How it scores

prompt = PREFIX + "<Instruct>:… <Query>:… <Document>:…" + SUFFIX     (Qwen3-Reranker template)
       →[host embed lookup]→ inputs_embeds[1,256,1024]
       →[GPU: 28-layer decoder + 2-logit head]→ logits[1,256,2]
       →[softmax over (no,yes) at the last token]→ P(yes) = relevance

The 2-logit head bakes the tied-embedding rows for "no" (2152) and "yes" (9693), so the graph emits [no,yes] directly. The host right-pads and pools the last real token (causal ⇒ it never sees the trailing pad, identical to the official left-pad + attention-mask). Token embedding is a GATHER (GPU-banned) so it is done host-side.

The GPU-clean re-authoring is the same as the embedder (host-embed, GQA cat-repeat to avoid BROADCAST_TO, max-normalized RMSNorm for the deep-stack fp16 overflow, baked RoPE / causal mask).

Minimal usage

Python (reference score with the original model):

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-Reranker-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-Reranker-0.6B").eval()
yes, no = tok.convert_tokens_to_ids("yes"), tok.convert_tokens_to_ids("no")
# … build the PREFIX/SUFFIX prompt, then:
logits = model(**inputs).logits[:, -1, :]
score = torch.softmax(torch.stack([logits[:, no], logits[:, yes]], 1), 1)[:, 1]  # P(yes)

Kotlin (on-device, LiteRT CompiledModel GPU):

val model = CompiledModel.create("qwen3rerank_gpu_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
// host: build prompt ids -> lookup embeddings_fp16.bin -> inputs_embeds[1,256,1024]
inputs[0].writeFloat(embedLookup(promptIds(query, doc)))
model.run(inputs, outputs)
val logits = outputs[0].readFloat()               // [256,2] = [no,yes] per position
val score = softmaxYes(logits, poolPos)           // P(yes) relevance

Full tokenizer + prompt template + reranking app: see the official LiteRT sample.

Conversion

Reproducible in the official sample's conversion/ (build_qwen3rerank.py, export_embeddings.py, device-parity harness).

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
TFLite `benchmark_model` (`TfLiteGpuDelegateV2`)	GPU (OpenCL)	3157 / 6425	9277.4 ms
TFLite `benchmark_model`	CPU (XNNPACK, 4 threads)	—	4066.1 ms

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

On this delegate the CPU is the faster choice for (4066.1 ms on CPU against 9277.4 ms on GPU) — worth knowing before you reach for the GPU on a mid-range phone.

Note that the GPU does not take the whole graph here (3157 of 6425); the remainder runs on the CPU and the split costs a per-partition round trip.