Skip to content

Implement GPU IVF add_preassigned via add_core (implement #3908)#5372

Open
weiwch wants to merge 4 commits into
facebookresearch:mainfrom
weiwch:main
Open

Implement GPU IVF add_preassigned via add_core (implement #3908)#5372
weiwch wants to merge 4 commits into
facebookresearch:mainfrom
weiwch:main

Conversation

@weiwch

@weiwch weiwch commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Summary

This implements GPU IVF support for ivf_tools.add_preassigned, addressing #3908.

GPU IVF already supports search_preassigned, but ivf_tools.add_preassigned previously only worked for CPU IVF indexes because GPU IVF did not expose the corresponding add_core path. This implements the follow-up I mentioned in #3908 by refactoring my earlier faiss_preassigned work into the mainline-compatible add_core path. This PR adds GpuIndexIVF::add_core(...), wires precomputed IVF list ids through the existing GPU add paging/copy path, and makes the existing ivf_tools.add_preassigned helper work for GPU IVF indexes without any change of ivf_tools.py.

The implementation follows the CPU IVF design:

  • add GpuIndexIVF::add_core(...) with the same naming and argument order as IndexIVF::add_core(...)
  • pass precomputed_idx through the existing GPU add paging/copy path
  • add IVFBase::addVectorsPreassigned(...) for appending vectors to precomputed IVF lists

contrib.ivf_tools.add_preassigned can now call index_ivf.add_core(...) for both CPU and GPU IVF indexes. Compared with my earlier fork implementation, this does not introduce separate GPU-specific add_preassigned* APIs; the naming and entry point are aligned with the existing CPU IVF API.

Implementation note

GpuIndex::addPaged_ and addPaged_ex_ were moved from private to protected so GpuIndexIVF::add_core(...) can reuse the existing GPU add paging path.

Tests

Added a GPU IVFFlat add-preassigned test that:

  • builds CPU and GPU IVF indexes
  • assigns vectors with an external quantizer
  • adds vectors through ivf_tools.add_preassigned
  • compares CPU/GPU results using search_preassigned

@meta-cla meta-cla Bot added the CLA Signed label Jul 3, 2026
@meta-codesync

meta-codesync Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

@mnorris11 has imported this pull request. If you are a Meta employee, you can view this in D111188416.

@mnorris11 mnorris11 self-assigned this Jul 16, 2026
@mnorris11

Copy link
Copy Markdown
Contributor

Thanks for this change!
Would you be able to add 1 small test that exercises the GpuIndexIVFPQ flow? Reason: it has by_residual = true by default.
Then it looks good to me to merge.

@weiwch

weiwch commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

I hadn’t used preassigned GpuIndexIVFPQ in my own project before, so I hadn’t tested it or added a test covering that flow. I’ll add a small test soon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants