Add CUDA-accelerated indexing and search

Cargo.toml

@@ -18,6 +18,12 @@ serde = { version = "1", features = ["derive"] }
 rust-stemmers = "1.2"
 unicode-normalization = "0.1"
 rayon = "1.10"
+rustc-hash = "2"
+cudarc = { version = "0.19", default-features = false, features = ["cuda-version-from-build-system", "driver", "dynamic-loading", "nvrtc"], optional = true }
+
+[features]
+default = []
+cuda = ["dep:cudarc"]
 
 [dev-dependencies]
 tempfile = "3"

Makefile

@@ -27,5 +27,11 @@ release: bump-version
 test:
 	cargo test --lib
 
+test-cuda:
+	cargo test --lib --features cuda
+
 build:
 	cd python && maturin develop --release
+
+build-gpu:
+	cd python && maturin develop --release --features cuda

README.md

@@ -157,7 +157,7 @@ fn is_empty(&self) -> bool
   <img src="assets/benchmarks.png" alt="bm25x benchmarks" width="100%" />
 </p>
 
-> [BEIR](https://github.com/beir-cellar/beir) datasets. Same or better NDCG@10 than bm25s, **3.5-6x faster** indexing, **1-4x faster** search.
+> [BEIR](https://github.com/beir-cellar/beir) datasets (log scale). Same or better NDCG@10 than bm25s. CPU: **3.5-6x faster** indexing, **1-4x faster** search. GPU (H100): up to **13x faster** indexing, **214x faster** search on MS MARCO 8.8M docs. GPU search has per-query kernel launch overhead, making it best suited for **batch querying** on **large datasets**.
 
 ---
 

benchmarks/bench_cuda.py

@@ -0,0 +1,122 @@
+"""
+Benchmark bm25x CUDA-accelerated indexing vs CPU (rayon) indexing.
+
+Uses BEIR SciFact dataset for realistic evaluation.
+"""
+
+import time
+
+import bm25x
+from beir import util
+from beir.datasets.data_loader import GenericDataLoader
+
+# ───────────────────────── load dataset ──────────────────────────
+
+print("Downloading SciFact dataset...")
+url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip"
+data_path = util.download_and_unzip(url, "benchmarks/data")
+
+corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="test")
+
+corpus_ids = list(corpus.keys())
+corpus_texts = [
+    (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "")
+    for did in corpus_ids
+]
+test_qids = sorted(qrels.keys())
+test_queries = [queries[qid] for qid in test_qids if qid in queries]
+
+gpu_available = bm25x.is_gpu_available()
+print(f"GPU available: {gpu_available}")
+print(f"Corpus: {len(corpus_texts)} docs | Queries: {len(test_queries)}")
+print("=" * 65)
+
+# ───────────────────────── CPU benchmark ─────────────────────────
+
+print("\n--- CPU indexing (rayon) ---")
+
+# Warmup
+warmup = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+warmup.add(corpus_texts[:100])
+del warmup
+
+cpu_times = []
+for trial in range(5):
+    idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+    t0 = time.perf_counter()
+    idx.add(corpus_texts)
+    t = time.perf_counter() - t0
+    cpu_times.append(t)
+    if trial < 4:
+        del idx
+
+cpu_index = idx  # keep last one for correctness check
+best_cpu = min(cpu_times)
+dps_cpu = len(corpus_texts) / best_cpu
+print(f"  Times: {[f'{t:.4f}s' for t in cpu_times]}")
+print(f"  Best:  {best_cpu:.4f}s  ({dps_cpu:,.0f} d/s)")
+
+# Search sanity check
+cpu_results = [cpu_index.search(q, 10) for q in test_queries[:10]]
+print(
+    f"  Search check: first 10 queries OK ({sum(len(r) for r in cpu_results)} results)"
+)
+
+# ───────────────────────── GPU benchmark ─────────────────────────
+
+if gpu_available:
+    print("\n--- GPU indexing (CUDA) ---")
+
+    # Warmup GPU (first call compiles NVRTC kernels)
+    warmup = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+    warmup.add(corpus_texts[:100])
+    del warmup
+
+    gpu_times = []
+    for trial in range(5):
+        idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+        t0 = time.perf_counter()
+        idx.add(corpus_texts)
+        t = time.perf_counter() - t0
+        gpu_times.append(t)
+        if trial < 4:
+            del idx
+
+    gpu_index = idx
+    best_gpu = min(gpu_times)
+    dps_gpu = len(corpus_texts) / best_gpu
+    print(f"  Times: {[f'{t:.4f}s' for t in gpu_times]}")
+    print(f"  Best:  {best_gpu:.4f}s  ({dps_gpu:,.0f} d/s)")
+
+    # Correctness: compare search results CPU vs GPU
+    print("\n--- Correctness verification ---")
+    mismatches = 0
+    for q in test_queries[:50]:
+        cpu_r = cpu_index.search(q, 10)
+        gpu_r = gpu_index.search(q, 10)
+        cpu_ids = [r[0] for r in cpu_r]
+        gpu_ids = [r[0] for r in gpu_r]
+        if cpu_ids != gpu_ids:
+            mismatches += 1
+    print(f"  Tested 50 queries: {mismatches} mismatches")
+    if mismatches == 0:
+        print("  PASS: GPU and CPU produce identical results")
+    else:
+        print("  WARN: Some results differ (may be due to tie-breaking)")
+
+    # Summary
+    print("\n" + "=" * 65)
+    speedup = best_cpu / best_gpu
+    print(f"CPU: {best_cpu:.4f}s  ({dps_cpu:,.0f} d/s)")
+    print(f"GPU: {best_gpu:.4f}s  ({dps_gpu:,.0f} d/s)")
+    print(f"Speedup: {speedup:.2f}x")
+else:
+    print(
+        "\nNOTE: GPU not available. Build with CUDA support to compare:\n"
+        "      make build-gpu\n"
+        "      # or: cd python && maturin develop --release --features cuda"
+    )
+    print("\n" + "=" * 65)
+    print(f"CPU only: {best_cpu:.4f}s  ({dps_cpu:,.0f} d/s)")
+
+print("Done.")

benchmarks/bench_msmarco_cuda.py

@@ -0,0 +1,134 @@
+"""
+Benchmark bm25x indexing on MS MARCO passage corpus (~8.8M passages).
+Compares CPU (rayon) vs GPU (CUDA) indexing throughput.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=3 python benchmarks/bench_msmarco_cuda.py
+"""
+
+import gc
+import os
+import time
+
+import bm25x
+
+# ───────────────────────── load MS MARCO ─────────────────────────
+
+
+def load_msmarco(max_docs=None):
+    """Load MS MARCO passage corpus via BEIR."""
+    from beir import util
+    from beir.datasets.data_loader import GenericDataLoader
+
+    dataset = "msmarco"
+    data_dir = "benchmarks/data"
+    data_path = os.path.join(data_dir, dataset)
+
+    if not os.path.exists(data_path):
+        print("Downloading MS MARCO dataset (this may take a while)...")
+        url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset}.zip"
+        data_path = util.download_and_unzip(url, data_dir)
+
+    print("Loading corpus...")
+    corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="dev")
+
+    corpus_ids = list(corpus.keys())
+    if max_docs:
+        corpus_ids = corpus_ids[:max_docs]
+
+    corpus_texts = [
+        (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "")
+        for did in corpus_ids
+    ]
+
+    test_qids = sorted(qrels.keys())[:200]
+    test_queries = [queries[qid] for qid in test_qids if qid in queries]
+
+    return corpus_texts, test_queries
+
+
+def bench_indexing(corpus_texts, label, trials=3):
+    """Benchmark indexing and return (best_time, index)."""
+    # Warmup
+    w = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+    w.add(corpus_texts[:500])
+    del w
+    gc.collect()
+
+    times = []
+    for trial in range(trials):
+        idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True)
+        t0 = time.perf_counter()
+        idx.add(corpus_texts)
+        t = time.perf_counter() - t0
+        times.append(t)
+        print(
+            f"  [{label}] trial {trial + 1}: {t:.3f}s  ({len(corpus_texts) / t:,.0f} d/s)"
+        )
+        if trial < trials - 1:
+            del idx
+            gc.collect()
+
+    best = min(times)
+    dps = len(corpus_texts) / best
+    print(f"  [{label}] best: {best:.3f}s  ({dps:,.0f} d/s)")
+    return best, idx
+
+
+# ───────────────────────── main ──────────────────────────────────
+
+
+def main():
+    gpu_available = bm25x.is_gpu_available()
+    print(f"GPU available: {gpu_available}")
+
+    # Load corpus
+    corpus_texts, test_queries = load_msmarco()
+    n = len(corpus_texts)
+    print(f"Corpus: {n:,} docs | Queries: {len(test_queries)}")
+    print("=" * 65)
+
+    # ── CPU baseline ──
+    print(f"\n--- CPU indexing (rayon) on {n:,} docs ---")
+    cpu_time, cpu_index = bench_indexing(corpus_texts, "CPU", trials=3)
+
+    # Quick search sanity check
+    t0 = time.perf_counter()
+    for q in test_queries:
+        cpu_index.search(q, 10)
+    t_search = time.perf_counter() - t0
+    print(
+        f"  [CPU] search {len(test_queries)} queries: {t_search:.3f}s ({len(test_queries) / t_search:,.0f} q/s)"
+    )
+
+    if not gpu_available:
+        print("\nNo GPU — done.")
+        return
+
+    # ── GPU ──
+    del cpu_index
+    gc.collect()
+
+    print(f"\n--- GPU indexing (CUDA) on {n:,} docs ---")
+    gpu_time, gpu_index = bench_indexing(corpus_texts, "GPU", trials=3)
+
+    # Search check
+    t0 = time.perf_counter()
+    for q in test_queries:
+        gpu_index.search(q, 10)
+    t_search_gpu = time.perf_counter() - t0
+    print(
+        f"  [GPU] search {len(test_queries)} queries: {t_search_gpu:.3f}s ({len(test_queries) / t_search_gpu:,.0f} q/s)"
+    )
+
+    # ── Summary ──
+    print("\n" + "=" * 65)
+    speedup = cpu_time / gpu_time
+    print(f"{'':>20} {'Time':>10} {'Throughput':>15}")
+    print(f"{'CPU (rayon)':>20} {cpu_time:>9.3f}s {n / cpu_time:>14,.0f} d/s")
+    print(f"{'GPU (CUDA)':>20} {gpu_time:>9.3f}s {n / gpu_time:>14,.0f} d/s")
+    print(f"{'Speedup':>20} {speedup:>9.2f}x")
+
+
+if __name__ == "__main__":
+    main()