fix: stabilize IVF-PQ k-means training

[jerry-024]

Summary

Fixes correctness and termination edge cases in IVF-PQ hierarchical k-means while preserving the historical strict largest-first split policy. The PR also parallelizes row-independent cluster assignment with bounded, deterministic blocks; the measured speedup is modest now that hierarchy scheduling is intentionally serial again.

Changes

• Preserve strict pop/split/reinsert largest-first hierarchy scheduling, avoiding the recall and skewed-distribution quality regression of batched breadth-first splitting.
• Terminate degenerate splits and pad exhausted hierarchies with valid centroids instead of looping or fabricating zero centroids.
• Parallelize only unbalanced cluster assignment; keep balanced assignment and zero-dimensional behavior compatible.
• Preserve historical objective summation boundaries so results are bitwise reproducible across Rayon pool sizes.
• Require at least 4M FLOPs per parallel block so the common 512-point k=2, d=768 split collapses to one SGEMM.
• Use the historical serial SGEMM chunks when Rayon has one worker, avoiding fork/join and per-row objective-buffer overhead.
• Validate KMeans matrix shapes and SGEMM backing lengths and strides before unsafe matrix multiplication, rejecting invalid or overflowing inputs.
• Bound aggregate concurrent ip_matrix scratch across Rayon workers to about 16 MiB, with a serial fallback when per-worker blocks are too small.
• Keep the benchmark bounded by default; set TRAIN_TARGET=1 to run the target 244606 x 768 workload.

Target benchmark

Command: TRAIN_TARGET=1 cargo +1.97.0 bench -p paimon-vindex-core --bench ivfpq_train_bench

Deterministic target-768 input: n=244606, d=768, nlist=1024, pq_m=96, inner product, OPQ disabled. Results are three-run medians on an Apple M3 Pro with 11 CPU cores and 36 GiB RAM. The 11-thread measurements alternated the base and PR binaries to reduce system-load drift.

version	threads	coarse_s	pq_s	train_total_s	coarse vs base	total vs base
base 6374d05	1	2.824	55.423	58.849	—	—
this PR 731d93e	1	2.780	54.908	58.459	1.02x	1.01x
base 6374d05	11	2.750	8.875	11.510	—	—
this PR 731d93e	11	2.387	7.728	9.878	1.15x	1.17x

train_total_s is the authoritative end-to-end measurement; phase timings come from separate stage replays and are not additive. The earlier 4.5x coarse result depended on the reverted batched hierarchy scheduler and is not claimed by this version. No 16-core Xeon result is available; 11 threads matches the test machine's physical CPU count.

Testing

• cargo +1.97.0 test -p paimon-vindex-core --lib — 439 passed, 1 ignored
• Cross-pool bitwise assignment/objective tests, including historical serial objective grouping
• Small split single-SGEMM and strict largest-first hierarchy regression tests
• cargo +1.97.0 fmt --all -- --check
• git diff --check
• TRAIN_TARGET=1 benchmark at 1 and 11 threads, three runs per configuration

Notes

• No public API signatures or serialized-format changes. KMeans input validation is stricter: data.len() and initial_centroids.len() must exactly match their declared shapes; invalid or overflowing shapes now panic instead of silently ignoring trailing elements or risking out-of-bounds SGEMM access.
• Target-scale benchmark execution remains opt-in because it can consume several GiB.

[leaves12138]

I think the batched hierarchical splitting changes the clustering policy in a way that can significantly hurt quality on skewed datasets.

At core/src/kmeans.rs:145, max_new is the entire remaining centroid count, and the loop then pops as many eligible clusters as possible before any children are reinserted. For target_k = 1024, this effectively splits all 16, then all 32, 64, 128, 256, and 512 clusters level by level. Therefore, the heap's “largest cluster first” ordering no longer influences most of the allocation. Small clusters receive roughly the same split depth as large clusters until they become too small, whereas the previous best-first policy allocated more centroids to the high-population clusters.

I isolated the scheduling difference while keeping the PR's current assignment implementation and the new degenerate-split handling. On a deterministic skewed dataset (n=21,500, d=2, k=512; 20,000 points in one broad population and 15 remote low-volume populations of 100 points each), I measured:

• Current batched schedule: SSE 339,858.785, 100 centroids assigned to the low-volume populations
• Best-first schedule: SSE 309,751.377, 38 centroids assigned to the low-volume populations

That is about a 9.7% increase in quantization error. The added recall scenario uses fairly uniform generated populations, so it does not cover this regression mode.

Could we keep the parallel assignment and degenerate-split fixes, but avoid replacing best-first splitting with a breadth-first schedule? Alternatively, the parallel split selection needs to preserve the population-aware allocation property, with a skewed-distribution quality regression test.

[shyjsarah]

Thanks for the update. I re-reviewed the latest head 60eecb1. The strict largest-first hierarchy behavior and concurrent sub_data memory regression from the previous revision are now fixed.

There is still one blocking memory-safety issue:

[Blocking] Validate tensor shapes before entering unsafe SGEMM

kmeans_train_with_init accepts data, n, d, k, and initial_centroids independently, but does not validate:

• data.len() == n * d
• initial_centroids.len() == k * d

In assign_clusters_fast/assign_block, the claimed row count reaches sgemm_a_bt before the first bounds-checked row access. For example, three rows of data with n=4, d=2, k=2, and valid initial centroids causes SGEMM to read a fourth row beyond the allocation. ASan reproduces a heap-buffer-overflow inside matrixmultiply.

Since this is reachable through a safe public API, malformed input may return an error or panic, but must not cause undefined behavior.

Please use checked multiplication to validate shapes at the public boundary and add defensive assertions at the unsafe SGEMM boundary. Tests should cover short/long data, invalid initial-centroid lengths, and dimension overflow.

Non-blocking: parallel assignment can retain one ip_matrix per Rayon worker. With a large custom max_points_per_centroid, peak scratch memory therefore scales with the ambient worker count. An aggregate scratch budget or worker-buffer reuse would make the memory behavior more predictable.

Local formatting, Clippy, core tests, 1/2/4/8-thread K-means tests, and the latest GitHub CI all pass.

[shyjsarah]

+1

[JingsongLi]

+1