Serialize Vamana index with SSD sector alignment per MSFT DiskANN format, generate quantized dataset for integration with DiskANN (#846)

(This supersedes [PR!703](#703))

Added an optional input flag to `cuvs::neighbors::vamana::serialize` to dump an input cuvs Vamana index to file with SSD sector alignment. File format follows [MSFT DiskANN](https://github.com/microsoft/DiskANN/blob/main/src/disk_utils.cpp).

Using the sector-aligned option also writes out the quantized dataset computed using user-supplied PQ codebooks file and rotation matrix file.

Authors:
  - James Xia (https://github.com/jamxia155)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Tarang Jain (https://github.com/tarang-jain)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Gil Forsyth (https://github.com/gforsyth)

URL: #846

Author: jamxia155

ci/get_test_data.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+# Copyright (c) 2025, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+NUMARGS=$#
+ARGS=$*
+function hasArg {
+    (( NUMARGS != 0 )) && (echo " ${ARGS} " | grep -q " $1 ")
+}
+
+if hasArg -h || hasArg --help; then
+    echo "Usage:"
+    echo "    $0 [--NEIGHBORS_ANN_VAMANA_TEST]"
+    exit 0
+fi
+
+DEST=${RAPIDS_DATASET_ROOT_DIR:-"$PWD"}
+
+# get test data for NEIGHBORS_ANN_VAMANA_TEST
+if hasArg "--NEIGHBORS_ANN_VAMANA_TEST"; then
+    echo "Downloading test data for NEIGHBORS_ANN_VAMANA_TEST"
+    echo "Destination: ${DEST}"
+    URL_PREFIX=https://data.rapids.ai/cuvs/tests/data
+    SUBDIR=neighbors/ann_vamana/randomized_codebooks
+    FILE_LIST=(
+        384_int8_pq_pivots.bin
+        384_int8_pq_pivots.bin_rotation_matrix.bin
+        64_float_pq_pivots.bin
+        64_float_pq_pivots.bin_rotation_matrix.bin
+    )
+    for f in "${FILE_LIST[@]}"; do
+        wget --no-verbose --directory-prefix="${DEST}/${SUBDIR}" "${URL_PREFIX}/${SUBDIR}/$f"
+    done
+fi

ci/test_cpp.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION.
 
 set -euo pipefail
 
@@ -30,6 +30,11 @@ rapids-print-env
 rapids-logger "Check GPU usage"
 nvidia-smi
 
+# RAPIDS_DATASET_ROOT_DIR is used by test scripts
+RAPIDS_DATASET_ROOT_DIR=${RAPIDS_TESTS_DIR}/dataset
+export RAPIDS_DATASET_ROOT_DIR
+./ci/get_test_data.sh --NEIGHBORS_ANN_VAMANA_TEST
+
 EXITCODE=0
 trap "EXITCODE=1" ERR
 set +e

cpp/CMakeLists.txt

@@ -484,6 +484,7 @@ if(BUILD_SHARED_LIBS)
     src/neighbors/vamana_build_float.cu
     src/neighbors/vamana_build_uint8.cu
     src/neighbors/vamana_build_int8.cu
+    src/neighbors/vamana_codebooks_float.cu
     src/neighbors/vamana_serialize_float.cu
     src/neighbors/vamana_serialize_uint8.cu
     src/neighbors/vamana_serialize_int8.cu

cpp/include/cuvs/neighbors/vamana.hpp

@@ -53,6 +53,17 @@ namespace cuvs::neighbors::vamana {
  */
 
 struct index_params : cuvs::neighbors::index_params {
+  /**
+   * @brief Parameters used to build quantized DiskANN index; to be generated using
+   * deserialize_codebooks()
+   */
+  template <typename T = float>
+  struct codebook_params {
+    int pq_codebook_size;
+    int pq_dim;
+    std::vector<T> pq_encoding_table;
+    std::vector<T> rotation_matrix;
+  };
   /** Maximum degree of output graph corresponds to the R parameter in the original Vamana
    * literature. */
   uint32_t graph_degree = 32;
@@ -72,6 +83,8 @@ struct index_params : cuvs::neighbors::index_params {
   uint32_t queue_size = 127;
   /** Max batchsize of reverse edge processing (reduces memory footprint) */
   uint32_t reverse_batchsize = 1000000;
+  /** Codebooks and related parameters */
+  std::optional<codebook_params<float>> codebooks = std::nullopt;
 };
 
 /**
@@ -127,6 +140,13 @@ struct index : cuvs::neighbors::index {
     return *dataset_;
   }
 
+  /** Quantized dataset [size, codes_rowlen] */
+  [[nodiscard]] inline auto quantized_data() const noexcept
+    -> raft::device_matrix_view<const uint8_t, int64_t, raft::row_major>
+  {
+    return quantized_dataset_.view();
+  }
+
   /** vamana graph [size, graph-degree] */
   [[nodiscard]] inline auto graph() const noexcept
     -> raft::device_matrix_view<const IdxT, int64_t, raft::row_major>
@@ -150,7 +170,8 @@ struct index : cuvs::neighbors::index {
     : cuvs::neighbors::index(),
       metric_(metric),
       graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
-      dataset_(new cuvs::neighbors::empty_dataset<int64_t>(0))
+      dataset_(new cuvs::neighbors::empty_dataset<int64_t>(0)),
+      quantized_dataset_(raft::make_device_matrix<uint8_t, int64_t>(res, 0, 0))
   {
   }
 
@@ -168,6 +189,7 @@ struct index : cuvs::neighbors::index {
       metric_(metric),
       graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
       dataset_(make_aligned_dataset(res, dataset, 16)),
+      quantized_dataset_(raft::make_device_matrix<uint8_t, int64_t>(res, 0, 0)),
       medoid_id_(medoid_id)
   {
     RAFT_EXPECTS(dataset.extent(0) == vamana_graph.extent(0),
@@ -212,11 +234,42 @@ struct index : cuvs::neighbors::index {
     graph_view_ = graph_.view();
   }
 
+  /**
+   * @brief Replace the current quantized dataset with a new quantized dataset.
+   *
+   * We create a copy of the quantized dataset on the device. The index manages the lifetime of this
+   * copy.
+   *
+   * @param[in] res
+   * @param[in] new_quantized_dataset the new quantized dataset for the index
+   *
+   */
+  void update_quantized_dataset(
+    raft::resources const& res,
+    raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_quantized_dataset)
+  {
+    RAFT_LOG_DEBUG("Creating device copy of Vamana quantized dataset");
+    if ((quantized_dataset_.extent(0) != new_quantized_dataset.extent(0)) ||
+        (quantized_dataset_.extent(1) != new_quantized_dataset.extent(1))) {
+      // clear existing memory before allocating to prevent OOM errors on large datasets
+      if (quantized_dataset_.size()) {
+        quantized_dataset_ = raft::make_device_matrix<uint8_t, int64_t>(res, 0, 0);
+      }
+      quantized_dataset_ = raft::make_device_matrix<uint8_t, int64_t>(
+        res, new_quantized_dataset.extent(0), new_quantized_dataset.extent(1));
+    }
+    raft::copy(quantized_dataset_.data_handle(),
+               new_quantized_dataset.data_handle(),
+               new_quantized_dataset.size(),
+               raft::resource::get_cuda_stream(res));
+  }
+
  private:
   cuvs::distance::DistanceType metric_;
   raft::device_matrix<IdxT, int64_t, raft::row_major> graph_;
   raft::device_matrix_view<const IdxT, int64_t, raft::row_major> graph_view_;
   std::unique_ptr<neighbors::dataset<int64_t>> dataset_;
+  raft::device_matrix<uint8_t, int64_t, raft::row_major> quantized_dataset_;
   IdxT medoid_id_;
 };
 /**
@@ -457,13 +510,15 @@ auto build(raft::resources const& res,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<float, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
 
 /**
  * Save the index to file.
@@ -486,12 +541,14 @@ void serialize(raft::resources const& handle,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<int8_t, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
 
 /**
  * Save the index to file.
@@ -514,12 +571,48 @@ void serialize(raft::resources const& handle,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<uint8_t, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
+
+/**
+ * @}
+ */
+
+/**
+ * @defgroup vamana_cpp_codebook Vamana codebook functions
+ * @{
+ */
+
+/**
+ * @brief Construct codebook parameters from input codebook files
+ *
+ * Expects pq pivots file at
+ * "${codebook_prefix}_pq_pivots.bin" and rotation matrix file at
+ * "${codebook_prefix}_pq_pivots.bin_rotation_matrix.bin".
+ *
+ * @code{.cpp}
+ *   #include <cuvs/neighbors/vamana.hpp>
+ *
+ *   // create a string with a filepath
+ *   std::string codebook_prefix("/path/to/index/prefix");
+ *   // define dimension of vectors in dataset
+ *   int dim = 64;
+ *   // construct codebook parameters from input codebook files
+ *   auto codebooks = cuvs::neighbors::vamana::deserialize_codebooks(codebook_prefix, dim);
+ * @endcode
+ *
+ * @param[in] codebook_prefix path prefix to pq pivots and rotation matrix files
+ * @param[in] dim dimension of vectors in dataset
+ *
+ */
+auto deserialize_codebooks(const std::string& codebook_prefix, const int dim)
+  -> index_params::codebook_params<float>;
 
 /**
  * @}