bench: publish search_batch threading curve (FIQA, M4 Pro)

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

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()