From 30983555573524436edc8c55960e18606965ad89 Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Sun, 18 Oct 2026 14:33:44 +0300 Subject: [PATCH] [MOD-17189] Share SQ8 training state and threshold --- src/VecSim/algorithms/hnsw/hnsw_tiered.h | 51 +++--------- src/VecSim/quantization/sq8_training_state.h | 68 ++++++++++++++++ src/VecSim/vec_sim_common.h | 3 +- tests/unit/test_components.cpp | 83 ++++++++++++++++++++ tests/unit/test_hnsw_tiered.cpp | 4 +- 5 files changed, 167 insertions(+), 42 deletions(-) create mode 100644 src/VecSim/quantization/sq8_training_state.h diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered.h b/src/VecSim/algorithms/hnsw/hnsw_tiered.h index 21ebaa342..1ec096fa6 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered.h @@ -18,7 +18,7 @@ #include #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" @@ -70,7 +70,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 @@ -150,17 +149,7 @@ class TieredHNSWIndex : public VecSimTieredIndex { // 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 runningSumVec; - - SQAccumulationState(const std::shared_ptr &allocator, size_t dim, - size_t normalizationSetSize) - : normalizationSetSize( - std::min(normalizationSetSize, MAX_QUANT_NORMALIZATION_SET_SIZE)), - runningSumVec(dim, 0.0, allocator) {} - }; - std::optional sqAccumulationState; + std::optional sqAccumulationState; #ifdef BUILD_TESTS std::function beforeQuantizationFinalization; @@ -171,25 +160,16 @@ class TieredHNSWIndex : public VecSimTieredIndex { void addToSum(std::span vector) { if constexpr (QuantInput) { - 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 vector) { if constexpr (QuantInput) { - 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 calculateQuantizationMean(size_t count) const; int addVectorDuringAccumulation(const void *blob, labelType label); int deleteVectorDuringAccumulation(labelType label); void finalizeQuantizationAndSubmitJobs(size_t count); @@ -1130,7 +1110,8 @@ int TieredHNSWIndex::addVectorDuringAccumulation(const void } } 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}; }(); @@ -1158,25 +1139,17 @@ int TieredHNSWIndex::deleteVectorDuringAccumulation(labelTyp for (const auto &[new_id, previous] : updated_ids) { updateInsertJobInternalId(previous.first, new_id, previous.second); } + assert(sqAccumulationState->getCount() == this->frontendIndex->indexSize()); return removed; } -template -vecsim_stl::vector -TieredHNSWIndex::calculateQuantizationMean(size_t count) const { - assert(sqAccumulationState); - assert(count >= sqAccumulationState->normalizationSetSize); - const auto &sum = sqAccumulationState->runningSumVec; - vecsim_stl::vector mean(sum.size(), this->allocator); - for (size_t i = 0; i < sum.size(); ++i) { - mean[i] = static_cast(sum[i] / static_cast(count)); - } - return mean; -} - template void TieredHNSWIndex::finalizeQuantizationAndSubmitJobs(size_t count) { - const auto mean = calculateQuantizationMean(count); + assert(sqAccumulationState); + assert(sqAccumulationState->getCount() == count); + assert(sqAccumulationState->getCount() == this->frontendIndex->indexSize()); + assert(sqAccumulationState->shouldFinalize()); + 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 jobs(this->allocator); diff --git a/src/VecSim/quantization/sq8_training_state.h b/src/VecSim/quantization/sq8_training_state.h new file mode 100644 index 000000000..0da9279d4 --- /dev/null +++ b/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 +#include +#include + +#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 runningSumVec; + size_t count{0}; + +public: + SQ8TrainingState(const std::shared_ptr &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 + void add(std::span vector) { + assert(vector.size() == runningSumVec.size()); + for (size_t i = 0; i < vector.size(); ++i) { + runningSumVec[i] += to_fp32(vector[i]); + } + ++count; + } + + template + void subtract(std::span 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 calculateMean() const { + assert(count > 0); + vecsim_stl::vector mean(runningSumVec.size(), runningSumVec.getAllocator()); + for (size_t i = 0; i < runningSumVec.size(); ++i) { + mean[i] = static_cast(runningSumVec[i] / static_cast(count)); + } + return mean; + } + + size_t getCount() const { return count; } + size_t getNormalizationSetSize() const { return normalizationSetSize; } + const vecsim_stl::vector &getRunningSum() const { return runningSumVec; } +}; diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h index 7d0676fb8..ad4a18cbf 100644 --- a/src/VecSim/vec_sim_common.h +++ b/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. diff --git a/tests/unit/test_components.cpp b/tests/unit/test_components.cpp index 7ef911762..375ee52b1 100644 --- a/tests/unit/test_components.cpp +++ b/tests/unit/test_components.cpp @@ -12,6 +12,7 @@ #include #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 +class SQ8TrainingStateTest : public ::testing::Test {}; + +using SQ8TrainingInputTypes = ::testing::Types; +TYPED_TEST_SUITE(SQ8TrainingStateTest, SQ8TrainingInputTypes); + +TYPED_TEST(SQ8TrainingStateTest, SumMeanAndRemoval) { + auto allocator = VecSimAllocator::newVecsimAllocator(); + SQ8TrainingState state(allocator, 2, 2); + const TypeParam first[] = {from_fp32(1.5f), from_fp32(-2.0f)}; + const TypeParam second[] = {from_fp32(-0.5f), from_fp32(4.0f)}; + EXPECT_EQ(state.getCount(), 0); + EXPECT_FALSE(state.shouldFinalize()); + state.add(std::span(first)); + EXPECT_EQ(state.getCount(), 1); + EXPECT_FALSE(state.shouldFinalize()); + state.add(std::span(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(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(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(2.0f)}; + const TypeParam replacement[] = {from_fp32(8.0f)}; + state.add(std::span(original)); + state.subtract(std::span(original)); + state.add(std::span(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::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(vector)); + } + EXPECT_FALSE(state.shouldFinalize()); + state.add(std::span(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(large)); + state.add(std::span(small)); + EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 16777217.0); + state.subtract(std::span(large)); + EXPECT_EQ(state.getCount(), 1); + EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 1.0); + EXPECT_FLOAT_EQ(state.calculateMean()[0], 1.0f); +} diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index 539f239d2..e1e3da09b 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -6784,11 +6784,11 @@ class HNSWTieredIndexTestSQ8 : public ::testing::Test { } const vecsim_stl::vector &getRunningSumVec(TieredHNSWIndex *idx) { - return idx->sqAccumulationState->runningSumVec; + return idx->sqAccumulationState->getRunningSum(); } size_t getNormalizationSetSize(TieredHNSWIndex *idx) { - return idx->sqAccumulationState->normalizationSetSize; + return idx->sqAccumulationState->getNormalizationSetSize(); } // Compare the actual cached dispatches and preprocessing with a backend constructed with