Support inner-product distance metric in IVF-RaBitQ - #2291
Open
jamxia155 wants to merge 7 commits into
Open
Conversation
IVF-RaBitQ previously supported L2 only. This adds InnerProduct support,
initially only for the bitwise (QUANT4/QUANT8) search paths. Support for LUT16
and LUT32 search modes are a natural follow-up.
Approach: rather than re-deriving the tuned RaBitQ per-vector factors, the
existing squared-L2 estimator is reused as-is and the inner product is
recovered via the identity <q,x> = (||q||^2 + ||x||^2 - ||q-x||^2) / 2. The
kernels emit the negated inner product (a "pseudo-distance"), so the existing
min-selection, block-sort queue, and per-query threshold are reused unchanged;
the result is negated after select_k. Probe selection for InnerProduct picks
clusters by argmax <q,c> while the centroid-distance buffer still carries
||q-c||^2 for the estimator (mirroring ivf_pq); k-means clustering stays L2.
The signed-distance behavior (the pseudo-distance transform plus a sign-aware
atomic threshold min) is selected at compile time via a new Signed axis on the
bitwise JIT-LTO fragments, so the L2 path's codegen is unchanged and only the
small entrypoint/emit fragments fan out.
Details:
- metric plumbed through index_params -> index<IdxT> -> IVFGPU (with accessor);
build validates {L2Expanded, InnerProduct}; LUT search modes reject IP.
- per-vector ||x||^2 stored (InnerProduct only) in cluster-permuted order,
computed in the quantizer and passed to the search kernels.
- serialization gains a leading version field and the metric enum, plus the
per-vector norm blob for InnerProduct (clean break, validated on load).
- tests: var_metric() adds InnerProduct cases over QUANT4/QUANT8 x
block-sort/non-block-sort x with_ex/no_ex under a dedicated
IvfRabitqInnerProduct instantiation; a signed_data flag feeds mean-zero data
to those cases so inner products take both signs and exercise the sign-aware
atomic threshold. Existing cases are unchanged.
Follow-up (not in this change):
- LUT16 and LUT32 search modes do not yet support InnerProduct; only the
bitwise QUANT4/QUANT8 path is implemented, and search() rejects IP for the
LUT modes. Extending them means adding the same compile-time Signed axis to
the LUT emit fragments.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
jamxia155
added
feature request
The initial InnerProduct support (earlier commit) covered only the bitwise QUANT4/QUANT8 search paths. This extends it to LUT16 and LUT32 so InnerProduct works across all four search modes. Apply the same compile-time Signed axis + pseudo-distance transform used for the bitwise path to the LUT emit fragments (lut_emit_distances, lut_block_sort_emit_topk, lut16_opt_emit_distances) and to the LUT16 block-sort kernel, which emits inline. The Signed axis flows through the fragment tags, planners, launchers (now take is_inner_product), the per-fragment matrix.json, and the CMake registrations. The non-block LUT entrypoints stay metric-agnostic (they call the emit symbol, whose Signed variant the launcher links). The LUT searchers (searcher_gpu.cu, searcher_gpu_shared_mem_opt.cu) now populate the per-query/per-vector norm pointers, pass is_inner_product to the launchers, and negate the emitted pseudo-distances back to true inner products after select_k. The search()-time guard that rejected InnerProduct for LUT modes is removed. Tests: var_metric() now sweeps all four search modes (LUT16/LUT32/QUANT4/QUANT8) under the IvfRabitqInnerProduct instantiation.
cuvs_ivf_rabitq_wrapper ignored the Metric passed to its constructor, so the ANN benchmark always built an L2 index regardless of the config's "distance". Apply index_params_.metric = parse_metric_type(metric) in the constructor, matching the ivf_pq/ivf_flat/ivf_sq wrappers, so "distance": "inner_product" in a config builds an InnerProduct index.
jamxia155
marked this pull request as ready for review
July 27, 2026 16:48
Replace the squared-L2 reconstruction with RaBitQ's native InnerProduct factorization. The previous approach recovered <q,x> from the L2 estimate via <q,x> = (||q||^2 + ||x||^2 - ||q-x||^2)/2, which required storing structural, not numerical: the metric now lives entirely in build-time constants plus one query-side term. - quantizer: the three factor kernels switch on the metric, using <o-c, c> from the existing reduction loop. IP halves the rescale and error factors and replaces the ||o-c||^2 offset with -<o-c, c>. The extra block reduction is guarded on the metric so L2 does not pay for it; the branch is uniform across the block, so skipping the reduction cannot desynchronize its __syncthreads(). - params gain d_g_add, separating the estimator's additive term from d_centroid_distances, which still supplies g_error = ||q-c|| for both metrics. For L2 the two alias; for IP, d_g_add holds -<q,c>, reusing the probe-selection GEMM that already computed it. - drops the per-vector norm array (4 B/vector), its serialization blob, and the two per-candidate loads it cost in the block-sort inner loop. - drops the Signed axis from the three non-block-sort emit fragments, which are now metric-agnostic. The block-sort fragments keep it: the sign-aware threshold atomics are required regardless of factorization, since -<q,x> is negative under both formulations. - serialization bumped to v2; a v1 InnerProduct index would be silently misinterpreted by the new estimator. Probe selection continues to pick clusters by argmax <q,c>. Verified against the pre-refactor build across all 48 InnerProduct cases: 30 are bit-identical and the other 18 differ by at most two neighbours out of 131072 (max delta recall 1.6e-5), all in the k=128 non-block-sort path where top-k boundary ties are resolved differently by the changed arithmetic order.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
IVF-RaBitQ previously supported L2 only. This PR adds InnerProduct support across all four search modes (LUT16, LUT32, QUANT4, QUANT8).
Approach: The per-vector RaBitQ factors are derived directly for the inner product rather than reconstructed from a squared-L2 estimate. Starting from
-<o,q> = -<q,c> - <o-c,c> - <o-c,q-c>, everything that does not depend on the code folds into the per-vector factors at build time, and the query side contributes a single additive termg_add = -<q,c>. Concretely,F_addbecomes-<o-c,c> + ||o-c||*<o_bar,c>/<o_bar,o>, and the rescale and error factors lose the factor of 2 that the squared-L2 expansion carries. The estimator kernel itself (est = f_add + g_add + f_rescale * (ip + kBxSumq)) is then byte-identical for both metrics.Kernels emit the negated inner product, so the existing min-selection, block-sort queue, and per-query threshold are reused unchanged; the result is negated after select_k.
Note: an earlier revision of this PR recovered
<q,x>from the squared-L2 estimate via<q,x> = (||q||^2 + ||x||^2 - ||q-x||^2)/2. That is algebraically the same estimator (expanding it cancels||q||^2and||x||^2against the terms inside||q-c||^2and||o-c||^2, and the error bounds coincide) but it required storing||x||^2per vector and applying a per-candidate correction at search time. The formulation above removes both. Thanks to @Stardust-SJF for pointing at RaBitQ-Library, whose estimator derivation this follows.Impact:
InnerProduct's marginal search-time cost over L2 is now zero: the same per-candidate loads, the same arithmetic, no extra per-vector storage. Index construction is roughly neutral: one FMA per dimension folded into an existing reduction loop, against a dropped full row-norm pass over the dataset.Details:
index_paramstoindex<IdxT>toIVFGPU; build validates{L2Expanded, InnerProduct}. All four search modes support both.cuvs::distance::DistanceType, so adding a metric means adding a case rather than editing a branch. The default arm emits NaN rather than silently falling back to L2, so an unhandled metric fails loudly. The extra<o-c,c>reduction is guarded on the metric (uniform across the block) so L2 does not pay for it.d_g_add, separating the estimator's additive term fromd_centroid_distances, which continues to supplyg_error = ||q-c||for both metrics. For L2 the two alias; for InnerProduct,d_g_add holds -<q,c>.argmax <q,c>(mirroring ivf_pq's select_clusters withalpha=-1); that same GEMM result isg_add, so it is computed once and reused rather than recomputed. k-means clustering stays L2.Signedaxis, retained only on the four block-sort fragments that update the threshold.-<q,x>is negative regardless of formulation, so the int-reinterpretatomicMintrick is not monotone there. The three non-block-sort emit fragments are metric-agnostic and carry noSignedaxis, so L2 codegen is unchanged and total fragment count decreases.Testing
var_metric()sweeps all four search modes × k ∈ {10, 128} (block-sort and non-block-sort) × bits_per_dim ∈ {1, 4} (no-ex and with-ex), under a dedicated IvfRabitqInnerProduct instantiation: 48 cases across the three test bodies, all of which round-trip through serialize/deserialize.signed_dataflag feeds mean-zero data to those cases so inner products take both signs and exercise the sign-aware atomic threshold. Existing L2 cases are unchanged.