[FEA] IVF-PQ to Write Flat PQ Codes

c/include/cuvs/neighbors/ivf_pq.h

@@ -28,6 +28,15 @@ enum codebook_gen {  // NOLINT
   PER_CLUSTER  = 1,  // NOLINT
 };
 
+/**
+ * @brief A type for specifying the memory layout of IVF-PQ list data
+ *
+ */
+enum list_layout {   // NOLINT
+  FLAT        = 0,   // NOLINT
+  INTERLEAVED = 1,   // NOLINT
+};
+
 /**
  * @brief Supplemental parameters to build IVF-PQ Index
  *
@@ -114,6 +123,14 @@ struct cuvsIvfPqIndexParams {
    * points to train each codebook.
    */
   uint32_t max_train_points_per_pq_code;
+  /**
+   * Memory layout of the IVF-PQ list data.
+   *
+   * - FLAT: Codes are stored contiguously, one vector's codes after another.
+   * - INTERLEAVED: Codes are interleaved for optimized search performance.
+   *   This is the default and recommended for search workloads.
+   */
+  enum list_layout codes_layout;
 };
 
 typedef struct cuvsIvfPqIndexParams* cuvsIvfPqIndexParams_t;

c/src/neighbors/ivf_pq.cpp

@@ -34,6 +34,7 @@ void convert_c_index_params(cuvsIvfPqIndexParams params, cuvs::neighbors::ivf_pq
   out->force_random_rotation          = params.force_random_rotation;
   out->conservative_memory_allocation = params.conservative_memory_allocation;
   out->max_train_points_per_pq_code   = params.max_train_points_per_pq_code;
+  out->codes_layout = static_cast<cuvs::neighbors::ivf_pq::list_layout>((int)params.codes_layout);
 }
 void convert_c_search_params(cuvsIvfPqSearchParams params,
                              cuvs::neighbors::ivf_pq::search_params* out)
@@ -218,8 +219,16 @@ void _get_list_indices(cuvsIvfPqIndex index,
                        uint32_t label,
                        DLManagedTensor* out_labels)
 {
-  auto index_ptr    = reinterpret_cast<cuvs::neighbors::ivf_pq::index<IdxT>*>(index.addr);
-  cuvs::core::to_dlpack(index_ptr->lists()[label]->indices.view(), out_labels);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_pq::index<IdxT>*>(index.addr);
+  if (index_ptr->codes_layout() == cuvs::neighbors::ivf_pq::list_layout::FLAT) {
+    auto& list =
+      static_cast<cuvs::neighbors::ivf_pq::list_data_flat<IdxT>&>(*index_ptr->lists()[label]);
+    cuvs::core::to_dlpack(list.indices.view(), out_labels);
+  } else {
+    auto& list = static_cast<cuvs::neighbors::ivf_pq::list_data_interleaved<IdxT>&>(
+      *index_ptr->lists()[label]);
+    cuvs::core::to_dlpack(list.indices.view(), out_labels);
+  }
 }
 }  // namespace
 
@@ -328,7 +337,8 @@ extern "C" cuvsError_t cuvsIvfPqIndexParamsCreate(cuvsIvfPqIndexParams_t* params
                                        .codebook_kind                  = codebook_gen::PER_SUBSPACE,
                                        .force_random_rotation          = false,
                                        .conservative_memory_allocation = false,
-                                       .max_train_points_per_pq_code   = 256};
+                                       .max_train_points_per_pq_code   = 256,
+                                       .codes_layout                   = list_layout::INTERLEAVED};
   });
 }
 

cpp/cmake/patches/faiss-1.13-cuvs-26.02.diff

@@ -1,5 +1,5 @@
 diff --git a/faiss/gpu/impl/CuvsIVFPQ.cu b/faiss/gpu/impl/CuvsIVFPQ.cu
-index 1e2fef225..35b388147 100644
+index 1e2fef225..2ee40da46 100644
 --- a/faiss/gpu/impl/CuvsIVFPQ.cu
 +++ b/faiss/gpu/impl/CuvsIVFPQ.cu
 @@ -129,8 +129,14 @@ void CuvsIVFPQ::updateQuantizer(Index* quantizer) {
@@ -122,7 +122,35 @@ index 1e2fef225..35b388147 100644
      }
 
      setPQCentroids_();
-@@ -520,7 +583,7 @@ void CuvsIVFPQ::setPQCentroids_() {
+@@ -404,10 +467,11 @@ void CuvsIVFPQ::copyInvertedListsFrom(const InvertedLists* ivf) {
+     auto& cuvs_index_lists = cuvs_index->lists();
+
+     // conservative memory alloc for cloning cpu inverted lists
+-    cuvs::neighbors::ivf_pq::list_spec<uint32_t, idx_t> ivf_list_spec{
+-            static_cast<uint32_t>(bitsPerSubQuantizer_),
+-            static_cast<uint32_t>(numSubQuantizers_),
+-            true};
++    cuvs::neighbors::ivf_pq::list_spec_interleaved<uint32_t, idx_t>
++            ivf_list_spec{
++                    static_cast<uint32_t>(bitsPerSubQuantizer_),
++                    static_cast<uint32_t>(numSubQuantizers_),
++                    true};
+
+     for (size_t i = 0; i < nlist; ++i) {
+         size_t listSize = ivf->list_size(i);
+@@ -426,9 +490,9 @@ void CuvsIVFPQ::copyInvertedListsFrom(const InvertedLists* ivf) {
+         // This cuVS list must currently be empty
+         FAISS_ASSERT(getListLength(i) == 0);
+
+-        cuvs::neighbors::ivf::resize_list(
++        cuvs::neighbors::ivf_pq::helpers::resize_list(
+                 raft_handle,
+-                cuvs_index_lists[i],
++                cuvs_index_lists[i],
+                 ivf_list_spec,
+                 static_cast<uint32_t>(listSize),
+                 static_cast<uint32_t>(0));
+@@ -520,7 +587,7 @@ void CuvsIVFPQ::setPQCentroids_() {
      auto stream = resources_->getDefaultStreamCurrentDevice();
 
      raft::copy(

cpp/include/cuvs/neighbors/common.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -683,11 +683,50 @@ template <typename IdxT>
 constexpr static IdxT kInvalidRecord =
   (std::is_signed_v<IdxT> ? IdxT{0} : std::numeric_limits<IdxT>::max()) - 1;
 
+/**
+ * Abstract base class for IVF list data.
+ * This allows polymorphic access to list data regardless of the underlying layout.
+ *
+ * @tparam ValueT The data element type (e.g., uint8_t for PQ codes, float for raw vectors)
+ * @tparam IdxT The index type for source indices
+ * @tparam SizeT The size type
+ */
+template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
+struct list_base {
+  using value_type = ValueT;
+  using index_type = IdxT;
+  using size_type  = SizeT;
+
+  virtual ~list_base() = default;
+
+  /** Get the raw data pointer. */
+  virtual value_type* data_ptr() noexcept             = 0;
+  virtual const value_type* data_ptr() const noexcept = 0;
+
+  /** Get the indices pointer. */
+  virtual index_type* indices_ptr() noexcept             = 0;
+  virtual const index_type* indices_ptr() const noexcept = 0;
+
+  /** Get the current size (number of records). */
+  virtual size_type get_size() const noexcept = 0;
+
+  /** Set the current size (number of records). */
+  virtual void set_size(size_type new_size) noexcept = 0;
+
+  /** Get the total size of the data array in bytes. */
+  virtual size_t data_byte_size() const noexcept = 0;
+
+  /** Get the capacity (number of indices that can be stored). */
+  virtual size_type indices_capacity() const noexcept = 0;
+};
+
 /** The data for a single IVF list. */
 template <template <typename, typename...> typename SpecT,
           typename SizeT,
           typename... SpecExtraArgs>
-struct list {
+struct list : public list_base<typename SpecT<SizeT, SpecExtraArgs...>::value_type,
+                               typename SpecT<SizeT, SpecExtraArgs...>::index_type,
+                               SizeT> {
   using size_type    = SizeT;
   using spec_type    = SpecT<size_type, SpecExtraArgs...>;
   using value_type   = typename spec_type::value_type;
@@ -703,6 +742,18 @@ struct list {
 
   /** Allocate a new list capable of holding at least `n_rows` data records and indices. */
   list(raft::resources const& res, const spec_type& spec, size_type n_rows);
+
+  value_type* data_ptr() noexcept override { return data.data_handle(); }
+  const value_type* data_ptr() const noexcept override { return data.data_handle(); }
+
+  index_type* indices_ptr() noexcept override { return indices.data_handle(); }
+  const index_type* indices_ptr() const noexcept override { return indices.data_handle(); }
+
+  size_type get_size() const noexcept override { return size.load(); }
+  void set_size(size_type new_size) noexcept override { size.store(new_size); }
+
+  size_t data_byte_size() const noexcept override { return data.size() * sizeof(value_type); }
+  size_type indices_capacity() const noexcept override { return indices.extent(0); }
 };
 
 template <typename ListT, class T = void>
@@ -727,6 +778,10 @@ using enable_if_valid_list_t = typename enable_if_valid_list<ListT, T>::type;
 /**
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
+ *
+ * @note This is an internal function that requires the concrete list type.
+ *       For IVF-PQ indexes, prefer using the helper functions in
+ *       `cuvs::neighbors::ivf_pq::helpers::resize_list` which handle type casting internally.
  */
 template <typename ListT>
 void resize_list(raft::resources const& res,
@@ -735,13 +790,23 @@ void resize_list(raft::resources const& res,
                  typename ListT::size_type new_used_size,
                  typename ListT::size_type old_used_size);
 
+/**
+ * Serialize a list to an output stream.
+ *
+ * @note This function requires the concrete list type (not the base class) because:
+ *       1. It needs access to the spec_type to determine the data layout for serialization
+ *       2. The serialized format depends on the spec's make_list_extents() method
+ *       When calling from code that only has a base class pointer, use std::static_pointer_cast
+ *       to obtain the typed pointer first.
+ */
 template <typename ListT>
 enable_if_valid_list_t<ListT> serialize_list(
   const raft::resources& handle,
   std::ostream& os,
   const ListT& ld,
   const typename ListT::spec_type& store_spec,
   std::optional<typename ListT::size_type> size_override = std::nullopt);
+
 template <typename ListT>
 enable_if_valid_list_t<ListT> serialize_list(
   const raft::resources& handle,

cpp/include/cuvs/neighbors/ivf_flat.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -2933,17 +2933,17 @@ void reset_index(const raft::resources& res, index<uint8_t, int64_t>* index);
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, float, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<float, int64_t>* index);
 
@@ -2961,17 +2961,17 @@ void recompute_internal_state(const raft::resources& res, index<float, int64_t>*
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, half, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<half, int64_t>* index);
 
@@ -2989,17 +2989,17 @@ void recompute_internal_state(const raft::resources& res, index<half, int64_t>*
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, int8_t, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<int8_t, int64_t>* index);
 
@@ -3018,9 +3018,9 @@ void recompute_internal_state(const raft::resources& res, index<int8_t, int64_t>
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
  *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);