Support inner-product distance metric in IVF-RaBitQ

cpp/CMakeLists.txt

@@ -698,12 +698,13 @@ if(NOT BUILD_CPU_ONLY)
   )
   generate_jit_lto_kernels(
     jit_lto_files
-    NAME_FORMAT "ivf_rabitq_lut_emit_distances_@with_ex_descriptor@"
+    NAME_FORMAT "ivf_rabitq_lut_emit_distances_@with_ex_descriptor@_@signed_descriptor@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in"
-    FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_emit_distances<@with_ex_value@>"
+    FRAGMENT_TAG_FORMAT
+      "${ivf_rabitq_ns}::fragment_tag_lut_emit_distances<@with_ex_value@, @signed_value@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_emit_distances"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
@@ -723,12 +724,13 @@ if(NOT BUILD_CPU_ONLY)
   )
   generate_jit_lto_kernels(
     jit_lto_files
-    NAME_FORMAT "ivf_rabitq_lut_block_sort_emit_topk_@with_ex_descriptor@"
+    NAME_FORMAT "ivf_rabitq_lut_block_sort_emit_topk_@with_ex_descriptor@_@signed_descriptor@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in"
-    FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut_block_sort_emit_topk<@with_ex_value@>"
+    FRAGMENT_TAG_FORMAT
+      "${ivf_rabitq_ns}::fragment_tag_lut_block_sort_emit_topk<@with_ex_value@, @signed_value@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut_block_sort_emit_topk"
@@ -789,26 +791,28 @@ if(NOT BUILD_CPU_ONLY)
   )
   generate_jit_lto_kernels(
     jit_lto_files
-    NAME_FORMAT "ivf_rabitq_lut16_opt_emit_distances_@with_ex_descriptor@"
+    NAME_FORMAT "ivf_rabitq_lut16_opt_emit_distances_@with_ex_descriptor@_@signed_descriptor@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in"
-    FRAGMENT_TAG_FORMAT "${ivf_rabitq_ns}::fragment_tag_lut16_opt_emit_distances<@with_ex_value@>"
+    FRAGMENT_TAG_FORMAT
+      "${ivf_rabitq_ns}::fragment_tag_lut16_opt_emit_distances<@with_ex_value@, @signed_value@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/lut16_opt_emit_distances"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
   generate_jit_lto_kernels(
     jit_lto_files
-    NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut16_opt_block_sort_@with_ex_descriptor@"
+    NAME_FORMAT
+      "ivf_rabitq_compute_inner_products_with_lut16_opt_block_sort_@with_ex_descriptor@_@signed_descriptor@"
     MATRIX_JSON_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_matrix.json"
     KERNEL_INPUT_FILE
       "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in"
     FRAGMENT_TAG_FORMAT
-      "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt_block_sort<@with_ex_value@>"
+      "${ivf_rabitq_ns}::fragment_tag_compute_inner_products_with_lut16_opt_block_sort<@with_ex_value@, @signed_value@>"
     FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_lut16_opt_block_sort"
@@ -829,40 +833,41 @@ 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 "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_emit_distances"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
   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 "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/compute_inner_products_with_bitwise_block_sort"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
   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 "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
     OUTPUT_DIRECTORY
       "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/bitwise_block_sort_emit_topk"

cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h

@@ -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
@@ -41,6 +41,7 @@ class cuvs_ivf_rabitq : public algo<T>, public algo_gpu {
   cuvs_ivf_rabitq(Metric metric, int dim, const build_param& param)
     : algo<T>(metric, dim), index_params_(param), dimension_(dim)
   {
+    index_params_.metric = parse_metric_type(metric);
   }
 
   void build(const T* dataset, size_t nrow) final;

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
  */
 
@@ -9,30 +9,31 @@ namespace cuvs::neighbors::ivf_rabitq::detail {
 
 struct fragment_tag_compute_inner_products_with_lut {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_lut_emit_distances {};
 
 struct fragment_tag_compute_inner_products_with_lut_block_sort {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_lut_block_sort_emit_topk {};
 
 struct fragment_tag_compute_inner_products_with_lut16_opt {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_lut16_opt_emit_distances {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_compute_inner_products_with_lut16_opt_block_sort {};
 
 struct fragment_tag_compute_inner_products_with_bitwise {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_bitwise_emit_distances {};
 
+template <bool Signed>
 struct fragment_tag_compute_inner_products_with_bitwise_block_sort {};
 
-template <bool WithEx>
+template <bool WithEx, bool Signed>
 struct fragment_tag_bitwise_block_sort_emit_topk {};
 
 template <int EX_BITS>

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 <cuda_fp16.h>
 
 #include <cuvs/core/export.hpp>
+#include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/common.hpp>
 
 #include <raft/core/device_mdarray.hpp>
@@ -149,14 +150,18 @@ 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;
 
   /** 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;
 

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<IdxT> index(handle, n_rows, dim, params.n_lists, params.bits_per_dim);
+  index<IdxT> 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) {
@@ -314,10 +317,12 @@ index<IdxT>::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<detail::IVFGPU>(handle, n_rows, dim, n_lists, bits_per_dim);
+  rabitq_index_ =
+    std::make_unique<detail::IVFGPU>(handle, n_rows, dim, n_lists, bits_per_dim, metric);
 }
 
 template <typename IdxT>
@@ -349,6 +354,12 @@ IdxT index<IdxT>::size() const noexcept
   return rabitq_index_->get_num_vectors();
 }
 
+template <typename IdxT>
+cuvs::distance::DistanceType index<IdxT>::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<const float, int64_t, raft::row_major> dataset)