Skip to content

Add native QGKMeans and RaBitQKMeans clustering - #120

Merged
gouyt13 merged 3 commits into
mainfrom
nativekmeans
Sep 29, 2026
Merged

gouyt13 merged 3 commits into
mainfrom
nativekmeans

Conversation

@gouyt13

@gouyt13 gouyt13 commented Sep 29, 2026

Copy link
Copy Markdown
Member

Summary

Add native C++ and Python k-means clustering so IVF and HNSW workflows no longer require FAISS for clustering.

  • Add RaBitQKMeans with flat RaBitQ assignment for small cluster counts, inspired by SuperKMeans, and QGKMeans with SymphonyQG graph assignment for larger counts. Both share Lloyd training, float64 centroid accumulation, spherical support, and result semantics.
  • Support approximate or exact final assignment. Exact assignment uses blocked GEMM and float64 refinement without allocating a full point-by-centroid distance matrix. Reuse graph, query, and scratch storage across iterations.
  • Add C++ file-based clustering, switch the C++ and Python examples to native clustering, and compare both methods against FAISS in compare_with_faiss.py.
  • Consolidate usage, parameters, memory costs, and assignment semantics in docs/docs/clustering.md, with regression coverage for both APIs and the example workflows.

No new production dependency or saved-index format change. FAISS remains an optional comparison dependency.

Testing

  • Core validation: 244 C++ tests and 565 Python tests passed; two platform-specific Python tests skipped. Portable build used RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF. All 72 before/after split parity cases matched bit-for-bit.
  • Final example changes: 16 C++ CLI/integration tests and four FAISS comparison tests passed, including L2/IP, exact saved labels, index-builder compatibility, repeated evaluations, and native-runtime isolation.
  • C++ examples and installed CMake consumer built and ran. Formatting, Python checks, ShellCheck, focused clang-tidy, strict MkDocs build, and Markdown link/anchor checks passed.
  • Benchmark: 36 fits across full SIFT, GIST, MSong, Contriever, MXBAI, and OpenAI datasets (769,382–1,000,000 vectors; 128–1,536 dimensions), with 4,096 clusters, seed 42, and exactly 25 actual iterations. Median of three runs per method; timings include exact final assignment and exclude loading/evaluation outside training. On an Intel Xeon Gold 6418H with 48 threads, GCC 14.3.0, native optimization OFF, and runtime AVX2/AVX-512 dispatch, QGKMeans was 1.55–2.03× faster than RaBitQKMeans with lower final WCSS on all six datasets and 19–112 MiB higher median peak RSS. WCSS was identical across repeats for each method/dataset.

Full repository clang-tidy remains to be run before merging. ARM64, macOS, and Windows were not tested locally.

- Add SymphonyQG graph assignment and flat RaBitQ assignment with shared
  Lloyd training, spherical support, and optional exact final assignment.
- Reuse graph and query scratch storage, dispatch centroid accumulation,
  and use blocked GEMM for exact assignment without a full distance matrix.
- Replace FAISS clustering in examples with native clustering methods
  and add C++ file-based clustering tools.
- Add comparisons against FAISS, regression tests, and a unified
  clustering guide with a SuperKMeans reference.
@gouyt13
gouyt13 merged commit 13ae120 into main Sep 29, 2026
36 checks passed
@gouyt13
gouyt13 deleted the nativekmeans branch September 29, 2026 06:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant