JIT LTO Cagra Search (#1807)

CUDA 13 binary size reduction from 282 MB to 257 MB (-8.86%).

Benchmark:
<img width="1137" height="776" alt="image" src="https://github.com/user-attachments/assets/81b4f25a-c999-423d-a2c9-3e45566904bf" />

Apply updates from 
CAGRA related PRs:

- [x] #1800
- [ ] #1781
- [x] #1780
- [x] #1771
- [x] #1834
- [x] #1851
- [x] #1831

JIT related PRs:
- [x] #1812

Authors:
  - Divye Gala (https://github.com/divyegala)
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #1807

Author: divyegala

conda/recipes/libcuvs/recipe.yaml

@@ -399,13 +399,13 @@ outputs:
         - librmm =${{ minor_version }}
         - nccl ${{ nccl_version }}
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
         - cuda-profiler-api
         - libcublas-dev
         - libcurand-dev
         - libcusolver-dev
         - libcusparse-dev
         - libnvjitlink-dev
-        - cuda-nvrtc-dev
       run:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}
         - ${{ pin_subpackage("libcuvs", exact=True) }}

cpp/CMakeLists.txt

@@ -27,7 +27,8 @@
         "kwargs": {
           "KERNEL_FILE": 1,
           "FATBIN_HEADER_FILE": 1,
-          "LINK_LIBRARIES": "*"
+          "LINK_LIBRARIES": "*",
+          "EXTRA_COMPILE_OPTIONS": "*"
         }
       },
       "process_jit_lto_matrix_entry": {
@@ -41,7 +42,8 @@
           "FRAGMENT_TAG_HEADER_FILES": "*",
           "OUTPUT_DIRECTORY": 1,
           "MATRIX_JSON_ENTRY": 1,
-          "KERNEL_LINK_LIBRARIES": "*"
+          "KERNEL_LINK_LIBRARIES": "*",
+          "KERNEL_EXTRA_COMPILE_OPTIONS": "*"
         }
       },
       "generate_jit_lto_kernels": {
@@ -56,7 +58,8 @@
           "FRAGMENT_TAG_FORMAT": 1,
           "FRAGMENT_TAG_HEADER_FILES": "*",
           "OUTPUT_DIRECTORY": 1,
-          "KERNEL_LINK_LIBRARIES": "*"
+          "KERNEL_LINK_LIBRARIES": "*",
+          "KERNEL_EXTRA_COMPILE_OPTIONS": "*"
         }
       },
       "generate_inst_matrix": {

cpp/cmake/config.json

@@ -12,7 +12,7 @@ include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake)
 function(add_jit_lto_kernel kernel_target)
   set(options)
   set(one_value KERNEL_FILE FATBIN_HEADER_FILE)
-  set(multi_value LINK_LIBRARIES)
+  set(multi_value LINK_LIBRARIES EXTRA_COMPILE_OPTIONS)
 
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
@@ -22,6 +22,9 @@ function(add_jit_lto_kernel kernel_target)
   # through the LINK_LIBRARIES argument.
   target_link_libraries(${kernel_target} PRIVATE ${_JIT_LTO_LINK_LIBRARIES})
   target_compile_options(${kernel_target} PRIVATE -Xfatbin=--compress-all --compress-mode=size)
+  if(_JIT_LTO_EXTRA_COMPILE_OPTIONS)
+    target_compile_options(${kernel_target} PRIVATE ${_JIT_LTO_EXTRA_COMPILE_OPTIONS})
+  endif()
   set_target_properties(
     ${kernel_target}
     PROPERTIES CUDA_SEPARABLE_COMPILATION ON
@@ -43,7 +46,7 @@ function(process_jit_lto_matrix_entry source_list_var)
   set(one_value NAME_FORMAT KERNEL_INPUT_FILE OUTPUT_DIRECTORY FRAGMENT_TAG_FORMAT
                 MATRIX_JSON_ENTRY
   )
-  set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES)
+  set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES KERNEL_EXTRA_COMPILE_OPTIONS)
 
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
@@ -71,6 +74,7 @@ function(process_jit_lto_matrix_entry source_list_var)
     KERNEL_FILE "${kernel_file}"
     FATBIN_HEADER_FILE "${fatbin_header_file}"
     LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES}
+    EXTRA_COMPILE_OPTIONS ${_JIT_LTO_KERNEL_EXTRA_COMPILE_OPTIONS}
   )
   list(APPEND ${source_list_var} "${fatbin_header_file}" "${fatbin_file}")
   set(${source_list_var}
@@ -84,7 +88,7 @@ function(generate_jit_lto_kernels source_list_var)
   set(one_value NAME_FORMAT MATRIX_JSON_FILE MATRIX_JSON_STRING KERNEL_INPUT_FILE
                 FRAGMENT_TAG_FORMAT OUTPUT_DIRECTORY
   )
-  set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES)
+  set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES KERNEL_EXTRA_COMPILE_OPTIONS)
 
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
@@ -121,6 +125,7 @@ function(generate_jit_lto_kernels source_list_var)
       OUTPUT_DIRECTORY "${_JIT_LTO_OUTPUT_DIRECTORY}"
       MATRIX_JSON_ENTRY "${matrix_json_entry}"
       KERNEL_LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES}
+      KERNEL_EXTRA_COMPILE_OPTIONS ${_JIT_LTO_KERNEL_EXTRA_COMPILE_OPTIONS}
     )
   endforeach()
 

cpp/cmake/modules/generate_jit_lto_kernels.cmake

@@ -37,10 +37,24 @@ struct AlgorithmLauncher {
     this->call(stream, grid, block, shared_mem, kernel_args);
   }
 
+  template <typename FuncT, typename... Args>
+  void dispatch_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
+  {
+    static_assert(
+      std::is_same_v<FuncT, void(std::remove_reference_t<Args>...)>,
+      "dispatch_cooperative() argument types do not match the kernel function signature FuncT");
+
+    void* kernel_args[] = {const_cast<void*>(static_cast<void const*>(&args))...};
+    this->call_cooperative(stream, grid, block, shared_mem, kernel_args);
+  }
+
   cudaKernel_t get_kernel() { return this->kernel; }
 
  private:
   void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
+  void call_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
   cudaKernel_t kernel;
   cudaLibrary_t library;
 };

cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp

@@ -44,6 +44,12 @@ struct AlgorithmPlanner {
     add_fragment(std::make_unique<StaticFatbinFragmentEntry<FragmentTag>>());
   }
 
+ protected:
+  /** Extra link-time option strings passed to nvJitLink. Base build()
+   *  always passes "-lto" and "-arch=sm_XX" first; derived planners may append here in their
+   *  constructor body. */
+  std::vector<std::string> linktime_extra_options;
+
  private:
   std::string get_fragments_key() const;
   std::shared_ptr<AlgorithmLauncher> build();

cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp

@@ -0,0 +1,93 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace cuvs::neighbors::cagra::detail {
+
+struct tag_dist_f {};
+struct tag_metric_l2 {};
+struct tag_metric_inner_product {};
+struct tag_metric_cosine {};
+struct tag_metric_hamming {};
+struct tag_codebook_none {};
+struct tag_codebook_half {};
+struct tag_metric_l1 {};
+struct tag_norm_noop {};
+struct tag_norm_cosine {};
+
+/// Multi-kernel planners that do not link `sample_filter` into the JIT link (e.g.
+/// `random_pickup`). Real filters use `cuvs::neighbors::detail::tag_filter_*` on
+/// `CagraPlannerBase`.
+struct tag_cagra_jit_sample_filter_link_absent {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          typename CodebookTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          uint32_t PqBits,
+          uint32_t PqLen>
+struct fragment_tag_setup_workspace {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          typename CodebookTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          uint32_t PqBits,
+          uint32_t PqLen>
+struct fragment_tag_compute_distance {};
+
+template <typename QueryTag, typename DistanceTag, typename MetricTag>
+struct fragment_tag_dist_op {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          typename NormTag>
+struct fragment_tag_apply_normalization_standard {};
+
+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,
+          bool TopkByBitonicSort,
+          bool BitonicSortAndMergeMultiWarps>
+struct fragment_tag_search_single_cta_p {};
+
+template <typename DataTag, typename SourceIndexTag, typename IndexTag, typename DistanceTag>
+struct fragment_tag_search_multi_cta {};
+
+template <typename DataTag, typename IndexTag, typename DistanceTag>
+struct fragment_tag_random_pickup {};
+
+template <typename DataTag, typename IndexTag, typename DistanceTag, typename SourceIndexTag>
+struct fragment_tag_compute_distance_to_child_nodes {};
+
+template <typename IndexTag, typename DistanceTag, typename SourceIndexTag>
+struct fragment_tag_apply_filter_kernel {};
+
+template <typename BitsetTag, typename SourceIndexTag, typename FilterTag>
+struct fragment_tag_sample_filter {};
+
+}  // namespace cuvs::neighbors::cagra::detail

cpp/include/cuvs/detail/jit_lto/cagra/cagra_fragments.hpp

@@ -7,11 +7,16 @@
 
 namespace cuvs::neighbors::detail {
 
+struct tag_f {};
+struct tag_h {};
+struct tag_i8 {};
+struct tag_u8 {};
 struct tag_filter_none {};
 struct tag_filter_bitset {};
 
 struct tag_bitset_u32 {};
 
+struct tag_index_u32 {};
 struct tag_index_i64 {};
 
 template <typename BitsetTag, typename IndexTag, typename FilterTag>

cpp/include/cuvs/detail/jit_lto/common_fragments.hpp

@@ -7,19 +7,13 @@
 
 namespace cuvs::neighbors::ivf_flat::detail {
 
-// Tag types for data types
-struct tag_f {};
-struct tag_h {};
-struct tag_i8 {};
-struct tag_u8 {};
-
 // Tag types for accumulator types
 struct tag_acc_f {};
 struct tag_acc_h {};
 struct tag_acc_i32 {};
 struct tag_acc_u32 {};
 
-// Tag types for distance metrics with full template info
+// Tag types for distance metrics
 struct tag_metric_euclidean {};
 struct tag_metric_inner_product {};
 struct tag_metric_custom_udf {};

cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp

@@ -24,7 +24,6 @@ AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept
 AlgorithmLauncher& AlgorithmLauncher::operator=(AlgorithmLauncher&& other) noexcept
 {
   if (this != &other) {
-    // Unload current library if it exists
     if (library != nullptr) { cudaLibraryUnload(library); }
     kernel        = other.kernel;
     library       = other.library;
@@ -47,3 +46,21 @@ void AlgorithmLauncher::call(
 
   RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
 }
+
+void AlgorithmLauncher::call_cooperative(
+  cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args)
+{
+  cudaLaunchAttribute attribute[1];
+  attribute[0].id              = cudaLaunchAttributeCooperative;
+  attribute[0].val.cooperative = 1;
+
+  cudaLaunchConfig_t config{};
+  config.gridDim          = grid;
+  config.blockDim         = block;
+  config.stream           = stream;
+  config.dynamicSmemBytes = shared_mem;
+  config.numAttrs         = 1;
+  config.attrs            = attribute;
+
+  RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
+}