[FlowANN step 2/5] Add experimental tiered CAGRA storage with FlowANN

cpp/CMakeLists.txt

@@ -84,7 +84,16 @@ 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)
 
+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 +122,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 +135,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 +180,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 +260,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 +990,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
@@ -1459,6 +1528,10 @@ 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_SEARCH}>:src/neighbors/cagra_tiered.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
@@ -1571,10 +1644,12 @@ if(NOT BUILD_CPU_ONLY)
 
   target_compile_definitions(
     cuvs_objs
-    PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}
-            $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:CUVS_BUILD_CAGRA_HNSWLIB>
-            $<$<BOOL:${BUILD_MG_ALGOS}>:CUVS_BUILD_MG_ALGOS>
-            $<$<BOOL:${CUVS_NVTX}>:NVTX_ENABLED>
+    PRIVATE
+      $<$<AND:$<BOOL:${CUVS_ENABLE_FLOWANN_SEARCH}>,$<BOOL:${CUVS_FLOWANN_USE_GDRCOPY}>>:CUVS_FLOWANN_USE_GDRCOPY>
+      CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}
+      $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:CUVS_BUILD_CAGRA_HNSWLIB>
+      $<$<BOOL:${BUILD_MG_ALGOS}>:CUVS_BUILD_MG_ALGOS>
+      $<$<BOOL:${CUVS_NVTX}>:NVTX_ENABLED>
   )
 
   target_link_libraries(
@@ -1584,6 +1659,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
@@ -1663,10 +1739,26 @@ 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)
+      # 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()
 
     # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
     file(
@@ -1724,6 +1816,7 @@ SECTIONS
               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