faiss 0.6.3 → 0.6.4

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/IndexIVF.h>
+#include <faiss/SuperKMeans.h>
 
 #include <omp.h>
 #include <atomic>
@@ -78,13 +79,27 @@ void Level1Quantizer::train_q1(
             printf("Training level-1 quantizer on %zd vectors in %zdD\n", n, d);
         }
 
-        Clustering clus(static_cast<int>(d), static_cast<int>(nlist), cp);
+        FAISS_THROW_IF_MSG(
+                cp.use_super_kmeans && clustering_index,
+                "cp.use_super_kmeans is incompatible with a user-provided "
+                "clustering_index: SuperKMeans assigns with its own index");
+
         quantizer->reset();
-        if (clustering_index) {
-            clus.train(n, x, *clustering_index);
+        if (cp.use_super_kmeans) {
+            SuperKMeansParameters super_cp;
+            static_cast<ClusteringParameters&>(super_cp) = cp;
+            SuperKMeans clus(
+                    static_cast<int>(d), static_cast<int>(nlist), super_cp);
+            clus.train(n, x);
             quantizer->add(nlist, clus.centroids.data());
         } else {
-            clus.train(n, x, *quantizer);
+            Clustering clus(static_cast<int>(d), static_cast<int>(nlist), cp);
+            if (clustering_index) {
+                clus.train(n, x, *clustering_index);
+                quantizer->add(nlist, clus.centroids.data());
+            } else {
+                clus.train(n, x, *quantizer);
+            }
         }
         quantizer->is_trained = true;
     } else if (quantizer_trains_alone == 2) {

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -526,9 +526,17 @@ InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         const IVFSearchParameters*) const {
+    return get_InvertedListScanner(
+            store_pairs, sel, IndexIVFPQ::ScannerMode::Precomputed);
+}
+
+InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        IndexIVFPQ::ScannerMode mode) const {
     return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
         return pq_code_distance::make_IVFPQInvertedListScanner<SL>(
-                *this, store_pairs, sel);
+                *this, store_pairs, sel, mode);
     });
 }
 

data/vendor/faiss/faiss/IndexIVFPQ.h

@@ -32,6 +32,11 @@ FAISS_API extern size_t precomputed_table_max_bytes;
  * vector is encoded as a product quantizer code.
  */
 struct IndexIVFPQ : IndexIVF {
+    enum class ScannerMode {
+        Precomputed,
+        OnTheFly,
+    };
+
     ProductQuantizer pq; ///< produces the codes
 
     bool do_polysemous_training; ///< reorder PQ centroids after training?
@@ -144,6 +149,12 @@ struct IndexIVFPQ : IndexIVF {
             const IDSelector* sel,
             const IVFSearchParameters* params) const override;
 
+    /** Returns a scanner with a fixed PQ distance-computation strategy. */
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            ScannerMode mode) const;
+
     /// build precomputed table
     void precompute_table();
 

data/vendor/faiss/faiss/IndexShards.cpp

@@ -243,8 +243,8 @@ void IndexShardsTemplate<IndexT>::search(
 
     this->runOnIndex(fn);
 
-    if (this->metric_type == METRIC_L2) {
-        merge_knn_results<idx_t, CMin<distance_t, int>>(
+    if (is_similarity_metric(this->metric_type)) {
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
                 n,
                 k,
                 nshard,
@@ -253,7 +253,7 @@ void IndexShardsTemplate<IndexT>::search(
                 distances,
                 labels);
     } else {
-        merge_knn_results<idx_t, CMax<distance_t, int>>(
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
                 n,
                 k,
                 nshard,

data/vendor/faiss/faiss/IndexShardsIVF.cpp

@@ -13,6 +13,7 @@
 #include <cstdio>
 #include <functional>
 
+#include <faiss/MetricType.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/WorkerThread.h>
@@ -227,8 +228,8 @@ void IndexShardsIVF::search(
 
     this->runOnIndex(fn);
 
-    if (this->metric_type == METRIC_L2) {
-        merge_knn_results<idx_t, CMin<distance_t, int>>(
+    if (is_similarity_metric(metric_type)) {
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
                 n,
                 k,
                 nshard,
@@ -237,7 +238,7 @@ void IndexShardsIVF::search(
                 distances,
                 labels);
     } else {
-        merge_knn_results<idx_t, CMax<distance_t, int>>(
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
                 n,
                 k,
                 nshard,

data/vendor/faiss/faiss/SuperKMeans.cpp

@@ -59,7 +59,7 @@ namespace {
 struct TrainState {
     /// Orthogonal rotation. Train in rotated space (X_tilde = X * R);
     /// un-rotate centroids before return.
-    faiss::RandomRotationMatrix R;
+    std::unique_ptr<faiss::VectorTransform> R;
 
     std::vector<float> X_tilde; // (n, d) row-major
     int n = 0;
@@ -77,7 +77,17 @@ struct TrainState {
     int low_pruning_streak = 0;
     bool low_pruning_warning_printed = false;
 
-    explicit TrainState(int d) : R(d, d) {}
+    explicit TrainState(int d, bool spherical)
+            : R([d, spherical]() -> std::unique_ptr<faiss::VectorTransform> {
+                  // Spherical (inner-product) clustering only: a power-of-two
+                  // dimension can use the fast Hadamard rotation instead of
+                  // the generic random rotation. L2 training keeps the
+                  // original RandomRotationMatrix path unchanged.
+                  if (spherical && d > 0 && (d & (d - 1)) == 0) {
+                      return std::make_unique<faiss::HadamardRotation>(d);
+                  }
+                  return std::make_unique<faiss::RandomRotationMatrix>(d, d);
+              }()) {}
 };
 
 /// PDX block layout for the trailing pruning sweep: block b covers original
@@ -295,10 +305,17 @@ std::unique_ptr<uint8_t[]> setup_train_state(
             "SuperKMeans: training set size exceeds INT_MAX after sampling");
     state.n = static_cast<int>(nx);
 
-    state.R.init(cp.seed);
+    if (auto* R = dynamic_cast<HadamardRotation*>(state.R.get())) {
+        R->init(cp.seed);
+    } else {
+        auto* dense_rotation =
+                dynamic_cast<RandomRotationMatrix*>(state.R.get());
+        FAISS_ASSERT(dense_rotation != nullptr);
+        dense_rotation->init(cp.seed);
+    }
 
     state.X_tilde.resize(static_cast<size_t>(state.n) * d);
-    state.R.apply_noalloc(state.n, x_sampled, state.X_tilde.data());
+    state.R->apply_noalloc(state.n, x_sampled, state.X_tilde.data());
 
     // Forgy init: pick k random rows from the rotated pool as initial
     // centroids. These remain in rotated space; un-rotation happens
@@ -313,6 +330,9 @@ std::unique_ptr<uint8_t[]> setup_train_state(
                     state.X_tilde.data() + static_cast<size_t>(perm[j]) * d,
                     sizeof(float) * d);
         }
+        if (cp.spherical) {
+            fvec_renorm_L2(d, k, state.Y_tilde.data());
+        }
     }
 
     state.d_prime =
@@ -339,7 +359,7 @@ std::unique_ptr<uint8_t[]> setup_train_state(
 /// reverse_transform applies R^T = R^-1.
 void untransform_centroids(
         std::vector<float>& centroids,
-        const RandomRotationMatrix& R,
+        const VectorTransform& R,
         int d,
         int k,
         const float* Y_tilde) {
@@ -411,7 +431,7 @@ void SuperKMeans::train(idx_t n, const float* x) {
                static_cast<idx_t>(k) * cp.min_points_per_centroid);
     }
 
-    TrainState state(d);
+    TrainState state(d, cp.spherical);
     std::vector<int64_t> labels64;
     SuperKMeansAssignScratch assign_scratch;
     std::vector<float> hassign;
@@ -446,6 +466,9 @@ void SuperKMeans::train(idx_t n, const float* x) {
 
         const int nsplit =
                 update_centroids_and_split(d, k, state, labels64, hassign);
+        if (cp.spherical) {
+            fvec_renorm_L2(d, k, state.Y_tilde.data());
+        }
         const float pruning_rate = (iter == 0)
                 ? 0.0f
                 : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm);
@@ -492,7 +515,7 @@ void SuperKMeans::train(idx_t n, const float* x) {
                (getmillisecs() - t_train_start) / 1000.0);
     }
 
-    untransform_centroids(centroids, state.R, d, k, state.Y_tilde.data());
+    untransform_centroids(centroids, *state.R, d, k, state.Y_tilde.data());
 }
 
 void super_kmeans_assign_iteration(
@@ -598,7 +621,7 @@ void super_kmeans_assign_iteration(
             }
 
             // One SIMD dispatch per (xi, yj) tile.
-            with_simd_level([&]<SIMDLevel SL>() {
+            with_simd_level_with_sve([&]<SIMDLevel SL>() {
                 [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk;
                 int64_t tile_total = 0;
                 int64_t tile_pruned = 0;

data/vendor/faiss/faiss/SuperKMeans.h

@@ -13,8 +13,9 @@
 //   "A Super Fast K-means for Indexing Vector Embeddings."
 //   arXiv preprint arXiv:2603.20009.
 //
-// Use when: L2 metric, k >= 1024, d >= 128, dense float embeddings.
-// Do not use for: IP/cosine (use Clustering with cp.spherical=true), small k,
+// Use when: L2 metric, or spherical inner-product clustering with
+// cp.spherical=true, k >= 1024, d >= 128, dense float embeddings.
+// Do not use for: small k,
 // binary data (use IndexBinaryIVF), or near-unit-sphere embeddings with
 // k < 4096 (chi-squared assumption breaks down).
 //

data/vendor/faiss/faiss/VectorTransform.cpp

@@ -513,6 +513,44 @@ void HadamardRotation::apply_noalloc(idx_t n, const float* x, float* xt) const {
     }
 }
 
+void HadamardRotation::reverse_transform(idx_t n, const float* xt, float* x)
+        const {
+    FAISS_THROW_IF_NOT_MSG(is_trained, "Transformation not trained yet");
+    FAISS_THROW_IF_NOT_MSG(
+            d_in == d_out,
+            "HadamardRotation inverse requires equal input/output dimensions");
+
+    const size_t p = d_out;
+    // Reverse of apply_noalloc: three unnormalized FWHT rounds scale norms
+    // by (sqrt(p))^3 = p*sqrt(p); the forward pass cancels this with
+    // total_scale = 1/(p*sqrt(p)), so the inverse applies the same factor.
+    const float inverse_scale = 1.0f / (p * std::sqrt(static_cast<float>(p)));
+
+#pragma omp parallel for schedule(dynamic)
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = xt + i * p;
+        float* xo = x + i * p;
+
+        // The inverse reverses the three sign-flip/Hadamard factors.
+        std::memcpy(xo, xi, p * sizeof(float));
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs3[j];
+        }
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs2[j];
+        }
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs1[j] * inverse_scale;
+        }
+    }
+}
+
 void HadamardRotation::check_identical(const VectorTransform& other) const {
     auto* hr = dynamic_cast<const HadamardRotation*>(&other);
     FAISS_THROW_IF_NOT_MSG(hr, "failed to cast to HadamardRotation");

data/vendor/faiss/faiss/VectorTransform.h

@@ -144,6 +144,9 @@ struct HadamardRotation : VectorTransform {
 
     void apply_noalloc(idx_t n, const float* x, float* xt) const override;
 
+    /// Apply the inverse transform when d_in == d_out.
+    void reverse_transform(idx_t n, const float* xt, float* x) const override;
+
     void check_identical(const VectorTransform& other) const override;
 
     HadamardRotation() {}

data/vendor/faiss/faiss/clone_index.cpp

@@ -38,6 +38,7 @@
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexPQFastScan.h>
 #include <faiss/IndexPreTransform.h>
+#include <faiss/IndexRaBitQ.h>
 #include <faiss/IndexRefine.h>
 #include <faiss/IndexRowwiseMinMax.h>
 #include <faiss/IndexScalarQuantizer.h>
@@ -139,6 +140,7 @@ IndexIDMap* clone_IndexIDMap(const IndexIDMap* im) {
 IndexHNSW* clone_IndexHNSW(const IndexHNSW* ihnsw) {
     TRYCLONE(IndexHNSW2Level, ihnsw)
     TRYCLONE(IndexHNSWFlatPanorama, ihnsw)
+    TRYCLONE(IndexHNSWRaBitQ, ihnsw)
     TRYCLONE(IndexHNSWFlat, ihnsw)
     TRYCLONE(IndexHNSWPQ, ihnsw)
     TRYCLONE(IndexHNSWSQ, ihnsw)
@@ -294,6 +296,7 @@ Index* Cloner::clone_Index(const Index* index) {
     TRYCLONE(IndexEDEN, index)
 
     TRYCLONE(IndexScalarQuantizer, index)
+    TRYCLONE(IndexRaBitQ, index)
     TRYCLONE(MultiIndexQuantizer, index)
 
     if (const IndexIVF* ivf = dynamic_cast<const IndexIVF*>(index)) {

data/vendor/faiss/faiss/gpu/GpuClonerOptions.h

@@ -39,12 +39,8 @@ struct GpuClonerOptions {
     /// Set verbose options on the index
     bool verbose = false;
 
-    /// use the cuVS implementation
-#if defined USE_NVIDIA_CUVS
-    bool use_cuvs = true;
-#else
+    /// Use the cuVS implementation. Opt-in: see GpuIndexConfig::use_cuvs.
     bool use_cuvs = false;
-#endif
 
     /// This flag controls the CPU fallback logic for coarse quantizer
     /// component of the index. When set to false (default), the cloner will

data/vendor/faiss/faiss/gpu/GpuDistance.h

@@ -107,12 +107,9 @@ struct GpuDistanceParams {
     /// execution
     int device = -1;
 
-    /// Should the index dispatch down to cuVS?
-#if defined USE_NVIDIA_CUVS
-    bool use_cuvs = true;
-#else
+    /// Should the search dispatch down to cuVS? Opt-in: see
+    /// GpuIndexConfig::use_cuvs.
     bool use_cuvs = false;
-#endif
 };
 
 /// A function that determines whether cuVS should be used based on various

data/vendor/faiss/faiss/gpu/GpuIndex.h

@@ -38,12 +38,11 @@ struct GpuIndexConfig {
     /// more memory than is available on the GPU.
     MemorySpace memorySpace = MemorySpace::Device;
 
-    /// Should the index dispatch down to cuVS?
-#if defined USE_NVIDIA_CUVS
-    bool use_cuvs = true;
-#else
+    /// Should the index dispatch down to cuVS? Opt-in: dispatching to cuVS
+    /// selects a different implementation with its own numerical behaviour, so
+    /// it must be an explicit choice by the caller rather than a consequence of
+    /// how the binary happened to be built.
     bool use_cuvs = false;
-#endif
 };
 
 /// A centralized function that determines whether cuVS should

data/vendor/faiss/faiss/gpu/GpuIndexCagra.h

@@ -47,7 +47,9 @@ enum class graph_build_algo {
     /// Use NN-Descent to build all-neighbors knn graph
     NN_DESCENT,
     /// Use iterative search to build knn graph
-    ITERATIVE_SEARCH
+    ITERATIVE_SEARCH,
+    /// Exact knn graph via tiled brute force
+    BRUTE_FORCE
 };
 
 /// A type for specifying how PQ codebooks are created.
@@ -173,6 +175,34 @@ struct IVFPQSearchCagraConfig {
     uint32_t max_internal_batch_size = 4096;
 };
 
+/// Knobs for the multi-GPU build path, selected by listing more than one
+/// device in GpuIndexCagraConfig::devices. Degrees, build_algo and metric
+/// still come from GpuIndexCagraConfig.
+struct AllNeighborsCagraConfig {
+    /// Number of overlapping clusters the dataset is partitioned into.
+    /// 0 selects max(2 * number of devices, 4).
+    size_t n_clusters = 0;
+
+    /// Clusters each vector is assigned to. Must be >= 2: with 1 there are no
+    /// cross-cluster edges and recall collapses. 0 selects 2.
+    size_t overlap_factor = 2;
+
+    /// Bounds IVF-PQ search memory during the knn build. cuVS's default can
+    /// exhaust device memory around 100M vectors; 8192 avoids that and does
+    /// not change recall. 0 keeps cuVS's default.
+    uint32_t ivf_pq_search_batch_size = 0;
+
+    /// IVF-PQ refinement multiplier. The refine pass runs on the CPU, so cost
+    /// scales with this. Raising it above 1.0 measured both slower and less
+    /// accurate, so it is separate from the shared refine_rate.
+    float refinement_rate = 1.0f;
+
+    /// Size the IVF-PQ index for one cluster rather than the whole dataset.
+    /// cuVS reuses these params for every cluster without rescaling, so
+    /// sizing from the full dataset badly over-partitions each one.
+    bool ivf_pq_size_from_cluster = true;
+};
+
 struct GpuIndexCagraConfig : public GpuIndexConfig {
     /// Degree of input graph for pruning.
     size_t intermediate_graph_degree = 128;
@@ -190,6 +220,27 @@ struct GpuIndexCagraConfig : public GpuIndexConfig {
 
     /// Whether to use MST optimization to guarantee graph connectivity.
     bool guarantee_connectivity = false;
+
+    /// Devices to build on. More than one selects the multi-GPU build in
+    /// train(); see AllNeighborsCagraConfig and train() for its restrictions.
+    /// Empty or one device uses GpuIndexConfig::device as usual.
+    std::vector<int> devices;
+
+    AllNeighborsCagraConfig all_neighbors_params;
+
+    /// Build the HNSW upper levels on the GPU during copyTo() instead of by
+    /// CPU insertion. Roughly 10x faster, but measured worse recall than
+    /// base_level_only, which is cheaper still. Off by default.
+    bool gpu_hnsw_upper_levels = false;
+
+    /// intermediate_graph_degree for the per-level subgraph builds.
+    /// 0 = twice the upper-level degree.
+    size_t gpu_hnsw_intermediate_degree = 0;
+
+    /// MST connectivity pass on the per-level subgraphs. On by default:
+    /// upper levels are walked greedily with no backtracking, so a
+    /// disconnected component is a trap the descent cannot escape.
+    bool gpu_hnsw_guarantee_connectivity = true;
 };
 
 enum class search_algo {
@@ -284,46 +335,6 @@ struct GpuIndexCagra : public GpuIndex {
     /// in the index instance
     void copyTo(faiss::IndexHNSWCagra* index) const;
 
-    /// Train CAGRA using multiple GPUs by sharding the dataset.
-    /// Uses cuVS native SNMG (single-node multi-GPU) CAGRA build.
-    /// Each device builds one shard in parallel via OpenMP.
-    /// Float32 only. After training, call copyTo() to produce a CPU
-    /// IndexHNSWCagra with full HNSW upper levels.
-    /// The training data pointer must remain valid until copyTo() completes.
-    /// stitch_mode: 0=CPU HNSW (Approach C), 1=GPU brute-force (Approach B)
-    void trainMultiGpu(
-            idx_t n,
-            const float* x,
-            std::vector<GpuResourcesProvider*>& providers,
-            std::vector<int>& devices,
-            idx_t stitch_per_shard = 0,
-            int stitch_k = 2,
-            int stitch_mode = 0);
-
-    /// Build a unified CAGRA graph using cuVS all_neighbors
-    /// (multi-GPU kNN graph construction with overlapping clusters) followed
-    /// by cagra::optimize (graph pruning). Produces a single unified graph
-    /// without stitching. The training data pointer must remain valid until
-    /// copyTo() completes.
-    /// build_algo: 0=NN-descent (default), 1=brute-force, 2=IVF-PQ
-    /// refinement_rate: IVF-PQ refinement multiplier (only used when
-    /// build_algo==2). Higher values trade build time for recall; the cuVS
-    /// default is 2.0.
-    /// ivfpq_search_batch: cap the IVF-PQ search `max_internal_batch_size` used
-    /// during the all_neighbors kNN build (build_algo==2). 0 = cuVS default
-    /// (128*1024), which can OOM at 100M; a smaller value (e.g. 8192) bounds
-    /// the GPU search workspace with no effect on results (recall-neutral).
-    void trainAllNeighbors(
-            idx_t n,
-            const float* x,
-            std::vector<int>& devices,
-            int n_clusters = 0,
-            int overlap_factor = 0,
-            bool multi_gpu_optimize = false,
-            int build_algo = 0,
-            float refinement_rate = 2.0f,
-            int ivfpq_search_batch = 0);
-
     void reset() override;
 
     std::vector<idx_t> get_knngraph() const;
@@ -357,6 +368,18 @@ struct GpuIndexCagra : public GpuIndex {
             idx_t* labels,
             const SearchParameters* search_params) const override;
 
+    /// Multi-GPU build path taken by train() when cagraConfig_.devices lists
+    /// more than one device: cuVS `all_neighbors` knn graph construction over
+    /// overlapping clusters, followed by graph pruning. Leaves the result in
+    /// merged_knngraph_ for copyTo(); index_ stays empty.
+    void trainAllNeighbors_(idx_t n, const float* x);
+
+    /// Populate HNSW levels >= 1 of `index` by building a CAGRA graph over
+    /// each level's node subset on the GPU. Requires the level table to be
+    /// prepared and the storage populated. Returns the max level.
+    void buildHnswUpperLevelsGpu_(faiss::IndexHNSWCagra* index, int max_lvl)
+            const;
+
     void copyToMultiGpu_(faiss::IndexHNSWCagra* index) const;
 
     /// Our configuration options
@@ -372,7 +395,7 @@ struct GpuIndexCagra : public GpuIndex {
             std::shared_ptr<CuvsCagra<int8_t>>>
             index_;
 
-    /// Multi-GPU state: populated by trainMultiGpu(), used by copyTo()
+    /// Multi-GPU state: populated by trainAllNeighbors_(), used by copyTo()
     std::vector<idx_t> merged_knngraph_;
     idx_t merged_knngraph_degree_ = 0;
     const float* multi_gpu_dataset_ = nullptr;

data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp

@@ -376,7 +376,7 @@ void AdditiveQuantizer::compute_centroid_norms(float* norms) const {
     size_t ntotal = (size_t)1 << tot_bits;
     int64_t ntotal_signed = ntotal;
     // TODO: make tree of partial sums
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_with_sve([&]<SIMDLevel SL>() {
 #pragma omp parallel
         {
             std::vector<float> tmp(d);

data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp

@@ -235,8 +235,11 @@ void ClusteringInitialization::init_kmeans_plus_plus(
             float* new_centroid = centroids + c * d;
             std::memcpy(new_centroid, x + next_idx * d, d * sizeof(float));
 
-            // Update min distances incrementally
-            for (size_t i = 0; i < n; i++) {
+            // Update min distances incrementally. The writes are independent,
+            // so the loop parallelizes without changing the output.
+            const int64_t ni = static_cast<int64_t>(n);
+#pragma omp parallel for
+            for (int64_t i = 0; i < ni; i++) {
                 double dist = fvec_L2sqr<SL>(x + i * d, new_centroid, d);
                 min_distances[i] = std::min(min_distances[i], dist);
             }