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
fmopaouter 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
vdpbf16psandvpmaddwdconsume; 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
@expon 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.