From e7aea74b49aac0003bc129f5d4fe28511aaab45c Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 2 Jul 2026 06:59:48 -0700 Subject: [PATCH] Add InnerProduct distance metric to IVF-RaBitQ (QUANT4/QUANT8 path) 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||^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 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 -> 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. --- cpp/CMakeLists.txt | 15 +-- .../ivf_rabitq/ivf_rabitq_fragments.hpp | 7 +- cpp/include/cuvs/neighbors/ivf_rabitq.hpp | 9 +- cpp/src/neighbors/ivf_rabitq.cu | 24 ++++- .../neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 96 +++++++++++++++++-- .../ivf_rabitq/gpu_index/ivf_gpu.cuh | 38 +++++++- .../ivf_rabitq/gpu_index/quantizer_gpu.cu | 31 +++++- .../ivf_rabitq/gpu_index/quantizer_gpu.cuh | 8 +- .../gpu_index/searcher_gpu_common.cuh | 23 +++-- .../gpu_index/searcher_gpu_quantize_query.cu | 67 ++++++++----- .../bitwise_block_sort_emit_topk_kernel.cu.in | 32 ++++--- .../bitwise_block_sort_emit_topk_matrix.json | 10 ++ .../bitwise_emit_distances_kernel.cu.in | 25 +++-- .../bitwise_emit_distances_matrix.json | 10 ++ ..._products_with_bitwise_block_sort_impl.cuh | 13 ++- ...ducts_with_bitwise_block_sort_kernel.cu.in | 10 +- ...oducts_with_bitwise_block_sort_matrix.json | 13 ++- ...oducts_with_bitwise_block_sort_planner.hpp | 10 +- ...te_inner_products_with_bitwise_planner.hpp | 6 +- ...roducts_with_lut16_opt_block_sort_impl.cuh | 18 ++-- .../jit_lto_kernels/device_functions.cuh | 39 +++++++- .../jit_lto_kernels/launcher_factory.hpp | 37 +++++-- .../lut_block_sort_emit_topk_kernel.cu.in | 34 +++---- cpp/tests/neighbors/ann_ivf_rabitq.cuh | 41 +++++++- .../ann_ivf_rabitq/test_float_int64_t.cu | 8 +- 25 files changed, 488 insertions(+), 136 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 89cbadfcfc..39d0926a1a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on cmake_minimum_required(VERSION 4.0 FATAL_ERROR) @@ -831,12 +831,13 @@ if(NOT BUILD_CPU_ONLY) ) generate_jit_lto_kernels( jit_lto_files - NAME_FORMAT "ivf_rabitq_bitwise_emit_distances_@with_ex_descriptor@" + NAME_FORMAT "ivf_rabitq_bitwise_emit_distances_@with_ex_descriptor@_@signed_descriptor@" MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json" KERNEL_INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in" - FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_bitwise_emit_distances<@with_ex_value@>" + FRAGMENT_TAG_FORMAT + "${ivf_rabitq_ns}::fragment_tag_bitwise_emit_distances<@with_ex_value@, @signed_value@>" FRAGMENT_TAG_HEADER_FILES "" OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_emit_distances" @@ -844,13 +845,13 @@ if(NOT BUILD_CPU_ONLY) ) generate_jit_lto_kernels( jit_lto_files - NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise_block_sort" + NAME_FORMAT "ivf_rabitq_compute_inner_products_with_bitwise_block_sort_@signed_descriptor@" MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_matrix.json" KERNEL_INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in" FRAGMENT_TAG_FORMAT - "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise_block_sort" + "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_bitwise_block_sort<@signed_value@>" FRAGMENT_TAG_HEADER_FILES "" OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise_block_sort" @@ -858,13 +859,13 @@ if(NOT BUILD_CPU_ONLY) ) generate_jit_lto_kernels( jit_lto_files - NAME_FORMAT "ivf_rabitq_bitwise_block_sort_emit_topk_@with_ex_descriptor@" + NAME_FORMAT "ivf_rabitq_bitwise_block_sort_emit_topk_@with_ex_descriptor@_@signed_descriptor@" MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json" KERNEL_INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in" FRAGMENT_TAG_FORMAT - "${ivf_rabitq_ns}::fragment_tag_bitwise_block_sort_emit_topk<@with_ex_value@>" + "${ivf_rabitq_ns}::fragment_tag_bitwise_block_sort_emit_topk<@with_ex_value@, @signed_value@>" FRAGMENT_TAG_HEADER_FILES "" OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_block_sort_emit_topk" diff --git a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp index 30885947b5..a8ee729573 100644 --- a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -27,12 +27,13 @@ struct fragment_tag_compute_inner_products_with_lut16_opt_block_sort {}; struct fragment_tag_compute_inner_products_with_bitwise {}; -template +template struct fragment_tag_bitwise_emit_distances {}; +template struct fragment_tag_compute_inner_products_with_bitwise_block_sort {}; -template +template struct fragment_tag_bitwise_block_sort_emit_topk {}; template diff --git a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp index d26bc04022..659b95fcc3 100644 --- a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -149,7 +150,8 @@ struct index : cuvs::neighbors::index { size_t n_rows, uint32_t dim, uint32_t n_lists, - uint32_t bits_per_dim); + uint32_t bits_per_dim, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); /** Dimensionality of the input data. */ uint32_t dim() const noexcept; @@ -157,6 +159,9 @@ struct index : cuvs::neighbors::index { /** Total length of the index. */ IdxT size() const noexcept; + /** Distance metric used to build the index. */ + cuvs::distance::DistanceType metric() const noexcept; + /** Accessor for underlying RaBitQ index */ detail::IVFGPU& rabitq_index() noexcept; diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu index d5572d4039..80deae17fa 100644 --- a/cpp/src/neighbors/ivf_rabitq.cu +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -37,6 +37,9 @@ auto build(raft::resources const& handle, RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); + RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::InnerProduct, + "ivf_rabitq only supports L2Expanded and InnerProduct metrics"); // Calculate dataset size and available workspace once size_t dataset_bytes = sizeof(T) * n_rows * dim; @@ -165,7 +168,7 @@ auto build(raft::resources const& handle, } } - index index(handle, n_rows, dim, params.n_lists, params.bits_per_dim); + index index(handle, n_rows, dim, params.n_lists, params.bits_per_dim, params.metric); // Call RaBitQ index construct - use streaming if dataset doesn't fit in GPU memory if (use_streaming) { @@ -213,6 +216,11 @@ void search(raft::resources const& handle, "n_probes (%u) cannot exceed number of IVF lists (%zu)", params.n_probes, idx.rabitq_index().get_num_centroids()); + // InnerProduct is currently implemented only on the bitwise (QUANT4/QUANT8) search paths; the + // LUT paths would silently return squared-L2 results, so reject them explicitly. + RAFT_EXPECTS(idx.rabitq_index().metric() != cuvs::distance::DistanceType::InnerProduct || + params.mode == search_mode::QUANT4 || params.mode == search_mode::QUANT8, + "ivf_rabitq: InnerProduct search currently supports only QUANT4 and QUANT8 modes"); auto padded_dim = idx.rabitq_index().get_num_padded_dim(); auto rotated_queries = raft::make_device_matrix(handle, NQ, padded_dim); @@ -314,10 +322,12 @@ index::index(raft::resources const& handle, size_t n_rows, uint32_t dim, uint32_t n_lists, - uint32_t bits_per_dim) + uint32_t bits_per_dim, + cuvs::distance::DistanceType metric) { RAFT_EXPECTS(bits_per_dim >= 1 && bits_per_dim <= 9, "Unsupported bits_per_dim"); - rabitq_index_ = std::make_unique(handle, n_rows, dim, n_lists, bits_per_dim); + rabitq_index_ = + std::make_unique(handle, n_rows, dim, n_lists, bits_per_dim, metric); } template @@ -349,6 +359,12 @@ IdxT index::size() const noexcept return rabitq_index_->get_num_vectors(); } +template +cuvs::distance::DistanceType index::metric() const noexcept +{ + return rabitq_index_->metric(); +} + auto build(raft::resources const& handle, const cuvs::neighbors::ivf_rabitq::index_params& index_params, raft::device_matrix_view dataset) diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index 49070acbcf..d0b1fedd95 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -36,13 +36,19 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -IVFGPU::IVFGPU(raft::resources const& handle, size_t n, size_t dim, size_t k, size_t bits_per_dim) +IVFGPU::IVFGPU(raft::resources const& handle, + size_t n, + size_t dim, + size_t k, + size_t bits_per_dim, + cuvs::distance::DistanceType metric) : handle_(handle), num_vectors(n), num_dimensions(dim), num_padded_dim(raft::round_up_safe(dim, 64)), num_centroids(k), ex_bits(bits_per_dim - 1), + metric_(metric), initializer(nullptr), DQ(std::make_unique(handle_, dim, bits_per_dim - 1)), Rota(std::make_unique(handle_, dim)) @@ -70,6 +76,12 @@ void IVFGPU::AllocateDeviceMemory() raft::make_device_vector(handle_, ex_factor_size / sizeof(ExFactor)); ids_ = raft::make_device_vector(handle_, pids_size / sizeof(PID)); + // Per-vector squared norms are only needed to reconstruct inner products. + if (is_inner_product()) { + vec_sqr_norms_ = + raft::make_device_vector(handle_, GetVecSqrNormBytes() / sizeof(float)); + } + // Allocate memory for the per-cluster metadata and centroids. cluster_meta_ = raft::make_device_vector(handle_, num_centroids); raft::resource::sync_stream(handle_); @@ -92,6 +104,10 @@ void IVFGPU::AllocateHostMemory() this->ids_host_ = raft::make_host_vector(pids_size / sizeof(PID)); } +// Serialized-format version. The IVF-RaBitQ on-disk format is not backward compatible: bumping +// this constant (or reading a pre-versioning index) is a hard error at load time. +static constexpr std::uint32_t kSerializationVersion = 1; + // load transposed data for short codes void IVFGPU::load_transposed(const char* filename) { @@ -105,6 +121,16 @@ void IVFGPU::load_transposed(const char* filename) filename); }; + // Format version is the first field. + std::uint32_t version = 0; + read_exact(&version, sizeof(std::uint32_t)); + RAFT_EXPECTS(version == kSerializationVersion, + "ivf_rabitq::deserialize: serialization version mismatch (got %u, expected %u) " + "in: %s", + version, + kSerializationVersion, + filename); + // Load metadata. read_exact(&this->num_vectors, sizeof(size_t)); read_exact(&this->num_dimensions, sizeof(size_t)); @@ -113,6 +139,16 @@ void IVFGPU::load_transposed(const char* filename) read_exact(&this->num_centroids, sizeof(size_t)); read_exact(&this->ex_bits, sizeof(size_t)); + // Distance metric (must be set before AllocateDeviceMemory, which allocates the per-vector + // norm array only for InnerProduct). + read_exact(&this->metric_, sizeof(cuvs::distance::DistanceType)); + RAFT_EXPECTS(cuvs::util::is_valid_distance_type(metric_) && + (metric_ == cuvs::distance::DistanceType::L2Expanded || + metric_ == cuvs::distance::DistanceType::InnerProduct), + "ivf_rabitq::deserialize: unsupported metric %d in: %s", + static_cast(metric_), + filename); + // Skip legacy batch_flag field for backward compatibility bool legacy_batch_flag; read_exact(&legacy_batch_flag, sizeof(bool)); @@ -256,6 +292,8 @@ void IVFGPU::load_transposed(const char* filename) read_into_device_host( ex_factor_.data_handle(), ex_factor_host_.data_handle(), GetExFactorBytes()); read_into_device_host(ids_.data_handle(), ids_host_.data_handle(), GetPIDsBytes()); + // Per-vector squared norms (InnerProduct only), same permuted order as ids_. + if (is_inner_product()) { read_into_device(vec_sqr_norms_.data_handle(), GetVecSqrNormBytes()); } // Initialize cluster metadata (host side) based on the loaded cluster sizes. init_clusters(cluster_sizes); @@ -352,11 +390,16 @@ void IVFGPU::save(const char* filename) const RAFT_EXPECTS(static_cast(output), "write failed to: %s", filename); }; + // Format version is the first field. + write_exact(&kSerializationVersion, sizeof(std::uint32_t)); + // Save meta data. write_exact(&num_vectors, sizeof(size_t)); write_exact(&num_dimensions, sizeof(size_t)); write_exact(&num_centroids, sizeof(size_t)); write_exact(&ex_bits, sizeof(size_t)); + // Distance metric (mirrors the read order in load_transposed). + write_exact(&metric_, sizeof(cuvs::distance::DistanceType)); // Write legacy batch_flag=true for backward compatibility bool legacy_batch_flag = true; write_exact(&legacy_batch_flag, sizeof(bool)); @@ -418,6 +461,18 @@ void IVFGPU::save(const char* filename) const write_exact(h_ex_factor_buf.data_handle(), ex_factor_size); write_exact(h_ids_buf.data_handle(), ids_size); + // Per-vector squared norms (InnerProduct only), same permuted order as ids. + if (is_inner_product()) { + size_t norms_size = GetVecSqrNormBytes(); + auto h_norms_buf = raft::make_host_vector(norms_size); + raft::copy(h_norms_buf.data_handle(), + reinterpret_cast(vec_sqr_norms_.data_handle()), + norms_size, + stream_); + raft::resource::sync_stream(handle_); + write_exact(h_norms_buf.data_handle(), norms_size); + } + output.close(); } @@ -852,7 +907,8 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, cp.short_factor_batch(*this, 0), cp.long_code(*this, 0, DQ->long_code_length()), reinterpret_cast(cp.ex_factor(*this, 0)), - cur_rotated_c); + cur_rotated_c, + is_inner_product() ? cp.vec_sqr_norm(*this, 0) : nullptr); batch_offset += cluster_size; } @@ -886,7 +942,8 @@ void IVFGPU::quantize_cluster(GPUClusterMeta& cp, cp.short_factor_batch(*this, 0), cp.long_code(*this, 0, DQ->long_code_length()), reinterpret_cast(cp.ex_factor(*this, 0)), - d_rotated_c); + d_rotated_c, + is_inner_product() ? cp.vec_sqr_norm(*this, 0) : nullptr); } // Optimized kernel to prepare keys and values from d_raft_idx @@ -1094,8 +1151,35 @@ void IVFGPU::PrepareClusterSearchInputs( // Step 4: select top-nprobe clusters per query auto d_raft_vals = raft::make_device_matrix(searcher_handle, batch_size, nprobe); auto d_raft_idx = raft::make_device_matrix(searcher_handle, batch_size, nprobe); - auto in_view = raft::make_device_matrix_view( - searcher.get_centroid_distances(), batch_size, num_centroids); + + // Clusters are scored by ‖q−c‖² for L2 (minimized). For InnerProduct we instead pick the + // clusters maximizing ⟨q,c⟩, computed into a separate scratch buffer (mirroring ivf_pq's + // select_clusters with alpha=-1) so that centroid_distances is left as ‖q−c‖² — it still feeds + // q_g_add in the search kernels, which reconstruct the estimate in the q−c residual space. + auto qc_scores = raft::make_device_matrix(searcher_handle, 0, 0); + const float* select_in = searcher.get_centroid_distances(); + if (is_inner_product()) { + qc_scores = + raft::make_device_matrix(searcher_handle, batch_size, num_centroids); + const float alpha_ip = -1.f, beta_ip = 0.f; + raft::linalg::detail::matmul(searcher_handle, + /* trans_a = */ true, + /* trans_b = */ false, + num_centroids, + batch_size, + num_padded_dim, + &alpha_ip, + initializer->GetCentroid(0), + num_padded_dim, + queries.data_handle(), + num_padded_dim, + &beta_ip, + qc_scores.data_handle(), + num_centroids); + select_in = qc_scores.data_handle(); + } + auto in_view = raft::make_device_matrix_view( + select_in, batch_size, num_centroids); cuvs::selection::select_k(searcher_handle, in_view, std::nullopt, diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index 75ddaee865..f05ca871d1 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,6 +14,8 @@ #include "quantizer_gpu.cuh" #include "rotator_gpu.cuh" +#include + #include #include #include @@ -91,6 +93,12 @@ class IVFGPU { return parent.get_short_factors_batch_device() + (start_index + i) * 3; } + // Per-vector ‖x‖² for the i-th vector in this cluster (InnerProduct metric only). + float* vec_sqr_norm(const IVFGPU& parent, size_t i) const + { + return parent.get_vec_sqr_norms_device() + (start_index + i); + } + /** * @brief Get the pointer to the long code for the i-th vector in this cluster. * @@ -156,7 +164,12 @@ class IVFGPU { * @param k Num of centroids. * @param bits_per_dim totalbits = EX_BITS+1 */ - IVFGPU(raft::resources const& handle, size_t n, size_t dim, size_t k, size_t bits_per_dim); + IVFGPU(raft::resources const& handle, + size_t n, + size_t dim, + size_t k, + size_t bits_per_dim, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); IVFGPU(raft::resources const& handle) : handle_(handle), initializer(nullptr), Rota(std::make_unique(handle_, 128)) { @@ -249,6 +262,15 @@ class IVFGPU { size_t get_num_centroids() const { return num_centroids; } size_t get_max_cluster_length() const noexcept { return max_cluster_length; } size_t get_ex_bits() const noexcept { return ex_bits; } + cuvs::distance::DistanceType metric() const noexcept { return metric_; } + bool is_inner_product() const noexcept + { + return metric_ == cuvs::distance::DistanceType::InnerProduct; + } + __host__ __device__ float* get_vec_sqr_norms_device() const noexcept + { + return const_cast(this->vec_sqr_norms_.data_handle()); + } // member object getters DataQuantizerGPU& quantizer() const { return *(this->DQ); } @@ -326,6 +348,9 @@ class IVFGPU { size_t GetPIDsBytes() const { return sizeof(PID) * num_vectors; } + // Per-vector squared-norm array; only allocated/serialized for the InnerProduct metric. + size_t GetVecSqrNormBytes() const { return sizeof(float) * num_vectors; } + size_t GetLongCodeBytes() const { return sizeof(uint8_t) * quantizer().long_code_length() * num_vectors; @@ -364,6 +389,13 @@ class IVFGPU { raft::device_vector short_factors_batch_ = raft::make_device_vector(handle_, 0); // N * 3 float rabitq factors + // Per-vector squared L2 norm ‖x‖² of the original (rotation-invariant) vectors, in the same + // cluster-permuted order as the factors. Only populated for the InnerProduct metric, where the + // reconstructed squared-L2 estimate is converted to an inner product via + // ⟨q,x⟩ = (‖q‖² + ‖x‖² − ‖q−x‖²)/2. + raft::device_vector vec_sqr_norms_ = + raft::make_device_vector(handle_, 0); + // host-side copies raft::host_vector short_data_host_ = raft::make_host_vector( 0); // TODO: CPU side, we need on factors from short_data_host_, so no need to @@ -384,6 +416,8 @@ class IVFGPU { size_t num_centroids; // Centroids && Clusters size_t max_cluster_length; // Maximum length of clusters size_t ex_bits; // Extra bits parameter for quantization. + cuvs::distance::DistanceType metric_ = + cuvs::distance::DistanceType::L2Expanded; // Distance metric of the index. std::unique_ptr initializer; // Initializer, indicates which initializer to use (currently only FlatIVF) diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index af0876acca..cb055ef68f 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -618,7 +619,8 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, float* d_short_data_factors, uint8_t* d_long_code, float* d_ex_factor, - float* d_rotated_c) + float* d_rotated_c, + float* d_vec_sqr_norms) { // 1. Data Transformation: data_transformation_batch_opt(d_data, @@ -635,6 +637,18 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, D, handle_); + // For the InnerProduct metric, capture the squared L2 norm ‖x‖² of each (padded, unrotated) + // input vector. gatherAndPadKernel just wrote the first num_points rows of d_X_and_C_pad with + // the padded data and the zero-padding does not affect the norm. raft's L2Norm returns the sum + // of squares (no sqrt) with the default finalizer, i.e. exactly ‖x‖². + if (d_vec_sqr_norms != nullptr) { + raft::linalg::norm( + handle_, + raft::make_device_matrix_view( + d_X_and_C_pad.data_handle(), num_points, D), + raft::make_device_vector_view(d_vec_sqr_norms, num_points)); + } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), @@ -739,7 +753,8 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d float* d_short_data_factors, uint8_t* d_long_code, float* d_ex_factor, - float* d_rotated_c) + float* d_rotated_c, + float* d_vec_sqr_norms) { // 1. Data Transformation: data_transformation_batch_opt_contiguous(d_contiguous_data, @@ -755,6 +770,16 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d D, handle_); + // See quantize_batch_opt: capture ‖x‖² for the InnerProduct metric before the residual + // transform. d_X_and_C_pad holds the padded, unrotated cluster data. + if (d_vec_sqr_norms != nullptr) { + raft::linalg::norm( + handle_, + raft::make_device_matrix_view( + d_X_and_C_pad.data_handle(), num_points, D), + raft::make_device_vector_view(d_vec_sqr_norms, num_points)); + } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh index 90ef71ed12..8ef006b7bd 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -115,7 +115,8 @@ class DataQuantizerGPU { float* short_data_factors, uint8_t* d_long_code, float* d_ex_factor, - float* d_rotated_c); + float* d_rotated_c, + float* d_vec_sqr_norms = nullptr); /*! * @brief Quantize contiguous cluster data for batch layout (without PID gathering). @@ -138,7 +139,8 @@ class DataQuantizerGPU { float* short_data_factors, uint8_t* d_long_code, float* d_ex_factor, - float* d_rotated_c); + float* d_rotated_c, + float* d_vec_sqr_norms = nullptr); private: // Dimension and quantization parameters diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh index 1e004d9164..fc9aaf0189 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -41,14 +41,19 @@ struct ComputeInnerProductsKernelParams { const float* d_G_k1xSumq = nullptr; const float* d_G_kbxSumq = nullptr; const float* d_centroid_distances = nullptr; - uint32_t topk = 0; - uint32_t num_queries = 0; - uint32_t nprobe = 0; - uint32_t num_pairs = 0; - uint32_t num_centroids = 0; - uint32_t D = 0; - const float* d_threshold = nullptr; // threshold for each query - uint32_t max_candidates_per_pair = 0; // max storage per pair, 1000 suggested + // Squared L2 norms for the InnerProduct metric: ‖q‖² per query and ‖x‖² per data vector. + // Used to convert the reconstructed squared-L2 estimate into an inner product via + // ⟨q,x⟩ = (‖q‖² + ‖x‖² − ‖q−x‖²)/2. Null for the L2 metric. + const float* d_q_sqr_norms = nullptr; + const float* d_vec_sqr_norms = nullptr; + uint32_t topk = 0; + uint32_t num_queries = 0; + uint32_t nprobe = 0; + uint32_t num_pairs = 0; + uint32_t num_centroids = 0; + uint32_t D = 0; + const float* d_threshold = nullptr; // threshold for each query + uint32_t max_candidates_per_pair = 0; // max storage per pair, 1000 suggested uint32_t max_candidates_per_query = 0; // max number of vectors in probed clusters for any particular query uint32_t ex_bits = 0; // bits per dimension in ex codes diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index 92906fb0e6..747ec85f32 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -22,6 +22,8 @@ #include #include +#include +#include #include #include @@ -488,23 +490,29 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( max(packed_query_size + candidate_storage + query_storage, (size_t)queue_buffer_smem_bytes); ComputeInnerProductsKernelParams kernelParams; - kernelParams.d_sorted_pairs = d_sorted_pairs; - kernelParams.d_query = queries.data_handle(); - kernelParams.d_short_data = cur_ivf.get_short_data_device(); - kernelParams.d_cluster_meta = d_cluster_meta; - kernelParams.d_packed_queries = d_packed_queries.data_handle(); - kernelParams.d_widths = d_widths.data_handle(); - kernelParams.d_short_factors = cur_ivf.get_short_factors_batch_device(); - kernelParams.d_G_k1xSumq = d_G_k1xSumq; - kernelParams.d_G_kbxSumq = d_G_kbxSumq; - kernelParams.d_centroid_distances = get_centroid_distances(); - kernelParams.topk = topk; - kernelParams.num_queries = num_queries; - kernelParams.nprobe = nprobe; - kernelParams.num_pairs = num_pairs; - kernelParams.num_centroids = cur_ivf.get_num_centroids(); - kernelParams.D = D; - kernelParams.d_threshold = d_topk_threshold_batch.data_handle(); + kernelParams.d_sorted_pairs = d_sorted_pairs; + kernelParams.d_query = queries.data_handle(); + kernelParams.d_short_data = cur_ivf.get_short_data_device(); + kernelParams.d_cluster_meta = d_cluster_meta; + kernelParams.d_packed_queries = d_packed_queries.data_handle(); + kernelParams.d_widths = d_widths.data_handle(); + kernelParams.d_short_factors = cur_ivf.get_short_factors_batch_device(); + kernelParams.d_G_k1xSumq = d_G_k1xSumq; + kernelParams.d_G_kbxSumq = d_G_kbxSumq; + kernelParams.d_centroid_distances = get_centroid_distances(); + // InnerProduct reconstructs ⟨q,x⟩ = (‖q‖² + ‖x‖² − ‖q−x‖²)/2 from the squared-L2 estimate; the + // kernels emit the negated inner product (a "pseudo-distance") so the existing min-selection is + // reused, and the result is negated after select_k below. + const bool is_inner_product = cur_ivf.is_inner_product(); + kernelParams.d_q_sqr_norms = is_inner_product ? get_q_norms() : nullptr; + kernelParams.d_vec_sqr_norms = is_inner_product ? cur_ivf.get_vec_sqr_norms_device() : nullptr; + kernelParams.topk = topk; + kernelParams.num_queries = num_queries; + kernelParams.nprobe = nprobe; + kernelParams.num_pairs = num_pairs; + kernelParams.num_centroids = cur_ivf.get_num_centroids(); + kernelParams.D = D; + kernelParams.d_threshold = d_topk_threshold_batch.data_handle(); kernelParams.max_candidates_per_pair = max_cluster_size; kernelParams.max_candidates_per_query = use_block_sort ? 0 /* unused */ : max_probed_vectors_count.value(); @@ -520,11 +528,12 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( const int num_bits_for_dispatch = use_4bit ? 4 : 8; const bool with_ex = (cur_ivf.get_ex_bits() != 0); - auto jit_launcher = use_block_sort ? make_compute_inner_products_with_bitwise_block_sort_launcher( - num_bits_for_dispatch, cur_ivf.get_ex_bits(), with_ex) - : make_compute_inner_products_with_bitwise_launcher( - cur_ivf.get_ex_bits(), with_ex); - auto const& kernel_launcher = [&]() -> void { + auto jit_launcher = use_block_sort + ? make_compute_inner_products_with_bitwise_block_sort_launcher( + num_bits_for_dispatch, cur_ivf.get_ex_bits(), with_ex, is_inner_product) + : make_compute_inner_products_with_bitwise_launcher( + cur_ivf.get_ex_bits(), with_ex, is_inner_product); + auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( stream_, gridDim, blockDim, shared_mem_size, kernelParams); }; @@ -533,7 +542,8 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( static_cast(shared_mem_size), kernel_launcher, jit_launcher->get_kernel()); RAFT_CUDA_TRY(cudaPeekAtLastError()); - // Merge results + // Merge results. The emitted distances are squared-L2 estimates for L2, or negated inner + // products (pseudo-distances) for InnerProduct; both are minimized here. cuvs::selection::select_k( handle_, raft::make_device_matrix_view( @@ -544,6 +554,15 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( d_final_pids, /*select_min=*/true, /*sorted=*/false); + + // Convert the pseudo-distances (−⟨q,x⟩) back to true inner products for the user-visible output. + if (is_inner_product) { + thrust::transform(thrust::cuda::par.on(stream_), + d_final_dists.data_handle(), + d_final_dists.data_handle() + d_final_dists.size(), + d_final_dists.data_handle(), + thrust::negate()); + } } } // namespace cuvs::neighbors::ivf_rabitq::detail diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in index d8dcff324d..3f5617129c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,7 +12,8 @@ namespace { -constexpr bool with_ex = @with_ex_value@; +constexpr bool with_ex = @with_ex_value@; +constexpr bool is_signed = @signed_value@; } // namespace @@ -38,6 +39,9 @@ __device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelPar const int tid = threadIdx.x; const int num_threads = blockDim.x; + // ‖q‖² for the InnerProduct pseudo-distance transform pd = (‖q−x‖² − ‖q‖² − ‖x‖²)/2 = −⟨q,x⟩. + float q_sqr_norm = is_signed ? params.d_q_sqr_norms[query_idx] : 0.0f; + // Shared-memory layout matches the caller's setup extern __shared__ __align__(256) char shared_mem_raw_2[]; size_t packed_query_bytes = max(params.num_bits * params.num_words * sizeof(uint32_t), @@ -107,6 +111,9 @@ __device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelPar ex_dist = f_ex_add + q_g_add + f_ex_rescale * (static_cast(1 << params.ex_bits) * ip + ip2 + q_kbxsumq); + if constexpr (is_signed) { + ex_dist = (ex_dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; + } pid = (uint32_t)params.d_pids[global_vec_idx]; } else { ex_dist = INFINITY; @@ -152,7 +159,10 @@ __device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelPar short_code_length, params.D); final_dist = f_add + q_g_add + f_rescale * (exact_ip + q_k1xsumq); - final_pid = (uint32_t)params.d_pids[global_vec_idx]; + if constexpr (is_signed) { + final_dist = (final_dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; + } + final_pid = (uint32_t)params.d_pids[global_vec_idx]; } else { final_dist = INFINITY; final_pid = 0; @@ -171,14 +181,14 @@ __device__ void bitwise_block_sort_emit_topk(const ComputeInnerProductsKernelPar } if (num_candidates >= params.topk) { - update_threshold_atomicmin(params.d_topk_dists, - params.d_threshold, - query_idx, - params.topk, - params.nprobe, - probe_slot, - threshold, - tid); + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json index 7e0f0f3603..6d32bd63fc 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_matrix.json @@ -8,5 +8,15 @@ "with_ex_descriptor": "no_ex", "with_ex_value": "false" } + ], + "_signed": [ + { + "signed_descriptor": "signed", + "signed_value": "true" + }, + { + "signed_descriptor": "unsigned", + "signed_value": "false" + } ] } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in index 20c5aec6a1..906c07cd57 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -10,7 +10,8 @@ namespace { -constexpr bool with_ex = @with_ex_value@; +constexpr bool with_ex = @with_ex_value@; +constexpr bool is_signed = @signed_value@; } // namespace @@ -32,6 +33,9 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa const int tid = threadIdx.x; const int num_threads = blockDim.x; + // ‖q‖² for the InnerProduct pseudo-distance transform pd = (‖q−x‖² − ‖q‖² − ‖x‖²)/2 = −⟨q,x⟩. + float q_sqr_norm = is_signed ? params.d_q_sqr_norms[query_idx] : 0.0f; + extern __shared__ __align__(256) char shared_mem_raw[]; float* shared_query = reinterpret_cast(shared_mem_raw); @@ -81,10 +85,14 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa float f_ex_add = ex_factors.x; float f_ex_rescale = ex_factors.y; - params.d_topk_dists[output_offset + vec_idx] = + float dist = f_ex_add + q_g_add + f_ex_rescale * (static_cast(1 << params.ex_bits) * exact_ip + ip2 + q_kbxsumq); - params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + if constexpr (is_signed) { + dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; + } + params.d_topk_dists[output_offset + vec_idx] = dist; + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; } } else { float q_k1xsumq = params.d_G_k1xSumq[query_idx]; @@ -104,9 +112,12 @@ __device__ void bitwise_emit_distances(const ComputeInnerProductsKernelParams pa short_code_length, params.D); - params.d_topk_dists[output_offset + vec_idx] = - f_add + q_g_add + f_rescale * (exact_ip + q_k1xsumq); - params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; + float dist = f_add + q_g_add + f_rescale * (exact_ip + q_k1xsumq); + if constexpr (is_signed) { + dist = (dist - q_sqr_norm - params.d_vec_sqr_norms[global_vec_idx]) * 0.5f; + } + params.d_topk_dists[output_offset + vec_idx] = dist; + params.d_topk_pids[output_offset + vec_idx] = params.d_pids[global_vec_idx]; } } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json index 7e0f0f3603..6d32bd63fc 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_matrix.json @@ -8,5 +8,15 @@ "with_ex_descriptor": "no_ex", "with_ex_value": "false" } + ], + "_signed": [ + { + "signed_descriptor": "signed", + "signed_value": "true" + }, + { + "signed_descriptor": "unsigned", + "signed_value": "false" + } ] } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh index a577ec44ce..26411f4db7 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -23,6 +23,7 @@ namespace cuvs::neighbors::ivf_rabitq::detail { // inner-product loop is dispatched at runtime through the // compute_bitwise_quantized_ip_for_vec JIT-LTO fragment, so this kernel is // also not templated on num_bits. +template __device__ void compute_inner_products_with_bitwise_block_sort_impl( const ComputeInnerProductsKernelParams params) { @@ -57,6 +58,8 @@ __device__ void compute_inner_products_with_bitwise_block_sort_impl( float q_k1xsumq = params.d_G_k1xSumq[query_idx]; float q_g_error = sqrtf(q_g_add); float threshold = params.d_threshold[query_idx]; + // ‖q‖² for the InnerProduct pseudo-distance transform pd = (‖q−x‖² − ‖q‖² − ‖x‖²)/2 = −⟨q,x⟩. + float q_sqr_norm = Signed ? params.d_q_sqr_norms[query_idx] : 0.0f; if (tid == 0) { num_candidates = 0; } __syncthreads(); @@ -93,6 +96,14 @@ __device__ void compute_inner_products_with_bitwise_block_sort_impl( float est_dist = f_add + q_g_add + f_rescale * (ip + q_k1xsumq); float low_dist = est_dist - f_error * q_g_error; + if constexpr (Signed) { + // Convert the squared-L2 lower bound into the pseudo-distance lower bound. The map + // est → (est − ‖q‖² − ‖x‖²)/2 is affine with positive slope, so a lower bound on the + // L2 estimate is a valid lower bound on pd = −⟨q,x⟩; the `< threshold` prune stays correct. + float x_sqr = params.d_vec_sqr_norms[cluster_start_index + vec_idx]; + low_dist = (low_dist - q_sqr_norm - x_sqr) * 0.5f; + } + if (low_dist < threshold) { is_candidate = true; local_ip_quantized = ip; diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in index 7bbdc24139..8da495007a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in @@ -1,17 +1,23 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include +namespace { + +constexpr bool is_signed = @signed_value@; + +} // namespace + namespace cuvs::neighbors::ivf_rabitq::detail { extern "C" __global__ void compute_inner_products_with_bitwise_block_sort( const ComputeInnerProductsKernelParams params) { - compute_inner_products_with_bitwise_block_sort_impl(params); + compute_inner_products_with_bitwise_block_sort_impl(params); } static_assert(std::is_same_v void add_entrypoint() { - this->add_static_fragment(); + this + ->add_static_fragment>(); } - template + template void add_bitwise_block_sort_emit_topk_device_function() { - this->add_static_fragment>(); + this->add_static_fragment>(); } template diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index 3ec4809395..67bfd636d0 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,10 +24,10 @@ struct ComputeInnerProductsWithBitwisePlanner : AlgorithmPlanner { this->add_static_fragment(); } - template + template void add_bitwise_emit_distances_device_function() { - this->add_static_fragment>(); + this->add_static_fragment>(); } template diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh index 45c8793b59..d53f409603 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -214,14 +214,14 @@ __device__ void compute_inner_products_with_lut16_opt_block_sort_impl( } if (num_candidates >= params.topk) { - update_threshold_atomicmin(params.d_topk_dists, - params.d_threshold, - query_idx, - params.topk, - params.nprobe, - probe_slot, - threshold, - tid); + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); } } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh index 7939b8ec13..296060ffd4 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -74,7 +74,25 @@ __inline__ __device__ float compute_bitwise_1bit_ip_for_vec(const uint32_t* d_sh return exact_ip; } +// Atomic float minimum that is correct for both signs. The int-reinterpret atomicMin trick is +// only monotone for non-negative floats; for negative values the ordering inverts, so we route +// them through an unsigned atomicMax. Needed by the InnerProduct metric, whose pseudo-distances +// (−⟨q,x⟩) are routinely negative. +__inline__ __device__ void atomic_min_float(float* addr, float val) +{ + if (val >= 0.0f) { + atomicMin(reinterpret_cast(addr), __float_as_int(val)); + } else { + atomicMax(reinterpret_cast(addr), __float_as_uint(val)); + } +} + // Note: contains __syncthreads(); must be called by all threads of the block. +// Signed selects the InnerProduct pseudo-distance behavior at compile time: the values stored in +// d_topk_dists are −⟨q,x⟩ (can be negative), so the positivity guards and the int-atomicMin are +// replaced by sentinel-only guards and a sign-aware atomic min. For Signed == false (L2) the body +// is identical to the original squared-distance implementation. +template __inline__ __device__ void update_threshold_atomicmin(const float* d_topk_dists, const float* d_threshold, uint32_t query_idx, @@ -90,13 +108,24 @@ __inline__ __device__ void update_threshold_atomicmin(const float* d_topk_dists, uint32_t output_offset = query_idx * (topk * nprobe) + probe_slot * topk; for (uint32_t i = 0; i < topk; i++) { float dist = d_topk_dists[output_offset + i]; - if (dist > 0 && dist > max_topk_dist && dist < INFINITY) { max_topk_dist = dist; } + if constexpr (Signed) { + // Only the +INF fill sentinel (and unfilled slots) are invalid; negatives are real. + if (dist < INFINITY && dist > max_topk_dist) { max_topk_dist = dist; } + } else { + if (dist > 0 && dist > max_topk_dist && dist < INFINITY) { max_topk_dist = dist; } + } } } __syncthreads(); - if (tid == 0 && max_topk_dist > 0 && max_topk_dist < threshold) { - int* threshold_ptr = (int*)(d_threshold + query_idx); - atomicMin(threshold_ptr, __float_as_int(max_topk_dist)); + if constexpr (Signed) { + if (tid == 0 && max_topk_dist > -INFINITY && max_topk_dist < threshold) { + atomic_min_float(const_cast(d_threshold) + query_idx, max_topk_dist); + } + } else { + if (tid == 0 && max_topk_dist > 0 && max_topk_dist < threshold) { + int* threshold_ptr = (int*)(d_threshold + query_idx); + atomicMin(threshold_ptr, __float_as_int(max_topk_dist)); + } } } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp index d9c94821f1..aafb9f0d6c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -99,15 +99,23 @@ make_compute_inner_products_with_lut16_opt_block_sort_launcher(int ex_bits, bool } inline std::shared_ptr make_compute_inner_products_with_bitwise_launcher( - int ex_bits, bool with_ex) + int ex_bits, bool with_ex, bool is_inner_product) { ComputeInnerProductsWithBitwisePlanner planner; planner.add_entrypoint(); if (with_ex) { - planner.add_bitwise_emit_distances_device_function(); + if (is_inner_product) { + planner.add_bitwise_emit_distances_device_function(); + } else { + planner.add_bitwise_emit_distances_device_function(); + } add_ex_bits_device_functions(planner, ex_bits); } else { - planner.add_bitwise_emit_distances_device_function(); + if (is_inner_product) { + planner.add_bitwise_emit_distances_device_function(); + } else { + planner.add_bitwise_emit_distances_device_function(); + } } return planner.get_launcher(); } @@ -115,20 +123,33 @@ inline std::shared_ptr make_compute_inner_products_with_bitwi inline std::shared_ptr make_compute_inner_products_with_bitwise_block_sort_launcher(int num_bits, int ex_bits, - bool with_ex) + bool with_ex, + bool is_inner_product) { ComputeInnerProductsWithBitwiseBlockSortPlanner planner; - planner.add_entrypoint(); + if (is_inner_product) { + planner.add_entrypoint(); + } else { + planner.add_entrypoint(); + } if (num_bits == 4) { planner.add_compute_bitwise_quantized_ip_for_vec_device_function<4>(); } else { planner.add_compute_bitwise_quantized_ip_for_vec_device_function<8>(); } if (with_ex) { - planner.add_bitwise_block_sort_emit_topk_device_function(); + if (is_inner_product) { + planner.add_bitwise_block_sort_emit_topk_device_function(); + } else { + planner.add_bitwise_block_sort_emit_topk_device_function(); + } add_ex_bits_device_functions(planner, ex_bits); } else { - planner.add_bitwise_block_sort_emit_topk_device_function(); + if (is_inner_product) { + planner.add_bitwise_block_sort_emit_topk_device_function(); + } else { + planner.add_bitwise_block_sort_emit_topk_device_function(); + } } return planner.get_launcher(); } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in index 5828027b6d..f970be5f4d 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -178,14 +178,14 @@ __device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams (uint32_t*)(params.d_topk_pids + output_offset)); if (final_num_candidates >= params.topk) { - update_threshold_atomicmin(params.d_topk_dists, - params.d_threshold, - query_idx, - params.topk, - params.nprobe, - probe_slot, - threshold, - tid); + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); } } } else { @@ -270,14 +270,14 @@ __device__ void lut_block_sort_emit_topk(const ComputeInnerProductsKernelParams } if (num_candidates >= params.topk) { - update_threshold_atomicmin(params.d_topk_dists, - params.d_threshold, - query_idx, - params.topk, - params.nprobe, - probe_slot, - threshold, - tid); + update_threshold_atomicmin(params.d_topk_dists, + params.d_threshold, + query_idx, + params.topk, + params.nprobe, + probe_slot, + threshold, + tid); } } } diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 938f41f846..8cf47be99f 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -18,6 +18,10 @@ struct ivf_rabitq_inputs { uint32_t dim = 64; uint32_t k = 10; std::optional min_recall = std::nullopt; + // Generate signed, mean-zero data instead of the default positive-only data. Used by the + // InnerProduct cases to produce mixed-sign inner products (which exercise both branches of the + // sign-aware atomic threshold update). Only affects signed DataT (e.g. float). + bool signed_data = false; cuvs::neighbors::ivf_rabitq::index_params index_params; cuvs::neighbors::ivf_rabitq::search_params search_params; @@ -58,6 +62,7 @@ inline auto operator<<(std::ostream& os, const ivf_rabitq_inputs& p) -> std::ost PRINT_DIFF(.dim); PRINT_DIFF(.k); PRINT_DIFF_V(.min_recall, p.min_recall.value_or(0)); + PRINT_DIFF(.signed_data); PRINT_DIFF(.index_params.n_lists); PRINT_DIFF(.index_params.bits_per_dim); PRINT_DIFF(.index_params.kmeans_n_iters); @@ -86,10 +91,13 @@ class ivf_rabitq_test : public ::testing::TestWithParam { raft::random::RngState r(1234ULL); if constexpr (std::is_same{}) { + // Signed data (used only by InnerProduct cases) is mean-zero over [-2, 2]; the default is + // positive-only over [0.1, 2] to match the other ANN harnesses and keep existing recall + // thresholds valid. + const DataT lo = (ps.signed_data && std::is_signed_v) ? DataT(-2.0) : DataT(0.1); + raft::random::uniform(handle_, r, database.data(), ps.num_db_vecs * ps.dim, lo, DataT(2.0)); raft::random::uniform( - handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); - raft::random::uniform( - handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); + handle_, r, search_queries.data(), ps.num_queries * ps.dim, lo, DataT(2.0)); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); @@ -367,6 +375,31 @@ inline auto var_search_mode_1_bit() -> test_cases_t }); } +// InnerProduct metric. Only the bitwise (QUANT4/QUANT8) search paths implement it, so restrict to +// those modes. Cover both the block-sort (k <= kMaxTopKBlockSort == 64) and non-block-sort (k > +// 64) sub-paths, and both the no-ex (bits_per_dim == 1) and with-ex (bits_per_dim > 1) code paths. +inline auto var_metric() -> test_cases_t +{ + test_cases_t xs; + for (auto mode : {ivf_rabitq::search_mode::QUANT4, ivf_rabitq::search_mode::QUANT8}) { + for (uint32_t k : {uint32_t{10}, uint32_t{128}}) { + for (uint32_t bits : {uint32_t{1}, uint32_t{4}}) { + ivf_rabitq_inputs x; + x.index_params.metric = cuvs::distance::DistanceType::InnerProduct; + x.search_params.mode = mode; + x.k = k; + x.index_params.bits_per_dim = bits; + // Mean-zero data so inner products take both signs. + x.signed_data = true; + // 1-bit quantization is coarse; relax the recall bound as the other 1-bit tests do. + if (bits == 1) { x.min_recall = 0.3; } + xs.push_back(x); + } + } + } + return xs; +} + /* Test instantiations */ // Currently IVF-RaBitQ deserialization reorganizes data for efficient search and is required for diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 2b412f3401..3106be9f85 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,4 +16,10 @@ INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims() + var_n_probes() + var_k() + var_bits_per_dim() + var_search_mode() + var_search_mode_1_bit()); +// InnerProduct cases in their own instantiation so they can be run in isolation via +// --gtest_filter='IvfRabitqInnerProduct/*'. +INSTANTIATE_TEST_SUITE_P(IvfRabitqInnerProduct, // NOLINT + f32_f32_i64, + ::testing::ValuesIn(var_metric())); + } // namespace cuvs::neighbors::ivf_rabitq