Kernel bm25

Kernel: bm25

bm25-mojo is a drop-in faster replacement for rank_bm25 — the same classes, constructor arguments, index attributes, and methods — powered by a clean-room Mojo kernel. Package: python/bm25_mojo · Kernel: kernels/bm25 · pip install bm25-mojo

The algorithm

rank_bm25 scores each query by walking every (query term × document) pair through a Python dict lookup per pair — interpreter work proportional to corpus size × query length, in pure Python. The kernel inverts the work:

1. Index build (once): documents are converted to postings-oriented flat arrays — per-term document lists with term frequencies — plus per-document lengths and the idf table, laid out for sequential access.
2. Score (per call, batch ABI): one FFI call scores a whole query against the whole corpus. Only documents that actually contain a query term are touched (postings traversal), and the inner accumulation is SIMD float64 — instead of dict-chasing every pair.
3. Reference-order arithmetic: idf, TF saturation, length normalization, and (for BM25L/BM25Plus) the delta flooring follow rank_bm25's formulas and operation order exactly, so results match to the last ulp rather than "approximately".

FFI cost is per query call, not per document — the batch shape is why the speedup survives ctypes.

Parity proof

• Differential suite (tests/) asserts element-wise agreement with rank_bm25 within 1e-8 absolute, run twice: native backend, then forced fallback (BM25_MOJO_DISABLE_NATIVE=1).
• Measured agreement on the benchmark: max abs diff 3.6e-15 at 100k docs — the scores are the same numbers, not approximations. The gate runs before every timing pass.
• All three variants — BM25Okapi, BM25L, BM25Plus — are covered, including get_scores, get_batch_scores, get_top_n, and the documented index attributes.

Measured numbers

Apple M4 Max, macOS 26.6.2 arm64, Python 3.12.14, NumPy 2.5.3, Mojo 1.1.0, rank_bm25 0.2.2; seeded corpora (Zipf-ish 20k-term vocabulary, 30–60 tokens/doc), 20 queries/cell, median of 5. Reproduce: pixi run bench.

corpus	query terms	rank_bm25 ms/query	bm25_mojo ms/query	speedup
1,000	5	0.338	0.0029	116×
10,000	10	7.570	0.0050	1,504×
100,000	20	162.178	0.0185	8,769×

vs bm25s — the incumbent, beaten

bm25s (numba-JIT over a precomputed sparse score matrix) is not a rank_bm25 drop-in, so this is a workload-identical speed comparison, not a parity one. The autopsy (benchmarks/AUTOPSY-bm25s.md) showed our v1 lost on fixed O(n_docs) costs — so we rebuilt the query path. bm25-mojo v2 wins every cell of the 9-cell grid: 1.10×–1.95× warm steady state (1k/5t 1.71×, 10k/5t 1.95×, 10k/10t 1.84×, 100k/5t 1.43×, 100k/20t 1.10×, re-verified at 1.12× over 11 interleaved samples) — same harness, same seeds, median of 5, bm25s in its recommended numba configuration. On top of that:

• Cold start: 10–1000× — bm25s pays ~4 s numba JIT plus 200–320 ms for the first query on each fresh index; bm25-mojo is AOT.
• Semantics: bit-faithful to rank_bm25 at 3.6e-15; bm25s implements its own documented variants. float32 storage was measured and rejected (94× over the 1e-8 parity gate) — float64 stays.
• Design: idf-weighted CSR baked at index time, query batching (get_scores_batch), no-rezero accumulation, deterministic single-threaded (threading is the documented reserve). Index build 0.81 s at 100k docs.

The v1-era table that showed bm25s ahead was replaced after this work; current numbers live in the repo README and benchmarks/AUTOPSY-bm25s.md.

Options matrix

Surface	Status
BM25Okapi(corpus, tokenizer=None, k1, b, epsilon)	full parity
BM25L(…, delta), BM25Plus(…, delta)	full parity
get_scores, get_batch_scores, get_top_n	full parity
Index attributes (doc_len, idf, doc_freqs, …)	same names/values as rank_bm25
BM25_MOJO_NATIVE_LIB / BM25_MOJO_DISABLE_NATIVE=1	loader override / force fallback
bm25_mojo.backend_info()	active backend, ABI version, load errors

Non-goals (by design): no numba-style sparse pre-scoring (that's bm25s's design point), no GPU, no Windows native (the vendored fallback runs there instead).

Autopsy: where the time went

At 100k docs and 20 query terms, rank_bm25 executes ~2M dict lookups and the same number of Python-level float operations per query — 162 ms of interpreter dispatch. The kernel touches only posting-list members (a small fraction of the corpus for typical queries) with SIMD accumulation, hence 0.0185 ms: the gap is interpreter overhead × corpus size, removed. This is the factory's canonical "pure-interpreter loop: 100×+ available" shape — see Writing a Kernel.

---

Benchmarks · Kernels · How It Works · Apache-2.0, © 2026 Algenta