[MOD-17189] Share SQ8 training state for RAM and disk indexes (#1078)

* [MOD-17189] Share SQ8 training state and threshold

* Share SQ8 mean installation between index backends

* Clarify SQ8 mean helper name and remove redundant assertions

* Wait for in-flight jobs before stopping async test workers

Author: dor-forer

src/VecSim/algorithms/hnsw/hnsw.h

@@ -21,6 +21,7 @@
 #include "VecSim/vec_sim_common.h"
 #include "VecSim/vec_sim_index.h"
 #include "VecSim/spaces/computer/preprocessors.h"
+#include "VecSim/quantization/sq8_mean.h"
 #include <span>
 #include "VecSim/tombstone_interface.h"
 
@@ -94,28 +95,15 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
 private:
     template <VecSimMetric Metric>
     void setSQ8Mean(std::span<const float> mean) noexcept {
-        // SQ8-specific component access. Introduce a trainer interface when a second quantizer
-        // needs training, replacing these concrete component casts.
-        // Assumes exactly one preprocessor: SQ8 WithNorm in slot 0. Tiered cosine inputs
-        // are normalized by the frontend; a pipeline with a preceding normalizer is unsupported.
-        auto *container = static_cast<MultiPreprocessorsContainer<DataType, 1> *>(
-            this->getPreprocessorsContainer());
-        auto *preprocessor = dynamic_cast<QuantPreprocessor<DataType, Metric, true> *>(
-            container->getPreprocessors()[0]);
-        auto *calculator = dynamic_cast<DistanceCalculatorWithNorm<DataType, float, Metric> *>(
-            this->getIndexCalculator());
-        assert(preprocessor && calculator);
-        preprocessor->setMean(mean);
-        calculator->setMeanSumSquares(mean);
+        applySQ8Mean<DataType, Metric>(*this->getPreprocessorsContainer(),
+                                       *this->getIndexCalculator(), mean);
     }
 
 public:
     // Writer-only operation: install the mean once, before any stored vector, under the
     // exclusive tiered main lock and before submitting insertion jobs. The backend must have
     // its final SQ8 WithNorm components; the caller owns and enforces the one-time transition.
-    // SQ8-specific: assumes a single SQ8 preprocessor in slot 0, with cosine normalization
-    // performed by the frontend. Introduce a trainer interface when a second quantizer needs
-    // training. A quantizer that learns storedDataSize from data needs deferred backend creation.
+    // A quantizer that learns storedDataSize from data needs deferred backend creation.
     void setQuantizationMean(std::span<const float> mean) noexcept {
         assert(this->isQuantized && curElementCount == 0 && mean.size() == this->dim);
         if constexpr (QuantInput<DataType> && std::is_same_v<DistType, float>) {

src/VecSim/algorithms/hnsw/hnsw_tiered.h

@@ -18,7 +18,7 @@
 #include <utility>
 
 #include "VecSim/algorithms/brute_force/brute_force_single.h"
-#include "VecSim/spaces/computer/preprocessors.h"
+#include "VecSim/quantization/sq8_training_state.h"
 #include "VecSim/vec_sim_tiered_index.h"
 #include "hnsw.h"
 
@@ -51,7 +51,6 @@ struct HNSWSwapJob : public VecsimBaseObject {
 
 static const size_t DEFAULT_PENDING_SWAP_JOBS_THRESHOLD = DEFAULT_BLOCK_SIZE;
 static const size_t MAX_PENDING_SWAP_JOBS_THRESHOLD = 100000;
-static const size_t MAX_QUANT_NORMALIZATION_SET_SIZE = 100 * DEFAULT_BLOCK_SIZE;
 
 /**
  * Definition of a job that repairs a certain node's connection in HNSW Index after delete
@@ -106,17 +105,7 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
 
     // Writer-owned state exists only while SQ8 vectors are accumulating in FLAT.
     // Reads and maintenance always use the real, initially empty backend.
-    struct SQAccumulationState {
-        size_t normalizationSetSize;
-        vecsim_stl::vector<double> runningSumVec;
-
-        SQAccumulationState(const std::shared_ptr<VecSimAllocator> &allocator, size_t dim,
-                            size_t normalizationSetSize)
-            : normalizationSetSize(
-                  std::min(normalizationSetSize, MAX_QUANT_NORMALIZATION_SET_SIZE)),
-              runningSumVec(dim, 0.0, allocator) {}
-    };
-    std::optional<SQAccumulationState> sqAccumulationState;
+    std::optional<SQ8TrainingState> sqAccumulationState;
 
 #ifdef BUILD_TESTS
     std::function<void()> beforeQuantizationFinalization;
@@ -126,25 +115,16 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
 
     void addToSum(std::span<const DataType> vector) {
         if constexpr (QuantInput<DataType>) {
-            auto &sum = sqAccumulationState->runningSumVec;
-            assert(vector.size() == sum.size());
-            for (size_t i = 0; i < vector.size(); ++i) {
-                sum[i] += to_fp32(vector[i]);
-            }
+            sqAccumulationState->add(vector);
         }
     }
 
     void subtractFromSum(std::span<const DataType> vector) {
         if constexpr (QuantInput<DataType>) {
-            auto &sum = sqAccumulationState->runningSumVec;
-            assert(vector.size() == sum.size());
-            for (size_t i = 0; i < vector.size(); ++i) {
-                sum[i] -= to_fp32(vector[i]);
-            }
+            sqAccumulationState->subtract(vector);
         }
     }
 
-    [[nodiscard]] vecsim_stl::vector<float> calculateQuantizationMean(size_t count) const;
     int addVectorDuringAccumulation(const void *blob, labelType label);
     int deleteVectorDuringAccumulation(labelType label);
     void finalizeQuantizationAndSubmitJobs(size_t count);
@@ -814,7 +794,8 @@ int TieredHNSWIndex<DataType, DistType>::addVectorDuringAccumulation(const void
             this->registerInsertJob(label, job);
         }
         const size_t count = this->frontendIndex->indexSize();
-        const bool should_finalize = count >= sqAccumulationState->normalizationSetSize;
+        assert(sqAccumulationState->getCount() == count);
+        const bool should_finalize = sqAccumulationState->shouldFinalize();
         return std::tuple{result, count, should_finalize};
     }();
 
@@ -842,25 +823,13 @@ int TieredHNSWIndex<DataType, DistType>::deleteVectorDuringAccumulation(labelTyp
     for (const auto &[new_id, previous] : updated_ids) {
         this->updateInsertJobInternalId(previous.first, new_id, previous.second);
     }
+    assert(sqAccumulationState->getCount() == this->frontendIndex->indexSize());
     return removed;
 }
 
-template <typename DataType, typename DistType>
-vecsim_stl::vector<float>
-TieredHNSWIndex<DataType, DistType>::calculateQuantizationMean(size_t count) const {
-    assert(sqAccumulationState);
-    assert(count >= sqAccumulationState->normalizationSetSize);
-    const auto &sum = sqAccumulationState->runningSumVec;
-    vecsim_stl::vector<float> mean(sum.size(), this->allocator);
-    for (size_t i = 0; i < sum.size(); ++i) {
-        mean[i] = static_cast<float>(sum[i] / static_cast<double>(count));
-    }
-    return mean;
-}
-
 template <typename DataType, typename DistType>
 void TieredHNSWIndex<DataType, DistType>::finalizeQuantizationAndSubmitJobs(size_t count) {
-    const auto mean = calculateQuantizationMean(count);
+    const auto mean = sqAccumulationState->calculateMean();
     // Allocate the snapshot before finishing accumulation. Synchronous execution removes jobs
     // from labelToInsertJobs, so iterating the map while dispatching would invalidate iterators.
     vecsim_stl::vector<AsyncJob *> jobs(this->allocator);

src/VecSim/quantization/sq8_mean.h

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * Licensed under your choice of the Redis Source Available License 2.0
+ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
+ * GNU Affero General Public License v3 (AGPLv3).
+ */
+#pragma once
+
+#include <cassert>
+#include <span>
+
+#include "VecSim/spaces/computer/calculator.h"
+#include "VecSim/spaces/computer/preprocessor_container.h"
+
+// Install before storing vectors, with queries excluded by the caller. Requires exactly one
+// SQ8 WithNorm preprocessor in slot 0; cosine normalization must be performed by the frontend.
+template <QuantInput DataType, VecSimMetric Metric>
+void applySQ8Mean(PreprocessorsContainerAbstract &preprocessors,
+                  IndexCalculatorInterface<float> &index_calculator,
+                  std::span<const float> mean) noexcept {
+    auto &container = static_cast<MultiPreprocessorsContainer<DataType, 1> &>(preprocessors);
+    auto *preprocessor =
+        dynamic_cast<QuantPreprocessor<DataType, Metric, true> *>(container.getPreprocessors()[0]);
+    auto *calculator =
+        dynamic_cast<DistanceCalculatorWithNorm<DataType, float, Metric> *>(&index_calculator);
+    assert(preprocessor && calculator);
+    preprocessor->setMean(mean);
+    calculator->setMeanSumSquares(mean);
+}

src/VecSim/quantization/sq8_training_state.h

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * Licensed under your choice of the Redis Source Available License 2.0
+ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
+ * GNU Affero General Public License v3 (AGPLv3).
+ */
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <span>
+
+#include "VecSim/spaces/computer/preprocessors.h"
+#include "VecSim/utils/vecsim_stl.h"
+
+static const size_t MAX_QUANT_NORMALIZATION_SET_SIZE = 100 * DEFAULT_BLOCK_SIZE;
+
+/** Writer-owned SQ8 training state. Inputs must be the stored, preprocessed vectors.
+ * Callers synchronize mutations and subtract a vector before overwriting or deleting it.
+ */
+class SQ8TrainingState {
+    size_t normalizationSetSize;
+    vecsim_stl::vector<double> runningSumVec;
+    size_t count{0};
+
+public:
+    SQ8TrainingState(const std::shared_ptr<VecSimAllocator> &allocator, size_t dim,
+                     size_t normalization_set_size)
+        : normalizationSetSize(std::min(normalization_set_size, MAX_QUANT_NORMALIZATION_SET_SIZE)),
+          runningSumVec(dim, 0.0, allocator) {}
+
+    template <QuantInput T>
+    void add(std::span<const T> vector) {
+        assert(vector.size() == runningSumVec.size());
+        for (size_t i = 0; i < vector.size(); ++i) {
+            runningSumVec[i] += to_fp32(vector[i]);
+        }
+        ++count;
+    }
+
+    template <QuantInput T>
+    void subtract(std::span<const T> vector) {
+        assert(count > 0);
+        assert(vector.size() == runningSumVec.size());
+        for (size_t i = 0; i < vector.size(); ++i) {
+            runningSumVec[i] -= to_fp32(vector[i]);
+        }
+        --count;
+    }
+
+    bool shouldFinalize() const { return count > 0 && count >= normalizationSetSize; }
+
+    /// Requires at least one live training vector.
+    [[nodiscard]] vecsim_stl::vector<float> calculateMean() const {
+        assert(count > 0);
+        vecsim_stl::vector<float> mean(runningSumVec.size(), runningSumVec.getAllocator());
+        for (size_t i = 0; i < runningSumVec.size(); ++i) {
+            mean[i] = static_cast<float>(runningSumVec[i] / static_cast<double>(count));
+        }
+        return mean;
+    }
+
+    size_t getCount() const { return count; }
+    size_t getNormalizationSetSize() const { return normalizationSetSize; }
+    const vecsim_stl::vector<double> &getRunningSum() const { return runningSumVec; }
+};

src/VecSim/vec_sim_common.h

@@ -254,7 +254,8 @@ typedef struct {
 // A struct that contains HNSW Disk tiered index specific params.
 // Consider removing and use TieredHNSWParams instead if they both share swapJobThreshold
 typedef struct {
-    char _placeholder; // Reserved for future fields and avoid compiler errors
+    size_t QuantNormalizationSetSize; // Training threshold, with the semantics of
+                                      // TieredHNSWParams::QuantNormalizationSetSize.
 } TieredHNSWDiskParams;
 
 // A struct that contains SVS tiered index specific params.

tests/unit/test_components.cpp

@@ -12,6 +12,7 @@
 #include <cmath>
 #include "gtest/gtest.h"
 #include "VecSim/vec_sim.h"
+#include "VecSim/quantization/sq8_training_state.h"
 #include "VecSim/spaces/computer/preprocessor_container.h"
 #include "VecSim/spaces/computer/preprocessors.h"
 #include "VecSim/spaces/computer/calculator.h"
@@ -2644,3 +2645,85 @@ TEST(DistanceCalculatorWithNormTest, OddDimension_MatchesBruteForce) {
         delete calc;
     }
 }
+
+template <typename T>
+class SQ8TrainingStateTest : public ::testing::Test {};
+
+using SQ8TrainingInputTypes = ::testing::Types<float, vecsim_types::float16>;
+TYPED_TEST_SUITE(SQ8TrainingStateTest, SQ8TrainingInputTypes);
+
+TYPED_TEST(SQ8TrainingStateTest, SumMeanAndRemoval) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 2, 2);
+    const TypeParam first[] = {from_fp32<TypeParam>(1.5f), from_fp32<TypeParam>(-2.0f)};
+    const TypeParam second[] = {from_fp32<TypeParam>(-0.5f), from_fp32<TypeParam>(4.0f)};
+    EXPECT_EQ(state.getCount(), 0);
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const TypeParam>(first));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const TypeParam>(second));
+    EXPECT_EQ(state.getCount(), 2);
+    EXPECT_TRUE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 1.0);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[1], 2.0);
+    const auto mean = state.calculateMean();
+    EXPECT_EQ(mean.getAllocator(), allocator);
+    EXPECT_FLOAT_EQ(mean[0], 0.5f);
+    EXPECT_FLOAT_EQ(mean[1], 1.0f);
+
+    state.subtract(std::span<const TypeParam>(first));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    const auto remaining_mean = state.calculateMean();
+    EXPECT_FLOAT_EQ(remaining_mean[0], -0.5f);
+    EXPECT_FLOAT_EQ(remaining_mean[1], 4.0f);
+    state.subtract(std::span<const TypeParam>(second));
+    EXPECT_EQ(state.getCount(), 0);
+    EXPECT_FALSE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 0.0);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[1], 0.0);
+}
+
+TYPED_TEST(SQ8TrainingStateTest, ReplacementKeepsLiveCount) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, 2);
+    const TypeParam original[] = {from_fp32<TypeParam>(2.0f)};
+    const TypeParam replacement[] = {from_fp32<TypeParam>(8.0f)};
+    state.add(std::span<const TypeParam>(original));
+    state.subtract(std::span<const TypeParam>(original));
+    state.add(std::span<const TypeParam>(replacement));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 8.0);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 8.0f);
+}
+
+TEST(SQ8TrainingStateTest, ThresholdIsCapped) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, std::numeric_limits<size_t>::max());
+    EXPECT_EQ(state.getNormalizationSetSize(), MAX_QUANT_NORMALIZATION_SET_SIZE);
+    const float vector[] = {1.0f};
+    for (size_t i = 0; i < MAX_QUANT_NORMALIZATION_SET_SIZE - 1; ++i) {
+        state.add(std::span<const float>(vector));
+    }
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const float>(vector));
+    EXPECT_TRUE(state.shouldFinalize());
+    EXPECT_EQ(state.getCount(), MAX_QUANT_NORMALIZATION_SET_SIZE);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 1.0f);
+}
+
+TEST(SQ8TrainingStateTest, AccumulatesInFP64) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, 3);
+    const float large[] = {16777216.0f};
+    const float small[] = {1.0f};
+    state.add(std::span<const float>(large));
+    state.add(std::span<const float>(small));
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 16777217.0);
+    state.subtract(std::span<const float>(large));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 1.0);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 1.0f);
+}

tests/unit/test_hnsw_tiered.cpp

@@ -3787,6 +3787,7 @@ TYPED_TEST(HNSWTieredIndexTest, bufferLimitAsync) {
             /*do nothing*/
         }
     }
+    mock_thread_pool.thread_pool_wait();
     mock_thread_pool.thread_pool_join();
     // Multi-value appends on the second run. Single-value overwrites: an overwrite that goes
     // through the flat buffer reuses its old id, while one inserted straight into HNSW (the buffer
@@ -6012,11 +6013,11 @@ class HNSWTieredIndexTestSQ8 : public ::testing::Test {
     }
 
     const vecsim_stl::vector<double> &getRunningSumVec(TieredHNSWIndex<data_t, dist_t> *idx) {
-        return idx->sqAccumulationState->runningSumVec;
+        return idx->sqAccumulationState->getRunningSum();
     }
 
     size_t getNormalizationSetSize(TieredHNSWIndex<data_t, dist_t> *idx) {
-        return idx->sqAccumulationState->normalizationSetSize;
+        return idx->sqAccumulationState->getNormalizationSetSize();
     }
 
     // Compare the actual cached dispatches and preprocessing with a backend constructed with