Skip to content

perf(hgraph): fuse RaBitQ split storage and search for x=1..4 - #2596

Open
LightWant wants to merge 1 commit into
antgroup:mainfrom
LightWant:feat/hgraph-rabitq-fused-x1to4
Open

perf(hgraph): fuse RaBitQ split storage and search for x=1..4#2596
LightWant wants to merge 1 commit into
antgroup:mainfrom
LightWant:feat/hgraph-rabitq-fused-x1to4

Conversation

@LightWant

@LightWant LightWant commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Change Type

  • Bug fix
  • New feature
  • Improvement/Refactor
  • Documentation
  • CI/Build/Infra

Linked Issue

What Changed

  • Add the opt-in rabitq_fused_datacell HGraph layout for RaBitQ split 1+y, 2+y, 3+y, and 4+y (x+y <= 8); the default remains disabled.
  • Co-locate bottom-layer neighbors, cluster id, label, x-bit code, and y-bit supplement in a cache-line-aligned node slab, with 16 reproducibly trained residual clusters and single-copy fused serialization.
  • Add dedicated fused search paths and runtime-dispatched SIMD kernels: HNSW-compatible 4-bit query bit-plane/popcount traversal for 1+y, and exact x-bit filter-inner-product reuse for 2+y through 4+y so reorder scans only the supplement planes.
  • Preserve HGraph lifecycle and query behavior, including Add/Update, delete/restore, resize/move, ordinary/streaming serialization, filters, duplicate aliases, range search, iterators, timeout/hop limits, reasoning, and mutable graph locks.
  • Add explicit validation for fused-v1 constraints: memory IO, L2/IP, no MCI, no deduplicated storage, no force remove, no PCA, and no legacy v0.14 serialization. Existing split indexes and configurations keep their original paths.

Test Evidence

  • make fmt
  • make lint
  • make test
  • make cov, run tests, and collect coverage
  • Other (described below)

Test details:

clang-format-15 --dry-run --Werror <all changed C++ files>
# passed

cmake --build build-pr-tests --target unittests functests -j32
# Release/-Ofast build passed

./scripts/linters/run-clang-tidy-15.sh ... <all 20 changed production TUs>
# 20/20 passed with clang-tidy 15

./build-pr-tests/tests/unittests "[fused_zero_residual]"
# 224 assertions / 1 case passed (L2/IP x filter bits 1-4, exact coarse/full)

./build-pr-tests/tests/unittests "[HGraphRaBitQSearcher][route]"
# 147 assertions / 1 case passed (route switch, scalar/batch fallback, overflow salvage)

./build-pr-tests/tests/unittests "[RaBitQSplitDataCell]"
# 14,323 assertions / 11 cases passed

./build-pr-tests/tests/functests "HGraph fused RaBitQ split honors disabled reorder"
# 952 assertions / 1 case passed; normal full/fallback/reorder counters stay zero

./build-pr-tests/tests/functests "HGraph fused RaBitQ split keeps zero-residual nodes without reorder"
# 120 assertions / 1 case passed

./build-pr-tests/tests/functests "[rabitq_split]"
# 12,670 assertions / 14 cases passed

git diff --check
# passed

GIST1M benchmark: 960-d L2, HGraph M=32, efConstruction=200, 16 clusters, MCI disabled, topK=10, 32 query threads, efSearch=10,20,40,80,160,320,640,1280. Each point used 20 searches per process and 3 paired alternating runs.

split fused mean-QPS gain over original split worst Recall@10 delta points faster
1+7 +150.79% to +239.52% -0.0050 8/8
2+6 +136.02% to +164.22% -0.0038 8/8
3+5 +109.08% to +133.76% -0.0008 8/8
4+4 +87.49% to +101.95% -0.0040 8/8

All 32 measured efSearch/split points outperform the existing split implementation.

Compatibility Impact

  • API/ABI compatibility: additive configuration symbol only; existing parameter defaults and serialized index paths are unchanged.
  • Behavior changes: only when rabitq_fused_datacell=true; unsupported fused-v1 combinations fail validation explicitly.

Performance and Concurrency Impact

  • Performance impact: improved for all measured fused x=1..4 variants; node/code locality, joint prefetch, direct IP hints, and dedicated SIMD kernels remove split-storage reads and duplicate x-code scans.
  • Concurrency/thread-safety impact: fused direct record reads use the HGraph shared mutex array; parallel search safely falls back to the basic searcher if no thread pool exists.

Documentation Impact

  • No docs update needed
  • Updated docs:
    • README.md
    • DEVELOPMENT.md
    • CONTRIBUTING.md
    • Other: docs/docs/en/src/quantization/rabitq_split.md, docs/docs/zh/src/quantization/rabitq_split.md

Risk and Rollback

  • Risk level: medium; this adds a substantial opt-in storage/search implementation but leaves the default path untouched.
  • Rollback plan: disable rabitq_fused_datacell for immediate configuration rollback, or revert commit 3305f72bd.

Checklist

  • I have linked the relevant issue (required for kind/bug and kind/feature; see "Linked Issue" above)
  • I have added/updated tests for new behavior or bug fixes
  • I have considered API compatibility impact
  • I have updated docs if behavior/workflow changed
  • My commit messages follow project conventions (Conventional Commits, optional [skip ci] prefix)

@vsag-bot

vsag-bot commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator

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

@vsag-bot

vsag-bot commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Automated pull request review completed.

Review effort: high (9605 changed lines across 61 files).

Submitted 2 inline comments.
Review: #2596 (review)

@mergify

mergify Bot commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

Merge Protections

🟢 All 2 merge protections satisfied — ready to merge.

Show 2 satisfied protections

🟢 Require kind label

  • label~=^kind/

🟢 Require version label

  • label~=^version/

@LightWant LightWant added kind/improvement Optimizations, UX polish, or minor improvements 性能优化、体验打磨或细节改良 version/1.1 labels Aug 1, 2026
Comment thread src/utils/visited_list.h
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp
Comment thread src/quantization/rabitq_quantization/rabitq_quantizer.cpp
Comment thread src/quantization/rabitq_quantization/rabitq_quantizer.cpp
Comment thread src/quantization/rabitq_quantization/rabitq_quantizer.h
Comment thread src/algorithm/hgraph/hgraph_build.cpp
Comment thread src/algorithm/hgraph/hgraph_serialize.cpp
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp

@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.

PR Review Summary

This PR adds an opt-in rabitq_fused_datacell feature that co-locates bottom-layer HGraph neighbors, cluster id, label, x-bit filter code, and y-bit supplement code into a single cache-line-aligned node record. It supports x=1..4 filter bits with x+y <= 8, uses 16 reproducibly trained residual clusters, and adds dedicated SIMD kernels (AVX2, AVX512, AVX512VPOPCNTDQ, generic).

Overall Assessment

The implementation is well-structured and thorough. The fused layout eliminates an indirection from the search hot path, and the dedicated HGraphRaBitQSearcher with AffineScorer templates provides a clean separation from the existing search path. Test coverage is good with functional tests for round-trip serialization, duplicate label handling, and alias expansion.

Key Strengths

  • Clean opt-in design: The fused path is gated behind a single boolean parameter and does not affect the existing code path
  • Good test coverage: Integration tests cover serialization round-trips, duplicate labels, alias expansion, and iterator search for the fused path
  • SIMD specialization: Properly dispatches to AVX2/AVX512/AVX512VPOPCNTDQ kernels with a generic fallback
  • Safe feature gating: Correctly disables incompatible features (merge, tune, concurrent mutations) at the InitFeatures level
  • Fast-math safety: IsFiniteRaBitQValue uses bitwise checks to avoid -ffast-math issues with std::isfinite

Issues Found (see inline comments for details)

  1. Thread safety: VisitedList::TestAndSet uses non-atomic operations but is called under shared lock in the fused searcher
  2. Numerical robustness: filter_rescale computation in EncodeFusedAffineMetadata can overflow for edge-case inputs
  3. Code organization: search_one_graph is ~700 lines and would benefit from being split into smaller functions
  4. Serialization versioning: The new format lacks an embedded version tag for forward compatibility
  5. Performance: BoundedResults::Insert is O(n) per insert; acceptable for typical ef but worth documenting
  6. Encoding performance: EncodeFusedAffineMetadata decodes bit-planes per dimension, which could be slow for large datasets

Recommendation

Approve with suggestions. The issues identified are non-blocking for correctness — the code handles edge cases gracefully (NaN fallbacks, feature gating). The thread safety concern in TestAndSet is pre-existing in the Set/Get methods and not introduced by this PR, but worth addressing in a follow-up.

@vsag-bot vsag-bot 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.

Automated inline review completed.

Review effort: high (8779 changed lines across 55 files).
Submitted 3 inline comments.

Reviewed commit 2c8430b.

Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp Outdated
Comment thread src/datacell/rabitq_split_datacell.h Outdated
Comment thread src/algorithm/hgraph/hgraph.cpp
@LightWant
LightWant force-pushed the feat/hgraph-rabitq-fused-x1to4 branch from 2c8430b to 346efc7 Compare August 2, 2026 07:30
@vsag-bot
vsag-bot self-requested a review August 2, 2026 07:30
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp
Comment thread src/impl/filter/duplicate_group_filter.h
Comment thread src/datacell/rabitq_split_datacell.h
Comment thread src/datacell/rabitq_split_datacell.h
Comment thread src/impl/reorder/flatten_reorder.cpp
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp

@vsag-bot vsag-bot 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.

Automated inline review completed.

Review effort: high (9139 changed lines across 60 files).
Submitted 2 inline comments.

Reviewed commit 346efc7.

Comment thread src/quantization/rabitq_quantization/rabitq_quantizer.cpp Outdated
Comment thread src/impl/searcher/hgraph_rabitq_searcher.cpp Outdated
Signed-off-by: zhuangye.yxw <2510035537@qq.com>
Assisted-by: Codex:gpt-5
@LightWant
LightWant force-pushed the feat/hgraph-rabitq-fused-x1to4 branch from 346efc7 to 3305f72 Compare August 2, 2026 08:38
@vsag-bot
vsag-bot self-requested a review August 2, 2026 08:38
if (exact_filter_ip_hint and IsFiniteRaBitQValue(filter_inner_products[i])) {
++full_count;
used_hint =
split_codes->ComputeFusedFullWithFilterIP(computer,

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.

Bug: full_count double-counted when filter-IP hint path falls back to full compute

When exact_filter_ip_hint is true and the hint-based ComputeFusedFullWithFilterIP returns false, full_count is incremented both at line 79 (before the hint attempt) and at line 93 (in the fallback else branch). This double-counts every candidate that enters the hint path but falls back.

Suggested fix: remove the first ++full_count on line 79 and only increment in the success/fallback branches, or move it into the if (used_hint) branch.

Vector<uint8_t> replacement(
static_cast<uint64_t>(new_capacity) * record_size_ + K_CACHE_LINE_SIZE - 1, 0, allocator_);
const auto replacement_address = reinterpret_cast<uintptr_t>(replacement.data());
const uint64_t replacement_offset =

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.

Thread safety: GetNodeRecord reads storage_ and aligned_offset_ without synchronization

GetNodeRecord (called from the search hot path via GetNodeView, GetCodeView, GetFusedCodeView, Prefetch, etc.) reads storage_.data() + aligned_offset_ without holding storage_mutex_. Meanwhile, Reallocate (called from Resize, ShrinkToFit) replaces storage_ and updates aligned_offset_ under storage_mutex_, but the search path never acquires even a shared lock.

This is a data race: a concurrent resize could free the old storage_ buffer while a search thread is still reading from it. The InitFeatures disables concurrent add+search for fused graphs, but Resize can also be triggered by InsertNeighborsById during build and SetFusedCodes during sync_fused_node_codes.

Consider either:

  • Holding storage_mutex_ shared in the search path (may hurt performance), or
  • Ensuring Reallocate never frees the old buffer while readers exist (e.g., RCU-style deferred reclamation), or
  • Documenting that resize and search are externally serialized for fused graphs.

return true;
}
}
++rabitq_reorder_fallback_full_count;

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.

Stats: rabitq_reorder_fallback_full_count over-counts for x=1 (LegacyOneBitScorer)

When K_EXACT_FILTER_HINT is false (i.e., for LegacyOneBitScorer used with x=1), the if constexpr block at line 633 is skipped entirely, and every call to refine_node unconditionally increments rabitq_reorder_fallback_full_count. This means the metric counts all rerank computations as "fallbacks" even though there is no filter-IP hint to fall back from.

Consider guarding the increment with if constexpr (Scorer::K_EXACT_FILTER_HINT) so the fallback count is only incremented when a hint was actually available but not usable.

@vsag-bot vsag-bot 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.

Automated inline review completed.

Review effort: high (9605 changed lines across 61 files).
Submitted 2 inline comments.

Reviewed commit 3305f72.

query_sum = 0.0F;
for (uint64_t i = 0; i < this->dim_; ++i) {
auto magnitude =
static_cast<uint32_t>(query_scale * std::fabs(transformed_query[i]) + 1e-5F);

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.

[P2] [P2] Reject non-finite values before integer query quantization

A fused 1+y query containing NaN or infinity reaches this cast because public search validation checks only dimension and element count. The computed scale is zero for an infinite norm or the false branch for a NaN norm, but 0 * fabs(NaN/Inf) is NaN, and converting that value to uint32_t is undefined behavior. Reject non-finite query coordinates or take a defined fallback before performing the floating-to-integer conversion.

layout.remove_flag_bit == expected_layout.remove_flag_bit,
"fused graph remove parameters do not match construction parameters");

auto codec_model = StreamReader::ReadString(reader);

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.

[P2] [P2] Validate the serialized codec length before allocating it

ReadString reads the wire-supplied uint64_t length and constructs std::vector<char>(length) before its subsequent read can detect that the input is truncated. A corrupt fused index can therefore request an arbitrarily large allocation and cause bad_alloc, severe memory pressure, or process termination instead of being rejected by the surrounding layout checks. Read and validate the length against the reader's remaining bytes and a valid codec-model bound before allocating.

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

Labels

kind/improvement Optimizations, UX polish, or minor improvements 性能优化、体验打磨或细节改良 module/api module/docs module/simd module/testing size/XXL version/1.1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants