enhance: default SuperKMeans coarse quantizer training for SCANN

cmake/libs/libfaiss.cmake

@@ -187,6 +187,7 @@ knowhere_file_glob(
   FAISS_DD_SVE_SRCS
   thirdparty/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp
   thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp
+  thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp
 )
 # combine files
 list(APPEND FAISS_SVE_SRCS ${FAISS_DD_SVE_SRCS})
@@ -548,6 +549,7 @@ if(__AARCH64)
                                      knowhere_utils)
   if(SVE_AVAILABLE)
     target_link_libraries(faiss PUBLIC faiss_sve)
+    target_compile_definitions(faiss PRIVATE COMPILE_SIMD_ARM_SVE)
   endif()
   target_compile_definitions(faiss PRIVATE FINTEGER=int FAISS_ENABLE_DD COMPILE_SIMD_ARM_NEON)
 endif()

include/knowhere/comp/index_param.h

@@ -145,6 +145,7 @@ namespace indexparam {
 constexpr const char* NPROBE = "nprobe";
 constexpr const char* NLIST = "nlist";
 constexpr const char* USE_ELKAN = "use_elkan";
+constexpr const char* USE_SUPER_KMEANS = "use_super_kmeans";
 constexpr const char* NBITS = "nbits";          // PQ/SQ
 constexpr const char* M = "m";                  // PQ param for IVFPQ
 constexpr const char* IVF_SQ_TYPE = "sq_type";  // SQ param for IVFSQ

src/index/ivf/ivf.cc

@@ -19,6 +19,7 @@
 #include "faiss/IndexIVFRaBitQ.h"
 #include "faiss/IndexIVFRaBitQFastScan.h"
 #include "faiss/IndexRefine.h"
+#include "faiss/SuperKMeans.h"
 #include "faiss/VectorTransform.h"
 #include "faiss/cppcontrib/knowhere/IndexBinaryFlat.h"
 #include "faiss/cppcontrib/knowhere/IndexBinaryIVF.h"
@@ -599,6 +600,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         }
         // apply clustering config
         ApplyClusteringConfig(index->cp);
+        index->cp.use_super_kmeans = ivf_flat_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
         // transfer ownership of qzr to index
@@ -621,6 +623,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         }
         // apply clustering config
         ApplyClusteringConfig(index->cp);
+        index->cp.use_super_kmeans = ivf_flat_cc_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
         // transfer ownership of qzr to index
@@ -652,6 +655,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
 
         // apply clustering config
         ApplyClusteringConfig(index->get_base_ivf_index()->cp);
+        index->get_base_ivf_index()->cp.use_super_kmeans = ivf_pq_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
         // transfer ownership of qzr to index
@@ -678,6 +682,14 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         }
         // apply clustering config
         ApplyClusteringConfig(base_index->cp);
+        // SuperKMeans requires d >= 2 * d_prime_min (default 32). Fall back to
+        // Clustering below that hard limit; use_super_kmeans otherwise remains
+        // the user's choice, including for small nlist values.
+        bool use_super_kmeans = scann_cfg.use_super_kmeans.value();
+        if (use_super_kmeans && dim < 2 * faiss::SuperKMeansParameters{}.d_prime_min) {
+            use_super_kmeans = false;
+        }
+        base_index->cp.use_super_kmeans = use_super_kmeans;
         // create scann index, which does not base_index by default,
         //    but owns the refine index by default omg
         if (scann_cfg.with_raw_data.value()) {
@@ -713,6 +725,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
 
         // apply clustering config
         ApplyClusteringConfig(index->get_base_ivf_index()->cp);
+        index->get_base_ivf_index()->cp.use_super_kmeans = ivf_sq_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
         // transfer ownership of qzr to index
@@ -730,6 +743,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         index = std::make_unique<faiss::cppcontrib::knowhere::IndexBinaryIVF>(qzr.get(), dim, nlist, metric.value());
         // apply clustering config
         ApplyClusteringConfig(index->cp);
+        index->cp.use_super_kmeans = ivf_bin_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const uint8_t*>(data));
         // transfer ownership of qzr to index
@@ -760,6 +774,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         }
         // apply clustering config
         ApplyClusteringConfig(index->cp);
+        index->cp.use_super_kmeans = ivf_sq_cc_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
         // transfer ownership of qzr to index
@@ -782,6 +797,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         index = std::move(result.value());
         // apply clustering config
         ApplyClusteringConfig(index->get_ivfrabitq_index()->cp);
+        index->get_ivfrabitq_index()->cp.use_super_kmeans = ivf_rabitq_cfg.use_super_kmeans.value();
         // train
         index->train(rows, static_cast<const float*>(data));
     }
@@ -798,6 +814,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
         auto* fs_idx = index->get_fastscan_index();
         if (fs_idx) {
             ApplyClusteringConfig(fs_idx->cp);
+            fs_idx->cp.use_super_kmeans = fs_cfg.use_super_kmeans.value();
         }
         index->train(rows, static_cast<const float*>(data));
     }

src/index/ivf/ivf_config.h

@@ -27,6 +27,7 @@ class IvfConfig : public BaseConfig {
     CFG_INT nlist;
     CFG_INT nprobe;
     CFG_BOOL use_elkan;
+    CFG_BOOL use_super_kmeans;
     CFG_BOOL ensure_topk_full;  // internal config, used for temp index
     CFG_INT max_empty_result_buckets;
     KNOWHERE_DECLARE_CONFIG(IvfConfig) {
@@ -46,6 +47,10 @@ class IvfConfig : public BaseConfig {
             .set_default(true)
             .description("whether to use elkan algorithm")
             .for_train();
+        KNOWHERE_CONFIG_DECLARE_FIELD(use_super_kmeans)
+            .set_default(false)
+            .description("whether to use SuperKMeans for coarse quantizer training")
+            .for_train();
         KNOWHERE_CONFIG_DECLARE_FIELD(ensure_topk_full)
             .set_default(true)
             .description("whether to make sure topk results full")
@@ -196,6 +201,13 @@ class ScannConfig : public IvfFlatConfig {
             .set_default(false)
             .description("whether to make sure topk results full")
             .for_search();
+        // SCANN defaults to SuperKMeans for coarse quantizer training: the
+        // super-fast k-means variant is recall-equivalent to Clustering on
+        // inner-product data but trains significantly faster for large nlist.
+        KNOWHERE_CONFIG_DECLARE_FIELD(use_super_kmeans)
+            .set_default(true)
+            .description("whether to use SuperKMeans for coarse quantizer training")
+            .for_train();
     }
 
     Status

tests/ut/test_cluster.cc

@@ -18,6 +18,7 @@
 #include "catch2/generators/catch_generators.hpp"
 #include "faiss/Clustering.h"
 #include "faiss/IndexFlat.h"
+#include "faiss/SuperKMeans.h"
 #include "faiss/cppcontrib/knowhere/utils/binary_distances.h"
 #include "hnswlib/hnswalg.h"
 #include "knowhere/bitsetview.h"
@@ -143,3 +144,68 @@ TEST_CASE("Test Kmeans With Float Vector", "[float metrics]") {
         REQUIRE(recall > kKnnRecallThreshold);
     }
 }
+
+// SuperKMeans spherical (inner-product) support: unit-normalized centroids
+// make L2 assignment equivalent to IP argmax, so the final objective must
+// track vanilla spherical Clustering and centroids must be unit norm.
+TEST_CASE("Test SuperKMeans Spherical", "[cluster]") {
+    const int d = 64;
+    const int k = 16;
+    const size_t n = 2000;
+
+    std::mt19937 rng(42);
+    std::normal_distribution<float> dist(0.f, 1.f);
+    std::vector<float> x(n * d);
+    for (auto& v : x) {
+        v = dist(rng);
+    }
+
+    faiss::SuperKMeansParameters sp;
+    sp.seed = 42;
+    sp.niter = 10;
+    sp.spherical = true;
+    faiss::SuperKMeans sc(d, k, sp);
+    sc.train(n, x.data());
+
+    // Centroids must be unit norm under spherical clustering.
+    for (int j = 0; j < k; ++j) {
+        float norm = 0.f;
+        for (int i = 0; i < d; ++i) {
+            norm += sc.centroids[j * d + i] * sc.centroids[j * d + i];
+        }
+        norm = std::sqrt(norm);
+        REQUIRE(norm == Catch::Approx(1.f).margin(1e-4));
+    }
+
+    // Final objective must track vanilla spherical Clustering.
+    const float sc_final = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj;
+    faiss::ClusteringParameters vp;
+    vp.seed = 42;
+    vp.niter = 10;
+    vp.spherical = true;
+    faiss::Clustering vanilla(d, k, vp);
+    faiss::IndexFlatL2 quantizer(d);
+    vanilla.train(n, x.data(), quantizer);
+    const float v_final = vanilla.iteration_stats.at(vanilla.iteration_stats.size() - 1).obj;
+    REQUIRE(std::abs(sc_final - v_final) / v_final < 0.05f);
+}
+
+TEST_CASE("Test SuperKMeans with 256 centroids", "[cluster]") {
+    constexpr int d = 64;
+    constexpr int k = 256;
+    constexpr size_t n = 1024;
+
+    std::mt19937 rng(43);
+    std::normal_distribution<float> dist(0.f, 1.f);
+    std::vector<float> x(n * d);
+    for (auto& v : x) {
+        v = dist(rng);
+    }
+
+    faiss::SuperKMeansParameters sp;
+    sp.seed = 43;
+    sp.niter = 2;
+    faiss::SuperKMeans clustering(d, k, sp);
+    REQUIRE_NOTHROW(clustering.train(n, x.data()));
+    REQUIRE(clustering.centroids.size() == static_cast<size_t>(k * d));
+}

tests/ut/test_scann_superkmeans.cc

@@ -0,0 +1,134 @@
+// Copyright (C) 2019-2023 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+#include <string>
+#include <vector>
+
+#include "catch2/catch_approx.hpp"
+#include "catch2/catch_test_macros.hpp"
+#include "knowhere/comp/brute_force.h"
+#include "knowhere/comp/index_param.h"
+#include "knowhere/index/index_factory.h"
+#include "knowhere/version.h"
+#include "utils.h"
+
+namespace {
+
+// Build a SCANN index over train_ds with the given use_super_kmeans value
+// and return recall@k against brute-force ground truth.
+float
+BuildScannAndRecall(const knowhere::DataSetPtr& train_ds, const knowhere::DataSetPtr& query_ds, int64_t nlist,
+                    int64_t nprobe, int64_t topk, bool use_super_kmeans) {
+    const auto version = knowhere::Version::GetCurrentVersion().VersionNumber();
+    auto idx = knowhere::IndexFactory::Instance()
+                   .Create<knowhere::fp32>(knowhere::IndexEnum::INDEX_FAISS_SCANN, version)
+                   .value();
+
+    knowhere::Json cfg;
+    cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP;
+    cfg[knowhere::indexparam::NLIST] = nlist;
+    cfg[knowhere::indexparam::NPROBE] = nprobe;
+    cfg[knowhere::indexparam::SUB_DIM] = 4;
+    cfg[knowhere::indexparam::WITH_RAW_DATA] = false;
+    cfg[knowhere::indexparam::USE_SUPER_KMEANS] = use_super_kmeans;
+
+    REQUIRE(idx.Build(train_ds, cfg) == knowhere::Status::success);
+
+    knowhere::Json search_cfg;
+    search_cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP;
+    search_cfg[knowhere::meta::TOPK] = topk;
+    search_cfg[knowhere::indexparam::NPROBE] = nprobe;
+    auto results = idx.Search(query_ds, search_cfg, nullptr);
+    REQUIRE(results.has_value());
+
+    auto gt = knowhere::BruteForce::Search<knowhere::fp32>(
+        train_ds, query_ds,
+        knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, {knowhere::meta::TOPK, topk}}, nullptr);
+    REQUIRE(gt.has_value());
+
+    return GetKNNRecall(*gt.value(), *results.value());
+}
+
+}  // namespace
+
+TEST_CASE("SCANN use_super_kmeans default matches Clustering recall", "[scann]") {
+    constexpr int64_t nb = 2000;
+    constexpr int64_t nq = 100;
+    constexpr int64_t dim = 64;
+    constexpr int64_t topk = 10;
+    constexpr int64_t nlist = 128;
+    constexpr int64_t nprobe = 8;
+
+    const auto train_ds = GenDataSet(nb, dim, kSeed);
+    const auto query_ds = GenDataSet(nq, dim, kSeed);
+
+    const float super_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, true);
+    const float cluster_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, false);
+
+    CAPTURE(super_recall, cluster_recall);
+    // SuperKMeans coarse quantizer training is recall-equivalent to Clustering.
+    // MatchNlist shrinks nlist on this small synthetic set. The enabled build
+    // still honors SuperKMeans for the resulting small centroid count; recall
+    // equivalence between the two clustering implementations is the invariant.
+    REQUIRE(super_recall == Catch::Approx(cluster_recall).margin(0.02f));
+}
+
+TEST_CASE("SCANN use_super_kmeans field is honored", "[scann]") {
+    // Explicitly disabled must build and search successfully too.
+    constexpr int64_t nb = 1000;
+    constexpr int64_t nq = 50;
+    constexpr int64_t dim = 32;
+    constexpr int64_t topk = 10;
+    constexpr int64_t nlist = 64;
+    constexpr int64_t nprobe = 4;
+
+    const auto train_ds = GenDataSet(nb, dim, kSeed + 1);
+    const auto query_ds = GenDataSet(nq, dim, kSeed + 1);
+
+    const float cluster_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, false);
+    REQUIRE(cluster_recall > 0.0f);
+}
+
+// SCANN with default use_super_kmeans=true must not fail the build on
+// low-dimensional data (e.g. d=16 emb-list scenarios) where SuperKMeans is
+// not applicable; it should fall back to Clustering.
+TEST_CASE("SCANN low-dim build succeeds with default superkmeans", "[scann]") {
+    constexpr int64_t nb = 200;
+    constexpr int64_t nq = 20;
+    constexpr int64_t dim = 16;
+    constexpr int64_t topk = 5;
+    constexpr int64_t nlist = 16;
+    constexpr int64_t nprobe = 2;
+
+    const auto train_ds = GenDataSet(nb, dim, kSeed + 2);
+    const auto query_ds = GenDataSet(nq, dim, kSeed + 2);
+
+    const auto version = knowhere::Version::GetCurrentVersion().VersionNumber();
+    auto idx = knowhere::IndexFactory::Instance()
+                   .Create<knowhere::fp32>(knowhere::IndexEnum::INDEX_FAISS_SCANN, version)
+                   .value();
+
+    knowhere::Json cfg;
+    cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP;
+    cfg[knowhere::indexparam::NLIST] = nlist;
+    cfg[knowhere::indexparam::NPROBE] = nprobe;
+    cfg[knowhere::indexparam::SUB_DIM] = 2;
+    cfg[knowhere::indexparam::WITH_RAW_DATA] = false;
+    // Default use_super_kmeans=true; must fall back to Clustering for d=16.
+    REQUIRE(idx.Build(train_ds, cfg) == knowhere::Status::success);
+
+    knowhere::Json search_cfg;
+    search_cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP;
+    search_cfg[knowhere::meta::TOPK] = topk;
+    search_cfg[knowhere::indexparam::NPROBE] = nprobe;
+    auto results = idx.Search(query_ds, search_cfg, nullptr);
+    REQUIRE(results.has_value());
+}

thirdparty/faiss/faiss/CMakeLists.txt

@@ -64,6 +64,7 @@ set(FAISS_SIMD_NEON_SRC
 set(FAISS_SIMD_SVE_SRC
   impl/pq_code_distance/pq_code_distance-sve.cpp
   utils/simd_impl/distances_arm_sve.cpp
+  utils/simd_impl/super_kmeans_kernels_sve.cpp
 )
 set(FAISS_SIMD_RVV_SRC
   impl/fast_scan/impl-riscv.cpp

thirdparty/faiss/faiss/Clustering.h

@@ -75,6 +75,11 @@ struct ClusteringParameters {
     /// so the training process stops only if an error
     /// is unchanged from the previous iteration.
     double early_stop_threshold = 0.0;
+
+    /// Whether to use the SuperKMeans (super fast k-means) variant instead of
+    /// the vanilla Clustering implementation. Only honored by callers that
+    /// explicitly support it (e.g. IVF level-1 quantizer training).
+    bool use_super_kmeans = false;
 };
 
 struct ClusteringIterationStats {