feat(ivf): support RaBitQ x+y split storage and search - #2707
Conversation
|
/label status/waiting-for-review |
Merge Protections🟢 All 3 merge protections satisfied — ready to merge. Show 3 satisfied protections🟢 Require kind label
🟢 Require version label
🟢 Require linked issue for feature/bug PRs
|
|
Maintainers: GitHub rejected the fork author's attempt to set upstream labels. Please add |
5ff8e2c to
261a807
Compare
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Consolidate the validated packed layout, inner-product reuse, residual query path, optimized build, and regression coverage for issue antgroup#2443. Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.6
d66f371 to
04c4d08
Compare
| ctx.reasoning_ctx->RecordReorder(ids[i], hints[i], distances[i]); | ||
| } | ||
| if (distance_threshold.has_value() and | ||
| (not std::isfinite(distances[i]) or distances[i] > *distance_threshold)) { |
There was a problem hiding this comment.
[critical] std::isfinite is used directly in SplitBucketReorder::Reorder but the same PR replaces std::isfinite with IsFiniteFloatBits in flat_bucket_searcher.cpp and searcher_utils.h to avoid -Ofast compatibility issues. This location was missed and will break under -Ofast builds.
The searcher_utils.h change explicitly documents the motivation: -Ofast eliminates std::isfinite/std::isnan checks because they are considered "finite math only" operations. The SplitBucketReorder class uses std::isfinite(distances[i]) directly in the distance threshold check, which would be eliminated under -Ofast.
Suggested fix: Replace with IsFiniteFloatBits(distances[i]) or use the same IEEE-754 bit-check pattern from searcher_utils.h.
| namespace { | ||
|
|
||
| bool | ||
| IsFiniteFloat(float value) { |
There was a problem hiding this comment.
[suggestion] IsFiniteFloatBits (or IsFiniteFloat) is duplicated in at least 4 files:
src/datacell/rabitq_split_bucket_datacell.cpp:33(asIsFiniteFloat)src/algorithm/ivf/rabitq_split_bucket_searcher.cpp:33(asIsFiniteFloatBits)src/algorithm/ivf/flat_bucket_searcher.cpp:31(asIsFiniteFloatBits)src/quantization/rabitq_quantization/rabitq_quantizer.cpp:61(asIsFiniteFloatBits)
Each is a copy of the same IEEE-754 bit-check pattern. Consider centralizing this in a shared utility header (e.g., simd/fp32_simd.h or a new utils/float_utils.h) to avoid duplication and ensure consistent usage across the codebase. The searcher_utils.h already has a similar is_finite_distance_bits helper that could serve as the canonical version.
| namespace { | ||
|
|
||
| uint64_t& | ||
| FastScanFallbackRandomState() { |
There was a problem hiding this comment.
[note] The FastScanFallbackRandomState() uses thread_local state initialized with std::random_device. This is a splittable PRNG (likely SplitMix64 variant) used for stochastic rounding in the FastScan32 LUT quantization.
Two observations:
- The
thread_localinitialization via lambda withstd::random_deviceis correct and ensures each thread gets an independent random stream. - The PRNG passes the
random_statepointer through the call chain (PrepareFastScan32Query→FastScanUniformRandom), allowing deterministic replay when an explicit state is provided. This is a good design for reproducibility.
| std::shared_lock bucket_lock(this->bucket_mutexes_[bucket_id]); | ||
| std::lock_guard codes_lock(this->codes_insert_mutex_); | ||
| { | ||
| // Writers publish locations while holding the corresponding bucket lock and follow the |
There was a problem hiding this comment.
[note] The lock ordering in query_current_by_inner_id is carefully documented: bucket lock → codes lock → locations lock (re-check). The comment on lines 1067-1069 explains that writers publish locations while holding the bucket lock and follow bucket → codes → locations order. Taking the same locks before rechecking ensures a consistent snapshot of lane + supplement without lock inversion.
This is correct but subtle. Consider adding a static assertion or a centralized lock-ordering document if more lock-taking paths are added in the future.
LHT129
left a comment
There was a problem hiding this comment.
Overall Review
This is a substantial and well-engineered PR that adds RaBitQ x+y split storage and search to the IVF index. The architecture is sound and the implementation is thorough, with strong attention to parameter validation, concurrency, and SIMD optimization.
Summary of findings
Issues identified (4 inline comments):
-
[critical]
std::isfiniteinSplitBucketReorder::Reorder(src/algorithm/ivf/ivf.cpp:614): The PR correctly replacedstd::isfinite/std::isnanwith IEEE-754 bit-based checks insearcher_utils.handflat_bucket_searcher.cppto avoid-Ofastelimination. However,std::isfiniteis still used inSplitBucketReorder::Reorder, which will silently produce incorrect results under-Ofast. This should use the sameIsFiniteFloatBitshelper. -
[suggestion]
IsFiniteFloatBitsduplicated across 4 files (rabitq_split_bucket_datacell.cpp:33,rabitq_split_bucket_searcher.cpp:33,flat_bucket_searcher.cpp:31,rabitq_quantizer.cpp:61): The same IEEE-754 bit-checking helper is defined as an anonymous-namespace function in four separate translation units. This should be extracted into a shared utility header (e.g.,simd/fp32_simd.hor a newfp32_utils.h) to avoid drift and ensure consistency. -
[note]
FastScanFallbackRandomStatethread_local PRNG (rabitq_quantizer.cpp:39): The SplitMix64 PRNG usesthread_localstate for reproducible random quantization in FastScan32. This is fine for the current use case but is not seedable per-call — if reproducibility across runs is ever needed, consider making the seed a parameter. -
[note] Lock ordering in
query_current_by_inner_id(rabitq_split_bucket_datacell.cpp:1067): The lock ordering (bucket → codes → locations) is carefully documented and consistent with the writer path. This is correct but subtle — consider a centralized lock-ordering document if more lock-taking paths are added.
What is done well
- Parameter validation: Comprehensive checks in
ivf.cppensure split storage is only enabled with compatible settings (both quantizers are rabitq,use_reorder=true,buckets_per_data=1, no bucket graphs). This prevents silent misconfiguration. - Concurrency design: Per-bucket mutexes + version counters + lock ordering discipline provide correct concurrent read/write without global locks. The version-based cache invalidation for filter inner products is well thought out.
- SIMD coverage: FastScan32 byte-LUT scanning and packed supplement kernels are implemented for AVX2, AVX512, SSE, NEON, and SVE — full architecture coverage.
- Residual L2 encoding: The residual query transform with per-lane factors and
residual_adjustmentcorrectly handles the L2SQR residual bias, which is easy to get wrong. - Two search strategies: Candidate reorder (scan x-bits, reorder top candidates with y-bits) and heap search (lower-bound pruning) give users a choice between throughput and precision.
- Build lifecycle:
IVFBucketOptimizedBuildSessionRAII wrapper ensures proper cleanup on error paths. - Documentation: Both EN and ZH docs are updated for the new parameters.
Recommendation
The PR is in good shape. The one critical issue (missing -Ofast fix in SplitBucketReorder) should be addressed before merge. The duplicated IsFiniteFloatBits helper is a maintenance concern worth fixing but not a correctness issue. The other notes are informational.
04c4d08 to
94df919
Compare
| } | ||
| const float lower_bound = lower_bounds[j]; | ||
| const bool has_usable_lower_bound = | ||
| std::isfinite(lower_bound) and lower_bound < std::numeric_limits<float>::max(); |
There was a problem hiding this comment.
[critical] std::isfinite is used in the RaBitQ heap-search path (lines 179 and 205) but the same PR introduces IsFiniteFloatBits to avoid -Ofast compatibility issues elsewhere. Under -Ofast, std::isfinite is eliminated as a "finite math only" operation, causing the lower-bound filter (not(lower_bound < distance_limit)) and the exact-distance validity check to be silently removed.
Line 179: std::isfinite(lower_bound) and lower_bound < std::numeric_limits<float>::max()
Line 205: not std::isfinite(exact_distance)
Suggested fix: Replace both with IsFiniteFloatBits(...) (already defined in the anonymous namespace of this file at line 33).
|
|
||
| void | ||
| RaBitQSplitBucketDataCell::query_non_residual_by_inner_ids( | ||
| float* result_dists, |
There was a problem hiding this comment.
[suggestion] query_non_residual_by_inner_ids (~200 lines) and query_residual_by_inner_ids (~200 lines) share the same structure: sort entries by bucket, group by bucket, check staleness, classify into reusable/fallback, batch query with filter inner products, then handle stale entries via query_current_by_inner_id. The only differences are residual-specific handling (transformed centroids, original-query filter inner product conversion, full_add computation).
Consider extracting the common grouping/staleness/batching logic into a shared helper to reduce the ~400 lines of near-duplicate code. This would make bug fixes and future optimizations apply to both paths automatically.
| void | ||
| AuxiliaryStandardHeap<max_heap, fixed_size>::PushWithAuxiliary(float dist, | ||
| InnerIdType id, | ||
| float auxiliary, |
There was a problem hiding this comment.
[note] Push(dist, id) delegates to PushWithAuxiliary passing -1 for source_bucket_id. Since BucketIdType is unsigned (uint32_t in bucket_datacell.h), -1 wraps to UINT32_MAX, which serves as a sentinel. This is intentional but subtle — a named constant like kUnknownBucketId would make the intent clearer and avoid confusion about the sign mismatch.
| @@ -0,0 +1,2591 @@ | |||
| // Copyright 2024-present the vsag project | |||
There was a problem hiding this comment.
[note] The new rabitq_split_bucket_datacell.cpp (2591 lines) and rabitq_split_bucket_searcher.cpp (353 lines) introduce ~2900 lines of new core logic without a dedicated unit test file. The existing test changes in bucket_datacell_test.cpp (+1213 lines) and flat_bucket_searcher_test.cpp (+799 lines) cover the generic bucket/searcher interfaces, but the split-bucket-specific code paths (filter inner product caching, version-aware staleness checks, routed bucket computer management, residual vs non-residual query dispatch) would benefit from focused tests that exercise the split-storage edge cases directly.
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5.6
LHT129
left a comment
There was a problem hiding this comment.
[critical] std::isfinite regression in flat_bucket_searcher.cpp
This commit removed the local IsFiniteFloatBits implementation from src/algorithm/ivf/flat_bucket_searcher.cpp and replaced it with std::isfinite at line 89 ((not std::isfinite(dist[j])). However, the same PR introduces IsFiniteFloatBits in rabitq_split_bucket_searcher.cpp specifically to avoid -Ofast compatibility issues — -Ofast eliminates std::isfinite/std::isnan checks because they are treated as "finite math only" operations.
This is inconsistent: rabitq_split_bucket_searcher.cpp correctly retains IsFiniteFloatBits, but flat_bucket_searcher.cpp regresses to std::isfinite. The distance threshold check in FlatBucketSearcher::Search will be silently removed under -Ofast builds.
Suggested fix: Restore IsFiniteFloatBits (or use the centralized version from searcher_utils.h) and replace std::isfinite(dist[j]) with IsFiniteFloatBits(dist[j]).
LHT129
left a comment
There was a problem hiding this comment.
[suggestion] SIMD fast path removed from FlatBucketSearcher::Search
This commit removes the SIMD-accelerated candidate push path (~30 lines) from FlatBucketSearcher::Search that used FP32LessThan32Mask for batch distance comparison when no filters or reasoning context are active. This path was a general optimization for the flat bucket searcher — it is not specific to RaBitQ split storage.
The removed code provided a meaningful throughput improvement for the common case of unfiltered KNN search by processing 32 distances at a time with a single SIMD comparison, avoiding per-element scalar checks. Without it, every candidate goes through the slower per-element loop even when no filters are configured.
If this removal is intentional (e.g., the SIMD path is being moved elsewhere or the performance impact was measured and deemed acceptable), please clarify in the commit message. Otherwise, consider restoring this optimization for the non-split flat bucket searcher path.
LHT129
left a comment
There was a problem hiding this comment.
[note] FilterInnerProductMode::CACHED enum value is now dead code
The FilterInnerProductMode enum in src/datacell/rabitq_split_bucket_datacell.h still contains CACHED as a member, but this commit removes all code paths that use it:
claim_filter_inner_product_cacheget_cached_filter_inner_productpublish_candidate_scan_versionpublish_heap_scan_versioncandidate_scan_version_matchesheap_scan_version_matches- All
routed_filter_cache_*fields inSplitBucketComputer
The CACHED mode is no longer reachable. Consider removing the enum value and the associated SplitBucketComputer fields (routed_filter_inner_products_, routed_filter_cache_claimed_, routed_filter_cache_versions_, routed_candidate_scan_versions_, routed_heap_scan_versions_) if they are truly no longer needed, to keep the codebase clean and avoid confusion for future readers.
性能结果
本 PR 让 IVF 的 flat bucket 支持 RaBitQ x+y split 存储与检索;未配置 split 参数时继续使用原有 IVF RaBitQ,bucket graph 路径不接入 split searcher。
测试环境:GIST1M,1,000,000 条 960 维 float32 L2 向量,1,000 queries,1,024 个 IVF buckets,100,000 条训练样本;构建使用 32 线程、12 轮 fast encode;搜索使用
nprobe=32、factor=10、top-10、index 内部parallelism=1。PR 前传统 IVF RaBitQ 与最终 residual 1+7
最终 residual 1+7 相对传统 8-bit:构建时间减少约 3.4%,构建 TPS 提升约 3.5%;单线程搜索吞吐约为 34.3 倍,Recall@10 提升 0.0018。
提交前 5 轮搜索回归
复用最终格式的 GIST1M residual 1+7 索引,固定 CPU core,预热后运行 5 轮并取中位数:
最终工程范围
IVF 分层
RaBitQSplitBucketSearcher。FlatBucketSearcher保持原搜索判定;Graph bucket 只做必要的 workspace 参数转发,不接入 split scan,也不改变图搜索策略。保留的 500+ QPS 路径
q-ccomputer,只有 stale/invalid lane 进入完整 fallback。已删除的实验代码
清理提交删除 2,496 行、增加 191 行;PR Changes 从 14,246 行新增降为 11,940 行新增。剩余改动主要由 split bucket/DataCell、跨架构 SIMD、packed layout、并发/序列化正确性及测试组成。
测试与检查
git diff --check:通过。评测参数
Fixes: #2443