enhance: default SuperKMeans coarse quantizer training for SCANN - #1783
Open
marcelo-cjl wants to merge 1 commit into
Open
enhance: default SuperKMeans coarse quantizer training for SCANN#1783marcelo-cjl wants to merge 1 commit into
marcelo-cjl wants to merge 1 commit into
Conversation
Collaborator
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: marcelo-cjl The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
|
@marcelo-cjl 🔍 Important: PR Classification Needed! For efficient project management and a seamless review process, it's essential to classify your PR correctly. Here's how:
For any PR outside the kind/improvement category, ensure you link to the associated issue using the format: “issue: #”. Thanks for your efforts and contribution to the community!. |
marcelo-cjl
changed the title
marcelo-cjl
force-pushed
the
enhance-superkmeans-scann
branch
from
August 20, 2026 03:09
0953e9a to
d37fc0a
Compare
marcelo-cjl
force-pushed
the
enhance-superkmeans-scann
branch
from
August 20, 2026 03:37
d37fc0a to
a1a0446
Compare
marcelo-cjl
changed the title
marcelo-cjl
force-pushed
the
enhance-superkmeans-scann
branch
from
August 20, 2026 03:51
a1a0446 to
b1a0e2b
Compare
Collaborator
Author
|
@foxspy @alexanderguzhva Please review this when you have some free time. |
marcelo-cjl
force-pushed
the
enhance-superkmeans-scann
branch
from
August 20, 2026 05:57
b1a0e2b to
27ed806
Compare
Enable SuperKMeans (super fast k-means) for SCANN coarse quantizer training by default. SuperKMeans is a faster k-means variant using ADSampling + PDX progressive pruning; on inner-product data its centroids are unit-normalized so L2 assignment is equivalent to IP argmax. Vendored faiss changes (tracked in thirdparty/faiss, mirrored to the upstream faiss PR): - SuperKMeans: spherical (inner-product) support via cp.spherical. Centroid renormalization makes L2 assignment equivalent to IP. Power-of-two d uses the fast HadamardRotation; L2 keeps RandomRotationMatrix. - Fix HadamardRotation::reverse_transform scaling: the inverse applied p*sqrt(p) instead of 1/(p*sqrt(p)), blowing up centroid norms by p^3 under spherical clustering. Verified with a roundtrip test. - Fix low-dimensional L2sqr nearest SVE fast path data layout. The previous per-lane/component layout loaded y row-major contiguously instead of strided across centroids, producing wrong distances for all PQ-encoded indexes (IVF_PQ/SCANN/HNSW_PQ). Rewritten to load each centroid's D components with a single svld1 and reduce with svaddv; verified against the scalar reference on real hardware. - ClusteringParameters::use_super_kmeans (default false) lets Level1Quantizer::train_q1 route coarse quantizer training through SuperKMeans. - ProductQuantizer::compute_code uses AVAILABLE_SIMD_LEVELS_A1 so PQ encoding reaches the ARM_SVE nearest kernel. - block_l2<ARM_SVE> SVE kernel for the SuperKMeans pruning loop. Knowhere changes: - ScannConfig::use_super_kmeans defaults to true, but only takes effect when SuperKMeans is applicable (d >= 2*d_prime_min and nlist >= 1024); otherwise SCANN falls back to Clustering so low-dimensional or small datasets keep working. - IvfConfig::use_super_kmeans plumbing into all IVF-family build paths. - Tests: tests/ut/test_scann_superkmeans.cc (SCANN default SuperKMeans recall-equivalent to Clustering, low-dim fallback), tests/ut/test_cluster.cc (spherical SuperKMeans unit centroids + objective), vendored test_distances_dispatch.cpp (low-dim nearest across SIMD levels). Benchmark (qwen 4096-dim, IP, nlist=1024, sub_dim=4, 1 thread, 50k rows): - SCANN_DVR Train 67.73s -> 24.36s, Add 10.86s -> 7.17s, Build 78.59s -> 31.52s (-60%) - recall@10 unchanged vs Clustering baseline (diff <= 0.0002) Signed-off-by: marcelo-cjl <marcelo.chen@zilliz.com>
marcelo-cjl
force-pushed
the
enhance-superkmeans-scann
branch
from
August 20, 2026 11:59
27ed806 to
98834e6
Compare
| // large nlist (k >= 1024). Outside that range fall back to Clustering | ||
| // instead of failing the build or degrading recall. | ||
| bool use_super_kmeans = scann_cfg.use_super_kmeans.value(); | ||
| if (use_super_kmeans && (dim < 2 * faiss::SuperKMeansParameters{}.d_prime_min || nlist < 1024)) { |
Collaborator
There was a problem hiding this comment.
This is questionable, because I constantly use SuperKMeans for 256 centroids. What is the justification and why not allow a user to decide?
| svfloat32_t yv = svld1_f32(pg, y + c * D); | ||
| svfloat32_t diff = svsub_f32_x(pg, yv, xv); | ||
| svfloat32_t sq = svmul_f32_x(pg, diff, diff); | ||
| const float dis = svaddv_f32(pg, sq); |
Collaborator
There was a problem hiding this comment.
a faster way would be to maintain svfloat32_t global_mins and svuint32_t global_ids. But this one works as well, although it is performance-wise suboptimal
| namespace { | ||
|
|
||
| struct TrainState { | ||
| /// Orthogonal rotation. Train in rotated space (X_tilde = X * R); |
Collaborator
There was a problem hiding this comment.
please keep the comment
| std::vector<int64_t>& labels64, | ||
| std::vector<float>& hassign) { | ||
| std::vector<float>& hassign, | ||
| bool spherical) { |
Collaborator
There was a problem hiding this comment.
the question from the Faiss team would be the following: why do you need to have spherical here at all, instead of pre-normalizing the input data outside of this code?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related to #1769 (Benchmark IVF_FLAT_CC and SCANN_DVR train performance)
Enable SuperKMeans coarse quantizer training for SCANN by default.
Motivation
SCANN (and other IVF-family indexes) train their coarse quantizer with
faiss::Clustering.SuperKMeansis a faster k-means variant usingADSampling + PDX progressive pruning, but it previously only supported L2
clustering. This change adds inner-product (spherical) support to
SuperKMeans and routes SCANN coarse quantizer training through it.
Changes
Vendored faiss (
thirdparty/faiss, mirrored to upstream faiss PR)cp.spherical=true, centroidsare unit-normalized each iteration so L2 assignment is equivalent to IP
argmax (
||x-c||² = ||x||² + 1 - 2x·c). Power-of-twoduses the fastHadamardRotation; L2 training keepsRandomRotationMatrix.ClusteringParameters::use_super_kmeans(default false): letsLevel1Quantizer::train_q1route coarse quantizer training throughSuperKMeans when explicitly enabled.
d in {2,4,8}): one SVE lane percentroid, min index tracked in registers, scratch buffer not written
(matching the x86 AVX2/AVX512
D2/D4/D8implementations).ProductQuantizer::compute_codeswitches toAVAILABLE_SIMD_LEVELS_A1so PQ encoding reaches the SVE path.
HadamardRotation::reverse_transform: fills the missing inverse.block_l2<ARM_SVE>: SVE kernel for the SuperKMeans pruning loop.Knowhere
ScannConfig::use_super_kmeansdefaults to true.IvfConfig::use_super_kmeansplumbed into all IVF-family build paths(default false for non-SCANN indexes).
Benchmark
qwen 4096-dim, IP, nlist=1024, sub_dim=4, 1 thread, 50k rows, SCANN_DVR:
recall@10 unchanged vs Clustering baseline (diff <= 0.0002).
Notes
The vendored faiss changes are also submitted upstream
(facebookresearch/faiss#5530). This PR carries them in the vendored copy
so Knowhere can use them immediately; they will be dropped at the next
faiss snapshot upgrade once merged upstream.