litert-community/Qwen3-Reranker-0.6B-LiteRT

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