Skip to content

[MOD-14957] Add SQ8 quantization support for tiered HNSW index #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xiangze-arm wants to merge 1 commit into
base: sq8_rebase
Choose a base branch
from
Open

[MOD-14957] Add SQ8 quantization support for tiered HNSW index #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 188 additions & 11 deletions src/VecSim/algorithms/hnsw/hnsw_tiered.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/*
* Copyright (c) 2006-Present, Redis Ltd.
* All rights reserved.
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
*
* 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
Expand Down Expand Up @@ -46,6 +47,7 @@ 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
Expand Down Expand Up @@ -99,6 +101,16 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
// Not atomic since it's only accessed from the main thread.
size_t directHNSWInsertions{0};

bool isQuantized{false};

// SQ accumulation phase members, only used for quantized tiered index
vecsim_stl::vector<float> runningSumVec;
size_t quantNormalizationSetSize;
bool isInAccumulationPhase;
VecSimParams backendIndexParams; // Saved params for creating the SQ backend

void initializeQuantizedBackend();

void executeInsertJob(HNSWInsertJob *job);
void executeRepairJob(HNSWRepairJob *job);

Expand Down Expand Up @@ -143,6 +155,26 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
// Handle deletion of vector inplace considering that async deletion might occurred beforehand.
int deleteLabelFromHNSWInplace(labelType label);

void addToSum(const DataType *vector) {
for (size_t i = 0; i < this->runningSumVec.size(); i++) {
if constexpr (std::is_same_v<DataType, float>) {
this->runningSumVec[i] += vector[i];
} else if constexpr (std::is_same_v<DataType, vecsim_types::float16>) {
this->runningSumVec[i] += FP16_to_FP32(vector[i]);
}
}
}

void subtractFromSum(const DataType *vector) {
for (size_t i = 0; i < this->runningSumVec.size(); i++) {
if constexpr (std::is_same_v<DataType, float>) {
this->runningSumVec[i] -= vector[i];
} else if constexpr (std::is_same_v<DataType, vecsim_types::float16>) {
this->runningSumVec[i] -= FP16_to_FP32(vector[i]);
}
}
}

#ifdef BUILD_TESTS
#include "VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h"
#endif
Expand Down Expand Up @@ -204,10 +236,23 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
std::shared_ptr<VecSimAllocator> allocator);
virtual ~TieredHNSWIndex();

// Override query routing so that for quantized tiered indices the flat+backend
// merge always uses withSet=true (the flat and quantized backend scores are not
// directly comparable, so the withSet=false optimization is unsound).
// Also short-circuits to a flat-only query while in accumulation phase.
VecSimQueryReply *topKQuery(const void *queryBlob, size_t k,
VecSimQueryParams *queryParams) const override;
VecSimQueryReply *rangeQuery(const void *queryBlob, double radius,
VecSimQueryParams *queryParams,
VecSimQueryReply_Order order) const override;

int addVector(const void *blob, labelType label) override;
int deleteVector(labelType label) override;
size_t getNumMarkedDeleted() const override {
return this->getHNSWIndex()->getNumMarkedDeleted();
this->mainIndexGuard.lock_shared();
size_t num_marked_deleted = this->getHNSWIndex()->getNumMarkedDeleted();
this->mainIndexGuard.unlock_shared();
return num_marked_deleted;
}
size_t indexSize() const override;
size_t indexCapacity() const override;
Expand All @@ -220,12 +265,17 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
VecSimDebugInfoIterator *debugInfoIterator() const override;
VecSimBatchIterator *newBatchIterator(const void *queryBlob,
VecSimQueryParams *queryParams) const override {
if (isInAccumulationPhase) {
return this->frontendIndex->newBatchIterator(queryBlob, queryParams);
}
// The query blob will be processed and copied by the internal indexes's batch iterator.
return new (this->allocator)
TieredHNSW_BatchIterator(queryBlob, this, queryParams, this->allocator);
}
inline void setLastSearchMode(VecSearchMode mode) override {
return this->backendIndex->setLastSearchMode(mode);
this->mainIndexGuard.lock_shared();
this->backendIndex->setLastSearchMode(mode);
this->mainIndexGuard.unlock_shared();
}
void runGC() override {
// Run no more than pendingSwapJobsThreshold value jobs.
Expand Down Expand Up @@ -255,8 +305,10 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
#ifdef BUILD_TESTS
void getDataByLabel(labelType label, std::vector<std::vector<DataType>> &vectors_output) const;
size_t indexMetaDataCapacity() const override {
return this->backendIndex->indexMetaDataCapacity() +
this->frontendIndex->indexMetaDataCapacity();
this->mainIndexGuard.lock_shared();
size_t capacity = this->backendIndex->indexMetaDataCapacity();
this->mainIndexGuard.unlock_shared();
return capacity + this->frontendIndex->indexMetaDataCapacity();
}
#endif
};
Expand Down Expand Up @@ -661,14 +713,30 @@ TieredHNSWIndex<DataType, DistType>::TieredHNSWIndex(HNSWIndex<DataType, DistTyp
: VecSimTieredIndex<DataType, DistType>(hnsw_index, bf_index, tiered_index_params, allocator),
labelToInsertJobs(this->allocator), idToRepairJobs(this->allocator),
idToSwapJob(this->allocator), invalidJobs(this->allocator), currInvalidJobId(0),
readySwapJobs(0) {
readySwapJobs(0), runningSumVec(this->allocator), quantNormalizationSetSize(0),
isInAccumulationPhase(false) {
// If the param for swapJobThreshold is 0 use the default value, if it exceeds the maximum
// allowed, use the maximum value.
this->pendingSwapJobsThreshold =
tiered_index_params.specificParams.tieredHnswParams.swapJobThreshold == 0
? DEFAULT_PENDING_SWAP_JOBS_THRESHOLD
: std::min(tiered_index_params.specificParams.tieredHnswParams.swapJobThreshold,
MAX_PENDING_SWAP_JOBS_THRESHOLD);

// Initialize SQ accumulation phase if quantization is enabled.
auto &hnswParams = tiered_index_params.primaryIndexParams->algoParams.hnswParams;
if (hnswParams.quantType != VecSimQuant_NONE) {
isQuantized = true;
size_t normSize =
tiered_index_params.specificParams.tieredHnswParams.QuantNormalizationSetSize;
if (normSize > 0) {
this->quantNormalizationSetSize = std::min(normSize, MAX_QUANT_NORMALIZATION_SET_SIZE);
this->isInAccumulationPhase = true;
this->runningSumVec.resize(hnswParams.dim, 0.0f);
}
// Save tiered params for creating the SQ backend
this->backendIndexParams = *(tiered_index_params.primaryIndexParams);
}
}

template <typename DataType, typename DistType>
Expand All @@ -695,19 +763,88 @@ TieredHNSWIndex<DataType, DistType>::~TieredHNSWIndex() {
}
}

template <typename DataType, typename DistType>
VecSimQueryReply *
TieredHNSWIndex<DataType, DistType>::topKQuery(const void *queryBlob, size_t k,
VecSimQueryParams *queryParams) const {
// Accumulation phase: backend is an empty placeholder. We can short-circuit
// to a flat-only query.
if (this->isInAccumulationPhase) {
this->flatIndexGuard.lock_shared();
auto *res = this->frontendIndex->topKQuery(queryBlob, k, queryParams);
this->flatIndexGuard.unlock_shared();
return res;
}
// For quantized tiered, flat/backend scores are not directly comparable, so we
// must use withSet=true even for single-value indexes. Multi-value already
// uses withSet=true in the base.
if (isQuantized && !this->backendIndex->isMultiValue()) {
return this->template topKQueryImp<true>(queryBlob, k, queryParams);
}
return VecSimTieredIndex<DataType, DistType>::topKQuery(queryBlob, k, queryParams);
}

template <typename DataType, typename DistType>
VecSimQueryReply *
TieredHNSWIndex<DataType, DistType>::rangeQuery(const void *queryBlob, double radius,
VecSimQueryParams *queryParams,
VecSimQueryReply_Order order) const {
if (this->isInAccumulationPhase) {
this->flatIndexGuard.lock_shared();
auto *res = this->frontendIndex->rangeQuery(queryBlob, radius, queryParams);
this->flatIndexGuard.unlock_shared();
if (res) {
sort_results(res, order);
}
return res;
}
if (isQuantized && !this->backendIndex->isMultiValue()) {
return this->template rangeQueryImp<true>(queryBlob, radius, queryParams, order);
}
return VecSimTieredIndex<DataType, DistType>::rangeQuery(queryBlob, radius, queryParams, order);
}

template <typename DataType, typename DistType>
void TieredHNSWIndex<DataType, DistType>::initializeQuantizedBackend() {
// Compute mean and create the SQ backend.
size_t dim = this->backendIndexParams.algoParams.hnswParams.dim;
vecsim_stl::vector<float> mean(dim, this->allocator);
for (size_t i = 0; i < dim; i++) {
mean[i] = this->runningSumVec[i] / this->quantNormalizationSetSize;
}

// Build HNSWParams with SQ quantization and computed mean.
VecSimParams newParams = this->backendIndexParams;
newParams.algoParams.hnswParams.quantParams = mean.data();

// Exclude readers while replacing the empty backend with its SQ8 counterpart.
this->lockMainIndexGuard();
VecSimIndex_Free(this->backendIndex);

// Create new SQ HNSW backend. Normalization is done by the frontend index.
this->backendIndex = reinterpret_cast<VecSimIndexAbstract<DataType, DistType> *>(
HNSWFactory::NewIndex(&newParams, true));
this->unlockMainIndexGuard();
}

template <typename DataType, typename DistType>
size_t TieredHNSWIndex<DataType, DistType>::indexSize() const {
this->flatIndexGuard.lock_shared();
this->mainIndexGuard.lock_shared();
this->getHNSWIndex()->lockSharedIndexDataGuard();
size_t res = this->backendIndex->indexSize() + this->frontendIndex->indexSize();
this->getHNSWIndex()->unlockSharedIndexDataGuard();
this->mainIndexGuard.unlock_shared();
this->flatIndexGuard.unlock_shared();
return res;
}

template <typename DataType, typename DistType>
size_t TieredHNSWIndex<DataType, DistType>::indexCapacity() const {
return this->backendIndex->indexCapacity() + this->frontendIndex->indexCapacity();
this->mainIndexGuard.lock_shared();
size_t capacity = this->backendIndex->indexCapacity();
this->mainIndexGuard.unlock_shared();
return capacity + this->frontendIndex->indexCapacity();
}

// In the tiered index, we assume that the blobs are processed by the flat buffer
Expand All @@ -721,7 +858,8 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
auto hnsw_index = this->getHNSWIndex();
// writeMode is not protected since it is assumed to be called only from the "main thread"
// (that is the thread that is exclusively calling add/delete vector).
if (this->getWriteMode() == VecSim_WriteInPlace) {
// VecSim_WriteInPlace is ignored during the accumulation phase
if (!this->isInAccumulationPhase && this->getWriteMode() == VecSim_WriteInPlace) {
// First, check if we need to overwrite the vector in-place for single (from both indexes).
if (!this->backendIndex->isMultiValue()) {
ret -= this->deleteVector(label);
Expand All @@ -739,7 +877,11 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
++this->directHNSWInsertions;
return ret;
}
if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
if (this->isInAccumulationPhase) {
// Accumulate the vectors in the running sum vector
auto storage_blob = this->frontendIndex->preprocessForStorage(blob);
this->addToSum(reinterpret_cast<const DataType *>(storage_blob.get()));
} else if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
// Handle overwrite situation.
if (!this->backendIndex->isMultiValue()) {
// This will do nothing (and return 0) if this label doesn't exist. Otherwise, it may
Expand All @@ -766,6 +908,10 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
if (this->frontendIndex->isLabelExists(label) && !this->frontendIndex->isMultiValue()) {
// Overwrite the vector and invalidate its only pending job (since we are not in MULTI).
auto *old_job = this->labelToInsertJobs.at(label).at(0);
if (this->isInAccumulationPhase) {
const DataType *vector_data = this->frontendIndex->getDataByInternalId(old_job->id);
this->subtractFromSum(vector_data);
}
old_job->id = this->setAndSaveInvalidJob(old_job);
this->labelToInsertJobs.erase(label);
ret = 0;
Expand Down Expand Up @@ -798,7 +944,7 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
// Here, a worker might ingest the previous vector that was stored under "label"
// (in case of override in non-MULTI index) - so if it's there, we remove it (and create the
// required repair jobs), *before* we submit the insert job.
if (!this->backendIndex->isMultiValue()) {
if (!this->isInAccumulationPhase && !this->backendIndex->isMultiValue()) {
// If we removed the previous vector from both HNSW and flat in the overwrite process,
// we still return 0 (not -1).
ret = std::max(ret - this->deleteLabelFromHNSW(label), 0);
Expand All @@ -812,8 +958,27 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
this->executeReadySwapJobs(this->pendingSwapJobsThreshold);
}

// Insert job to the queue and signal the workers' updater.
this->submitSingleJob(new_insert_job);
if (!this->isInAccumulationPhase) {
// Insert job to the queue and signal the workers' updater.
this->submitSingleJob(new_insert_job);
} else if (this->frontendIndex->indexSize() >= this->quantNormalizationSetSize) {
// If we are in the accumulation phase and we just reached the quantization set size, we
// can initialize the backend index and transition to the regular mode.
this->initializeQuantizedBackend();

// Submit all pending insert jobs to the job queue.
vecsim_stl::vector<AsyncJob *> jobs(this->allocator);
jobs.reserve(this->labelToInsertJobs.size());
for (auto &entry : this->labelToInsertJobs) {
for (auto *job : entry.second) {
jobs.push_back(job);
}
}
this->submitJobs(jobs);

// Transition complete.
this->isInAccumulationPhase = false;
}
return ret;
}

Expand All @@ -829,6 +994,10 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
// Invalidate the pending insert job(s) into HNSW associated with this label
auto &insert_jobs = this->labelToInsertJobs.at(label);
for (auto *job : insert_jobs) {
if (this->isInAccumulationPhase) {
const DataType *vector_data = this->frontendIndex->getDataByInternalId(job->id);
this->subtractFromSum(vector_data);
}
job->id = this->setAndSaveInvalidJob(job);
}
num_deleted_vectors += insert_jobs.size();
Expand All @@ -850,6 +1019,10 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
this->flatIndexGuard.unlock_shared();
}

if (this->isInAccumulationPhase) {
return num_deleted_vectors;
}

// Next, check if there vector(s) stored under the given label in HNSW and delete them as well.
// Note that we may remove the same vector that has been removed from the flat index, if it was
// being ingested at that time.
Expand Down Expand Up @@ -901,6 +1074,10 @@ double TieredHNSWIndex<DataType, DistType>::getDistanceFrom_Unsafe(labelType lab
// If the label doesn't exist, the distance will be NaN.
auto flat_dist = this->frontendIndex->getDistanceFrom_Unsafe(label, blob);

if (this->isInAccumulationPhase) {
return flat_dist;
}

// Optimization. TODO: consider having different implementations for single and multi indexes,
// to avoid checking the index type on every query.
if (!this->backendIndex->isMultiValue() && !std::isnan(flat_dist)) {
Expand Down
4 changes: 4 additions & 0 deletions src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/*
* Copyright (c) 2006-Present, Redis Ltd.
* All rights reserved.
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
*
* 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
Expand Down Expand Up @@ -65,6 +66,9 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarked
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test)

INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestSQ8)
INDEX_TEST_FRIEND_CLASS(SQ8TieredHNSWTest)

friend class CommonAPITest_SearchDifferentScores_Test;
friend class BF16TieredTest;
friend class FP16TieredTest;
Expand Down
12 changes: 12 additions & 0 deletions src/VecSim/index_factories/tiered_factory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/*
* Copyright (c) 2006-Present, Redis Ltd.
* All rights reserved.
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
*
* 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
Expand Down Expand Up @@ -91,12 +92,23 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) {
throw std::invalid_argument("Invalid hnsw_params.type");
}

VecSimQuantType quantType = hnsw_params.quantType;
size_t normSize = params->specificParams.tieredHnswParams.QuantNormalizationSetSize;
if (quantType != VecSimQuant_NONE && normSize > 0) {
// Add size of SQ accumulation buffer
est += allocations_overhead + hnsw_params.dim * sizeof(float);
}

return est;
}

VecSimIndex *NewIndex(const TieredIndexParams *params) {
// Tiered index that contains HNSW index as primary index
VecSimType type = params->primaryIndexParams->algoParams.hnswParams.type;
VecSimQuantType quantType = params->primaryIndexParams->algoParams.hnswParams.quantType;
if (quantType != VecSimQuant_NONE && type != VecSimType_FLOAT32 && type != VecSimType_FLOAT16) {
return nullptr;
}
if (type == VecSimType_FLOAT32) {
return TieredHNSWFactory::NewIndex<float>(params);
} else if (type == VecSimType_FLOAT64) {
Expand Down
Loading