Add CUDA-accelerated indexing and search

[raphaelsty]

Summary

• Feature-gated cuda support for GPU-accelerated BM25 indexing and search
• Zero quality loss — identical NDCG@10 to CPU on all BEIR datasets
• Indexing: 6x faster end-to-end (845K d/s on MS MARCO 8.8M docs)
• Search: 38x faster than CPU, 214x faster than bm25s (2,970 q/s)

Indexing optimizations

• Zero-alloc ASCII tokenizer with stack-allocated token buffer
• FxHashMap stem cache (each unique token stemmed once per thread)
• Parallel vocab building via thread-local HashSets
• Stable counting sort preserving doc_id order (no per-term sort)
• Unsafe FFI PyO3 extraction (44x faster) + GIL release for Rust processing
• Parallel drop of 44M intermediate strings

Search optimizations

• GPU-resident index (~2GB, uploaded once)
• Fused multi-term BM25 scoring kernel (single launch for all query terms)
• Correct GPU top-k: per-thread heaps in registers → block merge in shared memory
• Downloads only k results per query (~80 bytes) instead of full 35MB scores array

Benchmarks (H100, all BEIR datasets)

Dataset	Index (d/s)	Search (q/s)	NDCG@10
NFCorpus	32,551	6,940	0.3287
SciFact	57,570	5,682	0.6904
SciDocs	68,360	6,197	0.1600
FiQA	205,115	5,935	0.2514
MS MARCO	295,458	3,430	0.2240

API

Rust:

// Cargo.toml: bm25x = { features = ["cuda"] }
let mut gpu = index.to_gpu_search_index()?;
let results = index.search_gpu(&mut gpu, query, k);

Python:

# Build: maturin develop --features cuda
index.upload_to_gpu()
results = index.search(query, k)  # auto-dispatches to GPU

Test plan

• All 47 existing tests pass with --features cuda
• All 47 existing tests pass without cuda feature
• NDCG@10 identical to CPU on all 5 BEIR datasets
• GPU search correctness verified against CPU search
• Graceful CPU fallback when GPU unavailable

🤖 Generated with Claude Code