[MOD-14956] Add SQ8 quantization support for HNSW index

src/VecSim/index_factories/hnsw_factory.cpp

@@ -17,6 +17,7 @@
 
 using bfloat16 = vecsim_types::bfloat16;
 using float16 = vecsim_types::float16;
+using sq8 = vecsim_types::sq8;
 
 namespace HNSWFactory {
 
@@ -34,11 +35,127 @@ NewIndex_ChooseMultiOrSingle(const HNSWParams *params,
             HNSWIndex_Single<DataType, DistType>(params, abstractInitParams, components);
 }
 
+template <VecSimMetric Metric>
+[[nodiscard]] constexpr size_t GetSQ8StoredDataSize(size_t dim, bool with_norm) {
+    static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP);
+
+    // WithNorm is a template parameter, so dispatch the runtime flag to the two instantiations.
+    return with_norm ? sq8::storage_bytes_count<Metric, true>(dim)
+                     : sq8::storage_bytes_count<Metric, false>(dim);
+}
+
+// Alignment required by a query blob of type DataType. Per the asymmetric-types contract in
+// spaces.h, the hint returned alongside an asymmetric distance function describes its first
+// (storage) operand, so the query side must be obtained from the symmetric dispatcher for the
+// query's own type. Only that hint is wanted here, never the function it returns, so the call is
+// contained in this adapter instead of leaving a discarded value at the call site.
+template <typename DataType>
+[[nodiscard]] unsigned char GetQueryAlignment(VecSimMetric metric, size_t dim) {
+    unsigned char alignment = 0;
+    spaces::GetDistFunc<DataType, float>(metric, dim, &alignment);
+    return alignment;
+}
+
+// Helper to build an SQ8-quantized HNSW index given compile-time DataType and Metric.
+template <typename DataType, VecSimMetric Metric>
+VecSimIndex *NewIndex_SQ8(const HNSWParams *hnswParams, AbstractIndexInitParams abstractInitParams,
+                          const float *mean_ptr) {
+    auto &allocator = abstractInitParams.allocator;
+    const size_t dim = abstractInitParams.dim;
+    const bool with_norm = mean_ptr != nullptr;
+    unsigned char storage_alignment = 0, asym_storage_alignment = 0;
+
+    // Override blob size for the SQ8 storage layout.
+    abstractInitParams.storedDataSize = GetSQ8StoredDataSize<Metric>(dim, with_norm);
+
+    // Symmetric: both stored vectors are SQ8 blobs.
+    auto sym_func = spaces::GetDistFunc<sq8, float>(Metric, dim, &storage_alignment);
+    // Asymmetric: stored vector is SQ8 blob, query is DataType.
+    auto asym_func =
+        spaces::GetDistFunc<sq8, float, DataType>(Metric, dim, &asym_storage_alignment);
+    // Both hints describe the same stored blob, so they must be combined rather than overwritten.
+    storage_alignment = spaces::combineAlignments(storage_alignment, asym_storage_alignment);
+    // Queries stay in DataType and are compared against stored blobs by asym_func.
+    const unsigned char query_alignment = GetQueryAlignment<DataType>(Metric, dim);
+
+    PreprocessorInterface *pp = nullptr;
+    IndexCalculatorInterface<float> *calc = nullptr;
+
+    if (with_norm) {
+        // Mean-centered SQ8 quantization with norm correction.
+        vecsim_stl::vector<float> mean_vec(allocator);
+        mean_vec.assign(mean_ptr, mean_ptr + dim);
+
+        float mean_sum_squares = 0.0f;
+        for (float v : mean_vec) {
+            mean_sum_squares += v * v;
+        }
+
+        pp = new (allocator) QuantPreprocessor<DataType, Metric, true>(allocator, dim, mean_vec);
+        calc = new (allocator) DistanceCalculatorWithNorm<DataType, float, Metric>(
+            allocator, asym_func, sym_func, mean_sum_squares);
+    } else {
+        // Plain SQ8 quantization without mean centering.
+        pp = new (allocator) QuantPreprocessor<DataType, Metric>(allocator, dim);
+        // sym_func for storage-storage; asym_func for query-storage.
+        calc = new (allocator) DistanceCalculatorCommon<float>(allocator, sym_func, asym_func);
+    }
+
+    auto *container = new (allocator)
+        MultiPreprocessorsContainer<DataType, 1>(allocator, query_alignment, storage_alignment);
+    [[maybe_unused]] const int ret = container->addPreprocessor(pp);
+    assert(ret != -1 && "SQ8 preprocessor was not added correctly");
+
+    IndexComponents<DataType, float> components{calc, container};
+    return NewIndex_ChooseMultiOrSingle<DataType, float>(hnswParams, abstractInitParams,
+                                                         components);
+}
+
 VecSimIndex *NewIndex(const VecSimParams *params, bool is_normalized) {
     const HNSWParams *hnswParams = &params->algoParams.hnswParams;
+
     AbstractIndexInitParams abstractInitParams =
         VecSimFactory::NewAbstractInitParams(hnswParams, params->logCtx, is_normalized);
 
+    if (hnswParams->quantType == VecSimQuant_SQ8) {
+        if (hnswParams->type != VecSimType_FLOAT32 && hnswParams->type != VecSimType_FLOAT16) {
+            return NULL; // SQ8 supports FP32 and FP16 only.
+        }
+
+        VecSimMetric metric = hnswParams->metric;
+        if (is_normalized && metric == VecSimMetric_Cosine) {
+            metric = VecSimMetric_IP;
+        }
+
+        if (metric == VecSimMetric_Cosine) {
+            return NULL; // SQ8 does not support cosine metric.
+        }
+
+        const float *mean_ptr = static_cast<const float *>(hnswParams->quantParams);
+
+        if (hnswParams->type == VecSimType_FLOAT32) {
+            if (metric == VecSimMetric_L2) {
+                return NewIndex_SQ8<float, VecSimMetric_L2>(hnswParams, abstractInitParams,
+                                                            mean_ptr);
+            } else if (metric == VecSimMetric_IP) {
+                return NewIndex_SQ8<float, VecSimMetric_IP>(hnswParams, abstractInitParams,
+                                                            mean_ptr);
+            }
+        } else if (hnswParams->type == VecSimType_FLOAT16) {
+            if (metric == VecSimMetric_L2) {
+                return NewIndex_SQ8<float16, VecSimMetric_L2>(hnswParams, abstractInitParams,
+                                                              mean_ptr);
+            } else if (metric == VecSimMetric_IP) {
+                return NewIndex_SQ8<float16, VecSimMetric_IP>(hnswParams, abstractInitParams,
+                                                              mean_ptr);
+            }
+        }
+
+        // Unreachable today: the checks above leave only FP32/FP16 x L2/IP. Kept so that adding a
+        // type or metric cannot silently fall through and build an unquantized index instead.
+        return NULL;
+    }
+
     if (hnswParams->type == VecSimType_FLOAT32) {
         IndexComponents<float, float> indexComponents = CreateIndexComponents<float, float>(
             abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
@@ -94,7 +211,27 @@ size_t EstimateInitialSize(const HNSWParams *params, bool is_normalized) {
     size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize();
 
     size_t est = sizeof(VecSimAllocator) + allocations_overhead;
-    if (params->type == VecSimType_FLOAT32) {
+
+    if (params->quantType == VecSimQuant_SQ8) {
+        if (params->type != VecSimType_FLOAT32 && params->type != VecSimType_FLOAT16) {
+            throw std::invalid_argument("Invalid params->type for VecSimQuant_SQ8");
+        }
+        // Calculator + preprocessor container + preprocessor.
+        // Use representative types; sizeof is independent of the template parameters.
+        if (params->quantParams) { // mean provided, WithNorm = true
+            est += allocations_overhead +
+                   sizeof(DistanceCalculatorWithNorm<float, float, VecSimMetric_L2>);
+            est += allocations_overhead + sizeof(MultiPreprocessorsContainer<float, 1>);
+            est += allocations_overhead + sizeof(QuantPreprocessor<float, VecSimMetric_L2, true>);
+            est += allocations_overhead +
+                   params->dim * sizeof(float); // mean vector in QuantPreprocessor
+        } else {
+            est += allocations_overhead + sizeof(DistanceCalculatorCommon<float>);
+            est += allocations_overhead + sizeof(MultiPreprocessorsContainer<float, 1>);
+            est += allocations_overhead + sizeof(QuantPreprocessor<float, VecSimMetric_L2>);
+        }
+        est += EstimateInitialSize_ChooseMultiOrSingle<float>(params->multi);
+    } else if (params->type == VecSimType_FLOAT32) {
         est += EstimateComponentsMemory<float, float>(params->metric, is_normalized);
         est += EstimateInitialSize_ChooseMultiOrSingle<float>(params->multi);
     } else if (params->type == VecSimType_FLOAT64) {
@@ -125,9 +262,20 @@ size_t EstimateElementSize(const HNSWParams *params) {
     size_t M = (params->M) ? params->M : HNSW_DEFAULT_M;
     size_t elementGraphDataSize = sizeof(ElementGraphData) + sizeof(idType) * M * 2;
 
-    size_t size_total_data_per_element =
-        elementGraphDataSize +
-        VecSimParams_GetStoredDataSize(params->type, params->dim, params->metric);
+    size_t stored_data_size;
+    if (params->quantType == VecSimQuant_SQ8) {
+        bool with_norm = params->quantParams != nullptr;
+        if (params->metric == VecSimMetric_L2) {
+            stored_data_size = GetSQ8StoredDataSize<VecSimMetric_L2>(params->dim, with_norm);
+        } else {
+            stored_data_size = GetSQ8StoredDataSize<VecSimMetric_IP>(params->dim, with_norm);
+        }
+    } else {
+        stored_data_size =
+            VecSimParams_GetStoredDataSize(params->type, params->dim, params->metric);
+    }
+
+    size_t size_total_data_per_element = elementGraphDataSize + stored_data_size;
 
     // when reserving space for new labels in the lookup hash table, each entry is a pointer to a
     // label node (bucket).

src/VecSim/index_factories/tiered_factory.cpp

@@ -95,6 +95,14 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) {
 }
 
 VecSimIndex *NewIndex(const TieredIndexParams *params) {
+    // Quantization is not wired into the tiered index yet (MOD-14957). Reject it here rather than
+    // let it through: the primary index would be built from these params and quantize its storage,
+    // while NewBFParams does not carry quantType, so the frontend would stay unquantized and the
+    // two would disagree on the stored blob layout.
+    if (params->primaryIndexParams->algoParams.hnswParams.quantType != VecSimQuant_NONE) {
+        return nullptr;
+    }
+
     // Tiered index that contains HNSW index as primary index
     VecSimType type = params->primaryIndexParams->algoParams.hnswParams.type;
     if (type == VecSimType_FLOAT32) {

src/VecSim/spaces/computer/preprocessors.h

@@ -426,8 +426,7 @@ class QuantPreprocessor : public PreprocessorInterface {
     QuantPreprocessor(std::shared_ptr<VecSimAllocator> allocator, size_t dim)
         requires(!WithNorm)
         : PreprocessorInterface(allocator), dim(dim),
-          storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
-                              sq8::storage_metadata_count<Metric>() * sizeof(MetadataType)),
+          storage_bytes_count(sq8::storage_bytes_count<Metric>(dim)),
           query_bytes_count(dim * sizeof(DataType) +
                             sq8::query_metadata_count<Metric>() * sizeof(MetadataType)) {}
 
@@ -436,9 +435,7 @@ class QuantPreprocessor : public PreprocessorInterface {
                       const vecsim_stl::vector<float> &mean_vec)
         requires(WithNorm)
         : PreprocessorInterface(allocator), mean(mean_vec), dim(dim),
-          storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
-                              sq8::storage_metadata_count<Metric, WithNorm>() *
-                                  sizeof(MetadataType)),
+          storage_bytes_count(sq8::storage_bytes_count<Metric, WithNorm>(dim)),
           query_bytes_count(dim * sizeof(DataType) +
                             sq8::query_metadata_count<Metric, WithNorm>() * sizeof(MetadataType)) {
         assert(this->mean.size() == dim && "mean vector size must equal dim");

src/VecSim/types/sq8.h

@@ -47,6 +47,15 @@ struct sq8 {
                ((WithNorm && Metric == VecSimMetric_IP) ? 1 : 0);
     }
 
+    // Size of a stored SQ8 blob: one byte per dimension, followed by FP32 metadata. Single source
+    // of truth for the storage layout: every caller that sizes or allocates a stored blob must use
+    // this, so the layout cannot drift between the preprocessor and the index factories.
+    template <VecSimMetric Metric, bool WithNorm = false>
+    static constexpr size_t storage_bytes_count(size_t dim) {
+        return dim * sizeof(value_type) +
+               storage_metadata_count<Metric, WithNorm>() * sizeof(float);
+    }
+
     // Index of x_mean_ip / y_mean_ip in the last slot in metadata array
     template <VecSimMetric Metric>
     static constexpr size_t mean_ip_index() {

src/VecSim/vec_sim_common.h

@@ -68,6 +68,13 @@ typedef enum {
     VecSimType_INT64
 } VecSimType;
 
+// Quantization type for HNSW indices.
+typedef enum {
+    VecSimQuant_NONE = 0, // No quantization (default).
+    // 8-bit scalar quantization. Mean normalization is optional, selected by quantParams below.
+    VecSimQuant_SQ8 = 1,
+} VecSimQuantType;
+
 // Algorithm type/library.
 typedef enum { VecSimAlgo_BF, VecSimAlgo_HNSWLIB, VecSimAlgo_TIERED, VecSimAlgo_SVS } VecSimAlgo;
 
@@ -156,6 +163,10 @@ typedef struct {
     size_t efConstruction;
     size_t efRuntime;
     double epsilon;
+    VecSimQuantType quantType; // Quantization type. Default: VecSimQuant_NONE.
+    // For VecSimQuant_SQ8: pointer to float mean[dim], or NULL for zero mean. Read only, never
+    // retained: the index copies the mean vector during construction.
+    const void *quantParams;
 } HNSWParams;
 
 typedef struct {

tests/unit/CMakeLists.txt

@@ -36,6 +36,7 @@ endif()
 
 add_executable(test_hnsw ../utils/test_main_with_timeout.cpp ../utils/mock_thread_pool.cpp test_hnsw.cpp test_hnsw_multi.cpp test_hnsw_tiered.cpp unit_test_utils.cpp)
 add_executable(test_hnsw_parallel ../utils/test_main_with_timeout.cpp test_hnsw_parallel.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp)
+add_executable(test_hnsw_sq8 ../utils/test_main_with_timeout.cpp ../utils/mock_thread_pool.cpp test_hnsw_sq8.cpp unit_test_utils.cpp)
 add_executable(test_bruteforce ../utils/test_main_with_timeout.cpp test_bruteforce.cpp test_bruteforce_multi.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp)
 add_executable(test_allocator ../utils/test_main_with_timeout.cpp test_allocator.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp)
 add_executable(test_spaces ../utils/test_main_with_timeout.cpp test_spaces.cpp)
@@ -51,6 +52,7 @@ add_executable(test_svs ../utils/test_main_with_timeout.cpp ../utils/mock_thread
 
 target_link_libraries(test_hnsw PUBLIC gtest VectorSimilarity)
 target_link_libraries(test_hnsw_parallel PUBLIC gtest VectorSimilarity)
+target_link_libraries(test_hnsw_sq8 PUBLIC gtest VectorSimilarity)
 target_link_libraries(test_bruteforce PUBLIC gtest VectorSimilarity)
 target_link_libraries(test_allocator PUBLIC gtest VectorSimilarity)
 target_link_libraries(test_spaces PUBLIC gtest VectorSimilarity)
@@ -68,6 +70,7 @@ include(GoogleTest)
 
 gtest_discover_tests(test_hnsw)
 gtest_discover_tests(test_hnsw_parallel)
+gtest_discover_tests(test_hnsw_sq8)
 gtest_discover_tests(test_bruteforce)
 gtest_discover_tests(test_allocator)
 gtest_discover_tests(test_spaces)