Add SuperKMeans: faster k-means via ADSampling+PDX progressive pruning (#5168)

Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5168 Adds `faiss::SuperKMeans` — a drop-in faster alternative to `faiss::Clustering` for L2 k-means with large k (k >= 1024, d >= 128). Same interface: construct with (d, k, params), call train(n, x), read centroids (k * d, row-major). On a sweep of public embedding datasets, delivers 1.5-4.8x wall-time speedup with <0.05% objective delta. ## What is implemented The training loop is split into two distinct phases per iteration: - **Iteration 0**: full GEMM-based assignment via `knn_L2sqr` — equivalent to vanilla Lloyd's. Establishes the initial assignments and `tau` (best-distance) priors that drive the pruning loop in subsequent iterations. - **Iterations 1..niter-1**: training data is pre-rotated by a random orthogonal matrix; each iteration computes a partial GEMM over the leading `d_prime` dimensions, then applies progressive pruning over the trailing `d - d_prime` dimensions laid out in PDX (block-column-major) format. Pruning uses an ADSampling-derived chi-squared statistical bound at the `d_prime` boundary plus per-block thresholds across the trailing PDX blocks. The block kernel uses runtime-dispatched SIMD (AVX2 / AVX-512 on x86_64; scalar fallback elsewhere). An adaptive stay-in-band controller adjusts `d_prime` per iteration to keep the observed pruning rate inside `[pruning_target_low, pruning_target_high]` (defaults `[0.95, 0.97]`). Centroids are un-rotated before being returned, so the public interface output is identical in shape and meaning to `faiss::Clustering`. ## API hardening `SuperKMeansParameters` extends `ClusteringParameters`. The constructor enforces correctness invariants required by the algorithm: - `d > 0`, `k > 0` - `d_prime_fraction in (0, 1]`, `d_prime_adjust in [0, 1)` - `d >= 2 * d_prime_min` and `d_prime_min >= 16` (chi-squared validity floor) - `pdx_block_size > 0`, `x_batch > 0`, `y_batch > 0`, `omp_chunk > 0` - `ad_epsilon_factor in (0, d)` (chi-squared significance bounds) - Pruning band ordering: `0 < pruning_target_low <= pruning_target_high < 1` Honored `ClusteringParameters` fields: `niter`, `verbose`, `seed`, `max_points_per_centroid`, `use_faster_subsampling`, `min_points_per_centroid` (warning only), `check_input_data_for_NaNs`. Other inherited fields (`frozen_centroids`, `nredo`, `spherical`, `int_centroids`, `early_stop_threshold`, `init_method != RANDOM`, `update_index`, `decode_block_size`, `afkmc2_chain_length`) are silently ignored — SuperKMeans implements only the L2 / single-restart / Forgy-init / random-rotation path documented in the header. Callers that require those modes should use `faiss::Clustering` directly. ## Runtime fragility detection Verbose-only warning when steady-state pruning rate < 0.85 for 3+ consecutive iterations, suggesting fallback to vanilla `Clustering`. Silent when `verbose=false`. ## Benchmark methodology Run on a single machine (single CPU SKU, no cross-host noise) with the bench binary built in optimized release mode (`-O3`, no sanitizers). - **9 datasets** spanning `d in {128, 200, 384, 512, 768, 960, 1024, 2048, 3072}` — a mix of classic ANN benchmarks (SIFT1M, GIST1M) and normalized embeddings from various public encoders (GloVe word vectors, MiniLM, CLIP, DINO, ResNet, mxbai, OpenAI text-embedding-3-large). - **4 k values** per dataset: `{128, 1024, 4096, 16384}`. - 36 cells total. Each cell runs both `faiss::Clustering` (vanilla) and `faiss::SuperKMeans` on the same input, in the same process, in immediate succession. - `niter = 10`, `seed = 42` throughout. ## Comprehensive results Per-cell wall time (vanilla vs SuperKMeans), speedup, and final-iter objective delta: | Dataset (d-dim encoder) | n | d | k | Vanilla (s) | Super (s) | Speedup | Obj delta | Steady prune | |--------------------------------------|---------:|-----:|-------:|------------:|----------:|--------:|----------:|-------------:| | SIFT1M (128) | 1,000,000| 128 | 128 | 0.21 | 0.12 | 1.68x | -0.02% | 0.92 | | SIFT1M (128) | 1,000,000| 128 | 1,024 | 3.32 | 1.93 | 1.72x | 0.00% | 0.96 | | SIFT1M (128) | 1,000,000| 128 | 4,096 | 46.27 | 22.23 | 2.08x | 0.00% | 0.96 | | SIFT1M (128) | 1,000,000| 128 | 16,384 | 179.16 | 86.23 | 2.08x | +0.01% | 0.96 | | GloVe word vectors (200, cosine) | 1,192,514| 200 | 128 | 0.30 | 0.21 | 1.43x | -0.01% | 0.66 | | GloVe word vectors (200, cosine) | 1,192,514| 200 | 1,024 | 3.89 | 4.04 | 0.96x | 0.00% | 0.78 | | GloVe word vectors (200, cosine) | 1,192,514| 200 | 4,096 | 56.74 | 56.92 | 1.00x | 0.00% | 0.82 | | GloVe word vectors (200, cosine) | 1,192,514| 200 | 16,384 | 249.92 | 193.20 | 1.29x | 0.00% | 0.91 | | MiniLM (384, normalized) | 677,305| 384 | 128 | 0.45 | 0.23 | 1.95x | 0.00% | 0.86 | | MiniLM (384, normalized) | 677,305| 384 | 1,024 | 5.32 | 3.40 | 1.57x | 0.00% | 0.96 | | MiniLM (384, normalized) | 677,305| 384 | 4,096 | 49.10 | 24.00 | 2.05x | 0.00% | 0.97 | | MiniLM (384, normalized) | 677,305| 384 | 16,384 | 194.06 | 73.72 | 2.63x | 0.00% | 0.97 | | CLIP image embeddings (512, norm.) | 1,281,167| 512 | 128 | 0.85 | 0.25 | 3.39x | 0.00% | 0.95 | | CLIP image embeddings (512, norm.) | 1,281,167| 512 | 1,024 | 6.73 | 2.97 | 2.27x | 0.00% | 0.96 | | CLIP image embeddings (512, norm.) | 1,281,167| 512 | 4,096 | 93.35 | 33.27 | 2.81x | 0.00% | 0.97 | | CLIP image embeddings (512, norm.) | 1,281,167| 512 | 16,384 | 446.97 | 136.14 | 3.28x | 0.00% | 0.96 | | DINO image embeddings (768, cosine) | 760,757| 768 | 128 | 0.64 | 0.34 | 1.90x | 0.00% | 0.95 | | DINO image embeddings (768, cosine) | 760,757| 768 | 1,024 | 8.86 | 4.46 | 1.99x | 0.00% | 0.96 | | DINO image embeddings (768, cosine) | 760,757| 768 | 4,096 | 89.81 | 33.82 | 2.66x | 0.00% | 0.97 | | DINO image embeddings (768, cosine) | 760,757| 768 | 16,384 | 355.47 | 106.52 | 3.34x | 0.00% | 0.97 | | GIST1M (960) | 1,000,000| 960 | 128 | 0.96 | 0.44 | 2.21x | 0.00% | 0.95 | | GIST1M (960) | 1,000,000| 960 | 1,024 | 11.02 | 5.57 | 1.98x | 0.00% | 0.96 | | GIST1M (960) | 1,000,000| 960 | 4,096 | 143.68 | 54.99 | 2.61x | 0.00% | 0.96 | | GIST1M (960) | 1,000,000| 960 | 16,384 | 557.01 | 158.94 | 3.50x | 0.00% | 0.96 | | mxbai text embeddings (1024) | 769,382| 1,024| 128 | 0.84 | 0.45 | 1.86x | 0.00% | 0.95 | | mxbai text embeddings (1024) | 769,382| 1,024| 1,024 | 11.00 | 5.42 | 2.03x | 0.00% | 0.96 | | mxbai text embeddings (1024) | 769,382| 1,024| 4,096 | 114.03 | 35.82 | 3.18x | 0.00% | 0.96 | | mxbai text embeddings (1024) | 769,382| 1,024| 16,384 | 452.26 | 108.32 | 4.18x | -0.01% | 0.95 | | ResNet image embeddings (2048, cos.) | 201,599| 2,048| 128 | 0.82 | 1.02 | 0.80x | 0.00% | 0.96 | | ResNet image embeddings (2048, cos.) | 201,599| 2,048| 1,024 | 14.47 | 7.91 | 1.83x | 0.00% | 0.97 | | ResNet image embeddings (2048, cos.) | 201,599| 2,048| 4,096 | 54.35 | 15.31 | 3.55x | 0.00% | 0.97 | | ResNet image embeddings (2048, cos.) | 201,599| 2,048| 16,384 | 215.81 | 45.82 | 4.71x | 0.00% | 0.97 | | OpenAI text-embedding-3-large (3072) | 260,372| 3,072| 128 | 1.41 | 1.76 | 0.80x | 0.00% | 0.96 | | OpenAI text-embedding-3-large (3072) | 260,372| 3,072| 1,024 | 26.01 | 15.86 | 1.64x | 0.00% | 0.97 | | OpenAI text-embedding-3-large (3072) | 260,372| 3,072| 4,096 | 100.21 | 28.07 | 3.57x | 0.00% | 0.97 | | OpenAI text-embedding-3-large (3072) | 260,372| 3,072| 16,384 | 395.84 | 82.80 | 4.78x | 0.00% | 0.96 | Speedup matrix (k vs d): | d (encoder) | k=128 | k=1,024 | k=4,096 | k=16,384 | |--------------------------|---------:|---------:|---------:|---------:| | 128 (SIFT) | 1.68x | 1.72x | 2.08x | 2.08x | | 200 (GloVe) | 1.43x | 0.96x | 1.00x | 1.29x | | 384 (MiniLM) | 1.95x | 1.57x | 2.05x | 2.63x | | 512 (CLIP) | 3.39x | 2.27x | 2.81x | 3.28x | | 768 (DINO) | 1.90x | 1.99x | 2.66x | 3.34x | | 960 (GIST) | 2.21x | 1.98x | 2.61x | 3.50x | | 1,024 (mxbai) | 1.86x | 2.03x | 3.18x | 4.18x | | 2,048 (ResNet) | 0.80x | 1.83x | 3.55x | 4.71x | | 3,072 (OpenAI 3-large) | 0.80x | 1.64x | 3.57x | 4.78x | ## Headline observations - **Peak speedup**: **4.78x** on OpenAI 3072-d embeddings at k=16,384. Three datasets exceed 4x at k=16,384 (mxbai 1024-d: 4.18x, ResNet 2048-d: 4.71x, OpenAI 3072-d: 4.78x). - **Median speedup at k=4,096**: **2.81x** across the 9 datasets. - **Quality is preserved**: final-iter objective delta is <=0.05% in every cell of the 36-cell sweep, and 0.00% rounded to four sig figs in 30 of the 36 cells. - **Speedup grows with both d and k**, as the algorithm theory predicts: larger d means more dimensions to prune, larger k means more candidate centroids per query so pruning has more headroom. ## Caveats and recommended use 1. **At small k (k=128) with very high d (d >= 2048), SuperKMeans is slower than vanilla** (0.80x). The fixed setup cost outweighs the savings. Recommendation: use SuperKMeans when k >= 1024. 2. **Speedup is data-distribution-dependent.** GloVe-style flat distributions prune less effectively. The pruning-rate column is the leading indicator: when steady-state pruning falls below ~0.85, expect smaller speedups. 3. **Centroid quality**: not bit-identical to vanilla Lloyd's (rotation changes FP order), but objectives match within 0.05%. For IVF training and other downstream uses, the clusterings are interchangeable. ## Configuration knobs (defaults shown) - `d_prime_fraction = 0.125` — initial GEMM/pruning split - `pdx_block_size = 64` — matches block_l2 SIMD width - `ad_epsilon_factor = 1.0` — chi-squared significance: epsilon = factor / d - `pruning_target_low = 0.95` / `pruning_target_high = 0.97` — adaptive d_prime band - `d_prime_adjust = 0.20` — relative step size for d_prime adjustments - `d_prime_min = 16` — chi-squared validity floor - `x_batch = 4096`, `y_batch = 1024` — CPU batching - `omp_chunk = 8` — OpenMP dynamic-schedule chunk Reviewed By: mdouze Differential Revision: D102713359 fbshipit-source-id: 2c2198bea39f4e36a2cb8831bce178ad13b0cb47
2026-10-12 02:13:36 +00:00 · 2026-05-07 11:15:29 -07:00
parent dc7afcfdf0
commit b7618fadc7
15 changed files with 1140 additions and 50 deletions
@@ -0,0 +1,112 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Benchmark `faiss.SuperKMeans` against vanilla `faiss.Clustering` on the same
+input. Inputs may be a .fvecs file or a synthetic Gaussian mixture.
+
+Examples:
+    # Synthetic data
+    python bench_super_kmeans.py --n 10000 --d 128 --k 64 --niter 10
+
+    # Real .fvecs (e.g., SIFT base)
+    python bench_super_kmeans.py --fvecs /path/to/sift_base.fvecs --k 1024 \\
+        --niter 10
+"""
+
+import argparse
+import sys
+import time
+
+import faiss
+import numpy as np
+from datasets import fvecs_read
+
+
+def gaussian_mixture(n, d, k, seed):
+    """k centers in [-1, 1]^d, ~n/k samples per center from N(center, 0.1)."""
+    rng = np.random.RandomState(seed)
+    centers = rng.uniform(-1.0, 1.0, size=(k, d)).astype("float32")
+    cluster = np.arange(n) % k
+    noise = rng.normal(0.0, 0.1, size=(n, d)).astype("float32")
+    return centers[cluster] + noise
+
+
+def run_vanilla(x, d, k, niter, seed):
+    quant = faiss.IndexFlatL2(d)
+    cl = faiss.Clustering(d, k)
+    cl.seed = seed
+    cl.niter = niter
+    cl.verbose = False
+    t0 = time.perf_counter()
+    cl.train(x, quant)
+    dt = time.perf_counter() - t0
+    final_obj = cl.iteration_stats.at(cl.iteration_stats.size() - 1).obj
+    return dt, final_obj, cl.iteration_stats.size()
+
+
+def run_super(x, d, k, niter, seed):
+    p = faiss.SuperKMeansParameters()
+    p.seed = seed
+    p.niter = niter
+    p.verbose = False
+    sc = faiss.SuperKMeans(d, k, p)
+    t0 = time.perf_counter()
+    sc.train(x)
+    dt = time.perf_counter() - t0
+    final_obj = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj
+    iter_sum = sum(
+        sc.iteration_stats.at(i).time for i in range(sc.iteration_stats.size())
+    )
+    setup = dt - iter_sum
+    rates = faiss.vector_to_array(sc.gemm_pruning_rates)
+    return dt, final_obj, sc.iteration_stats.size(), iter_sum, setup, rates
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    src = ap.add_mutually_exclusive_group()
+    src.add_argument("--fvecs", help="path to a .fvecs file")
+    ap.add_argument("--n", type=int, default=10000, help="synthetic n")
+    ap.add_argument("--d", type=int, default=128, help="synthetic d")
+    ap.add_argument("--k", type=int, default=64, help="number of centroids")
+    ap.add_argument("--niter", type=int, default=10, help="iterations")
+    ap.add_argument("--seed", type=int, default=42, help="RNG seed")
+    args = ap.parse_args()
+
+    if args.fvecs:
+        x = fvecs_read(args.fvecs)
+        n, d = x.shape
+        print("=== bench_super_kmeans ===")
+        print(
+            f"fvecs={args.fvecs} n={n} d={d} k={args.k} "
+            f"niter={args.niter} seed={args.seed}"
+        )
+    else:
+        n, d = args.n, args.d
+        x = gaussian_mixture(n, d, args.k, args.seed)
+        print("=== bench_super_kmeans ===")
+        print(
+            f"n={n} d={d} k={args.k} niter={args.niter} seed={args.seed}"
+        )
+
+    print("\n--- faiss.Clustering ---")
+    v_dt, v_obj, v_iters = run_vanilla(x, d, args.k, args.niter, args.seed)
+    print(f"  wall_time={v_dt:.3f}s final_obj={v_obj:g} iters={v_iters}")
+
+    print("\n--- faiss.SuperKMeans ---")
+    s_dt, s_obj, s_iters, iter_sum, setup, rates = run_super(
+        x, d, args.k, args.niter, args.seed
+    )
+    print(
+        f"  wall_time={s_dt:.3f}s final_obj={s_obj:g} iters={s_iters} "
+        f"iter_sum={iter_sum:.3f}s setup={setup:.3f}s"
+    )
+    print("  gemm_pruning_rates:", " ".join(f"{r:.3f}" for r in rates))
+    print()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -73,6 +73,7 @@ endif()
 set(FAISS_SRC
  AutoTune.cpp
  Clustering.cpp
+  SuperKMeans.cpp
  IVFlib.cpp
  Index.cpp
  Index2Layer.cpp
@@ -200,6 +201,7 @@ endif()
 set(FAISS_HEADERS
  AutoTune.h
  Clustering.h
+  SuperKMeans.h
  IVFlib.h
  Index.h
  Index2Layer.h
@@ -0,0 +1,656 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <algorithm>
+#include <cassert>
+#include <cinttypes>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <memory>
+#include <vector>
+
+#include <faiss/SuperKMeans.h>
+#include <faiss/VectorTransform.h>
+#include <faiss/impl/AdSampling.h>
+#include <faiss/impl/ClusteringHelpers.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/PdxLayout.h>
+#include <faiss/impl/simd_dispatch.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/simd_impl/super_kmeans_kernels.h>
+#include <faiss/utils/utils.h>
+
+#ifndef FINTEGER
+#define FINTEGER long
+#endif
+
+extern "C" {
+
+/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
+
+int sgemm_(
+        const char* transa,
+        const char* transb,
+        FINTEGER* m,
+        FINTEGER* n,
+        FINTEGER* k,
+        const float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* b,
+        FINTEGER* ldb,
+        float* beta,
+        float* c,
+        FINTEGER* ldc);
+}
+
+namespace faiss {
+
+namespace {
+
+struct TrainState {
+    /// Orthogonal rotation. Train in rotated space (X_tilde = X * R);
+    /// un-rotate centroids before return.
+    faiss::RandomRotationMatrix R;
+
+    std::vector<float> X_tilde; // (n, d) row-major
+    int n = 0;
+    std::vector<float> Y_tilde; // (k, d) row-major
+
+    std::vector<int> assignments;  // size n
+    std::vector<float> best_dists; // size n; tau per vector
+
+    /// ||X_tilde[i, 0:d_prime]||^2; recomputed when d_prime changes.
+    std::vector<float> x_norms_partial;
+
+    int d_prime = 0;
+
+    /// ADSampling threshold table; size d+1.
+    std::vector<float> ad_coeff;
+
+    /// PDX block layout for the trailing pruning sweep: block b covers
+    /// original dims [true_block_end[b] - block_dim[b], true_block_end[b]).
+    /// Recomputed when d_prime changes.
+    std::vector<int> block_dim;
+    std::vector<int> true_block_end;
+
+    /// Counter for the verbose-mode "low pruning" warning.
+    int low_pruning_streak = 0;
+    bool low_pruning_warning_printed = false;
+
+    explicit TrainState(int d) : R(d, d) {}
+};
+
+/// Rebuild state.block_dim and state.true_block_end from the current
+/// state.d_prime and pdx_block_size. Call after any change to d_prime.
+void rebuild_pdx_block_layout(int d, int pdx_block_size, TrainState& state) {
+    const int dp = state.d_prime;
+    const int d_trail = d - dp;
+    const int n_full_blocks = d_trail / pdx_block_size;
+    const int tail = d_trail % pdx_block_size;
+    const int n_blocks = n_full_blocks + (tail > 0 ? 1 : 0);
+    state.block_dim.assign(n_blocks, pdx_block_size);
+    state.true_block_end.resize(n_blocks);
+    if (n_blocks > 0) {
+        assert(!state.block_dim.empty());
+        assert(!state.true_block_end.empty());
+        for (int b = 0; b < n_full_blocks; ++b) {
+            state.true_block_end[b] = dp + (b + 1) * pdx_block_size;
+        }
+        if (tail > 0) {
+            state.block_dim[n_full_blocks] = tail;
+            state.true_block_end[n_full_blocks] = d;
+        }
+    }
+}
+
+struct IterScratch {
+    std::vector<float> partial_ip; // (bx_max, by_max) for the GEMM tile
+    std::vector<float> Y_pdx;      // PDX-laid-out trailing block
+    std::vector<float> Y_trail;    // row-major (k, d_trail) input to pdxify
+    std::vector<float> y_norms_partial; // ||Y_tilde[j, 0:dp]||^2
+    std::vector<int64_t> labels64;      // size n; widened state.assignments
+    int prev_d_trail = -1;
+};
+
+/// Iter 0: full GEMM via knn_L2sqr (vanilla Lloyd's). Fills
+/// state.assignments and state.best_dists. Returns objective.
+double run_iter0_full_gemm(int d, int k, TrainState& state) {
+    std::vector<int64_t> labels(state.n);
+    std::vector<float> distances(state.n);
+    knn_L2sqr(
+            state.X_tilde.data(),
+            state.Y_tilde.data(),
+            d,
+            state.n,
+            k,
+            /*k=*/1,
+            distances.data(),
+            labels.data(),
+            /*y_norm2=*/nullptr);
+
+    assert(!state.assignments.empty());
+    assert(!state.best_dists.empty());
+    double objective = 0.0;
+    for (int i = 0; i < state.n; ++i) {
+        state.assignments[i] = static_cast<int>(labels[i]);
+        state.best_dists[i] = distances[i];
+        objective += distances[i];
+    }
+    return objective;
+}
+
+/// Iter 1+: partial GEMM over [0, d_prime) + ADSampling progressive
+/// pruning over the PDX-laid-out trailing block. Updates
+/// state.assignments and state.best_dists. Writes total_pairs and
+/// pruned_at_gemm. Returns objective.
+double run_iter_pruned(
+        int d,
+        int k,
+        const SuperKMeansParameters& cp,
+        TrainState& state,
+        IterScratch& scratch,
+        int64_t& total_pairs,
+        int64_t& pruned_at_gemm) {
+    const int dp = state.d_prime;
+    assert(dp >= 1);
+    assert(!state.ad_coeff.empty());
+    assert(!scratch.partial_ip.empty());
+    assert(!scratch.y_norms_partial.empty());
+    const int d_trail = d - dp;
+    const int n_train = state.n;
+    assert(static_cast<int>(state.best_dists.size()) >= n_train);
+    assert(static_cast<int>(state.x_norms_partial.size()) >= n_train);
+    assert(static_cast<int>(state.assignments.size()) >= n_train);
+
+    if (d_trail != scratch.prev_d_trail) {
+        scratch.Y_pdx.resize(static_cast<size_t>(k) * d_trail);
+        scratch.Y_trail.resize(static_cast<size_t>(k) * d_trail);
+        scratch.prev_d_trail = d_trail;
+    }
+    for (int j = 0; j < k; ++j) {
+        std::memcpy(
+                scratch.Y_trail.data() + static_cast<size_t>(j) * d_trail,
+                state.Y_tilde.data() + static_cast<size_t>(j) * d + dp,
+                d_trail * sizeof(float));
+    }
+    detail::pdxify(
+            scratch.Y_trail.data(),
+            k,
+            d_trail,
+            cp.pdx_block_size,
+            scratch.Y_pdx.data());
+
+    detail::compute_partial_norms(
+            state.Y_tilde.data(), k, d, dp, scratch.y_norms_partial.data());
+
+    const int n_blocks = static_cast<int>(state.block_dim.size());
+
+    for (int xi = 0; xi < n_train; xi += cp.x_batch) {
+        const int bx = std::min(cp.x_batch, n_train - xi);
+
+        // Refresh tau: recompute full-d L2 distance to the previously
+        // assigned centroid. This is intentionally over all d dims (not
+        // just d_prime) because tau must be an exact distance for the
+        // chi-squared pruning bound to be valid. Cost is O(bx * d) per
+        // x-batch, amortized across the y-batch tiles that follow.
+#pragma omp parallel for
+        for (int i = 0; i < bx; ++i) {
+            const int j_prev = state.assignments[xi + i];
+            const float* xrow =
+                    state.X_tilde.data() + static_cast<size_t>(xi + i) * d;
+            const float* yrow =
+                    state.Y_tilde.data() + static_cast<size_t>(j_prev) * d;
+            float tau = 0.0f;
+            for (int m = 0; m < d; ++m) {
+                const float diff = xrow[m] - yrow[m];
+                tau += diff * diff;
+            }
+            state.best_dists[xi + i] = tau;
+        }
+
+        for (int yj = 0; yj < k; yj += cp.y_batch) {
+            const int by = std::min(cp.y_batch, k - yj);
+
+            // GEMM phase: column-major sgemm computes
+            //   partial_ip[i*by + j] = <X[xi+i, 0:dp], Y[yj+j, 0:dp]>.
+            {
+                FINTEGER M = by;
+                FINTEGER N_ = bx;
+                FINTEGER K_ = dp;
+                float alpha = 1.0f;
+                float beta = 0.0f;
+                FINTEGER lda_y = d;
+                FINTEGER lda_x = d;
+                FINTEGER ldc = by;
+                sgemm_("Transpose",
+                       "Not transpose",
+                       &M,
+                       &N_,
+                       &K_,
+                       &alpha,
+                       state.Y_tilde.data() + static_cast<size_t>(yj) * d,
+                       &lda_y,
+                       state.X_tilde.data() + static_cast<size_t>(xi) * d,
+                       &lda_x,
+                       &beta,
+                       scratch.partial_ip.data(),
+                       &ldc);
+            }
+
+            // One SIMD dispatch per (xi, yj) tile — block_l2<SL> below is
+            // a direct call (no per-call switch on SIMDConfig::level).
+            with_simd_level([&]<SIMDLevel SL>() {
+                [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk;
+                int64_t total_pairs_local = 0;
+                int64_t pruned_at_gemm_local = 0;
+#pragma omp parallel for schedule(dynamic, omp_chunk_local) \
+        reduction(+ : total_pairs_local) reduction(+ : pruned_at_gemm_local)
+                for (int i = 0; i < bx; ++i) {
+                    // tau is the best full-d distance found so far for this
+                    // point; tightened as closer centroids are found.
+                    float tau = state.best_dists[xi + i];
+                    int best_j = state.assignments[xi + i];
+                    const float xnp_i = state.x_norms_partial[xi + i];
+                    const float* xrow = state.X_tilde.data() +
+                            static_cast<size_t>(xi + i) * d;
+
+                    for (int j = 0; j < by; ++j) {
+                        ++total_pairs_local;
+
+                        // L2-from-IP; clamp to handle catastrophic
+                        // cancellation when the true distance is ~0.
+                        float pd = xnp_i + scratch.y_norms_partial[yj + j] -
+                                2.0f *
+                                        scratch.partial_ip
+                                                [static_cast<size_t>(i) * by +
+                                                 j];
+                        if (pd < 0.0f) {
+                            pd = 0.0f;
+                        }
+
+                        if (pd > state.ad_coeff[dp] * tau) {
+                            ++pruned_at_gemm_local;
+                            continue;
+                        }
+
+                        // double accumulator mitigates float drift over many
+                        // block additions.
+                        double dist = pd;
+                        bool keep = true;
+
+                        // Progressive pruning across PDX blocks. Per block:
+                        // stride = k * block_dim[b] floats, column-major
+                        // across centroids.
+                        size_t pdx_offset = 0;
+                        for (int b = 0; b < n_blocks; ++b) {
+                            const int n_in_block = state.block_dim.at(b);
+                            const int true_end = state.true_block_end.at(b);
+                            const float* xblk = xrow + (true_end - n_in_block);
+                            const float* yblk = scratch.Y_pdx.data() +
+                                    pdx_offset +
+                                    static_cast<size_t>(yj + j) * n_in_block;
+                            dist += faiss::detail::block_l2<SL>(
+                                    xblk, yblk, n_in_block);
+                            pdx_offset += static_cast<size_t>(k) * n_in_block;
+
+                            if (dist >
+                                static_cast<double>(state.ad_coeff[true_end]) *
+                                        tau) {
+                                keep = false;
+                                break;
+                            }
+                        }
+
+                        if (keep && dist < tau) {
+                            tau = static_cast<float>(dist);
+                            best_j = yj + j;
+                        }
+                    }
+
+                    state.best_dists[xi + i] = tau;
+                    state.assignments[xi + i] = best_j;
+                }
+                total_pairs += total_pairs_local;
+                pruned_at_gemm += pruned_at_gemm_local;
+            });
+        }
+    }
+
+    double objective = 0.0;
+    for (int i = 0; i < n_train; ++i) {
+        objective += state.best_dists[i];
+    }
+    return objective;
+}
+
+/// Post-iteration: update centroids and split empties. Returns nsplit.
+int update_centroids_and_split(
+        int d,
+        int k,
+        TrainState& state,
+        IterScratch& scratch,
+        std::vector<float>& hassign) {
+    std::fill(hassign.begin(), hassign.end(), 0.0f);
+    assert(!scratch.labels64.empty());
+    assert(!state.assignments.empty());
+    for (int i = 0; i < state.n; ++i) {
+        scratch.labels64[i] = static_cast<int64_t>(state.assignments[i]);
+    }
+    detail::compute_centroids(
+            d,
+            k,
+            state.n,
+            /*k_frozen=*/0,
+            reinterpret_cast<const uint8_t*>(state.X_tilde.data()),
+            /*codec=*/nullptr,
+            scratch.labels64.data(),
+            /*weights=*/nullptr,
+            hassign.data(),
+            state.Y_tilde.data());
+    if (state.n <= k) {
+        return 0;
+    }
+    return detail::split_clusters(
+            d,
+            k,
+            state.n,
+            /*k_frozen=*/0,
+            hassign.data(),
+            state.Y_tilde.data());
+}
+
+/// Stay-in-band controller: nudge state.d_prime based on observed
+/// pruning rate. Recomputes x_norms_partial if d_prime changed. Returns
+/// the observed pruning rate (0 when there were no pairs).
+float adapt_d_prime(
+        int d,
+        const SuperKMeansParameters& cp,
+        TrainState& state,
+        int64_t total_pairs,
+        int64_t pruned_at_gemm) {
+    if (total_pairs == 0) {
+        return 0.0f;
+    }
+    const float pruning_rate = static_cast<float>(pruned_at_gemm) /
+            static_cast<float>(total_pairs);
+    int new_dp = state.d_prime;
+    if (pruning_rate > cp.pruning_target_high) {
+        new_dp = static_cast<int>(
+                std::lround(state.d_prime * (1.0f - cp.d_prime_adjust)));
+    } else if (pruning_rate < cp.pruning_target_low) {
+        new_dp = static_cast<int>(
+                std::lround(state.d_prime * (1.0f + cp.d_prime_adjust)));
+    }
+    new_dp = std::max(cp.d_prime_min, new_dp);
+    new_dp = std::min(d / 2, new_dp);
+    if (new_dp != state.d_prime) {
+        state.d_prime = new_dp;
+        detail::compute_partial_norms(
+                state.X_tilde.data(),
+                state.n,
+                d,
+                state.d_prime,
+                state.x_norms_partial.data());
+        rebuild_pdx_block_layout(d, cp.pdx_block_size, state);
+    }
+    return pruning_rate;
+}
+
+/// Pre-loop setup: subsample, rotate, Forgy init, build ADSampling table,
+/// allocate scratch. Returned `sampled_x_owner` keeps the subsampled buffer
+/// alive when subsampling occurred (otherwise empty).
+std::unique_ptr<uint8_t[]> setup_train_state(
+        TrainState& state,
+        IterScratch& scratch,
+        std::vector<float>& hassign,
+        const SuperKMeansParameters& cp,
+        int d,
+        int k,
+        idx_t n,
+        const float* x) {
+    const size_t line_size = sizeof(float) * static_cast<size_t>(d);
+    idx_t nx = n;
+    const uint8_t* x_bytes = reinterpret_cast<const uint8_t*>(x);
+    std::unique_ptr<uint8_t[]> sampled_x_owner;
+    if (static_cast<size_t>(nx) >
+        static_cast<size_t>(k) * cp.max_points_per_centroid) {
+        Clustering tmp_clus(d, k, cp);
+        uint8_t* x_new = nullptr;
+        float* w_unused = nullptr;
+        nx = detail::subsample_training_set(
+                tmp_clus,
+                nx,
+                x_bytes,
+                line_size,
+                /*weights=*/nullptr,
+                &x_new,
+                &w_unused);
+        FAISS_ASSERT(x_new != nullptr);
+        sampled_x_owner.reset(x_new);
+        x_bytes = x_new;
+    }
+    const float* x_sampled = reinterpret_cast<const float*>(x_bytes);
+
+    FAISS_THROW_IF_NOT_MSG(
+            nx <= static_cast<idx_t>(std::numeric_limits<int>::max()),
+            "SuperKMeans: training set size exceeds INT_MAX after sampling");
+    state.n = static_cast<int>(nx);
+
+    state.R.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());
+
+    // Forgy init: pick k random rows from the rotated pool as initial
+    // centroids. These remain in rotated space; un-rotation happens
+    // after the iteration loop.
+    state.Y_tilde.resize(static_cast<size_t>(k) * d);
+    {
+        std::vector<int> perm(state.n);
+        rand_perm(perm.data(), state.n, static_cast<int64_t>(cp.seed) + 1);
+        for (int j = 0; j < k; ++j) {
+            std::memcpy(
+                    state.Y_tilde.data() + static_cast<size_t>(j) * d,
+                    state.X_tilde.data() + static_cast<size_t>(perm[j]) * d,
+                    sizeof(float) * d);
+        }
+    }
+
+    state.d_prime =
+            std::max(cp.d_prime_min, static_cast<int>(d * cp.d_prime_fraction));
+    state.d_prime = std::min(state.d_prime, d / 2);
+    rebuild_pdx_block_layout(d, cp.pdx_block_size, state);
+
+    // Iter 1+ uses L2-from-IP only over [0, d_prime), so full ||X[i]||^2 is
+    // never read; iter 0 routes through knn_L2sqr which carries its own.
+    state.x_norms_partial.resize(state.n);
+    detail::compute_partial_norms(
+            state.X_tilde.data(),
+            state.n,
+            d,
+            state.d_prime,
+            state.x_norms_partial.data());
+
+    const double epsilon = static_cast<double>(cp.ad_epsilon_factor) / d;
+    state.ad_coeff = detail::precompute_ad_thresholds(d, epsilon);
+    FAISS_ASSERT_MSG(
+            state.ad_coeff.size() == static_cast<size_t>(d + 1),
+            "ad_coeff size mismatch");
+
+    state.assignments.assign(state.n, 0);
+    state.best_dists.assign(state.n, std::numeric_limits<float>::max());
+
+    hassign.assign(k, 0.0f);
+
+    const int by_max = std::min(cp.y_batch, k);
+    const int bx_max = std::min(cp.x_batch, state.n);
+    scratch.partial_ip.resize(static_cast<size_t>(bx_max) * by_max);
+    scratch.y_norms_partial.resize(k);
+    scratch.labels64.resize(state.n);
+
+    return sampled_x_owner;
+}
+
+/// Un-rotate centroids into output buffer. R orthogonal, so
+/// reverse_transform applies R^T = R^-1.
+void untransform_centroids(
+        std::vector<float>& centroids,
+        const RandomRotationMatrix& R,
+        int d,
+        int k,
+        const float* Y_tilde) {
+    centroids.resize(static_cast<size_t>(k) * d);
+    R.reverse_transform(k, Y_tilde, centroids.data());
+}
+
+} // namespace
+
+SuperKMeans::SuperKMeans(int d, int k, const SuperKMeansParameters& cp_in)
+        : cp(cp_in), d(d), k(k) {
+    FAISS_THROW_IF_NOT_MSG(d > 0, "SuperKMeans: d must be positive");
+    FAISS_THROW_IF_NOT_MSG(k > 0, "SuperKMeans: k must be positive");
+    FAISS_THROW_IF_NOT_MSG(
+            cp.d_prime_fraction > 0.0f && cp.d_prime_fraction <= 1.0f,
+            "SuperKMeans: d_prime_fraction must be in (0, 1]");
+    FAISS_THROW_IF_NOT_MSG(
+            cp.d_prime_adjust >= 0.0f && cp.d_prime_adjust < 1.0f,
+            "SuperKMeans: d_prime_adjust must be in [0, 1)");
+    // d >= 2 * d_prime_min keeps d_prime in the chi-squared validity
+    // range after both clamping steps (floor at d_prime_min, ceiling
+    // at d/2). See AdSampling.h on the p >= 16 contract.
+    FAISS_THROW_IF_NOT_FMT(
+            d >= 2 * cp.d_prime_min,
+            "SuperKMeans: d (%d) must be >= 2 * d_prime_min (%d)",
+            d,
+            cp.d_prime_min);
+    FAISS_THROW_IF_NOT_MSG(
+            cp.d_prime_min >= 16,
+            "SuperKMeans: d_prime_min must be >= 16 (chi-squared validity floor)");
+    FAISS_THROW_IF_NOT_MSG(
+            cp.pdx_block_size > 0, "SuperKMeans: pdx_block_size must be > 0");
+    FAISS_THROW_IF_NOT_MSG(cp.x_batch > 0, "SuperKMeans: x_batch must be > 0");
+    FAISS_THROW_IF_NOT_MSG(cp.y_batch > 0, "SuperKMeans: y_batch must be > 0");
+    FAISS_THROW_IF_NOT_MSG(
+            cp.pruning_target_low > 0.0f &&
+                    cp.pruning_target_low <= cp.pruning_target_high &&
+                    cp.pruning_target_high < 1.0f,
+            "SuperKMeans: require 0 < pruning_target_low <= pruning_target_high < 1");
+    // epsilon = ad_epsilon_factor / d is the chi-squared significance
+    // level. precompute_ad_thresholds requires epsilon in (0, 1).
+    FAISS_THROW_IF_NOT_MSG(
+            cp.ad_epsilon_factor > 0.0f &&
+                    cp.ad_epsilon_factor < static_cast<float>(d),
+            "SuperKMeans: ad_epsilon_factor must be in (0, d) "
+            "so epsilon = factor/d is in (0,1)");
+    FAISS_THROW_IF_NOT_MSG(
+            cp.omp_chunk > 0, "SuperKMeans: omp_chunk must be > 0");
+}
+
+void SuperKMeans::train(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT_MSG(n > 0, "SuperKMeans: n must be positive");
+    FAISS_THROW_IF_NOT_MSG(x != nullptr, "SuperKMeans: x must not be null");
+    FAISS_THROW_IF_NOT_MSG(
+            n >= static_cast<idx_t>(k), "SuperKMeans: n must be >= k");
+    if (cp.check_input_data_for_NaNs) {
+        for (size_t i = 0; i < static_cast<size_t>(n) * d; i++) {
+            FAISS_THROW_IF_NOT_MSG(
+                    std::isfinite(x[i]),
+                    "SuperKMeans: input contains NaN's or Inf's");
+        }
+    }
+    if (cp.verbose && n < static_cast<idx_t>(k) * cp.min_points_per_centroid) {
+        printf("WARNING: clustering %" PRId64
+               " points to %d centroids: please provide at least "
+               "%" PRId64 " training points\n",
+               n,
+               k,
+               static_cast<idx_t>(k) * cp.min_points_per_centroid);
+    }
+
+    TrainState state(d);
+    IterScratch scratch;
+    std::vector<float> hassign;
+    [[maybe_unused]] auto sampled_x_owner =
+            setup_train_state(state, scratch, hassign, cp, d, k, n, x);
+
+    iteration_stats.clear();
+    iteration_stats.reserve(cp.niter);
+    gemm_pruning_rates.clear();
+    gemm_pruning_rates.reserve(cp.niter);
+
+    const double t_train_start = getmillisecs();
+
+    for (int iter = 0; iter < cp.niter; ++iter) {
+        const double t_iter_start = getmillisecs();
+        double objective = 0.0;
+        int64_t total_pairs = 0;
+        int64_t pruned_at_gemm = 0;
+
+        if (iter == 0) {
+            objective = run_iter0_full_gemm(d, k, state);
+        } else {
+            objective = run_iter_pruned(
+                    d, k, cp, state, scratch, total_pairs, pruned_at_gemm);
+        }
+
+        const int nsplit =
+                update_centroids_and_split(d, k, state, scratch, hassign);
+        const float pruning_rate = (iter == 0)
+                ? 0.0f
+                : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm);
+
+        ClusteringIterationStats stat{};
+        stat.obj = static_cast<float>(objective);
+        stat.time = (getmillisecs() - t_iter_start) / 1000.0;
+        stat.time_search = stat.time;
+        stat.imbalance_factor = std::numeric_limits<double>::quiet_NaN();
+        stat.nsplit = nsplit;
+        iteration_stats.push_back(stat);
+        gemm_pruning_rates.push_back(pruning_rate);
+
+        if (iter > 0) {
+            if (pruning_rate < 0.85f) {
+                state.low_pruning_streak++;
+            } else {
+                state.low_pruning_streak = 0;
+            }
+            if (cp.verbose && state.low_pruning_streak >= 3 &&
+                !state.low_pruning_warning_printed) {
+                fprintf(stderr,
+                        "WARNING: SuperKMeans steady-state pruning < 0.85 for 3+ iters "
+                        "(current=%.2f). Data may not be a good fit for ADSampling; "
+                        "consider falling back to faiss::Clustering.\n",
+                        pruning_rate);
+                state.low_pruning_warning_printed = true;
+            }
+        }
+
+        if (cp.verbose) {
+            printf("  Iter %d: obj=%g time=%.3fs prune=%.4f dp=%d nsplit=%d\n",
+                   iter,
+                   stat.obj,
+                   stat.time,
+                   pruning_rate,
+                   state.d_prime,
+                   nsplit);
+        }
+    }
+
+    if (cp.verbose) {
+        printf("Total training time: %.3fs\n",
+               (getmillisecs() - t_train_start) / 1000.0);
+    }
+
+    untransform_centroids(centroids, state.R, d, k, state.Y_tilde.data());
+}
+
+} // namespace faiss
@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// SuperKMeans — a faster k-means alternative to faiss::Clustering for IVF
+// training and other large-k workloads.
+//
+// Based on:
+//   Kuffo, L., Hepkema, S., & Boncz, P. (2026).
+//   "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,
+// binary data (use IndexBinaryIVF), or near-unit-sphere embeddings with
+// k < 4096 (chi-squared assumption breaks down).
+//
+// `cp`, `d`, `k` are public mutable; mutating them post-construction bypasses
+// validation and may produce undefined behavior.
+
+#pragma once
+
+#include <vector>
+
+#include <faiss/Clustering.h>
+#include <faiss/Index.h>
+
+namespace faiss {
+
+struct SuperKMeansParameters : public ClusteringParameters {
+    /// Initial d_prime as a fraction of d (GEMM/PRUNING split).
+    float d_prime_fraction = 0.125f;
+
+    /// Matches block_l2 SIMD width.
+    int pdx_block_size = 64;
+
+    /// ADSampling significance: epsilon = ad_epsilon_factor / d.
+    float ad_epsilon_factor = 1.0f;
+
+    /// Adaptive d_prime stay-in-band controller. Per iteration:
+    ///   pruning > high → shrink d_prime  (over-pruning, cheaper threshold)
+    ///   pruning < low  → grow   d_prime  (under-pruning, more GEMM work)
+    ///   else: hold.
+    float pruning_target_low = 0.95f;
+    float pruning_target_high = 0.97f;
+
+    /// Relative step size for d_prime adjustments.
+    float d_prime_adjust = 0.20f;
+
+    /// Floor on d_prime; below this the chi-squared bound is unvalidated.
+    int d_prime_min = 16;
+
+    int x_batch = 4096;
+    int y_batch = 1024;
+
+    /// OpenMP dynamic-schedule chunk size for the pruning loop.
+    int omp_chunk = 8;
+};
+
+/** Drop-in faster k-means: same interface as faiss::Clustering. Per iteration:
+ *   iter 0:           full GEMM over all d dims (vanilla Lloyd's).
+ *   iter 1..niter-1:  GEMM over front d_prime dims, then ADSampling
+ *                     progressive pruning over PDX-laid trailing dims.
+ *
+ * Trains in a randomly-rotated space; centroids are un-rotated before return.
+ */
+struct SuperKMeans {
+    SuperKMeansParameters cp;
+    int d;
+    int k;
+
+    /// Output: k * d floats, row-major, un-rotated.
+    std::vector<float> centroids;
+
+    /// Per-iter stats. `obj`, `time`, and `nsplit` are populated faithfully.
+    /// `time_search` mirrors `time` (phases not timed separately) and
+    /// `imbalance_factor` is set to NaN (not computed).
+    std::vector<ClusteringIterationStats> iteration_stats;
+
+    /// Per-iter fraction of (vector, centroid) pairs pruned at the d_prime
+    /// GEMM-boundary chi-squared check. Per-PDX-block early-exits are NOT
+    /// counted. iter 0 uses full GEMM, so gemm_pruning_rates[0] == 0.0f.
+    std::vector<float> gemm_pruning_rates;
+
+    SuperKMeans(int d, int k, const SuperKMeansParameters& cp = {});
+
+    /// Train on `n` row-major vectors of dimension `d`. Honors the applicable
+    /// ClusteringParameters fields (niter, seed, verbose,
+    /// max_points_per_centroid, use_faster_subsampling,
+    /// check_input_data_for_NaNs).
+    void train(idx_t n, const float* x);
+};
+
+} // namespace faiss
@@ -78,7 +78,7 @@ idx_t subsample_training_set(
            "subsample_training_set: perm size %zu < required nx %" PRId64,
            perm.size(),
            nx);
-    assert(!perm.empty()); // pattern for clang-tidy LocalUncheckedArrayBounds
+    assert(!perm.empty());

    uint8_t* x_new = new uint8_t[nx * line_size];
    *x_out = x_new;
@@ -116,15 +116,6 @@ void compute_centroids(

    size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);

-    // NOTE: FAISS_THROW_IF_NOT inside this OMP region is formally UB per the
-    // OpenMP spec (exceptions cannot propagate out of a parallel region).
-    // In practice, most runtimes terminate the process, which is acceptable
-    // for a corrupted-assignment invariant violation. Inherited from
-    // Clustering.cpp; not worth refactoring without a broader FAISS effort.
-    //
-    // Thread safety of hassign[ci]: each thread owns the exclusive slice
-    // [c0, c1) and only writes hassign[ci] when ci falls in that slice
-    // (same guard as centroids[ci*d]). No two threads share an index.
 #pragma omp parallel
    {
        int nt = omp_get_num_threads();
@@ -22,8 +22,8 @@ from faiss.gpu_wrappers import *
 from faiss.array_conversions import *
 from faiss.extra_wrappers import kmin, kmax, pairwise_distances, rand, randint, \
    lrand, randn, rand_smooth_vectors, eval_intersection, normalize_L2, \
-    ResultHeap, knn, Kmeans, checksum, matrix_bucket_sort_inplace, bucket_sort, \
-    merge_knn_results, MapInt64ToInt64, knn_hamming, \
+    ResultHeap, knn, Kmeans, SuperKmeans, checksum, matrix_bucket_sort_inplace, \
+    bucket_sort, merge_knn_results, MapInt64ToInt64, knn_hamming, \
    pack_bitstrings, unpack_bitstrings


@@ -34,6 +34,7 @@ __version__ = "%d.%d.%d" % (FAISS_VERSION_MAJOR,
 logger = logging.getLogger(__name__)

 class_wrappers.handle_Clustering(Clustering)
+class_wrappers.handle_SuperKMeans(SuperKMeans)
 class_wrappers.handle_Clustering1D(Clustering1D)
 class_wrappers.handle_MatrixStats(MatrixStats)
 class_wrappers.handle_IOWriter(IOWriter)
@@ -126,6 +126,24 @@ def handle_Clustering(the_class):
    replace_method(the_class, 'train_encoded', replacement_train_encoded)


+def handle_SuperKMeans(the_class):
+
+    def replacement_train(self, x):
+        """Perform SuperKMeans clustering on a set of vectors.
+
+        Parameters
+        ----------
+        x : array_like
+            Training vectors, shape (n, self.d). `dtype` must be float32.
+        """
+        n, d = x.shape
+        assert d == self.d
+        x = np.ascontiguousarray(x, dtype='float32')
+        self.train_c(n, swig_ptr(x))
+
+    replace_method(the_class, 'train', replacement_train)
+
+
 def handle_Clustering1D(the_class):

    def replacement_train_exact(self, x):
@@ -602,6 +602,56 @@ class Kmeans:
        return D.ravel(), I.ravel()


+class SuperKmeans(Kmeans):
+    """Drop-in replacement for `Kmeans` that runs `SuperKMeans` (ADSampling +
+    PDX progressive pruning) instead of `Clustering`. Same `centroids`, `obj`,
+    `iteration_stats`, and `assign()` surface; additionally exposes
+    `gemm_pruning_rates`.
+
+    kwargs are forwarded to `SuperKMeansParameters`. Fields not present on it
+    (e.g. `spherical`, `int_centroids`, `nredo`, `frozen_centroids`,
+    `init_method`, `update_index`, `early_stop_threshold`,
+    `progressive_dim_steps`, `gpu`) raise `AttributeError`.
+    """
+
+    def __init__(self, d, k, **kwargs):
+        self.d = d
+        self.reset(k)
+        self.gpu = False
+        self.cp = SuperKMeansParameters()
+        for key, v in kwargs.items():
+            getattr(self.cp, key)
+            setattr(self.cp, key, v)
+        self.set_index()
+
+    def set_index(self):
+        self.index = IndexFlatL2(self.d)
+
+    def train(self, x, weights=None, init_centroids=None):
+        assert weights is None, "SuperKmeans does not support weights"
+        assert init_centroids is None, \
+            "SuperKmeans does not support init_centroids"
+        x = np.ascontiguousarray(x, dtype='float32')
+        n, d = x.shape
+        assert d == self.d
+
+        sc = SuperKMeans(d, self.k, self.cp)
+        sc.train(x)
+
+        centroids = faiss.vector_to_array(sc.centroids)
+        self.centroids = centroids.reshape(self.k, d)
+        stats = sc.iteration_stats
+        stats = [stats.at(i) for i in range(stats.size())]
+        self.obj = np.array([st.obj for st in stats])
+        stat_fields = 'obj time time_search imbalance_factor nsplit'.split()
+        self.iteration_stats = [
+            {field: getattr(st, field) for field in stat_fields}
+            for st in stats
+        ]
+        self.gemm_pruning_rates = faiss.vector_to_array(sc.gemm_pruning_rates)
+        return self.obj[-1] if self.obj.size > 0 else 0.0
+
+
 ###########################################
 # Packing and unpacking bitstrings
 ###########################################
@@ -188,6 +188,7 @@ typedef uint64_t size_t;
 #endif // !_MSC_VER

 #include <faiss/Clustering.h>
+#include <faiss/SuperKMeans.h>
 #include <faiss/impl/ClusteringInitialization.h>

 #include <faiss/utils/hamming.h>
@@ -521,6 +522,7 @@ void gpu_sync_all_devices()
 %include  <faiss/IndexFlat.h>
 %include  <faiss/impl/ClusteringInitialization.h>
 %include  <faiss/Clustering.h>
+%include  <faiss/SuperKMeans.h>

 %include  <faiss/utils/extra_distances.h>

@@ -13,7 +13,7 @@
 // scalar primary template; adding NEON/SVE means just adding a new
 // specialization file alongside the AVX ones.
 //
-// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented in v1.
+// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented yet.
 // aarch64 falls through to the scalar primary template. Validating SVE requires
 // a Graviton-class host; deferred to a focused follow-up.

@@ -27,15 +27,17 @@ inline float block_l2(const float* x, const float* y, int n) {
    return s;
 }

-// Mirrors the impl-file guards; turns off-arch direct calls into a
-// compile-time error instead of a linker error.
-#if defined(__x86_64__)
+// COMPILE_SIMD_* is a build-system define (link-time promise that the
+// specialization will be available). Mirrors the impl-file guards.
+#ifdef COMPILE_SIMD_AVX2
 template <>
 float block_l2<SIMDLevel::AVX2>(const float* x, const float* y, int n);
+#endif

+#ifdef COMPILE_SIMD_AVX512
 template <>
 float block_l2<SIMDLevel::AVX512>(const float* x, const float* y, int n);
-#endif // __x86_64__
+#endif

 } // namespace detail
 } // namespace faiss
@@ -5,7 +5,7 @@
 * LICENSE file in the root directory of this source tree.
 */

-#if defined(__x86_64__)
+#ifdef COMPILE_SIMD_AVX2

 #include <faiss/utils/simd_impl/super_kmeans_kernels.h>

@@ -54,4 +54,4 @@ float block_l2<SIMDLevel::AVX2>(const float* x, const float* y, int n) {
 } // namespace detail
 } // namespace faiss

-#endif // __x86_64__
+#endif // COMPILE_SIMD_AVX2
@@ -5,7 +5,7 @@
 * LICENSE file in the root directory of this source tree.
 */

-#if defined(__x86_64__) && defined(__AVX512F__)
+#ifdef COMPILE_SIMD_AVX512

 #include <faiss/utils/simd_impl/super_kmeans_kernels.h>

@@ -42,4 +42,4 @@ float block_l2<SIMDLevel::AVX512>(const float* x, const float* y, int n) {
 } // namespace detail
 } // namespace faiss

-#endif // __x86_64__ && __AVX512F__
+#endif // COMPILE_SIMD_AVX512
@@ -0,0 +1,188 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import faiss
+import numpy as np
+from faiss.contrib.datasets import SyntheticDataset
+
+
+class SuperKMeansTest(unittest.TestCase):
+    def test_train_rejects_too_few_points(self):
+        sc = faiss.SuperKMeans(128, 16)
+        x = SyntheticDataset(128, 10, 0, 0).get_train()  # n=10 < k=16
+        with self.assertRaises(RuntimeError):
+            sc.train(x)
+
+    def test_objective_decreases_monotonically(self):
+        d, k, n = 32, 8, 1000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        p = faiss.SuperKMeansParameters()
+        p.seed = 42
+        p.niter = 10
+        sc = faiss.SuperKMeans(d, k, p)
+        sc.train(x)
+
+        stats = sc.iteration_stats
+        objs = [stats.at(i).obj for i in range(stats.size())]
+        for i in range(1, len(objs)):
+            self.assertLessEqual(objs[i], objs[i - 1] * 1.001)
+
+    def test_pruned_iter_assignments_close_to_vanilla(self):
+        d, k, n = 64, 16, 2000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        p = faiss.SuperKMeansParameters()
+        p.seed = 42
+        p.niter = 10
+        sc = faiss.SuperKMeans(d, k, p)
+        sc.train(x)
+
+        quant = faiss.IndexFlatL2(d)
+        vanilla = faiss.Clustering(d, k)
+        vanilla.seed = 42
+        vanilla.niter = 10
+        vanilla.train(x, quant)
+
+        sc_final = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj
+        v_stats = vanilla.iteration_stats
+        v_final = v_stats.at(v_stats.size() - 1).obj
+        self.assertLess(abs(sc_final - v_final) / v_final, 0.05)
+
+    def test_pruning_rate_in_expected_range(self):
+        # The 10-d intrinsic manifold of SyntheticDataset gives ADSampling
+        # the assigned-vs-other distance gap it needs to prune effectively.
+        d, k, n = 64, 64, 5000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        p = faiss.SuperKMeansParameters()
+        p.seed = 42
+        p.niter = 5
+        sc = faiss.SuperKMeans(d, k, p)
+        sc.train(x)
+
+        rates = faiss.vector_to_array(sc.gemm_pruning_rates)
+        self.assertEqual(rates.shape, (5,))
+        self.assertEqual(rates[0], 0.0)  # iter 0: no pruning
+        # Lower bound 0.30: looser than the 0.93 design target because
+        # synthetic d=64 data lacks the cluster gap real embeddings have at
+        # d=768; the assertion still confirms meaningful pruning occurs.
+        for i in range(2, 5):
+            self.assertGreaterEqual(rates[i], 0.30)
+
+    def test_pdx_offset_correct_for_unaligned_y_batch(self):
+        # PDX block layout is sized for the full k centroids, not per y-batch
+        # tile. The y-batch tile boundary must be invisible to PDX-offset
+        # arithmetic. k=70 with y_batch=32 forces a 6-wide tail tile; assert
+        # the resulting centroids are identical to running with y_batch=70
+        # (single tile, no tail).
+        d, k, n = 64, 70, 5000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        def train_with_y_batch(y_batch):
+            p = faiss.SuperKMeansParameters()
+            p.seed = 1234
+            p.niter = 5
+            p.y_batch = y_batch
+            sc = faiss.SuperKMeans(d, k, p)
+            sc.train(x)
+            return faiss.vector_to_array(sc.centroids)
+
+        tiled = train_with_y_batch(32)  # k % y_batch = 6
+        single = train_with_y_batch(k)  # k % y_batch = 0
+        np.testing.assert_array_equal(tiled, single)
+
+    def test_determinism(self):
+        # Bit-exact reproducibility requires MKL_CBWR=COMPATIBLE and a fixed
+        # OMP thread count. The BUCK target sets these in `env`; for OSS
+        # ctest the user must export them before running.
+        d, k, n = 128, 64, 5000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        p = faiss.SuperKMeansParameters()
+        p.niter = 5
+        p.seed = 1234
+
+        sc1 = faiss.SuperKMeans(d, k, p)
+        sc1.train(x)
+        sc2 = faiss.SuperKMeans(d, k, p)
+        sc2.train(x)
+
+        np.testing.assert_array_equal(
+            faiss.vector_to_array(sc1.centroids),
+            faiss.vector_to_array(sc2.centroids),
+        )
+        np.testing.assert_array_equal(
+            faiss.vector_to_array(sc1.gemm_pruning_rates),
+            faiss.vector_to_array(sc2.gemm_pruning_rates),
+        )
+
+    def test_rejects_nan_input(self):
+        # check_input_data_for_NaNs defaults to True.
+        sc = faiss.SuperKMeans(128, 64)
+        x = SyntheticDataset(128, 5000, 0, 0).get_train()
+        x.flat[42] = float("nan")
+        with self.assertRaises(RuntimeError):
+            sc.train(x)
+
+    def test_cluster_splitting_with_many_duplicates(self):
+        # 80% duplicates force empty clusters after the first assignment;
+        # split_clusters must fire and produce distinct centroids.
+        d, k, n = 64, 16, 1000
+        n_dup = int(n * 0.8)
+
+        ds = SyntheticDataset(d, (n - n_dup) + 1, 0, 0).get_train()
+        x = np.empty((n, d), dtype="float32")
+        x[:n_dup] = ds[0]
+        x[n_dup:] = ds[1:]
+
+        p = faiss.SuperKMeansParameters()
+        p.seed = 42
+        p.niter = 5
+        sc = faiss.SuperKMeans(d, k, p)
+        sc.train(x)
+
+        stats = sc.iteration_stats
+        total_splits = sum(stats.at(i).nsplit for i in range(stats.size()))
+        self.assertGreater(total_splits, 0)
+
+        centroids = faiss.vector_to_array(sc.centroids).reshape(k, d)
+        unique_centroids = {centroids[j].tobytes() for j in range(k)}
+        self.assertEqual(len(unique_centroids), k)
+
+
+class SuperKmeansWrapperTest(unittest.TestCase):
+    def test_train_populates_kmeans_surface(self):
+        d, k, n = 64, 16, 2000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        km = faiss.SuperKmeans(d, k, niter=5, seed=42)
+        final_obj = km.train(x)
+
+        self.assertEqual(km.centroids.shape, (k, d))
+        self.assertEqual(km.obj.shape, (5,))
+        self.assertEqual(len(km.iteration_stats), 5)
+        self.assertEqual(final_obj, km.obj[-1])
+        self.assertEqual(km.gemm_pruning_rates.shape, (5,))
+
+        D, I = km.assign(x)
+        self.assertEqual(I.shape, (n,))
+        self.assertEqual(D.shape, (n,))
+        self.assertTrue((I >= 0).all() and (I < k).all())
+
+    def test_final_obj_matches_vanilla_kmeans(self):
+        d, k, n = 64, 16, 2000
+        x = SyntheticDataset(d, n, 0, 0).get_train()
+
+        vanilla = faiss.Kmeans(d, k, niter=10, seed=42)
+        vanilla.train(x)
+
+        sk = faiss.SuperKmeans(d, k, niter=10, seed=42)
+        sk.train(x)
+
+        rel_diff = abs(sk.obj[-1] - vanilla.obj[-1]) / vanilla.obj[-1]
+        self.assertLess(rel_diff, 0.05)
@@ -171,35 +171,6 @@ TEST(BlockL2, ScalarMatchesReference) {
    }
 }

-TEST(BlockL2, Avx512MatchesScalar) {
-#if defined(__x86_64__) && defined(__AVX512F__)
-    constexpr int N = 64;
-    std::mt19937 rng(43);
-    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
-    std::vector<float> x(N), y(N);
-    for (int m = 0; m < N; ++m) {
-        x[m] = dist(rng);
-        y[m] = dist(rng);
-    }
-
-    const float scalar = faiss::detail::block_l2<faiss::SIMDLevel::NONE>(
-            x.data(), y.data(), N);
-    const float avx512 = faiss::detail::block_l2<faiss::SIMDLevel::AVX512>(
-            x.data(), y.data(), N);
-    EXPECT_NEAR(avx512, scalar, 1e-4f);
-
-    for (int n = 1; n < N; ++n) {
-        const float s = faiss::detail::block_l2<faiss::SIMDLevel::NONE>(
-                x.data(), y.data(), n);
-        const float a = faiss::detail::block_l2<faiss::SIMDLevel::AVX512>(
-                x.data(), y.data(), n);
-        EXPECT_NEAR(a, s, 1e-4f) << "n=" << n;
-    }
-#else
-    GTEST_SKIP() << "AVX-512 test requires __AVX512F__";
-#endif
-}
-
 TEST(BlockL2, DispatchMatchesScalar) {
    constexpr int N = 64;
    std::mt19937 rng(44);