feat: add optional cuVS training and construction for IVF-SQ

[TheR1sing3un]

Closes #100

This adds optional cuVS/CuPy acceleration for IVF-SQ training and construction, with CPU-readable v1 output and no mandatory CUDA dependency.

Performance at a glance

2.63× faster fixed-center index builds and 6.80× faster full-vector addition on real GIST1M, with unchanged Recall@10.

10M GIST-derived synthetic workload: 4.63× faster fixed-center builds (78.4% less time) and 9.25× faster full-vector addition.

One NVIDIA L20Z vs. 8 CPU threads; 1M vectors × 960 dimensions, L2, batch 8,192. Three-run medians on this PR's source:

Work measured (partitions)	CPU	GPU	Speedup	Less time
Index build with fixed centers (4,096)	26.453 s	10.055 s	2.63×	62.0%
Full-vector assignment + SQ8 encoding + append (4,096)	19.289 s	2.837 s	6.80×	85.3%
Training: sampling + fitting + SQ calibration (1,024)	4.385 s	1.996 s	2.20×	54.5%
Center fitting only, vs. CPU Lloyd (1,024)	7.699 s	0.219 s	35.1×	97.2%

The fixed-center build includes sampling, CPU SQ calibration, full-vector addition, and buffered serialization; it excludes center fitting. GPU builder setup is included. Training rows come from a separate experiment with a reused GPU worker; its first setup costs 0.639 s. Total training compares CPU Auto with cuVS, which use different clustering strategies; the fitting-only comparison uses CPU Lloyd and identical initial centers.

• Quality preserved in the fixed-center build comparison: Recall@10 remains 96.86% at 4,096 partitions (nprobe=256), and 95.75% at 1,024 partitions (nprobe=64). With the same GPU assignments, CPU and GPU SQ8 encoding produce byte-identical index files.
• GPU encoding has an independent benefit: after moving assignment to GPU, moving SQ8 encoding to GPU cuts add time by another 21.1% at 4,096 partitions (3.597 → 2.837 s), and 26.4% at 1,024 partitions. The ablations below measure each step separately.

Scaling to 10M vectors (GIST-derived synthetic)

At 10M vectors, fixed-center build time drops from 217.444 s to 47.002 s, saving 170.442 s per build.

One L20Z vs. 8 Rayon threads; 960 dimensions, L2, 4,096 fixed centers, exact assignment, batch 8,192. Three fresh-process repeats per case with rotated CPU/GPU order. Each CPU/GPU pair uses identical input and the same SQ encoding model. Times below are medians in seconds.

Vectors	CPU build	GPU build	Build speedup	CPU add	GPU add	Add speedup
1M	26.312	9.975	2.64×	18.918	2.839	6.66×
5M	111.720	26.127	4.28×	96.250	10.964	8.78×
10M	217.444	47.002	4.63×	190.970	20.655	9.25×

At 10M, build time is 78.4% lower and add time is 89.2% lower. GPU add throughput reaches approximately 484K vectors/s.

• Data: prefixes of one contiguous float32 .npy matrix. The first 1M rows are unchanged GIST1M; subsequent rows are independently perturbed copies (Gaussian noise stddev 0.0001, seed 20260911). This measures scale/capacity on a GIST-derived workload. The 1M row is this experiment's own matched baseline, measured separately from the .fvecs experiment above. The GPU adapter source matches this PR.
• Timing: build includes ingestion/sampling, CPU SQ calibration, add, and buffered serialization; center fitting is excluded. Add includes GPU builder setup/teardown, validation, transfers, and native append. Data generation, conversion, diagnostics, and read checks are outside these timers; writes do not call fsync.
• Validation: all 18 CPU/GPU builds in this scale table passed native-reader checks (16 queries per index, valid IDs and finite distances). Recall was not measured for the expanded dataset; the real GIST1M recall results above apply to their own experiment.
• Memory: sampled GPU add peak stays around 1.15 GiB across 1M/5M/10M at this fixed batch size. The 10M GPU process reaches approximately 49.22 GiB host RSS, including mapped input pages; the writer still holds the encoded index in host memory.

Changes

• Add owned prepared-training and encoding-model interfaces in Rust/C/Python, with validation before externally assigned or encoded batches are appended.
• Add reusable CuvsKMeans and CuvsIvfSqWriter context managers. The builder supports GPU assignment with either CPU or GPU SQ8 encoding and borrows the caller's native writer.
• Match native SQ8 arithmetic, including SIMD/scalar-tail rounding and subnormals. Explicit PTX is needed because CuPy's flush-to-zero default breaks the subnormal parity tests.
• Cover invalid inputs, handle ownership, device selection, and CUDA cleanup failures; preserve the original operation exception.
• Add usage documentation and a reproducible stage benchmark. Record CPU SQ calibration as a follow-up TODO at train_sq_from_processed.

GPU construction requires a model calibrated with exact assignment (ivf.coarse-assignment=exact). Default CPU paths, native dependencies, serialization, and readers remain unchanged. GPU execution is exposed through Python; native extension points are backend-independent.

Detailed measurements

Construction ablations: CPU vs. GPU assignment vs. GPU assignment + encoding (18 builds)

Measured on the submitted source: real GIST1M (1M × 960, L2), one L20Z, 8 Rayon threads, batch 8,192, cuVS 26.08.01 / CuPy 14.2.0. These are workload-specific observations against an 8-thread CPU baseline, not the host's full-core maximum.

For construction, all modes use identical sampled rows and fixed centers per nlist. Each case runs in a fresh process; three repetitions rotate mode order (18 builds). Values are medians, with min–max add times in parentheses.

nlist	Builder	Full-row add, s (range)	Fixed-center total, s	Recall@10
1,024	CPU	6.273 (6.100–6.594)	8.781	0.9575
1,024	GPU assignment / CPU encoding	3.605 (3.314–3.756)	6.076	0.9575
1,024	GPU assignment / GPU encoding	2.653 (2.612–2.713)	5.195	0.9575
4,096	CPU	19.289 (19.247–19.416)	26.453	0.9686
4,096	GPU assignment / CPU encoding	3.597 (3.554–4.589)	10.753	0.9686
4,096	GPU assignment / GPU encoding	2.837 (2.790–2.856)	10.055	0.9686

At nlist=4,096, GPU assignment alone reduces add time from 19.289 to 3.597 s. GPU encoding further reduces it by 21.1%, to 2.837 s. Combined add speedup is 6.80×; fixed-center total speedup is 2.63×. At nlist=1,024, GPU encoding independently reduces add time by 26.4%.

The add timer includes builder initialization/close, validation, transfers, and native append. Fixed-center total additionally includes sampling, CPU SQ calibration, and buffered serialization; centroid fitting is excluded. Source opening, diagnostics, and query evaluation are outside these timers; writes do not call fsync. Recall uses 1,000 held-out queries, nprobe=64/256 respectively. With the same GPU assignments, CPU and GPU encoding produce byte-identical whole-index files in every repeat; independent CPU/GPU assignments may differ near ties.

Training ablations: CPU Auto, CPU Lloyd, and cuVS (9 runs)

The separate training ablation uses nlist=1,024, the same 65,536-row sample, 25-iteration limit, seed 1234, and three rotated repetitions:

Backend	Center fitting, s	Preparation + fitting + SQ calibration, s
CPU Auto	2.620	4.385
CPU Lloyd	7.699	9.480
cuVS	0.219	1.996

CPU Lloyd and cuVS start from identical centers, but numerical convergence can differ. Center fitting is 35.1× faster than CPU Lloyd. CPU Auto uses hierarchical clustering; the 2.20× total-training comparison against Auto is a comparison of available strategies. GPU-fitted centers can change partition balance and CPU query latency; deployment should validate recall and query cost. Training reuses the GPU worker and reports its first initialization (0.639 s) separately. Fitting timings include required sample export and transfers.

Reproduction

Run tools/benchmark_gpu_training.py --preset baseline --base "$BASE" --backends cpu-auto cpu-lloyd cuvs --coarse-assignment exact --threads 8 --batch-rows 8192 --repeats 3 --output-dir "$NEW_DIR" for training.

For construction, use --backends external --centroids "$CENTERS" --build-index --build-backend MODE --queries "$QUERIES" --neighbors "$TRUTH" with the same remaining settings. Use baseline/medium for nlist 1,024/4,096 and nprobe 64/256. Rotate cpu, cuvs-assign, and cuvs-encode across fresh processes with --repeats 1, collecting three runs of each. Every output directory must be new. The tool records stage timings, sample/center hashes, native-library hash, versions, and recall.

Validation and limits

• Rust workspace: 532 passed, 2 ignored. C/C++ smoke: 9/7 passed.
• Python: 52 passed on CPU (18 GPU cases skipped); all 70 passed on the GPU host. GPU cases include L2/cosine/inner-product retrieval, subnormal encoding, and lifecycle faults.
• Formatting, Clippy, Python 3.9 syntax, release-helper tests, ASF headers, binary legal inventory, and reproducible source archive checks pass.
• The GPU adapter remains experimental, validated on one GPU/package combination. Ordinary CI has no GPU runner. Device scratch is bounded by batch size, but native writer memory grows with the full index. CPU SQ calibration, querying, multi-GPU construction, automatic batch sizing, and out-of-core writing are outside this change.

[TheR1sing3un]

@JingsongLi Hi, this is the first pr about cuda acceleration. Could you review it for me?

[JunRuiLee]

LGTM. The cuVS integration is optional and well isolated, preserves CPU compatibility and the existing index format, and shows strong real-GPU performance improvements without recall regression.