trit-embed-3l

trit-embed-3l is a small English sentence embedder built for the AMD Ryzen AI NPU. There it encodes about 51,500 sentences a second, 1.4 times as fast as bge-micro-v2 in fp16 on the GPU, with half the single-sentence latency. Its layers use just three weight values, −1, 0 and +1, so the whole model is a 13.9 MB file. It matches the quality of bge-micro-v2, a float model with the same parameter count, at less than half the size.

• 17.2M parameters, 3 transformer layers, 384-dimensional vectors
• Ternary (2-bit) weights in every encoder layer, int8 embeddings, int8 activations
• English, up to 128 tokens per input
• Runs on CPU or GPU with plain PyTorch, and natively on the AMD Ryzen AI (XDNA2) NPU

Quick start

import sys
from huggingface_hub import snapshot_download

path = snapshot_download("Cyronius/trit-embed-3l")
sys.path.insert(0, path)
from trit_embed import TritEmbed  # needs torch, transformers and safetensors

model = TritEmbed.from_pretrained(path)
vectors = model.encode(["how do index funds work", "Index funds track a market index."])
print(vectors @ vectors.T)  # cosine similarities

The vectors are unit length, so a dot product is the cosine similarity. Queries and passages need no prefix. To use a GPU, call model.to("cuda") first.

NPU without a toolchain

This repository also carries the NPU encoder already built, so on an XDNA2 machine it runs on the NPU straight from the download. There is nothing to compile, and you need no mlir-aie, Peano or phlegm checkout.

import sys
from huggingface_hub import snapshot_download

path = snapshot_download("Cyronius/trit-embed-3l")
sys.path.insert(0, path)
from trit_npu import TritNpu, NpuUnavailable

try:
    npu = TritNpu(path)
except NpuUnavailable as e:  # names the missing prerequisite
    sys.exit(e)
vectors = npu.encode(["how do index funds work", "Index funds track a market index."])
npu.close()

If you already run FastFlowLM or OpenFlowLM on Linux, your machine most likely has everything this needs: the amdxdna driver, XRT with its NPU plugin, and a raised locked-memory limit. You need only the Python packages above. The kernels are ordinary mlir-aie output (an xclbin with the MLIR_AIE kernel, plus its instruction stream) loaded through plain XRT, so nothing here depends on phlegm or replaces your runtime. They need the whole NPU array, though, so stop your FLM or OFLM server, or unload its model, before encoding. To call the NPU from your own code instead of Python, libtd_npu2.so exports a small C API (tdn_open, tdn_host_init, tdn_encode, tdn_close and a few lower-level calls), and trit_npu.py shows how to call it. Windows is not supported yet.

To check a machine first, run python <path>/trit_npu.py <path>. It prints the device, XRT and locked-memory limit it found, then encodes a few texts on the NPU and compares them with the CPU path.

Requirements

• Linux on x86-64.
• An XDNA2 NPU (PCI ID 1022:17f0): Ryzen AI 300 (Strix Point), Ryzen AI Max (Strix Halo) or Krackan Point. The design uses the whole 32-core array, so no other program may be holding NPU columns at the same time. The XDNA1 NPUs of Phoenix and Hawk Point (1022:1502) cannot run it.
• The amdxdna kernel driver. It is in mainline Linux from 6.14; for older kernels, use AMD's xdna-driver. Your user also needs read and write access to /dev/accel/accel0.
• XRT with its amdxdna plugin: libxrt_coreutil.so.2 and libxrt_driver_xdna.so.2. The library was built and tested against XRT 2.27.0 and needs 22 XRT C++ symbols, which npu/build.json lists. An older XRT that lacks one fails at load with the missing symbol named. On Fedora, install xrt-npu and xrt-plugin-amdxdna. Elsewhere, use AMD's packages under /opt/xilinx/xrt, or set XILINX_XRT.
• glibc 2.32 and libstdc++ from GCC 12 (GLIBCXX_3.4.30) or newer. Ubuntu 22.04, Debian 12 and Fedora 36 and later all have these.
• A locked-memory limit of at least 64 MiB, because XRT pins a 64 MiB heap when it opens the device. Desktop sessions often get 8 MiB. trit_npu.py raises its own soft limit when the hard limit allows; otherwise it stops and says how to raise the limit (ulimit -l, /etc/security/limits.d, or systemd-run --user -p LimitMEMLOCK=infinity).
• The Python packages of the CPU path: torch, transformers and safetensors.

If something is missing, TritNpu raises NpuUnavailable before it touches the device, and the message names the problem: no Linux, no amdxdna, a non-XDNA2 NPU, no device access, no XRT, an old libstdc++ or XRT, too little locked memory, or another program holding the NPU. The CPU/GPU path (TritEmbed) does not change.

Files

• final.xclbin holds the NPU program for all 32 cores. insts.bin is the runtime sequence that streams the weights and texts.
• libtd_npu2.so is the host library. It packs texts into dispatches and does the embedding lookup, mean pooling, output projection and normalization. It keeps several dispatches in flight through XRT.
• trit_npu.py turns model.safetensors into the NPU's weight stream at load time. It checks the prerequisites and the three hashes in npu/build.json before it loads anything.

How they were built. The files come from phlegm at commit c236878, on Fedora 44.

• The graph: mlir-aie 1.4.2 and Peano (llvm-aie) 21.0.0.2026080301+c9c5ecb7, packaged with xclbinutil from XRT 2.27.0. Each core's program takes 15,968 of its 16,384 bytes.
• The host library: GCC 16.2.1 against the headers of XRT 2.27.0 (a29919c), with -O3 -march=x86-64-v3 -ffp-contract=off -fno-math-errno and no rpath. It needs AVX2, which every Ryzen AI CPU has.

npu/build.json records the commands, source hashes, versions and the library's imported symbol versions. Two builds of the same sources differ only in the xclbin's timestamp and UUIDs: the core programs and insts.bin come out byte for byte the same.

To rebuild, for example against another XRT, run these from a phlegm checkout with the toolchain in target/npu-dev:

source target/npu-dev/env.sh
python open_kernels/designs/trit_embed/build.py                     # final.xclbin, insts.bin
python open_kernels/designs/trit_embed/bundle.py <model folder> --xrt /usr --cxx /usr/bin/g++

bundle.py fills <model folder>/npu/ and copies trit_npu.py next to trit_embed.py. It refuses to build a host library that has an rpath or needs newer glibc or libstdc++ symbols than those above.

Checked from a clean download layout, 2026-11-04. We laid these files out as snapshot_download does, with blobs and symlinks, and loaded them from a fresh Python 3.14 environment containing only torch 2.15 (CPU), transformers and safetensors. XRT 2.27.0 came from the Fedora packages, and no phlegm, mlir-aie or development-XRT paths were set. On the Strix Halo machine above (Linux 7.2.13, NPU firmware 1.1.2.65), the NPU vectors scored retrieval 33.91 and STSB 82.88. The CPU path scores 33.85 and 82.87 (paired 95% interval for the retrieval difference: −0.07 to +0.21). Each NPU vector's cosine to the CPU vector was at least 0.998. This is the same result as the phlegm build that produced the speed figures above. In alternating runs on the same busy machine, the prebuilt files ran as fast as a phlegm development build. The machine was not quiet, so these runs only check that nothing got slower; the table above gives the quiet-machine figures.

Windows is not supported yet: the host library is Linux-only.

Speed

These timings are from text to vectors, tokenization included, on one otherwise idle AMD Ryzen AI Max+ 395 (Strix Halo) machine. The three workloads:

• Sentences: the 2,758 STSB test sentences.
• Passages: the 5,183 SciFact abstracts, nearly all at the 128-token cap.
• Latency: one sentence at a time.

On the NPU, trit-embed-3l encodes sentences 1.4 times as fast as bge-micro-v2 does in fp16 on the GPU, matches it on passages, and takes half the time for a single sentence. It does this in under 15 MB, and it leaves the GPU and most of the CPU free for other work.

PyTorch has no ternary math. On CPU and GPU, trit_embed.py unpacks the weights to fp32 when it loads, which is why it needs 39.1 MB there, and imitates the int8 rounding. That makes it slower than a float model of the same shape. The NPU runs the 2-bit weights directly: all three layers in one pass per 256 tokens, with short texts packed together into a pass. The NPU code is in phlegm, built copies ship in this repository (NPU without a toolchain), and its vectors score the same as the CPU's (retrieval 33.91, STSB 82.88).

CPU runs used 16 threads. GPU runs used the Radeon 8060S with ROCm. bge-micro-v2's GPU figures come from runs made after the GPU was already busy; its first run after the GPU has idled is about 14% slower on passages. On the NPU, the CPU tokenizes, looks up the embeddings, and does the pooling and output projection, on three threads. Each figure is a median: of three timed runs for single measurements, and over repeated runs for the NPU (nine) and bge-micro-v2 on the GPU (five).

Quality

Higher is better for both scores:

• Retrieval: nDCG@10 averaged over SciFact, NFCorpus, FiQA and SCIDOCS.
• STSB: Spearman correlation on the STSBenchmark test set, ×100.

Every model ran through the same evaluator with a 128-token cap.

"Weights" is the size of each model's weights at its published default precision. all-MiniLM-L6-v2 retrieves better, but it is a 6-layer model with 1.3× the parameters and 6.5× the size.

Scores per task, and an MTEB subset

nDCG@10 per retrieval task, with STSB:

trit-embed-3l on 12 MTEB tasks (main score ×100, 128-token cap):

What's in the file

• Encoder layers (attention and feed-forward projections): 2-bit codes for −1, 0 and +1, plus one bf16 scale per output row.
• Word and position embeddings: int8, with one fp32 scale per row.
• Everything else: fp32. This covers the LayerNorms, biases, the token-type embedding and the 384×384 output projection.

Activations are rounded to int8 at the input of each ternary layer. The architecture is a 3-layer BERT encoder (12 heads, feed-forward width 1536). Its token vectors are mean-pooled, projected and normalized.

How it was trained

All training was quantization-aware, BitNet b1.58 style. The model keeps full-precision master weights during training, and each forward pass runs through the ternary rounding.

1. Starting point: an earlier 3-layer ternary embedder of the same shape, descended from bge-micro.
2. Ranking distillation, 768 steps of 512 text pairs. The model learned to rank candidates the way Qwen3-Embedding-4B does. This ran alongside standard in-batch contrastive training and a smaller replay of human similarity ratings.
3. Science pass: 512 more steps with the same recipe. It added citation triplets from SciRepEval and claim-to-abstract pairs from PubMedQA. SciRepEval overlaps the SCIDOCS benchmark, so every example that touched a benchmark paper was removed first.

The training pairs came from public retrieval data:

• Snowflake's arctic-embed-ft-v1
• MS MARCO
• Natural Questions
• SQuAD
• HotpotQA
• StackExchange
• GooAQ
• S2ORC
• a finance set

The human similarity ratings came from SemRel2024 and the STSB training split. The final checkpoint was chosen on held-out development data.

Limitations

• English only. Inputs longer than 128 tokens are cut off.
• It needs trit_embed.py from this repository. It is not a transformers or sentence-transformers model.
• The benchmark tasks above were also used to compare training variants during development, so they may flatter it slightly.

License

MIT, the same as the parent model, bge-micro-v2. The teacher, Qwen3-Embedding-4B, is Apache 2.0. The training datasets each have their own terms.