Introduce UDF Architecture

ci/build_wheel.sh

@@ -42,6 +42,7 @@ EXCLUDE_ARGS=(
   --exclude "libcusparse.so.*"
   --exclude "libnccl.so.*"
   --exclude "libnvJitLink.so.*"
+  --exclude "libnvrtc.so.*"
   --exclude "libraft.so"
   --exclude "librapids_logger.so"
   --exclude "librmm.so"

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - clang==20.1.8
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.9.2,<13.0

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - clang==20.1.8
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.9.2,<13.0

conda/environments/all_cuda-131_arch-aarch64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - clang==20.1.8
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=13.0.1,<14.0

conda/environments/all_cuda-131_arch-x86_64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - clang==20.1.8
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=13.0.1,<14.0

conda/environments/bench_ann_cuda-129_arch-aarch64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - click
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.9.2,<13.0

conda/environments/bench_ann_cuda-129_arch-x86_64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - click
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=12.9.2,<13.0

conda/environments/bench_ann_cuda-131_arch-aarch64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - click
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=13.0.1,<14.0

conda/environments/bench_ann_cuda-131_arch-x86_64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - click
 - cmake>=3.30.4
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-python>=13.0.1,<14.0

conda/environments/go_cuda-129_arch-aarch64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.9

conda/environments/go_cuda-129_arch-x86_64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.9

conda/environments/go_cuda-131_arch-aarch64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=13.1

conda/environments/go_cuda-131_arch-x86_64.yaml

@@ -12,6 +12,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=13.1

conda/environments/rust_cuda-129_arch-aarch64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.9

conda/environments/rust_cuda-129_arch-x86_64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.9

conda/environments/rust_cuda-131_arch-aarch64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=13.1

conda/environments/rust_cuda-131_arch-x86_64.yaml

@@ -11,6 +11,7 @@ dependencies:
 - cmake>=3.30.4
 - cuda-cudart-dev
 - cuda-nvcc
+- cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=13.1

conda/recipes/libcuvs/recipe.yaml

@@ -78,6 +78,7 @@ cache:
       - nccl ${{ nccl_version }}
       - cuda-version =${{ cuda_version }}
       - cuda-cudart-dev
+      - cuda-nvrtc-dev
       - cuda-profiler-api
       - libcublas-dev
       - libcurand-dev
@@ -114,6 +115,7 @@ outputs:
         - nccl ${{ nccl_version }}
         - cuda-version =${{ cuda_version }}
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
         - cuda-profiler-api
         - libcublas-dev
         - libcurand-dev
@@ -124,6 +126,7 @@ outputs:
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
         - libraft-headers =${{ minor_version }}
         - librmm =${{ minor_version }}
+        - cuda-nvrtc
         - nccl
         - libcublas
         - libcurand
@@ -132,6 +135,7 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libaio
           - libboost
@@ -173,6 +177,7 @@ outputs:
         - nccl ${{ nccl_version }}
         - cuda-version =${{ cuda_version }}
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
         - cuda-profiler-api
         - libcublas-dev
         - libcurand-dev
@@ -184,6 +189,7 @@ outputs:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}
         - libraft-headers =${{ minor_version }}
         - librmm =${{ minor_version }}
+        - cuda-nvrtc
         - nccl
         - libcublas
         - libcurand
@@ -192,6 +198,7 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libaio
           - libboost
@@ -231,6 +238,7 @@ outputs:
         - nccl ${{ nccl_version }}
         - cuda-version =${{ cuda_version }}
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
         - cuda-profiler-api
         - libcublas-dev
         - libcurand-dev
@@ -242,6 +250,7 @@ outputs:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}
         - libraft-headers =${{ minor_version }}
         - librmm =${{ minor_version }}
+        - cuda-nvrtc
         - nccl
         - libcublas
         - libcurand
@@ -250,6 +259,7 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libaio
           - libboost
@@ -306,6 +316,7 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libaio
           - libboost
@@ -394,10 +405,12 @@ outputs:
         - libcusolver-dev
         - libcusparse-dev
         - libnvjitlink-dev
+        - cuda-nvrtc-dev
       run:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}
         - ${{ pin_subpackage("libcuvs", exact=True) }}
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
+        - cuda-nvrtc
         - nccl
         - libcublas
         - libcurand
@@ -406,6 +419,7 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libaio
           - libboost
@@ -478,14 +492,14 @@ outputs:
       ignore_run_exports:
         by_name:
           - cuda-cudart
+          - cuda-nvrtc
           - cuda-version
           - libcublas
           - libcurand
           - libcusolver
           - libcusparse
           - librmm
           - mkl
-          - nccl
     about:
       homepage: ${{ load_from_file("python/cuvs_bench/pyproject.toml").project.urls.Homepage }}
       license: ${{ load_from_file("python/cuvs_bench/pyproject.toml").project.license }}

cpp/CMakeLists.txt

@@ -847,6 +847,7 @@ if(NOT BUILD_CPU_ONLY)
     src/detail/jit_lto/AlgorithmPlanner.cpp
     src/detail/jit_lto/FragmentEntry.cpp
     src/detail/jit_lto/nvjitlink_checker.cpp
+    src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp
     src/distance/detail/kernels/gram_matrix.cu
     src/distance/detail/kernels/kernel_factory.cu
     src/distance/detail/kernels/kernel_matrices.cu
@@ -1033,8 +1034,12 @@ 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 rmm::rmm $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
-              $<COMPILE_ONLY:nvidia::cutlass::cutlass> $<COMPILE_ONLY:cuco::cuco> CUDA::nvJitLink
+      PRIVATE rmm::rmm
+              $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
+              $<COMPILE_ONLY:nvidia::cutlass::cutlass>
+              $<COMPILE_ONLY:cuco::cuco>
+              CUDA::nvJitLink
+              CUDA::nvrtc
     )
     set_property(TARGET cuvs PROPERTY NO_CUDART_DEP ON)
 
@@ -1093,6 +1098,7 @@ SECTIONS
       PRIVATE rmm::rmm
               $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
               CUDA::nvJitLink
+              CUDA::nvrtc
               $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
               $<COMPILE_ONLY:nvidia::cutlass::cutlass>
               $<COMPILE_ONLY:cuco::cuco>

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

@@ -44,5 +44,3 @@ struct AlgorithmLauncher {
   cudaKernel_t kernel;
   cudaLibrary_t library;
 };
-
-std::unordered_map<std::string, std::shared_ptr<AlgorithmLauncher>>& get_cached_launchers();

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

@@ -6,16 +6,26 @@
 #pragma once
 
 #include <memory>
+#include <shared_mutex>
 #include <string>
 #include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
 #include "AlgorithmLauncher.hpp"
 #include "FragmentEntry.hpp"
 
+struct LauncherJitCache {
+  std::shared_mutex mutex;
+  std::unordered_map<std::string, std::shared_ptr<AlgorithmLauncher>> launchers;
+};
+
 struct AlgorithmPlanner {
-  AlgorithmPlanner(std::string entrypoint) : entrypoint(std::move(entrypoint)) {}
+  AlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache)
+    : entrypoint(std::move(entrypoint)), jit_cache_(jit_cache)
+  {
+  }
 
   std::shared_ptr<AlgorithmLauncher> get_launcher();
 
@@ -37,4 +47,8 @@ struct AlgorithmPlanner {
  private:
   std::string get_fragments_key() const;
   std::shared_ptr<AlgorithmLauncher> build();
+
+  std::shared_ptr<AlgorithmLauncher> read_cache(std::string const& launch_key) const;
+
+  LauncherJitCache& jit_cache_;
 };

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

@@ -45,3 +45,20 @@ struct StaticFatbinFragmentEntry final : FatbinFragmentEntry {
   static const uint8_t* const data;
   static const size_t length;
 };
+
+struct UDFFatbinFragment final : FatbinFragmentEntry {
+  UDFFatbinFragment(std::string key, std::vector<uint8_t> bytes)
+    : key_(std::move(key)), bytes_(std::move(bytes))
+  {
+  }
+
+  const uint8_t* get_data() const override { return bytes_.data(); }
+
+  size_t get_length() const override { return bytes_.size(); }
+
+  const char* get_key() const override { return key_.c_str(); }
+
+ private:
+  std::string key_;
+  std::vector<uint8_t> bytes_;
+};

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

@@ -0,0 +1,29 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuvs/detail/jit_lto/FragmentEntry.hpp>
+
+#include <memory>
+#include <shared_mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+struct NVRTCLTOFragmentCompiler {
+  NVRTCLTOFragmentCompiler();
+
+  std::vector<std::string> standard_compile_opts;
+  std::unordered_map<std::string, std::vector<uint8_t>> cache;
+  mutable std::shared_mutex cache_mutex_;
+
+  std::unique_ptr<UDFFatbinFragment> compile(std::string const& key, std::string const& code);
+
+ private:
+  std::unique_ptr<UDFFatbinFragment> read_cache(std::string const& key) const;
+};
+
+NVRTCLTOFragmentCompiler& nvrtc_compiler();

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

@@ -22,6 +22,7 @@ struct tag_acc_u32 {};
 // Tag types for distance metrics with full template info
 struct tag_metric_euclidean {};
 struct tag_metric_inner_product {};
+struct tag_metric_custom_udf {};
 
 // Tag types for post-processing
 struct tag_post_process_identity {};