diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 23b4d0a..6b9ea83 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -32,6 +32,42 @@ Same corpus, same hardware. `snapvec` at `nprobe=64` + rerank. Disk footprint is 2-8x smaller across the range. +## Batched search threading curve + +`IVFPQSnapIndex.search_batch` fans out per-query scoring across worker +threads. Threading is a throughput knob, not a per-call latency knob: +the single `search()` API is serial on purpose (it competes with +NumPy's internal BLAS pool; see the docstring). + +Same FIQA corpus as above (N = 57,638, dim = 384), batch_size = 128, +measured on an Apple M4 Pro (12 cores), NumPy 2.4.3, Python 3.12. + +| `nprobe` | t=1 ms/q | t=2 ms/q | t=4 ms/q | t=8 ms/q | best speedup | +|---------:|---------:|---------:|---------:|---------:|:------------:| +| 4 | 0.09 | 0.06 | **0.05** | 0.06 | 1.69x | +| 8 | 0.15 | 0.09 | **0.07** | 0.09 | 2.31x | +| 16 | 0.29 | 0.16 | **0.10** | 0.13 | 2.92x | +| 32 | 0.50 | 0.28 | **0.16** | 0.20 | 3.06x | +| 64 | 0.95 | 0.51 | **0.29** | 0.33 | 3.31x | +| 128 | 1.84 | 0.98 | **0.54** | 0.57 | 3.43x | +| 256 | 3.70 | 1.91 | **1.01** | 1.09 | 3.67x | + +Observations: + +- **`num_threads=4` is the sweet spot** on this machine across every + nprobe. At `num_threads=8` the curve regresses; the executor + over-subscribes the efficiency cores and starts fighting BLAS. +- **Scaling improves with `nprobe`** because per-query work grows: + at `nprobe=4`, threading overhead caps speedup at 1.7x; at + `nprobe=256` it reaches 3.7x. +- **Sub-millisecond at 4 threads** for `nprobe <= 64`, which spans + the 0.85 to 0.977 recall range from the headline table above. + That is 3,400 - 20,000 queries per second per process. + +Reproduce with `python experiments/bench_ivf_pq_threading.py` after +caching both the FIQA corpus (`experiments/.cache_fiqa_bge_small.npy`) +and the FIQA queries (`experiments/.cache_fiqa_queries_bge_small.npy`). + ## Compression ratios For BGE-small (dim=384, float32 baseline = 1536 B/vec): @@ -50,8 +86,9 @@ For BGE-small (dim=384, float32 baseline = 1536 B/vec): ```bash pip install -e ".[dev]" -python experiments/bench_v090_fiqa.py # FIQA recall / latency -python experiments/bench_sqlite_vec_baseline.py # sqlite-vec comparison +python experiments/bench_v090_fiqa.py # FIQA recall / latency +python experiments/bench_ivf_pq_threading.py # search_batch threading curve +python experiments/bench_sqlite_vec_baseline.py # sqlite-vec comparison ``` The `experiments/` folder is WIP; expect rough edges. A first-class diff --git a/experiments/bench_ivf_pq_threading.py b/experiments/bench_ivf_pq_threading.py new file mode 100644 index 0000000..f6867b1 --- /dev/null +++ b/experiments/bench_ivf_pq_threading.py @@ -0,0 +1,113 @@ +"""A/B for ``num_threads`` on IVFPQSnapIndex.search_batch. + +Reuses the FIQA recall harness; reports each nprobe at num_threads +in {1, 2, 4, 8} so the speedup curve and the small-nprobe overhead +crossover can be read off a single table. + +Threading lives on ``search_batch`` (not ``search``) because batch-level +fan-out amortises Python overhead over a whole query's worth of work. +The single-query ``search()`` API is serial by design; that is where +latency per call is minimised against NumPy's BLAS pool. +""" +from __future__ import annotations + +from pathlib import Path +from time import perf_counter + +import numpy as np +from numpy.typing import NDArray + +from snapvec import IVFPQSnapIndex + + +CORPUS_CANDIDATES = [ + Path("experiments/.cache_fiqa_corpus_bge_small.npy"), + Path("experiments/.cache_fiqa_bge_small.npy"), +] +QUERIES_PATH = Path("experiments/.cache_fiqa_queries_bge_small.npy") + +NLIST = 512 +M = 192 +K = 256 +KK = 10 +NPROBES = [4, 8, 16, 32, 64, 128, 256] +THREAD_COUNTS = [1, 2, 4, 8] +BATCH_SIZE = 128 +SEED = 0 + + +def find_corpus() -> Path: + for p in CORPUS_CANDIDATES: + if p.exists(): + return p + raise SystemExit("missing FIQA corpus cache (run _colab_embed_corpus.py).") + + +def time_batch( + idx: IVFPQSnapIndex, + queries: NDArray[np.float32], + nprobe: int, + num_threads: int, + batch_size: int = BATCH_SIZE, +) -> tuple[float, float]: + """Return (ms_per_query, total_elapsed_s) averaged over all batches.""" + # Warm-up: first batch warms LUT + gather caches. The thread-pool + # executor itself is only created when num_threads > 1. + _ = idx.search_batch( + queries[:batch_size], k=KK, nprobe=nprobe, num_threads=num_threads, + ) + + n = len(queries) + t0 = perf_counter() + for start in range(0, n, batch_size): + chunk = queries[start : start + batch_size] + idx.search_batch(chunk, k=KK, nprobe=nprobe, num_threads=num_threads) + elapsed = perf_counter() - t0 + ms_per_q = elapsed / n * 1e3 + return ms_per_q, elapsed + + +def main() -> None: + corpus = np.load(find_corpus()).astype(np.float32) + queries = np.load(QUERIES_PATH).astype(np.float32) + corpus /= np.linalg.norm(corpus, axis=1, keepdims=True) + 1e-12 + queries /= np.linalg.norm(queries, axis=1, keepdims=True) + 1e-12 + queries = queries[:512] # enough batches to amortise scheduling + + print(f"corpus : {corpus.shape}") + print(f"queries : {queries.shape} (batch_size={BATCH_SIZE})") + + idx = IVFPQSnapIndex( + dim=corpus.shape[1], nlist=NLIST, M=M, K=K, + normalized=True, seed=SEED, + ) + print(f"\nfit({len(corpus)}) + add_batch...") + t0 = perf_counter() + idx.fit(corpus, kmeans_iters=15) + idx.add_batch(list(range(len(corpus))), corpus) + print(f" done in {perf_counter() - t0:.1f}s") + + header = "nprobe" + "".join(f" t={t} (ms/q)" for t in THREAD_COUNTS) + " best-speedup" + print() + print(header) + print("-" * len(header)) + for nprobe in NPROBES: + row = [] + for t in THREAD_COUNTS: + # The executor is lazily created and locked to the first + # num_threads value. Call close() between t settings to + # release the executor so it can be recreated for the next t. + idx.close() + ms, _ = time_batch(idx, queries, nprobe, t) + row.append(ms) + base = row[0] + best = min(row) + speedup = base / best if best > 0 else 0.0 + cells = "".join(f" {ms:5.2f}" for ms in row) + print(f"{nprobe:>5}{cells} {speedup:4.2f}x") + + idx.close() + + +if __name__ == "__main__": + main()