SuperKMeans retains the sampled training buffer after rotation

Author: marcelo-cjlCreated Sep 16, 2026Updated Sep 16, 2026

Summary

SuperKMeans::train retains its temporary subsampled input buffer after rotation has already materialized X_tilde. For large, high-dimensional training sets, this unnecessarily increases average build memory throughout centroid initialization and the iterative clustering loop.

Current behavior

setup_train_state returns sampled_x_owner, and SuperKMeans::train keeps that owner alive until training returns. The sampled input is no longer read after RandomRotationMatrix::apply_noalloc completes.

Expected behavior

Release the temporary sample buffer immediately after rotation populates X_tilde, while preserving the existing algorithm, output, and rotation-time peak-memory behavior.

Environment

FAISS main at d63faeaf5a9b5133c1a500d6c0cf9d934c075cf8
Linux x86_64, Release/AVX2 CPU build, 16 threads
Downstream Knowhere SCANN build using Cohere 1M x 768 FP32, nlist=1024, sub_dim=2, with_raw_data=false

Reproduction and evidence

Run one full index build while sampling process RSS every 20 ms:

SuperKMeans disabled: 6369.7 MiB average RSS
SuperKMeans enabled before early release: 6749.1 MiB average RSS
SuperKMeans enabled with early release: 6425.4 MiB average RSS
Optimized versus unoptimized: -323.7 MiB average RSS (-4.8% overall; -8.5% of build RSS above baseline)

With k=1024 and the default max_points_per_centroid=256, the sampled FP32 buffer contains 262,144 x 768 values, or 768 MiB. The full-build average reduction is smaller because the buffer remains necessary during sampling and rotation.

Root cause

The sample owner's lifetime follows the entire train call even though the last consumer of the sampled input is the rotation call. X_tilde owns the transformed data needed by all later phases.

Proposed direction

Keep ownership local to setup_train_state and reset it immediately after apply_noalloc returns. Change the helper to return void so the caller cannot accidentally extend the lifetime.

Acceptance criteria

The sampled input remains alive through rotation and is released before centroid initialization and iterative clustering.
SuperKMeans training results and existing tests remain unchanged.
The FAISS CPU test target builds and the related SuperKMeans foundation tests pass.