Skip to content

feat(ivf): support RaBitQ x+y split storage and search - #2707

Open
LightWant wants to merge 12 commits into
antgroup:mainfrom
LightWant:feat/issue-2443
Open

feat(ivf): support RaBitQ x+y split storage and search#2707
LightWant wants to merge 12 commits into
antgroup:mainfrom
LightWant:feat/issue-2443

Conversation

@LightWant

@LightWant LightWant commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

性能结果

本 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=32factor=10、top-10、index 内部 parallelism=1

PR 前传统 IVF RaBitQ 与最终 residual 1+7

版本 布局 构建时间 构建 TPS 索引内存 索引文件 单线程 QPS Avg Recall@10
PR 前 传统 8-bit RaBitQ 149.369 s 6,694.851 1,284,595,232 B 946,484,662 B 16.784 59.576 ms 0.9166
PR 最终 residual 1-bit filter + 7-bit supplement 144.253 s 6,932.263 1,340,849,360 B 1,026,371,655 B 576.11 1.733 ms 0.9184

最终 residual 1+7 相对传统 8-bit:构建时间减少约 3.4%,构建 TPS 提升约 3.5%;单线程搜索吞吐约为 34.3 倍,Recall@10 提升 0.0018。

提交前 5 轮搜索回归

复用最终格式的 GIST1M residual 1+7 索引,固定 CPU core,预热后运行 5 轮并取中位数:

QPS Avg P50 P95 Recall@10 扫描/query 精排/query x-bit IP 命中 fallback
576.11 1.733 ms 1.652 ms 2.475 ms 0.9184 60,274.60 100 100 0

最终工程范围

IVF 分层

  • 仅当 IVF flat bucket 使用 RaBitQ split storage 时选择专用 RaBitQSplitBucketSearcher
  • 普通 FlatBucketSearcher 保持原搜索判定;Graph bucket 只做必要的 workspace 参数转发,不接入 split scan,也不改变图搜索策略。
  • IVF searcher 负责过滤、threshold、reasoning 和候选 heap;Bucket/DataCell 负责 split block 扫描、版本校验及 supplement 读取。

保留的 500+ QPS 路径

  • filter code 按 32-vector FastScan block 打包,1+7、2+6、3+5 共用 HighAcc query LUT 与 bit-plane 加权。
  • supplement 使用紧凑 5/6/7-bit 布局,并提供 generic、AVX2、AVX512 运行时分派。
  • candidate heap 携带量化 x-bit 内积以及 source bucket、offset、version;正常精排只读取 ybits,不重新读取 xbits。
  • residual candidate 使用 original-query full factors;正常路径不创建 bucket-specific q-c computer,只有 stale/invalid lane 进入完整 fallback。
  • optimized build、Add/Update/Delete、relocation、Merge 和 Serialize/Deserialize 都维护相同 packed layout 与版本语义。

已删除的实验代码

  • 可选 lower-bound heap 搜索策略及 YAML 参数。
  • 低精度 uint8 FastScan query LUT 和重复 lower-bound/filter-IP 接口。
  • dense whole-bucket filter-IP cache、旧 centroid-IP conversion、旧 residual computer factory。
  • 通用 FlatBucketSearcher 中与 split 无关的 SIMD candidate 快路径。
  • 对应的实验文档、性能表和重复测试。

清理提交删除 2,496 行、增加 191 行;PR Changes 从 14,246 行新增降为 11,940 行新增。剩余改动主要由 split bucket/DataCell、跨架构 SIMD、packed layout、并发/序列化正确性及测试组成。

测试与检查

  • clang-format 15 dry-run、git diff --check:通过。
  • 完整 Release 构建(library、unittests、functests、eval_performance):通过。
  • split/IVF/RaBitQ 定向测试:43 cases / 3,669,967 assertions 通过。
  • 专用 split searcher 与参数测试:3 cases / 145 assertions 通过。
  • GIST1M residual 1+7:5 轮中位数 576.11 QPS,Recall@10 0.9184。

评测参数

global:
  num_threads_searching: 1

residual_1_7_candidate:
  datapath: /root/data/gist1M/gist-960-euclidean.hdf5
  type: search
  index_name: ivf
  index_path: /tmp/vsag-pr-2443-consolidated-gist/residual-1-7-v4.index
  create_params: >-
    {"dtype":"float32","metric_type":"l2","dim":960,"index_param":{
      "buckets_count":1024,"base_quantization_type":"rabitq",
      "precise_quantization_type":"rabitq","rabitq_bits_per_dim_query":32,
      "rabitq_bits_per_dim_base":1,"rabitq_bits_per_dim_precise":7,
      "rabitq_error_rate":1.9,"use_reorder":true,"use_residual":true,
      "partition_strategy_type":"ivf","ivf_train_type":"kmeans",
      "train_sample_count":100000,"thread_count":32,
      "fast_encode_rabitq":true,"fast_encode_rabitq_rounds":12}}
  search_params: >-
    {"ivf":{"scan_buckets_count":32,"factor":10.0,"parallelism":1}}
  search_mode: knn
  topk: 10
  search_query_count: 1000
  delete_index_after_search: false

Fixes: #2443

@vsag-bot

vsag-bot commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

/label status/waiting-for-review
/waiting-on reviewer
/request-review @jiaweizone
/request-review @wxyucs
/request-review @inabao

@mergify

mergify Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

🟢 All 3 merge protections satisfied — ready to merge.

Show 3 satisfied protections

🟢 Require kind label

  • label~=^kind/

🟢 Require version label

  • label~=^version/

🟢 Require linked issue for feature/bug PRs

  • body~=(?im)(?:^|[\s\-\*])(?:close[sd]?|fix(?:e[sd])?|resolve[sd]?)\s*:?\s+(?:#\d+|[\w.\-]+/[\w.\-]+#\d+|https?://github\.com/[\w.\-]+/[\w.\-]+/issues/\d+)

@LightWant

Copy link
Copy Markdown
Collaborator Author

Maintainers: GitHub rejected the fork author's attempt to set upstream labels. Please add kind/feature and version/1.1, matching linked issue #2443.

@LightWant LightWant added kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 version/1.0 labels Aug 18, 2026
@mergify mergify Bot added the module/simd label Aug 18, 2026
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
Comment thread src/algorithm/ivf/ivf.cpp
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)) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[suggestion] IsFiniteFloatBits (or IsFiniteFloat) is duplicated in at least 4 files:

  • src/datacell/rabitq_split_bucket_datacell.cpp:33 (as IsFiniteFloat)
  • src/algorithm/ivf/rabitq_split_bucket_searcher.cpp:33 (as IsFiniteFloatBits)
  • src/algorithm/ivf/flat_bucket_searcher.cpp:31 (as IsFiniteFloatBits)
  • src/quantization/rabitq_quantization/rabitq_quantizer.cpp:61 (as IsFiniteFloatBits)

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() {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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:

  1. The thread_local initialization via lambda with std::random_device is correct and ensures each thread gets an independent random stream.
  2. The PRNG passes the random_state pointer through the call chain (PrepareFastScan32QueryFastScanUniformRandom), 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

  1. [critical] std::isfinite in SplitBucketReorder::Reorder (src/algorithm/ivf/ivf.cpp:614): The PR correctly replaced std::isfinite/std::isnan with IEEE-754 bit-based checks in searcher_utils.h and flat_bucket_searcher.cpp to avoid -Ofast elimination. However, std::isfinite is still used in SplitBucketReorder::Reorder, which will silently produce incorrect results under -Ofast. This should use the same IsFiniteFloatBits helper.

  2. [suggestion] IsFiniteFloatBits duplicated 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.h or a new fp32_utils.h) to avoid drift and ensure consistency.

  3. [note] FastScanFallbackRandomState thread_local PRNG (rabitq_quantizer.cpp:39): The SplitMix64 PRNG uses thread_local state 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.

  4. [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.cpp ensure 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_adjustment correctly 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: IVFBucketOptimizedBuildSession RAII 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.

}
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();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 LHT129 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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_cache
  • get_cached_filter_inner_product
  • publish_candidate_scan_version
  • publish_heap_scan_version
  • candidate_scan_version_matches
  • heap_scan_version_matches
  • All routed_filter_cache_* fields in SplitBucketComputer

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/feature Brand-new functionality or capabilities 引入全新的功能、新特性或新能力 module/docs module/simd module/testing size/XXL version/1.1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[feat](ivf): support RaBitQ x+y split storage and search

4 participants