apache / apache/paimon-vector-index

Add optional GPU acceleration for IVF-SQ training and construction

Open
#100 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21
Forks
21
Avg merge
2d 10h
Merged PRs (30d)
10

Description

IVF-SQ currently fits centers, assigns vectors to partitions, and encodes residual SQ8 values on CPU. Accelerating center fitting alone leaves full-corpus assignment and encoding on the construction critical path.

I propose an optional Python cuVS/CuPy adapter, backed by backend-independent native interfaces, to accelerate both training and construction while keeping the v1 index format and existing CPU readers.

### Scope

- Prepare an owned, bounded training sample using native sampling and metric preprocessing; expose the resolved parameters.
- Accept validated external IVF centers and finish residual SQ calibration in Rust.
- Export an owned encoding-model snapshot and accept validated preassigned vectors or encoded SQ8 batches.
- Provide reusable GPU workers for center fitting, exact partition assignment, and optional SQ8 encoding. Keep serialization and querying on CPU.

GPU construction should use a model calibrated with exact assignment. CPU and GPU distance calculations may choose different partitions near ties, so validation must check retrieval quality. For identical assignments, GPU encoding should reproduce native SQ8 arithmetic, including rounding, clipping, scalar tails, constant bounds, and subnormal values.

### Acceptance criteria

- No mandatory CUDA dependency, change to default CPU behavior, or file-format change.
- Explicit ownership and batch validation at the Rust/C/Python boundaries; invalid batches are rejected before append.
- GPU errors propagate without a silent CPU fallback; cleanup preserves the original operation error.
- CPU CI covers native interfaces and lifecycle/error contracts. Explicit GPU tests cover L2, cosine, inner product, numerical parity, and held-out retrieval.
- Controlled ablations compare CPU construction, GPU assignment with CPU encoding, and GPU assignment with GPU encoding using the same sample and centers. Report repeated stage timings, recall, hardware, and timing boundaries.

CPU SQ calibration remains a follow-up optimization, recorded as a TODO at its implementation entry point. GPU querying, multi-GPU execution, automatic batch sizing, and an out-of-core writer are outside this proposal. Device scratch is batch-sized; the native writer still retains the complete encoded index in host memory.

The GPU adapter would initially be experimental, with broader GPU/CUDA/cuVS version coverage left for follow-up.

Contributor guide

Open the contributing guide

Research direction

The issue names no files or tests; first map the IVF-SQ construction path and its native Rust/C/Python boundaries, then identify the interfaces for sampling, center fitting, assignment, encoding, and serialization. Use the acceptance criteria to define ownership, validation, error cleanup, numerical parity, retrieval validation, and CPU/GPU ablation reporting as done.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
machine-learning, search
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.