Proposal: Python GPU-accelerated indexing with Collection API integration

[cluster2600]

Motivation

Issues #100 and #147 request GPU-accelerated indexing. zvec already has standalone GPU backends (FAISS GPU, cuVS CAGRA/IVF-PQ, Apple MPS) but they're not integrated with the Collection API — users can't do collection.query() with GPU acceleration.

Proposed approach

1. UnifiedGpuIndex ABC (python/zvec/backends/unified.py)

Abstract base with train() + add() + search() + size() + backend_name, plus 6 adapters:

Adapter	Wraps	Priority
CppCuvsAdapter	Native _zvec pybind11 (zero-copy)	1 (highest)
CuvsCAGRAAdapter	cuvs.neighbors.cagra	2
CuvsIvfPqAdapter	cuvs.neighbors.ivf_pq	3
FaissGpuAdapter	backends/gpu.py::GPUIndex	4
AppleMpsAdapter	backends/apple_silicon.py	5
FaissCpuAdapter	FAISS CPU fallback	6

C++ native path is preferred (avoids Python→GPU data copies).

2. GpuIndex bridge (python/zvec/gpu_index.py)

Connects GPU search to Collection:

gpu = collection.gpu_index("embedding")
gpu.build(vectors, ids)
docs = gpu.query(query_vector, topk=10, output_fields=["title"])
# Returns list[Doc] — same format as collection.query()

Flow: GPU search → map indices to doc IDs → collection.fetch() → attach scores → return list[Doc].

3. Backend detection (python/zvec/backends/detect.py)

Adds C++ cuVS and Python cuVS detection with proper priority chain.

Tested on RTX 4090

Backend	QPS (50K vectors, dim=128)
FAISS GPU (flat)	529,316
cuVS IVF-PQ	45,771
cuVS CAGRA	43,711

Questions for maintainers

1. Is collection.gpu_index(field_name) the right API surface, or would you prefer a different entry point?
2. The current design requires explicit gpu.build(vectors, ids) since _Collection has no scan/iterate API. Would you consider exposing a bulk vector extraction API from C++?
3. Any preference on how the C++ cuVS priority should interact with the existing backend detection?

Draft implementation: #176 (addresses #100 and #147)

[xXMrNidaXx]

Nice proposal. A few thoughts on the API design and implementation:

1. On collection.gpu_index(field_name) vs alternatives:

The proposed API is intuitive, but I'd consider collection.index(field_name, device="gpu") for consistency with PyTorch/JAX conventions and future flexibility (multi-GPU, specific device selection). You could then have:

# GPU with auto-detection
idx = collection.index("embedding", device="gpu")
# Specific device
idx = collection.index("embedding", device="cuda:0")
# Explicit backend
idx = collection.index("embedding", device="gpu", backend="cagra")

2. On the bulk vector extraction API:

Rather than exposing a full scan API (which has its own complexity around batching, memory, streaming), consider a callback-based builder pattern:

gpu = collection.gpu_index("embedding")
# Collection internally streams vectors in batches to the builder
gpu.build_from_collection(batch_size=10000)

This keeps the streaming/batching logic internal where you can optimize it, and avoids exposing implementation details that might change.

3. On backend priority:

The priority chain makes sense, but consider making it configurable via environment variable or config file. Production deployments often have specific backend preferences (e.g., FAISS GPU might be preferred over cuVS in some environments for stability/compatibility reasons, even when cuVS is available).

Something like:

ZVEC_GPU_BACKEND_PRIORITY=faiss_gpu,cuvs_cagra,cuvs_ivfpq

4. Consideration: hybrid CPU/GPU search

At 50K vectors, the overhead of moving data GPU→CPU for result mapping might dominate. Have you measured end-to-end latency including the collection.fetch() step? For small collections, pure CPU might actually win.

A nice addition would be a threshold-based auto-selector that picks CPU for small collections and GPU for larger ones, with the threshold configurable.

Looking forward to seeing this land — GPU-accelerated vector search in the Collection API would be a great addition.

[cluster2600]

@xXMrNidaXx Thanks for the detailed feedback — really appreciate the thoughtful suggestions!

On collection.index(field_name, device="gpu"): Great idea. The device= pattern is familiar from the PyTorch ecosystem and opens up multi-GPU targeting cleanly. I like that it also leaves room for device="cpu" as an explicit choice. Will adopt this in the next revision.

On build_from_collection(batch_size=): This is a much better API than making users extract vectors manually. Keeping the streaming/batching logic internal is the right call — it avoids leaking implementation details and lets us optimize the pipeline later (e.g. prefetching batches while the GPU builds). The callback-based builder pattern is the way to go here.

On configurable backend priority: Agreed — hardcoded priority doesn't account for the reality of production environments. The ZVEC_GPU_BACKEND_PRIORITY env var is a clean solution. Will add this alongside the backend= parameter for programmatic control.

On hybrid CPU/GPU thresholds: Spot on. I've seen this in practice — at 50K vectors the collection.fetch() round-trip can dominate latency. An auto-selector with a configurable threshold would make the default experience much better. The draft PR (#176) already has benchmark infrastructure, so measuring the crossover point should be straightforward.

I'll incorporate all four points into the next iteration of #176. Thanks again for engaging with the proposal!