sentence embeddings in pure zig: bge-small with an HF-exact tokenizer and an SME matmul
README.md

embedz #

sentence embeddings in pure zig from sentence-transformers BERT snapshots, with a tokenizer that matches Hugging Face's exactly. no python, no runtime dependencies. developed against bge-small-en-v1.5, also verified with all-MiniLM-L6-v2.

the matmuls pick a kernel for the CPU they run on:

CPU kernel numerics
Apple M4 and later SME fmopa outer products fp32, exact
x86 with AVX-512 BF16 (e.g. Zen 4) vdpbf16ps, Precision.bf16 bf16 inputs, fp32 accumulation
x86 with AVX2 (e.g. Zen 3) vpmaddwd, Precision.int8 int8 weights, 12-bit activations, exact int32 accumulation
anything else portable @Vector fp32, exact

the default is exact fp32. Precision.fastest takes bf16, then int8, then fp32, whichever the running CPU supports and passes a load-time self test (the asm kernel must match a plain-Zig emulation of the same numerics). the x86 kernels are only compiled under the LLVM backend; zig's self-hosted x86 backend (the Debug default on x86_64 Linux) cannot assemble them, so there the fp32 path runs.

install #

zig fetch --save https://tangled.sh/@zzstoatzz.io/embedz/archive/main
const embedz = b.dependency("embedz", .{ .target = target, .optimize = optimize }).module("embedz");
exe.root_module.addImport("embedz", embedz);

usage #

Model.load reads a Hugging Face snapshot directory (see models):

const embedz = @import("embedz");

var model = try embedz.Model.load(gpa, io, "/path/to/bge-small-en-v1.5", .{});
defer model.deinit(gpa);

const texts = [_][]const u8{ "the stoat changes its coat in winter", "ermine" };
const out = try gpa.alloc(f32, texts.len * model.dim());
defer gpa.free(out);
try model.embed(gpa, &texts, out); // pooled and normalized as the model specifies

embed is thread-safe on a shared Model, so you can run one call per thread.

models #

everything that differs between models of this family is read from the snapshot, not configured in code:

file gives
config.json layers, hidden size, heads, FFN width, vocab and position counts, LayerNorm eps
1_Pooling/config.json CLS or mean pooling
modules.json whether outputs are L2-normalized
sentence_bert_config.json the token cap (e.g. 256 for MiniLM)
tokenizer.json checked against what the tokenizer implements
vocab.txt, model.safetensors vocabulary and weights (a bert. tensor prefix is detected)

supported: BertModel, GELU, absolute positions, head width 32 or 64 (e.g. 384/12 or 768/12), CLS or mean pooling, and an uncased BertNormalizer + BertPreTokenizer + ## WordPiece tokenizer. anything else fails at load with a named error (UnsupportedArchitecture, UnsupportedTokenizer, UnsupportedPooling, UnsupportedHeadDim, ...) instead of producing plausible wrong vectors. multilingual-e5, for instance, is a BertModel but tokenizes with SentencePiece, so it is refused. the next boundaries are cased tokenizers (the unicode tables are generated for an uncased normalizer), SentencePiece/Unigram, and RoPE architectures (ModernBERT, Gemma).

a model some texts must be prefixed for (e5's query: / passage: ) needs the caller to add the prefix.

to try a new model, dump its reference and check it:

just reference sentence-transformers/all-MiniLM-L6-v2 data/ref-minilm
just xval 6 data/ref-minilm
model layers × width pooling token ids min cosine vs sentence-transformers
bge-small-en-v1.5 12 × 384 CLS 5000 / 5000 0.9999994
all-MiniLM-L6-v2 6 × 384 mean, cap 256 5000 / 5000 0.9999994

fidelity #

scripts/reference.py runs sentence-transformers on a corpus and dumps its token ids and embeddings. embedz-xval checks this library against that dump. For bge-small on 5,000 real Bluesky posts (multilingual, with emoji, 164k tokens), token ids match HF tokenizers for every post, and cosine against sentence-transformers (fp32) is:

precision min mean
fp32 (portable, SME, x86) 0.9999994 1.00000003
bf16 0.99996 0.999992
int8 0.9969 0.99979

the normalizer and pre-tokenizer tables (src/unicode_table.zig) are not hand-derived. scripts/gen_unicode.py generates them by asking HF's own BertNormalizer and BertPreTokenizer about every Unicode scalar value.

speed #

Apple M5 Pro, same corpus, batches of 64, ranges over repeated runs. torch 2.14 uses Accelerate, so it runs on SME too.

1 thread 6 threads
sentence-transformers (torch CPU) 78–84 posts/s 148–157 posts/s
embedz, SME (-Dcpu=native+sme+sme2) 193 posts/s 473–624 posts/s
embedz, portable @Vector 47–48 posts/s 196–199 posts/s

on macOS the benchmark threads request QOS_CLASS_USER_INTERACTIVE. at default QoS the scheduler puts them on efficiency cores, and single-thread SME runs then vary from 50 to 190 posts/s. the library leaves scheduling to its caller.

fly.io, one thread, built -Dcpu=x86_64_v3, posts of about 55 tokens (coral's traffic averages 53):

host precision cost per post
Zen 3 (Milan), performance-1x fp32 ~30 ms
Zen 3 (Milan), performance-1x int8 21.1 ms
Zen 4 (Genoa), coral's shared-cpu-2x bf16 14.0 ms

new fly machines have landed on Zen 3 in every region tried (ord, iad, dfw, sjc), so a host with AVX-512 BF16 cannot be assumed; runtime dispatch matters.

zig's CPU detection does not enable SME on Apple silicon yet (it reports apple_m1), so the SME path needs -Dcpu=native+sme+sme2. at runtime it also checks the streaming vector length.

how it works #

  • packed batches. all posts in a batch are concatenated with no padding. the dense layers see one [tokens × 384] matrix, and only attention runs per post.
  • fused qkv. query, key and value are a single 384×1152 matmul.
  • SME. each call computes a 16×64 output block. per K step, one packed column of activations and four weight panels feed four fmopa outer products into the four ZA tiles. streaming mode is entered and left inside one asm block (src/sme.zig).
  • x86 kernels. bf16 and int8 pack weights and activations as K pairs, the layout vdpbf16ps and vpmaddwd consume; each K step broadcasts one pair per row against a 16-lane weight panel. int8 activations are 12-bit because that is the widest that cannot overflow int32 over K = 1536. activation quantization is vectorized; it was 10% of int8 time when scalar.
  • attention. the head dimension is a comptime 32 and K is transposed per head, so scores accumulate as vectors along the keys. four queries share each K and V load, with scalars broadcast from memory so a tile fits AVX2's 16 registers. exp is a vectorized Cephes polynomial, because @exp on vectors lowers to scalar libcalls.
  • last layer. with CLS pooling only row 0 of the final layer is read, so after its QKV everything runs on one row per text.

on a Zen 3 box with int8, time goes to the matmuls (~79%), attention (~13%) and GELU (~6%). those last two are the next targets.

develop #

just test         # unit tests, portable and SME builds
just corpus       # sample 5000 posts from a jetstream into data/
just reference    # sentence-transformers ids + embeddings + torch timing
just xval 6       # fidelity + throughput at 6 threads
just bench        # kernel GFLOP/s

x86 timing without deploying anything: a throwaway fly machine (fly apps create embedz-bench, then fly machine run alpine:3.19 sleep infinity --app embedz-bench --region ord --vm-size performance-1x --name bench). just remote-stage [data/ref] uploads the reference dump's model snapshot, the corpus and the dump; just remote-xval int8 and just remote-bench cross-compile, copy and run in about 30 seconds. perf record -e cpu-clock works there even though the VM hides hardware counters. destroy the app when done.