[FlowANN step 7/12] Add VPQ vector compression to tiered construction and search

c/CMakeLists.txt

@@ -32,6 +32,8 @@ option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON)
 option(BUILD_MG_ALGOS "Build with multi-GPU support" ON)
 
 option(CUVSC_STATIC_CUVS_LIBRARY "Link against statical version of libcuvs" OFF)
+set(CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH "$ORIGIN" CACHE STRING
+    "Install RPATH for external GDRCopy runtime libraries")
 
 # Check if cuVS is already available. If so, it is the user's responsibility to ensure that the
 # CMake package is also available at build time of the Python cuvs package.
@@ -123,6 +125,13 @@ set_target_properties(
              VISIBILITY_INLINES_HIDDEN ON
 )
 
+if(CUVSC_STATIC_CUVS_LIBRARY)
+  # Imported cuvs_static can carry GDRCopy without exporting its build-time feature flags.
+  # Keep the runtime deployment setting available to standalone C builds too.
+  set_property(TARGET cuvs_c APPEND PROPERTY INSTALL_RPATH
+               "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}")
+endif()
+
 # Setup the VERSION and SOVERSION of our library
 #
 # Long term goal is that every release that breaks the C ABI will update `cuvs_abi_version` allowing

cpp/CMakeLists.txt

@@ -84,7 +84,18 @@ option(DISABLE_OPENMP "Disable OpenMP" OFF)
 option(CUVS_COMPILE_DYNAMIC_ONLY "Only build the shared library and skip the static library." OFF)
 option(CUVS_NVTX "Enable nvtx markers" OFF)
 option(CUVS_RAFT_CLONE_ON_PIN "Explicitly clone RAFT branch when pinned to non-feature branch" ON)
-
+option(CUVS_ENABLE_FLOWANN_SEARCH "Enable experimental FlowANN tiered-graph search" OFF)
+option(CUVS_ENABLE_FLOWANN_BUILD "Enable the experimental integrated FlowANN graph builder" OFF)
+option(CUVS_FLOWANN_USE_GDRCOPY "Use GDRCopy for FlowANN GPU-to-CPU command queues" ON)
+set(CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH "\$ORIGIN" CACHE STRING
+    "Runtime search paths for externally deployed GDRCopy; never inferred from the build host")
+
+if(CUVS_ENABLE_FLOWANN_SEARCH AND BUILD_CPU_ONLY)
+  message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_SEARCH requires a CUDA-enabled build")
+endif()
+if(CUVS_ENABLE_FLOWANN_BUILD AND NOT CUVS_ENABLE_FLOWANN_SEARCH)
+  message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_BUILD requires CUVS_ENABLE_FLOWANN_SEARCH")
+endif()
 if(BUILD_CPU_ONLY)
   set(BUILD_SHARED_LIBS OFF)
   set(BUILD_TESTS OFF)
@@ -113,7 +124,6 @@ elseif(NOT BUILD_SHARED_LIBS)
   set(cuvs_compile_mode "static_only")
   set(_cuvs_lib_targets cuvs_static)
 endif()
-
 # Needed because GoogleBenchmark changes the state of FindThreads.cmake, causing subsequent runs to
 # have different values for the `Threads::Threads` target. Setting this flag ensures
 # `Threads::Threads` is the same value across all builds so that cache hits occur
@@ -127,6 +137,9 @@ message(VERBOSE "cuVS: Build cuVS unit-tests: ${BUILD_TESTS}")
 message(VERBOSE "cuVS: Build CPU only components: ${BUILD_CPU_ONLY}")
 message(VERBOSE "cuVS: Build ANN benchmarks: ${BUILD_CUVS_BENCH}")
 message(VERBOSE "cuVS: Build only the shared library: ${CUVS_COMPILE_DYNAMIC_ONLY}")
+message(VERBOSE "cuVS: Enable FlowANN tiered-graph search: ${CUVS_ENABLE_FLOWANN_SEARCH}")
+message(VERBOSE "cuVS: Enable FlowANN integrated build: ${CUVS_ENABLE_FLOWANN_BUILD}")
+message(VERBOSE "cuVS: Use GDRCopy for FlowANN queues: ${CUVS_FLOWANN_USE_GDRCOPY}")
 message(VERBOSE "cuVS: Enable detection of conda environment for dependencies: ${DETECT_CONDA_ENV}")
 message(VERBOSE "cuVS: Disable deprecation warnings " ${DISABLE_DEPRECATION_WARNINGS})
 message(VERBOSE "cuVS: Disable OpenMP: ${DISABLE_OPENMP}")
@@ -169,6 +182,22 @@ if(NOT BUILD_CPU_ONLY)
     BUILD_EXPORT_SET cuvs-exports
     INSTALL_EXPORT_SET cuvs-exports
   )
+
+  if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY)
+    rapids_find_generate_module(
+      GDRCopy
+      HEADER_NAMES gdrapi.h
+      LIBRARY_NAMES gdrapi
+      BUILD_EXPORT_SET cuvs-exports
+      INSTALL_EXPORT_SET cuvs-exports
+    )
+    rapids_find_package(
+      GDRCopy REQUIRED
+      BUILD_EXPORT_SET cuvs-exports
+      INSTALL_EXPORT_SET cuvs-exports
+    )
+    set(CUVS_FLOWANN_GDRCOPY_DEPENDENCIES GDRCopy::GDRCopy CUDA::cuda_driver)
+  endif()
 else()
   add_compile_definitions(BUILD_CPU_ONLY)
 endif()
@@ -233,6 +262,10 @@ if(NOT BUILD_CPU_ONLY)
               "$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/include>"
               "$<INSTALL_INTERFACE:include>"
   )
+  target_compile_definitions(
+    cuvs_cpp_headers INTERFACE $<$<BOOL:${CUVS_ENABLE_FLOWANN_SEARCH}>:CUVS_ENABLE_FLOWANN_SEARCH>
+                               $<$<BOOL:${CUVS_ENABLE_FLOWANN_BUILD}>:CUVS_ENABLE_FLOWANN_BUILD>
+  )
   target_link_libraries(cuvs_cpp_headers INTERFACE raft::raft rmm::rmm)
 
   generate_inst_matrix(
@@ -959,6 +992,44 @@ if(NOT BUILD_CPU_ONLY)
     OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/cagra/search_single_cta"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
+  if(CUVS_ENABLE_FLOWANN_SEARCH)
+    set(flowann_ns "cuvs::neighbors::cagra::experimental::flowann::detail")
+    generate_jit_lto_kernels(
+      jit_lto_files
+      NAME_FORMAT
+        "flowann_search_single_cta_@topk_by_bitonic_sort_str@_@bitonic_sort_and_merge_multi_warps_str@_data_@data_abbrev@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_single_cta_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in"
+      FRAGMENT_TAG_FORMAT
+        "${flowann_ns}::fragment_tag_search_single_cta<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@source_index_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, @topk_by_bitonic_sort@, @bitonic_sort_and_merge_multi_warps@>"
+      FRAGMENT_TAG_HEADER_FILES
+        "<cuvs/detail/jit_lto/flowann/flowann_fragments.hpp>"
+        "<cuvs/detail/jit_lto/common_fragments.hpp>"
+        "<cuvs/detail/jit_lto/cagra/cagra_fragments.hpp>"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/flowann/search_single_cta"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+      KERNEL_EXTRA_COMPILE_OPTIONS --maxrregcount=64
+    )
+    generate_jit_lto_kernels(
+      jit_lto_files
+      NAME_FORMAT "flowann_search_multi_cta_data_@data_abbrev@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_multi_cta_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in"
+      FRAGMENT_TAG_FORMAT
+        "${flowann_ns}::fragment_tag_search_multi_cta<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@source_index_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@>"
+      FRAGMENT_TAG_HEADER_FILES
+        "<cuvs/detail/jit_lto/flowann/flowann_fragments.hpp>"
+        "<cuvs/detail/jit_lto/common_fragments.hpp>"
+        "<cuvs/detail/jit_lto/cagra/cagra_fragments.hpp>"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/flowann/search_multi_cta"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+      KERNEL_EXTRA_COMPILE_OPTIONS --maxrregcount=64
+    )
+  endif()
   generate_jit_lto_kernels(
     jit_lto_files
     NAME_FORMAT
@@ -1486,6 +1557,9 @@ if(NOT BUILD_CPU_ONLY)
     src/neighbors/brute_force.cu
     src/neighbors/brute_force_serialize.cu
     src/neighbors/cagra.cpp
+    $<$<BOOL:${CUVS_ENABLE_FLOWANN_SEARCH}>:src/neighbors/flowann.cu>
+    $<$<BOOL:${CUVS_ENABLE_FLOWANN_BUILD}>:src/neighbors/flowann_build.cu>
+    $<$<BOOL:${CUVS_ENABLE_FLOWANN_SEARCH}>:src/neighbors/flowann_queue.cpp>
     ${cagra_build_inst_files}
     ${cagra_extend_inst_files}
     src/neighbors/cagra_optimize.cu
@@ -1592,6 +1666,12 @@ if(NOT BUILD_CPU_ONLY)
             "$<$<AND:$<COMPILE_LANGUAGE:CUDA>,$<CONFIG:Debug>>:${CUVS_DEBUG_CUDA_FLAGS}>"
   )
 
+  target_compile_definitions(
+    cuvs_objs
+    PRIVATE
+      $<$<AND:$<BOOL:${CUVS_ENABLE_FLOWANN_SEARCH}>,$<BOOL:${CUVS_FLOWANN_USE_GDRCOPY}>>:CUVS_FLOWANN_USE_GDRCOPY>
+  )
+
   target_compile_definitions(
     cuvs_objs
     PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}
@@ -1608,6 +1688,7 @@ if(NOT BUILD_CPU_ONLY)
             nvidia::cutlass::cutlass
             rtcx::rtcx
             ${CUVS_CTK_MATH_DEPENDENCIES}
+            ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES}
             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
             $<TARGET_NAME_IF_EXISTS:hnswlib::hnswlib>
             kvikio::kvikio
@@ -1687,10 +1768,30 @@ if(NOT BUILD_CPU_ONLY)
              $<BUILD_LOCAL_INTERFACE:$<TARGET_NAME_IF_EXISTS:NCCL::NCCL>>
              $<BUILD_LOCAL_INTERFACE:$<TARGET_NAME_IF_EXISTS:hnswlib::hnswlib>>
              $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
-      PRIVATE $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX> $<COMPILE_ONLY:nvidia::cutlass::cutlass>
-              $<COMPILE_ONLY:cuco::cuco> CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx kvikio::kvikio
+      PRIVATE $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
+              $<COMPILE_ONLY:nvidia::cutlass::cutlass>
+              $<COMPILE_ONLY:cuco::cuco>
+              CUDA::nvJitLink
+              CUDA::nvrtc
+              rtcx::rtcx
+              kvikio::kvikio
+              ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES}
       INTERFACE $<LINK_ONLY:kvikio::kvikio>
     )
+    if(CUVS_ENABLE_FLOWANN_SEARCH)
+      target_link_libraries(cuvs INTERFACE $<LINK_ONLY:$<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>>)
+    endif()
+    if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY)
+      # GDRCopy is an external runtime dependency. Default to an install-relative path;
+      # packagers may supply an explicit deployment path, independent of the build prefix.
+      set_property(TARGET cuvs APPEND PROPERTY INSTALL_RPATH
+                   "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}")
+      # Consumers must resolve this runtime dependency even when it is installed outside the
+      # linker's default search paths. Its headers remain private to the queue implementation.
+      target_link_libraries(
+        cuvs INTERFACE $<LINK_ONLY:GDRCopy::GDRCopy> $<LINK_ONLY:CUDA::cuda_driver>
+      )
+    endif()
 
   endif()
 
@@ -1735,6 +1836,7 @@ if(NOT BUILD_CPU_ONLY)
               CUDA::nvJitLink
               CUDA::nvrtc
               rtcx::rtcx
+              ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES}
               $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
               $<COMPILE_ONLY:nvidia::cutlass::cutlass>
               $<COMPILE_ONLY:cuco::cuco>

cpp/include/cuvs/detail/jit_lto/flowann/flowann_fragments.hpp

@@ -0,0 +1,24 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * FlowANN is described in "Disentangling Graph Dependencies for Efficient
+ * Billion-Scale GPU Vector Search" (OSDI 2026).
+ */
+
+#pragma once
+
+namespace cuvs::neighbors::cagra::experimental::flowann::detail {
+
+template <typename DataTag,
+          typename SourceIndexTag,
+          typename IndexTag,
+          typename DistanceTag,
+          bool TopkByBitonicSort,
+          bool BitonicSortAndMergeMultiWarps>
+struct fragment_tag_search_single_cta {};
+
+template <typename DataTag, typename SourceIndexTag, typename IndexTag, typename DistanceTag>
+struct fragment_tag_search_multi_cta {};
+
+}  // namespace cuvs::neighbors::cagra::experimental::flowann::detail

cpp/src/neighbors/detail/cagra/cagra_build.cuh

@@ -719,6 +719,8 @@ auto iterative_build_graph(raft::resources const& res,
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
+  // Other callers also pass the original parameters after checking local degree copies.
+  check_graph_degree<T, IdxT>(intermediate_degree, graph_degree, dataset.n_rows());
 
   const auto& iter_params =
     std::get<cagra::graph_build_params::iterative_search_params>(params.graph_build_params);
@@ -755,11 +757,15 @@ auto iterative_build_graph(raft::resources const& res,
     final_graph_size = static_cast<uint64_t>(dataset.n_rows());
   }
 
+  // Every initial row must have enough distinct non-self neighbors for the
+  // first search-and-optimize iteration, including very small datasets.
+  auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24));
   // Determine initial graph size.
   uint64_t initial_graph_size = (final_graph_size + 1) / 2;
   while (initial_graph_size > graph_degree * 64) {
     initial_graph_size = (initial_graph_size + 1) / 2;
   }
+  initial_graph_size = std::max(initial_graph_size, small_graph_degree + 1);
   RAFT_LOG_DEBUG("# initial graph size = %lu", (uint64_t)initial_graph_size);
 
   // Preallocate the kNN graph at last-iteration size (N × (intermediate_degree+1)) from the large
@@ -793,7 +799,6 @@ auto iterative_build_graph(raft::resources const& res,
 
   // Determine graph degree and number of search results while increasing
   // graph size.
-  auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24));
   RAFT_LOG_DEBUG("# small_graph_degree = %lu", (uint64_t)small_graph_degree);
   RAFT_LOG_DEBUG("# graph_degree = %lu", (uint64_t)graph_degree);
   RAFT_LOG_DEBUG("# topk = %lu", (uint64_t)topk);
@@ -1062,6 +1067,61 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res,
   return cagra_graph;
 }
 
+/** Build only the final optimized host graph from a dense dataset view. */
+template <typename T, typename IdxT = uint32_t, typename DatasetViewT>
+  requires cuvs::neighbors::is_dense_row_major_dataset_view_v<DatasetViewT>
+auto build_graph(raft::resources const& res,
+                 index_params const& params,
+                 DatasetViewT const& dataset) -> raft::host_matrix<IdxT, int64_t>
+{
+  auto const n_rows        = static_cast<size_t>(dataset.n_rows());
+  auto const dim           = static_cast<size_t>(dataset.dim());
+  auto intermediate_degree = params.intermediate_graph_degree;
+  auto graph_degree        = params.graph_degree;
+  common::nvtx::range<common::nvtx::domain::cuvs> function_scope(
+    "cagra::detail::build_graph(%zu, %zu)", intermediate_degree, graph_degree);
+  check_graph_degree<T, IdxT>(intermediate_degree, graph_degree, n_rows);
+
+  auto const dataset_extents =
+    raft::matrix_extent<int64_t>(static_cast<int64_t>(n_rows), static_cast<int64_t>(dim));
+  auto knn_build_params = resolve_cagra_default_knn_graph_build_params<IdxT>(
+    res, params, dataset_extents, intermediate_degree);
+  validate_cagra_knn_graph_build_constraints<T>(params, knn_build_params);
+
+  if (std::holds_alternative<cagra::graph_build_params::iterative_search_params>(
+        knn_build_params)) {
+    auto adjusted_params                      = params;
+    adjusted_params.intermediate_graph_degree = intermediate_degree;
+    adjusted_params.graph_degree              = graph_degree;
+    auto device_graph                         = [&]() {
+      if constexpr (cuvs::neighbors::is_host_dataset_view_v<DatasetViewT>) {
+        auto padded = cuvs::neighbors::make_device_padded_dataset(res, dataset.data_view());
+        auto graph  = iterative_build_graph<T, IdxT>(res, adjusted_params, padded->view());
+        raft::resource::sync_stream(res);
+        return graph;
+      } else {
+        return iterative_build_graph<T, IdxT>(res, adjusted_params, dataset);
+      }
+    }();
+    auto host_graph =
+      raft::make_host_matrix<IdxT, int64_t>(device_graph.extent(0), device_graph.extent(1));
+    raft::copy(host_graph.data_handle(),
+               device_graph.data_handle(),
+               host_graph.size(),
+               raft::resource::get_cuda_stream(res));
+    raft::resource::sync_stream(res);
+    return host_graph;
+  }
+  return build_cagra_host_graph_from_knn_params<T, IdxT>(
+    res,
+    params,
+    knn_build_params,
+    static_cast<int64_t>(n_rows),
+    intermediate_degree,
+    graph_degree,
+    cuvs::neighbors::dense_row_major_view(dataset));
+}
+
 /**
  * Build from a host row-major matrix without uploading the full dataset early when IVF-PQ graph
  * construction can consume host batches directly. Iterative CAGRA search needs the rows on device