Add optional GPU acceleration for IVF-SQ training and construction

Open
#100 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

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

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.

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.

Dominant language
Rust
Stars
21
Forks
22
Avg merge
2d 1h
Merged PRs (30d)
12

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from apache/paimon-vector-index

All issues in apache/paimon-vector-index

Similar issues

More Rust issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.