xagent2025/velavec-rust-inference

VelaVec-infer: 9.8M-parameter retrieval embeddings, Rust engine for Apple Silicon

The inference engine behind VelaVec — a 9.8M-parameter hybrid retrieval encoder (frozen n-gram table + 3 layers of bidirectional attention + dual pooling heads).

This engine encodes text 22.5–53.9× faster than the 33M teacher (bge-small) on a single Apple M4 core, with a 5.4 ms cold start and 39 MB of weights. No Python runtime, no transformers, no torch — pure Rust on top of Accelerate.

Benchmark (Apple M4, 1 thread)	VelaVec (9.8M)	Teacher bge-small (33M)
Query encode (len≈15)	103 µs	~5.5 ms (53.9×)
Document encode (len≈128)	579 µs	~13 ms (22.5×)
Cold start (load weights + graph)	5.4 ms	~29 s (PyTorch init)
Weights on disk	39 MB	~130 MB
NLI r@1 / MARCO r@1	0.912 / 0.694	0.909 / 0.689

The point is not the model architecture itself (that is described in the model card) — it is how the same model runs on an Apple CPU:

• Small GEMMs (15 graphs, shapes l×256 → l×512 etc.) are compiled by CoreML into a BNNSGraph with pre-packed weights (AMX-native layout). Measured 2.5–3.1× faster than cblas_sgemm at these shapes (e.g. l=15: 159 µs → 51 µs for 15 GEMMs).
• Elementwise ops (RMSNorm, SwiGLU, RoPE, attention softmax) are hand-written 4-wide NEON kernels (fused dual-chain reductions, a degree-9 exp2 polynomial with max rel. error 5.5e-14 — numerically identical to vvexpf).
• Zero steady-state heap allocation: all intermediate buffers are pre-allocated and reused.
• Dynamic-length encoding (no padding to 64) is 2.0–2.7× faster at zero quality loss.

Requirements

• macOS 13+ on Apple Silicon (the BNNS graph path is macOS-only).
• Current Rust toolchain (stable).
• macOS ships the Accelerate framework; no crates are required (cargo pulls nothing).

On other platforms the engine still compiles and runs correctness checks in the BLAS fallback mode, but the speed claims are Apple Silicon / macOS specific. The pre-compiled gemms/*.mlmodelc graphs are CoreML artifacts; if they fail to load on your macOS version, regenerate them with export_bnns_graphs.py (see below).

Build

cargo build --release

That's it — no external crates, no network access needed beyond a normal toolchain.

Run

# correctness check (single hard-coded sentence, layer-by-layer debug print)
./target/release/hybrid_infer --debug

# full eval: NLI retrieval + MARCO retrieval + single-item latency benchmarks
./target/release/hybrid_infer weights.bin eval_ids.bin

# BLAS fallback (A/B reference. Everything else identical)
./target/release/hybrid_infer weights.bin eval_ids.bin --blas

Expected output (Apple M4, low load; latency is machine-state sensitive):

NLI  r@1=0.912 r@5=0.978 r@10=0.982（pot3v2 PyTorch 对照 0.940/0.978/0.982）
MARCO r@1=0.694 r@5=0.938 r@10=0.956
GEMM 后端: BNNSGraph（CoreML 预打包 AMX 权重）（图编译 ~2-20ms）
加载权重: ~5ms（PyTorch 冷启动 ~29s）
查询延迟 (len=15): ~100µs（PyTorch: 动态544µs / 固定843µs / 优化前 Rust 881µs）
文档延迟 (len=128): ~580-780µs（单线程，负载敏感）

Files

File	Purpose
src/main.rs	Full engine: model, NEON kernels, BNNS wiring, eval & bench harness
src/bnns.rs	BNNSGraph FFI (CoreML pre-packed linear layers, dynamic batch via BNNSGraphContextSetBatchSize)
gemms/	15 CoreML mlmodelc graphs (3 blocks × qkv / o_proj / w13 / w2 / fused-ffn), pre-packed AMX layouts
export_bnns_graphs.py	(Re)generate the graphs from weights.bin via coremltools
weights.bin	9.8M-parameter fp32 weights (same model as xagent2025/VelaVec)
eval_ids.bin, marco_eval_ids.bin	Pre-tokenized eval sets (NLI 500q/500p, MARCO dev) for rerunning the reported numbers

To regenerate the graphs (e.g. for a different macOS/CoreML version):

python3 export_bnns_graphs.py        # needs coremltools; writes gemms/*.mlmodelc

Notes

• fp16 is slower than fp32 on AMX at these shapes — Apple's AMX has no fp16 path for small GEMMs, so fp16 incurs convert+bf16x2 emulation overhead. We verified this empirically (0.74–0.90× of fp32 speed).
• int8 gives no speed-up on Apple CPU either: BNNSMatMul/filter reject Int8 outright, and via CoreML+BNNSGraph the int8 weights version equals fp32 speed while int8 weights+activations is ~2× slower. There is no public AMX int8 GEMM path.
• These findings hold at the model's GEMM shapes (l=15 and l=256, compute-bound); they should not be over-generalized to large-batch server workloads.
• Latency is machine-state sensitive (load, core scheduling). Always benchmark with --blas A/B in the same session; the engine pins interactive QoS on the bench thread.

License and links

• MIT.
• Models: xagent2025/VelaVec · xagent2025/VelaVec-T2I
• Issues/PRs: this repo mirrors the source of record (ask to open issues here or on the model repos).