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

Qwen3-Embedding-0.6B — LiteRT on-device text embeddings (fully GPU)

Qwen3-Embedding-0.6B (Apache-2.0), the 2025 state-of-the-art small text-embedding model, re-authored to run entirely on the LiteRT CompiledModel GPU (ML Drift). Embed a query and a set of documents on-device and rank by cosine similarity — the retrieval half of a RAG pipeline, no server, no NPU, no CPU fallback.

Because sentence embedding uses last-token pooling of a single forward pass (no generation, no KV cache), the model is a plain single-graph .tflite — not a .litertlm — and runs on the same GPU path as any CNN/ViT.

Verified on a Pixel 8a / Tensor G3: CompiledModel GPU compile OK, all 3264/3264 nodes on the GPU delegate (zero CPU fallback), ~390 ms per embedding, fp16 881 MB, output cosine 0.9997 vs the HF fp32 reference. Semantic ranking is correct — query "What is the capital of China?" → "The capital of China is Beijing" at 0.77, unrelated docs at <0.1.

Files

file	purpose	runs on
qwen3emb_gpu_fp16.tflite	28-layer Qwen3 transformer, inputs_embeds[1,128,1024] → hidden[1,128,1024]	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

Pipeline

text →[BPE tokenize]→ ids →[host embed lookup]→ inputs_embeds[1,128,1024]
     →[GPU: 28-layer Qwen3 decoder]→ hidden[1,128,1024]
     →[pool last token + L2-normalize (+ optional Matryoshka 1024→N)]→ embedding
     →[cosine]→ ranked documents

The token-embedding lookup is a GATHER (GPU-banned), so it is done on the host and fed in as inputs_embeds, exactly like mel/log-mel preprocessing in the audio samples.

Why it runs fully on GPU — a Mali fp16 finding

A 28-layer decoder is the first decoder-transformer verified end-to-end on this GPU path. The one device-only fix: the residual stream grows across depth, so the deep RMSNorm's mean(x²) overflows fp16 (>65504) → rsqrt(inf)=0 → the whole output collapses to 0 (even though Qwen3 already RMSNorm+qk-norms every sub-layer input — it is the residual that overflows). The fix is a per-row max-normalized RMSNorm, mathematically identical to the original:

m  = max(|x|).clamp_min(1e-4)          # per-row scale
xs = x / m                              # x² now in [0,1] — the 1024-term sum never overflows
y  = xs * rsqrt(mean(xs²) + eps/m²) * w

GQA heads are cat-repeated to 16 (a broadcast matmul would emit the Mali-rejected BROADCAST_TO), RoPE and the causal mask are baked constants, and every tensor stays ≤4D.

Minimal usage

Python (reference embeddings with the original model):

from sentence_transformers import SentenceTransformer
m = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")
q = m.encode("What is the capital of China?", prompt_name="query")
docs = m.encode(["The capital of China is Beijing.", "Paris is the capital of France."])
print((q @ docs.T))          # cosine similarity — the .tflite reproduces this at corr 0.9997

Kotlin (on-device, LiteRT CompiledModel GPU):

val model = CompiledModel.create("qwen3emb_gpu_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers(); val outputs = model.createOutputBuffers()

// host: BPE tokenize -> lookup embeddings_fp16.bin -> inputs_embeds[1,128,1024]
inputs[0].writeFloat(embedLookup(tokenize(text)))
model.run(inputs, outputs)
val hidden = outputs[0].readFloat()          // [128,1024]
val emb = l2normalize(hidden, poolPos)       // last real token -> 1024-d embedding

Full tokenizer + embedding-lookup + semantic-search app: see the official LiteRT sample.

Conversion

The GPU-clean re-authoring is fully reproducible — conversion scripts (build_qwen3emb.py, export_embeddings.py) and a device-parity harness are provided with the official sample.