[MOD-14957] Add SQ8 quantization support for tiered HNSW index - #5
[MOD-14957] Add SQ8 quantization support for tiered HNSW index#5xiangze-arm wants to merge 1 commit into
Conversation
|
|
||
| // Exclude readers while replacing the empty backend with its SQ8 counterpart. | ||
| this->lockMainIndexGuard(); | ||
| VecSimIndex_Free(this->backendIndex); |
There was a problem hiding this comment.
Freeing and replacing backendIndex here is not safe with the current locking. mainIndexGuard only excludes readers that take it, and several existing readers of backendIndex do not take any lock, because before this PR the pointer was immutable after construction:
vec_sim_tiered_index.h:160 preferAdHocSearch(), which runs on the query-planning pathvec_sim_tiered_index.h:149 getAllocationSize()vec_sim_tiered_index.h:100 getLabelsSet()vec_sim_tiered_index.h:173 fitMemory()vec_sim_tiered_index.h:357 debugInfo()vec_sim_tiered_index.h:22 the TIERED_LOG macro
A query thread in VecSimIndex_PreferAdHocSearch dereferences the index that line 822 just destroyed, so this is a use-after-free rather than a torn read.
Rather than adding lock_shared() to each of those call sites (which leaves the same trap for the next reader added), could we keep the backend object identity stable and re-initialize its SQ8 preprocessor and DistanceCalculatorWithNorm in place once the mean is known? The pointer then never changes and never becomes null, which also matches the intent of HLD 5.2 ("thebackend pointer transitions ... exactly once ... Once non-null, it stays non-null forever").
That would additionally let us drop the lock_shared() calls added to getNumMarkedDeleted, indexSize,indexCapacity, setLastSearchMode and indexMetaDataCapacity, and it removes the need for the accumulation-phase special case in newBatchIterator.
| // SQ accumulation phase members, only used for quantized tiered index | ||
| vecsim_stl::vector<float> runningSumVec; | ||
| size_t quantNormalizationSetSize; | ||
| bool isInAccumulationPhase; |
There was a problem hiding this comment.
isInAccumulationPhase is written by the main thread at line 980 and read concurrently by query paths at lines 268, 772, 792, and 1077 without synchronization. This is a C++ data race and undefined behavior.
Compare this with directHNSWInsertions above, which explicitly documents that it is non-atomic because it is accessed only from the main thread. That justification does not apply to this phase state.
The flag is also acting as a publication boundary: readers that observe the living phase must see the fully initialized backend and all initialization performed before the transition. A plain bool provides no such happens-before relationship.
If we implement HLD §5.2 literally, the phase should instead be represented by a one-time atomic nullptr → backend publication. If backend identity remains stable, we still need an atomic readiness/phase state, with a release transition after initialization and acquire loads in query paths.
This only orders readers that consult the phase state; it does not fix the unguarded backend readers described at line 822.
HLD Risk 4 also calls for ThreadSanitizer coverage of Phase 1 under concurrent queries. Could we add a synchronized TSan test that crosses the threshold while threads execute topKQuery()/rangeQuery() and an otherwise unguarded backend reader such as preferAdHocSearch()? Using a barrier or transition test hook would make the overlap deterministic rather than relying on a very small race window.
| // 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); |
There was a problem hiding this comment.
submitJobs() makes the pending jobs runnable and may wake workers before it returns. executeInsertJob() first inserts a vector into HNSW and then removes its frontend copy. Therefore, a worker can complete that migration before line 980 publishes the living phase.
During this window, topKQuery() and rangeQuery() still observe the accumulation phase and search only the frontend. newBatchIterator() can likewise return a frontend-only iterator. Any already-migrated vector is consequently omitted, including vectors whose earlier addVector() calls completed before the query was submitted.
This remains incorrect even if isInAccumulationPhase is made atomic: the publication is simply too late. HLD §3.2 requires queries to search and merge both tiers while the pending insert jobs drain.
The query-visible living state must be published before the jobs become runnable. With an explicit phase state, perform a release store immediately after successful backend initialization and before submitJobs(), with acquire loads in query paths. With the HLD’s pointer-based design, atomically publish the initialized backend before submitting the jobs.
Keeping backend identity stable does not by itself fix the ordering while query paths continue to skip the backend based on the phase state.
Could we add a deterministic concurrency test whose submission callback allows one migration job to finish, then runs a query before returning? That would exercise this exact window without relying on scheduler timing.
| // 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()) { |
There was a problem hiding this comment.
This fix is correct, and the same reasoning applies to rangeQuery() below, but the batch iterator was not updated and still uses the withSet=false path. Its relevant code is outside the diff, so I’m raising it here.
TieredHNSW_BatchIterator::getNextResults() derives isMulti only from backendIndex->isMultiValue() and dispatches single-value indexes to compute_current_batch<false>() at lines 1226–1229. That optimization is safe only when a shared label has an identical score in both tiers. SQ8 breaks this invariant because the frontend uses a full-precision score while the backend uses a quantized score.
During migration, a label can appear in both snapshots after its backend insertion but before its frontend copy is removed. This can produce duplicates:
Within one batch, merge_results<false> can emit both differently scored copies.
Across batches, one copy can be returned while the other remains buffered or appears later. The single-value path tracks only IDs consumed from the flat results and does not filter both remaining lists, so the second copy can be returned by a later getNextResults() call.
This violates the batch-iterator contract that a result is not returned twice. Could we use the same condition here, e.g. needsDedup = isMulti || index->isQuantized, both for the filtering decision around line 1187 and for selecting compute_current_batch<true>() at line 1226?
Please also add a test that creates the backend-inserted/frontend-not-yet-removed overlap and drains multiple batches while asserting that every label is returned at most once.
| if (isInAccumulationPhase) { | ||
| return this->frontendIndex->newBatchIterator(queryBlob, queryParams); | ||
| } |
There was a problem hiding this comment.
Returning the frontend’s own BF_BatchIterator here drops two guarantees required by a tiered index.
First, synchronization: BF_BatchIterator::calculateScores() iterates the frontend’s vectors and label mappings without taking flatIndexGuard. That is appropriate for a standalone BF index, but not here, where the main thread can concurrently add/delete vectors and workers can remove migrated vectors. TieredHNSW_BatchIterator::getNextResults() takes flatIndexGuard around this scan for exactly that reason.
Second, the iterator type is selected permanently at construction. An iterator created during accumulation remains frontend-only even if its first getNextResults() call occurs after the transition. By then workers may have moved vectors that existed when the iterator was created into the backend, but this iterator can never see them. Calling reset() after transition has the same problem.
A deterministic example is:
Create the iterator immediately below the training threshold.
Cross the threshold and drain one or more pending insert jobs.
Call getNextResults() for the first time.
It scans only the remaining frontend and omits every previously existing vector already migrated to HNSW.
With the current non-null empty placeholder, could we always return TieredHNSW_BatchIterator and let it query the empty backend normally? Its construction should also occur while holding flatIndexGuard, because constructing the internal BF iterator reads the frontend label count.
If the backend instead follows HLD §5.2 and is null during Phase 0, the tiered iterator should be made phase/null-aware rather than exposing a raw frontend iterator. This keeps locking and phase-transition behavior inside one iterator implementation.
Please add a cross-transition test that creates the iterator before the threshold, migrates vectors before its first Next(), and verifies that all labels existing at iterator creation remain visible. A concurrent Next()/frontend-mutation test under TSan would cover the locking issue separately.
| if (this->frontendIndex->indexSize() >= this->flatBufferLimit) { | ||
| if (this->isInAccumulationPhase) { | ||
| // Accumulate the vectors in the running sum vector | ||
| auto storage_blob = this->frontendIndex->preprocessForStorage(blob); |
There was a problem hiding this comment.
During Phase 0 this preprocesses the vector solely to update the running sum, and a new frontend insertion preprocesses it again through BruteForceIndex::appendVector(). For cosine indexes that means two allocations, copies, and normalization passes for every newly appended vector during training.
Could we update the sum from the frontend’s canonical stored vector after insertion, using new_flat_id, or otherwise add a preprocess-once insertion path that lets both storage and accumulation consume the same result?
Reading back getDataByInternalId(new_flat_id) looks suitable for newly appended vectors and matches the existing subtraction paths. Please check the single-value overwrite case separately, because its BF path currently updates the stored element directly rather than going through appendVector().
|
|
||
| bool isQuantized{false}; | ||
|
|
||
| // SQ accumulation phase members, only used for quantized tiered index |
There was a problem hiding this comment.
Could the main-thread-only Phase-0 data be grouped into an SQAccumulationState, containing runningSumVec and the saved backend parameters? That would make their temporary lifetime explicit and allow releasing the running-sum allocation once backend initialization completes.
I would keep the query-visible readiness state separate and atomic; using the presence of an optional<SQAccumulationState> as the phase indicator would not be safe across threads. isQuantized and the configured training threshold should also remain index-lifetime configuration, since they are needed after training for query behavior and reporting.
Describe the changes in the pull request
Enable SQ8 quantization for tiered HNSW index. The index now accumulates a configurable initial vector set to calculate the mean, builds the quantized backend, and then migrates pending vectors through the normal tiered flow.
Which issues this PR fixes
Main objects this PR modified
TieredHNSWIndex— Add support for SQ8 and accumulation phase etc.TieredFactory— Consider quantization inEstimateInitialSize()andNewIndex()Mark if applicable