From 51032ae970e7f9a3ef69bba9c38632dfcb466b80 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Thu, 18 Dec 2025 20:38:25 +0000
Subject: [PATCH 01/20] Refactor JIT LTO kernel generation

Add an algorithm that computes a matrix product, and add a generic
CMake function that uses this algorithm to generate a matrix of kernels
with desired parameters
---
 cpp/CMakeLists.txt                            |  79 ++++-
 cpp/cmake/config.json                         |  49 ++-
 cpp/cmake/modules/compute_matrix_product.py   |  73 +++++
 .../modules/generate_jit_lto_kernels.cmake    | 299 +++++++-----------
 .../jit_lto_kernels/filter_embedded.cpp.in    |  21 ++
 .../{filter.cu.in => filter_kernel.cu.in}     |  19 +-
 .../jit_lto_kernels/filter_matrix.json        |  12 +
 .../interleaved_scan_embedded.cpp.in          |  26 ++
 .../interleaved_scan_kernel.cu.in             |  26 +-
 .../interleaved_scan_matrix.json              |  61 ++++
 .../interleaved_scan_planner.hpp              |  10 +-
 .../ivf_flat/jit_lto_kernels/metric.cu.in     |  36 ---
 .../jit_lto_kernels/metric_embedded.cpp.in    |  25 ++
 .../jit_lto_kernels/metric_kernel.cu.in       |  15 +
 .../jit_lto_kernels/metric_matrix.json        |  48 +++
 .../jit_lto_kernels/post_lambda.cu.in         |  32 --
 .../post_lambda_embedded.cpp.in               |  21 ++
 .../jit_lto_kernels/post_lambda_kernel.cu.in  |  15 +
 .../jit_lto_kernels/post_lambda_matrix.json   |  16 +
 19 files changed, 582 insertions(+), 301 deletions(-)
 create mode 100644 cpp/cmake/modules/compute_matrix_product.py
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
 rename cpp/src/neighbors/ivf_flat/jit_lto_kernels/{filter.cu.in => filter_kernel.cu.in} (52%)
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_matrix.json
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_matrix.json
 delete mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric.cu.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
 delete mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda.cu.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
 create mode 100644 cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_matrix.json

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index a75890737e..b4ede2b42c 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -354,7 +354,84 @@ if(NOT BUILD_CPU_ONLY)
   if(JIT_LTO_COMPILATION)
     # Generate interleaved scan kernel files at build time
     include(cmake/modules/generate_jit_lto_kernels.cmake)
-    generate_jit_lto_kernels(cuvs_jit_lto_kernels)
+
+    add_library(jit_lto_kernel_usage_requirements INTERFACE)
+    target_include_directories(
+      jit_lto_kernel_usage_requirements
+      INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"
+                "${CMAKE_CURRENT_SOURCE_DIR}/../c/include"
+    )
+    target_compile_options(
+      jit_lto_kernel_usage_requirements INTERFACE "$<$<COMPILE_LANGUAGE:CXX>:${CUVS_CXX_FLAGS}>"
+                                                  "$<$<COMPILE_LANGUAGE:CUDA>:${CUVS_CUDA_FLAGS}>"
+    )
+    target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_17)
+    target_link_libraries(
+      jit_lto_kernel_usage_requirements INTERFACE rmm::rmm raft::raft CCCL::CCCL
+    )
+
+    set(cuda_architectures_old ${CMAKE_CUDA_ARCHITECTURES})
+    set(CMAKE_CUDA_ARCHITECTURES ${JIT_LTO_TARGET_ARCHITECTURE})
+    generate_jit_lto_kernels(
+      interleaved_scan_files
+      NAME_FORMAT
+        "interleaved_scan_capacity_@capacity@_veclen_@veclen@_@ascending_descending@_@compute_norm_name@_data_@type_abbrev@_acc_@acc_abbrev@_idx_@idx_abbrev@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in"
+      EMBEDDED_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/interleaved_scan"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+    )
+    generate_jit_lto_kernels(
+      metric_files
+      NAME_FORMAT "metric_@metric_name@_veclen_@veclen@_data_@type_abbrev@_acc_@acc_abbrev@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in"
+      EMBEDDED_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/metric"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+    )
+    generate_jit_lto_kernels(
+      filter_files
+      NAME_FORMAT "@filter_name@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/filter_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in"
+      EMBEDDED_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/filter"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+    )
+    generate_jit_lto_kernels(
+      post_lambda_files
+      NAME_FORMAT "@post_lambda_name@"
+      MATRIX_JSON_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_matrix.json"
+      KERNEL_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in"
+      EMBEDDED_INPUT_FILE
+        "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in"
+      OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/post_lambda"
+      KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+    )
+    set(CMAKE_CUDA_ARCHITECTURES ${cuda_architectures_old})
+
+    add_library(
+      cuvs_jit_lto_kernels STATIC ${interleaved_scan_files} ${metric_files} ${filter_files}
+                                  ${post_lambda_files}
+    )
+    target_include_directories(
+      cuvs_jit_lto_kernels
+      PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"
+              "${CMAKE_CURRENT_SOURCE_DIR}/../c/include"
+    )
     add_library(cuvs::cuvs_jit_lto_kernels ALIAS cuvs_jit_lto_kernels)
 
     set(JIT_LTO_FILES
diff --git a/cpp/cmake/config.json b/cpp/cmake/config.json
index aa46006a44..a3ab8cb48f 100644
--- a/cpp/cmake/config.json
+++ b/cpp/cmake/config.json
@@ -11,13 +11,50 @@
           "PATH": "*"
         }
       },
-      "embed_jit_lto_fatbin": {
+      "compute_matrix_product": {
+        "pargs": {
+          "nargs": 1
+        },
         "kwargs": {
-          "FATBIN_TARGET": 1,
-          "FATBIN_SOURCE": 1,
-          "EMBEDDED_TARGET": 1,
-          "EMBEDDED_HEADER": 1,
-          "EMBEDDED_ARRAY": 1
+          "MATRIX_JSON_FILE": "?",
+          "MATRIX_JSON_STRING": "?"
+        }
+      },
+      "add_jit_lto_kernel": {
+        "pargs": {
+          "nargs": 1
+        },
+        "kwargs": {
+          "KERNEL_FILE": 1,
+          "EMBEDDED_HEADER_FILE": 1,
+          "LINK_LIBRARIES": "*"
+        }
+      },
+      "process_matrix_entry": {
+        "pargs": {
+          "nargs": 1
+        },
+        "kwargs": {
+          "NAME_FORMAT": 1,
+          "KERNEL_INPUT_FILE": 1,
+          "EMBEDDED_INPUT_FILE": 1,
+          "OUTPUT_DIRECTORY": 1,
+          "MATRIX_JSON_ENTRY": 1,
+          "KERNEL_LINK_LIBRARIES": "*"
+        }
+      },
+      "generate_jit_lto_kernels": {
+        "pargs": {
+          "nargs": 1
+        },
+        "kwargs": {
+          "NAME_FORMAT": 1,
+          "MATRIX_JSON_FILE": "?",
+          "MATRIX_JSON_STRING": "?",
+          "KERNEL_INPUT_FILE": 1,
+          "EMBEDDED_INPUT_FILE": 1,
+          "OUTPUT_DIRECTORY": 1,
+          "KERNEL_LINK_LIBRARIES": "*"
         }
       }
     }
diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
new file mode 100644
index 0000000000..296a03f68f
--- /dev/null
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# This algorithm takes a JSON dictionary and computes a matrix product of all
+# of its arrays. We use this to compute all matrix combinations for kernel
+# generation. We *could* write this in CMake with `string(JSON)`, but writing
+# it in Python is much easier. Once we have a version of CMake that has
+# https://gitlab.kitware.com/cmake/cmake/-/merge_requests/11516, we may be able
+# to port the algorithm to CMake script and use it in other RAPIDS projects.
+
+import json
+import sys
+from itertools import chain
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterator
+
+    MatrixValue = None | bool | int | float | str
+    Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"]
+
+
+def iterate_matrix_product(
+    matrix: "Matrix",
+) -> "Generator[dict[str, MatrixValue]]":
+    def iterate_next_dimension(
+        queue: "Iterator[tuple[str | None, Matrix]]",
+        entry: "dict[str, MatrixValue]",
+    ) -> "Generator[dict[str, MatrixValue]]":
+        try:
+            key, matrix = next(queue)
+        except StopIteration:
+            yield entry
+        else:
+            yield from iterate_impl(matrix, key, queue, entry)
+
+    def iterate_impl(
+        matrix: "Matrix",
+        key: str | None,
+        queue: "Iterator[tuple[str | None, Matrix]]",
+        entry: "dict[str, MatrixValue]",
+    ) -> "Generator[dict[str, MatrixValue]]":
+        if isinstance(matrix, dict):
+            yield from iterate_next_dimension(
+                chain(sorted(matrix.items()), queue),
+                entry,
+            )
+        elif isinstance(matrix, list):
+            queue_list = list(queue)
+            for v in matrix:
+                yield from iterate_next_dimension(
+                    chain([(key, v)], queue_list), entry
+                )
+        else:
+            assert key is not None
+            yield from iterate_next_dimension(queue, {**entry, key: matrix})
+
+    return iterate_impl(matrix, None, chain(), {})
+
+
+try:
+    filename = sys.argv[1]
+except IndexError:
+    filename = "-"
+with sys.stdin if filename == "-" else open(filename) as f:
+    matrix: "Matrix" = json.load(f)
+
+json.dump(
+    list(iterate_matrix_product(matrix)),
+    sys.stdout,
+    indent=2,
+    sort_keys=True,
+)
diff --git a/cpp/cmake/modules/generate_jit_lto_kernels.cmake b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
index 671b08321c..5d30eda6b2 100644
--- a/cpp/cmake/modules/generate_jit_lto_kernels.cmake
+++ b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
@@ -7,214 +7,157 @@
 
 include_guard(GLOBAL)
 
-function(embed_jit_lto_fatbin)
+function(compute_matrix_product output_var)
   set(options)
-  set(one_value FATBIN_TARGET FATBIN_SOURCE EMBEDDED_TARGET EMBEDDED_HEADER EMBEDDED_ARRAY)
+  set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING)
   set(multi_value)
 
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
-  find_package(CUDAToolkit REQUIRED)
-  find_program(
-    bin_to_c
-    NAMES bin2c
-    PATHS ${CUDAToolkit_BIN_DIR}
-  )
+  if(_JIT_LTO_MATRIX_JSON_FILE)
+    set(commands #
+        COMMAND
+        "${Python3_EXECUTABLE}"
+        "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
+        "${_JIT_LTO_MATRIX_JSON_FILE}" #
+    )
+  else()
+    set(commands #
+        COMMAND
+        "${CMAKE_COMMAND}"
+        -E
+        echo
+        "${_JIT_LTO_MATRIX_JSON_STRING}" #
+        COMMAND
+        "${Python3_EXECUTABLE}"
+        "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
+        - #
+    )
+  endif()
 
-  add_library(${_JIT_LTO_FATBIN_TARGET} OBJECT "${_JIT_LTO_FATBIN_SOURCE}")
-  target_compile_definitions(${_JIT_LTO_FATBIN_TARGET} PRIVATE BUILD_KERNEL)
-  target_include_directories(
-    ${_JIT_LTO_FATBIN_TARGET}
-    PRIVATE "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
-            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/src>"
-            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../c/include>"
-  )
-  target_compile_options(
-    ${_JIT_LTO_FATBIN_TARGET}
-    PRIVATE -Xfatbin=--compress-all
-            --compress-mode=size
-            "$<$<COMPILE_LANGUAGE:CXX>:${CUVS_CXX_FLAGS}>"
-            "$<$<COMPILE_LANGUAGE:CUDA>:${CUVS_CUDA_FLAGS}>"
+  execute_process(${commands} OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY)
+  set(${output_var}
+      "${output}"
+      PARENT_SCOPE
   )
+endfunction()
+
+function(populate_matrix_variables matrix_json_entry)
+  string(JSON len LENGTH "${matrix_json_entry}")
+  math(EXPR last "${len} - 1")
+
+  # cmake-lint: disable=C0103,E1120
+  foreach(i RANGE "${last}")
+    string(JSON key MEMBER "${matrix_json_entry}" "${i}")
+    string(JSON value GET "${matrix_json_entry}" "${key}")
+    set(${key}
+        "${value}"
+        PARENT_SCOPE
+    )
+  endforeach()
+endfunction()
+
+function(add_jit_lto_kernel kernel_target)
+  set(options)
+  set(one_value KERNEL_FILE EMBEDDED_HEADER_FILE)
+  set(multi_value LINK_LIBRARIES)
+
+  cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
+
+  add_library(${kernel_target} OBJECT EXCLUDE_FROM_ALL "${_JIT_LTO_KERNEL_FILE}")
+  target_link_libraries(${kernel_target} PRIVATE ${_JIT_LTO_LINK_LIBRARIES})
+  target_compile_options(${kernel_target} PRIVATE -Xfatbin=--compress-all --compress-mode=size)
   set_target_properties(
-    ${_JIT_LTO_FATBIN_TARGET}
-    PROPERTIES CUDA_ARCHITECTURES ${JIT_LTO_TARGET_ARCHITECTURE}
-               CUDA_STANDARD 17
-               CUDA_STANDARD_REQUIRED ON
-               CUDA_SEPARABLE_COMPILATION ON
+    ${kernel_target}
+    PROPERTIES CUDA_SEPARABLE_COMPILATION ON
                CUDA_FATBIN_COMPILATION ON
                POSITION_INDEPENDENT_CODE ON
                INTERPROCEDURAL_OPTIMIZATION ON
   )
-  target_link_libraries(${_JIT_LTO_FATBIN_TARGET} PRIVATE rmm::rmm raft::raft CCCL::CCCL)
 
   add_custom_command(
-    OUTPUT "${_JIT_LTO_EMBEDDED_HEADER}"
-    COMMAND "${bin_to_c}" -c -p 0x0 --name "${_JIT_LTO_EMBEDDED_ARRAY}" --static
-            $<TARGET_OBJECTS:${_JIT_LTO_FATBIN_TARGET}> > "${_JIT_LTO_EMBEDDED_HEADER}"
-    DEPENDS $<TARGET_OBJECTS:${_JIT_LTO_FATBIN_TARGET}>
+    OUTPUT "${_JIT_LTO_EMBEDDED_HEADER_FILE}"
+    COMMAND "${bin_to_c}" --const --name embedded_fatbin --static $<TARGET_OBJECTS:${kernel_target}>
+            > "${_JIT_LTO_EMBEDDED_HEADER_FILE}"
+    DEPENDS $<TARGET_OBJECTS:${kernel_target}>
   )
-  target_sources(
-    ${_JIT_LTO_EMBEDDED_TARGET} PRIVATE "${_JIT_LTO_FATBIN_SOURCE}" "${_JIT_LTO_EMBEDDED_HEADER}"
-  )
-  cmake_path(GET _JIT_LTO_EMBEDDED_HEADER PARENT_PATH header_dir)
-  target_include_directories(${_JIT_LTO_EMBEDDED_TARGET} PRIVATE "${header_dir}")
 endfunction()
 
-function(parse_jit_lto_data_type_configs config)
+function(process_matrix_entry source_list_var)
   set(options)
-  set(one_value DATA_TYPE ACC_TYPE VECLENS TYPE_ABBREV ACC_ABBREV)
-  set(multi_value)
+  set(one_value NAME_FORMAT KERNEL_INPUT_FILE EMBEDDED_INPUT_FILE OUTPUT_DIRECTORY
+                MATRIX_JSON_ENTRY
+  )
+  set(multi_value KERNEL_LINK_LIBRARIES)
 
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
-  if(config MATCHES [==[^([^,]+),([^,]+),\[([0-9]+(,[0-9]+)*)?\],([^,]+),([^,]+)$]==])
-    if(_JIT_LTO_DATA_TYPE)
-      set(${_JIT_LTO_DATA_TYPE}
-          "${CMAKE_MATCH_1}"
-          PARENT_SCOPE
-      )
-    endif()
-    if(_JIT_LTO_ACC_TYPE)
-      set(${_JIT_LTO_ACC_TYPE}
-          "${CMAKE_MATCH_2}"
-          PARENT_SCOPE
-      )
-    endif()
-    if(_JIT_LTO_VECLENS)
-      string(REPLACE "," ";" veclens_value "${CMAKE_MATCH_3}")
-      set(${_JIT_LTO_VECLENS}
-          "${veclens_value}"
-          PARENT_SCOPE
-      )
-    endif()
-    if(_JIT_LTO_TYPE_ABBREV)
-      set(${_JIT_LTO_TYPE_ABBREV}
-          "${CMAKE_MATCH_5}"
-          PARENT_SCOPE
-      )
-    endif()
-    if(_JIT_LTO_ACC_ABBREV)
-      set(${_JIT_LTO_ACC_ABBREV}
-          "${CMAKE_MATCH_6}"
-          PARENT_SCOPE
-      )
-    endif()
-  else()
-    message(FATAL_ERROR "Invalid data type config: ${config}")
-  endif()
+  populate_matrix_variables("${_JIT_LTO_MATRIX_JSON_ENTRY}")
+  string(CONFIGURE "${_JIT_LTO_NAME_FORMAT}" kernel_name @ONLY)
+
+  set(kernel_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_kernel.cu")
+  set(kernel_target "${kernel_name}_kernel")
+  set(embedded_header_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_embedded.h")
+  set(embedded_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_embedded.cpp")
+  configure_file("${_JIT_LTO_KERNEL_INPUT_FILE}" "${kernel_file}" @ONLY)
+  configure_file("${_JIT_LTO_EMBEDDED_INPUT_FILE}" "${embedded_file}" @ONLY)
+
+  add_jit_lto_kernel(
+    ${kernel_target}
+    KERNEL_FILE "${kernel_file}"
+    EMBEDDED_HEADER_FILE "${embedded_header_file}"
+    LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES}
+  )
+  list(APPEND ${source_list_var} "${embedded_header_file}" "${embedded_file}")
+  set(${source_list_var}
+      "${${source_list_var}}"
+      PARENT_SCOPE
+  )
 endfunction()
 
-# cmake-lint: disable=R0915
-function(generate_jit_lto_kernels target)
-  add_library(${target} STATIC)
-  target_include_directories(
-    ${target}
-    PRIVATE "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
-            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/src>"
-            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../c/include>"
+function(generate_jit_lto_kernels source_list_var)
+  set(options)
+  set(one_value NAME_FORMAT MATRIX_JSON_FILE MATRIX_JSON_STRING KERNEL_INPUT_FILE
+                EMBEDDED_INPUT_FILE OUTPUT_DIRECTORY
   )
-  set_target_properties(${target} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+  set(multi_value KERNEL_LINK_LIBRARIES)
 
-  set(generated_kernels_dir "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels")
-  string(TIMESTAMP year "%Y")
+  cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
-  set(capacities 0 1 2 4 8 16 32 64 128 256)
-  set(ascending_values true false)
-  set(compute_norm_values true false)
-  set(data_type_configs "float,float,[1,4],f,f" "__half,__half,[1,8],h,h"
-                        "uint8_t,uint32_t,[1,16],uc,ui" "int8_t,int32_t,[1,16],sc,i"
+  find_package(CUDAToolkit REQUIRED)
+  find_program(
+    bin_to_c
+    NAMES bin2c
+    PATHS ${CUDAToolkit_BIN_DIR}
   )
-  set(idx_type int64_t)
-  set(idx_abbrev l)
-  set(metric_configs euclidean inner_prod)
-  set(filter_configs filter_none filter_bitset)
-  set(post_lambda_configs post_identity post_sqrt post_compose)
-
-  foreach(config IN LISTS data_type_configs)
-    parse_jit_lto_data_type_configs(
-      "${config}" DATA_TYPE data_type ACC_TYPE acc_type VECLENS veclens TYPE_ABBREV type_abbrev
-      ACC_ABBREV acc_abbrev
-    )
-    foreach(veclen IN LISTS veclens)
-      foreach(capacity IN LISTS capacities)
-        foreach(ascending IN LISTS ascending_values)
-          foreach(compute_norm IN LISTS compute_norm_values)
-            set(kernel_name
-                "interleaved_scan_kernel_${capacity}_${veclen}_${ascending}_${compute_norm}_${type_abbrev}_${acc_abbrev}_${idx_abbrev}"
-            )
-            set(filename
-                "${generated_kernels_dir}/interleaved_scan_kernels/fatbin_${kernel_name}.cu"
-            )
-            configure_file(
-              "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in"
-              "${filename}"
-              @ONLY
-            )
-            embed_jit_lto_fatbin(
-              FATBIN_TARGET "fatbin_${kernel_name}"
-              FATBIN_SOURCE "${filename}"
-              EMBEDDED_TARGET "${target}"
-              EMBEDDED_HEADER "${generated_kernels_dir}/interleaved_scan_kernels/${kernel_name}.h"
-              EMBEDDED_ARRAY "embedded_${kernel_name}"
-            )
-          endforeach()
-        endforeach()
-      endforeach()
-
-      foreach(metric_name IN LISTS metric_configs)
-        set(header_file "neighbors/ivf_flat/jit_lto_kernels/metric_${metric_name}.cuh")
-
-        set(kernel_name "metric_${metric_name}_${veclen}_${type_abbrev}_${acc_abbrev}")
-        set(filename "${generated_kernels_dir}/metric_device_functions/fatbin_${kernel_name}.cu")
-        configure_file(
-          "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/metric.cu.in"
-          "${filename}" @ONLY
-        )
-        embed_jit_lto_fatbin(
-          FATBIN_TARGET "fatbin_${kernel_name}"
-          FATBIN_SOURCE "${filename}"
-          EMBEDDED_TARGET "${target}"
-          EMBEDDED_HEADER "${generated_kernels_dir}/metric_device_functions/${kernel_name}.h"
-          EMBEDDED_ARRAY "embedded_${kernel_name}"
-        )
-      endforeach()
-    endforeach()
-  endforeach()
+  find_package(Python3 REQUIRED COMPONENTS Interpreter)
 
-  foreach(filter_name IN LISTS filter_configs)
-    set(header_file "neighbors/ivf_flat/jit_lto_kernels/${filter_name}.cuh")
+  if(_JIT_LTO_MATRIX_JSON_FILE)
+    compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_JIT_LTO_MATRIX_JSON_FILE}")
+  else()
+    compute_matrix_product(matrix_product MATRIX_JSON_STRING "${_JIT_LTO_MATRIX_JSON_STRING}")
+  endif()
 
-    set(kernel_name "${filter_name}")
-    set(filename "${generated_kernels_dir}/filter_device_functions/fatbin_${kernel_name}.cu")
-    configure_file(
-      "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/filter.cu.in"
-      "${filename}" @ONLY
-    )
-    embed_jit_lto_fatbin(
-      FATBIN_TARGET "fatbin_${kernel_name}"
-      FATBIN_SOURCE "${filename}"
-      EMBEDDED_TARGET "${target}"
-      EMBEDDED_HEADER "${generated_kernels_dir}/filter_device_functions/${kernel_name}.h"
-      EMBEDDED_ARRAY "embedded_${kernel_name}"
+  string(JSON len LENGTH "${matrix_product}")
+  math(EXPR last "${len} - 1")
+
+  # cmake-lint: disable=C0103,E1120
+  foreach(i RANGE "${last}")
+    string(JSON matrix_json_entry GET "${matrix_product}" "${i}")
+    process_matrix_entry(
+      "${source_list_var}"
+      NAME_FORMAT "${_JIT_LTO_NAME_FORMAT}"
+      KERNEL_INPUT_FILE "${_JIT_LTO_KERNEL_INPUT_FILE}"
+      EMBEDDED_INPUT_FILE "${_JIT_LTO_EMBEDDED_INPUT_FILE}"
+      OUTPUT_DIRECTORY "${_JIT_LTO_OUTPUT_DIRECTORY}"
+      MATRIX_JSON_ENTRY "${matrix_json_entry}"
+      KERNEL_LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES}
     )
+    list(APPEND ${source_list_var} "${embedded_file}")
   endforeach()
 
-  foreach(post_lambda_name IN LISTS post_lambda_configs)
-    set(header_file "neighbors/ivf_flat/jit_lto_kernels/${post_lambda_name}.cuh")
-
-    set(kernel_name "${post_lambda_name}")
-    set(filename "${generated_kernels_dir}/post_lambda_device_functions/${post_lambda_name}.cu")
-    configure_file(
-      "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda.cu.in"
-      "${filename}" @ONLY
-    )
-    embed_jit_lto_fatbin(
-      FATBIN_TARGET "fatbin_${kernel_name}"
-      FATBIN_SOURCE "${filename}"
-      EMBEDDED_TARGET "${target}"
-      EMBEDDED_HEADER "${generated_kernels_dir}/post_lambda_device_functions/${kernel_name}.h"
-      EMBEDDED_ARRAY "embedded_${kernel_name}"
-    )
-  endforeach()
+  set(${source_list_var}
+      "${${source_list_var}}"
+      PARENT_SCOPE
+  )
 endfunction()
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
new file mode 100644
index 0000000000..9a1be3bf43
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
@@ -0,0 +1,21 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
+#include "@embedded_header_file@"
+
+namespace {
+
+__attribute__((__constructor__)) void register_kernel()
+{
+  registerAlgorithm(
+    "@filter_name@",
+    embedded_fatbin,
+    sizeof(embedded_fatbin));
+}
+
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
similarity index 52%
rename from cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter.cu.in
rename to cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
index 934e36dba7..c6f671f9f2 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
@@ -1,12 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) @year@, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 // This file is auto-generated. Do not edit manually.
 
-#ifdef BUILD_KERNEL
-
 #include <@header_file@>
 
 namespace cuvs::neighbors::ivf_flat::detail {
@@ -15,18 +13,3 @@ namespace cuvs::neighbors::ivf_flat::detail {
 template __device__ bool sample_filter<int64_t>(int64_t* const* const, const uint32_t, const uint32_t, const uint32_t, uint32_t*, int64_t, int64_t);
 
 }  // namespace cuvs::neighbors::ivf_flat::detail
-
-#else
-
-#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
-#include "@filter_name@.h"
-
-__attribute__((__constructor__)) static void register_@filter_name@()
-{
-  registerAlgorithm(
-    "@filter_name@",
-    embedded_@filter_name@,
-    sizeof(embedded_@filter_name@));
-}
-
-#endif
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_matrix.json b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_matrix.json
new file mode 100644
index 0000000000..6ceebe78c3
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_matrix.json
@@ -0,0 +1,12 @@
+{
+  "_filter": [
+    {
+      "filter_name": "filter_none",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/filter_none.cuh"
+    },
+    {
+      "filter_name": "filter_bitset",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/filter_bitset.cuh"
+    }
+  ]
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
new file mode 100644
index 0000000000..7620880925
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
@@ -0,0 +1,26 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
+#include <cuvs/detail/jit_lto/ivf_flat/interleaved_scan_tags.hpp>
+#include "@embedded_header_file@"
+
+using namespace cuvs::neighbors::ivf_flat::detail;
+
+namespace {
+
+__attribute__((__constructor__)) void register_kernel()
+{
+  registerAlgorithm<tag_@type_abbrev@,
+    tag_acc_@acc_abbrev@,
+    tag_idx_@idx_abbrev@>(
+    "interleaved_scan_kernel_capacity_@capacity@_veclen_@veclen@_@ascending_descending@_@compute_norm_name@",
+    embedded_fatbin,
+    sizeof(embedded_fatbin));
+}
+
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
index 5e75253939..ed39e7eb25 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
@@ -1,40 +1,18 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) @year@, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 // This file is auto-generated. Do not edit manually.
 
-#ifdef BUILD_KERNEL
-
 #include <neighbors/ivf_flat/jit_lto_kernels/ivf_flat_interleaved_scan_kernel.cuh>
 
 namespace cuvs::neighbors::ivf_flat::detail {
 
 // Instantiate the kernel template
-template __global__ void interleaved_scan_kernel<@capacity@, @veclen@, @ascending@, @compute_norm@, @data_type@, @acc_type@, @idx_type@>(
+template __global__ void interleaved_scan_kernel<@capacity@, @veclen@, @ascending_value@, @compute_norm_value@, @data_type@, @acc_type@, @idx_type@>(
     const uint32_t, const @data_type@*, const uint32_t*, const @data_type@* const*, const uint32_t*,
     const uint32_t, const uint32_t, const uint32_t, const uint32_t, const uint32_t*, const uint32_t,
     @idx_type@* const* const, uint32_t*, @idx_type@, @idx_type@, uint32_t*, float*);
 
 }  // namespace cuvs::neighbors::ivf_flat::detail
-
-#else
-
-#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
-#include <cuvs/detail/jit_lto/ivf_flat/interleaved_scan_tags.hpp>
-#include "interleaved_scan_kernel_@capacity@_@veclen@_@ascending@_@compute_norm@_@type_abbrev@_@acc_abbrev@_@idx_abbrev@.h"
-
-using namespace cuvs::neighbors::ivf_flat::detail;
-
-__attribute__((__constructor__)) static void register_kernel_@capacity@_@veclen@_@ascending@_@compute_norm@_@type_abbrev@_@acc_abbrev@_@idx_abbrev@()
-{
-  registerAlgorithm<tag_@type_abbrev@,
-    tag_acc_@acc_abbrev@,
-    tag_idx_@idx_abbrev@>(
-    "interleaved_scan_kernel_@capacity@_@veclen@_@ascending@_@compute_norm@",
-    embedded_interleaved_scan_kernel_@capacity@_@veclen@_@ascending@_@compute_norm@_@type_abbrev@_@acc_abbrev@_@idx_abbrev@,
-    sizeof(embedded_interleaved_scan_kernel_@capacity@_@veclen@_@ascending@_@compute_norm@_@type_abbrev@_@acc_abbrev@_@idx_abbrev@));
-}
-
-#endif
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_matrix.json b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_matrix.json
new file mode 100644
index 0000000000..41409ab7ad
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_matrix.json
@@ -0,0 +1,61 @@
+{
+  "capacity": [
+    "0", "1", "2", "4", "8", "16", "32", "64", "128", "256"
+  ],
+  "_ascending": [
+    {
+      "ascending_descending": "ascending",
+      "ascending_value": "true"
+    },
+    {
+      "ascending_descending": "descending",
+      "ascending_value": "false"
+    }
+  ],
+  "_compute_norm": [
+    {
+      "compute_norm_name": "compute_norm",
+      "compute_norm_value": "true"
+    },
+    {
+      "compute_norm_name": "no_compute_norm",
+      "compute_norm_value": "false"
+    }
+  ],
+  "_data_type": [
+    {
+      "data_type": "float",
+      "acc_type": "float",
+      "veclen": ["1", "4"],
+      "type_abbrev": "f",
+      "acc_abbrev": "f"
+    },
+    {
+      "data_type": "__half",
+      "acc_type": "__half",
+      "veclen": ["1", "8"],
+      "type_abbrev": "h",
+      "acc_abbrev": "h"
+    },
+    {
+      "data_type": "uint8_t",
+      "acc_type": "uint32_t",
+      "veclen": ["1", "16"],
+      "type_abbrev": "uc",
+      "acc_abbrev": "ui"
+    },
+    {
+      "data_type": "int8_t",
+      "acc_type": "int32_t",
+      "veclen": ["1", "16"],
+      "type_abbrev": "sc",
+      "acc_abbrev": "i"
+    }
+  ],
+  "_index": [
+    {
+      "idx_type": "int64_t",
+      "idx_abbrev": "l"
+    }
+  ]
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_planner.hpp b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_planner.hpp
index 792c64f39a..77592e8fc3 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_planner.hpp
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_planner.hpp
@@ -11,14 +11,12 @@
 #include <iostream>
 #include <string>
 
-inline std::string bool_to_string(bool b) { return b ? "true" : "false"; }
-
 template <typename... Args>
 struct InterleavedScanPlanner : AlgorithmPlanner {
   InterleavedScanPlanner(int Capacity, int Veclen, bool Ascending, bool ComputeNorm)
-    : AlgorithmPlanner("interleaved_scan_kernel_" + std::to_string(Capacity) + "_" +
-                         std::to_string(Veclen) + "_" + bool_to_string(Ascending) + "_" +
-                         bool_to_string(ComputeNorm),
+    : AlgorithmPlanner("interleaved_scan_kernel_capacity_" + std::to_string(Capacity) + "_veclen_" +
+                         std::to_string(Veclen) + "_" + (Ascending ? "ascending" : "descending") +
+                         "_" + (ComputeNorm ? "compute_norm" : "no_compute_norm"),
                        make_fragment_key<Args...>())
   {
   }
@@ -26,7 +24,7 @@ struct InterleavedScanPlanner : AlgorithmPlanner {
   template <typename... FuncTags>
   void add_metric_device_function(std::string metric_name, int Veclen)
   {
-    auto key    = metric_name + "_" + std::to_string(Veclen);
+    auto key    = metric_name + "_veclen_" + std::to_string(Veclen);
     auto params = make_fragment_key<FuncTags...>();
     this->device_functions.push_back(key + "_" + params);
   }
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric.cu.in
deleted file mode 100644
index 0f6bb904d1..0000000000
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric.cu.in
+++ /dev/null
@@ -1,36 +0,0 @@
-/*
- * SPDX-FileCopyrightText: Copyright (c) @year@, NVIDIA CORPORATION.
- * SPDX-License-Identifier: Apache-2.0
- */
-
-// This file is auto-generated. Do not edit manually.
-
-#ifdef BUILD_KERNEL
-
-#include <@header_file@>
-
-namespace cuvs::neighbors::ivf_flat::detail {
-
-// Instantiate the device function template
-template __device__ void compute_dist<@veclen@, @data_type@, @acc_type@>(@acc_type@&, @acc_type@, @acc_type@);
-
-}  // namespace cuvs::neighbors::ivf_flat::detail
-
-#else
-
-#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
-#include <cuvs/detail/jit_lto/ivf_flat/interleaved_scan_tags.hpp>
-#include "metric_@metric_name@_@veclen@_@type_abbrev@_@acc_abbrev@.h"
-
-using namespace cuvs::neighbors::ivf_flat::detail;
-
-__attribute__((__constructor__)) static void register_metric_@metric_name@_@veclen@_@type_abbrev@_@acc_abbrev@()
-{
-  registerAlgorithm<tag_@type_abbrev@,
-    tag_acc_@acc_abbrev@>(
-    "@metric_name@_@veclen@",
-    embedded_metric_@metric_name@_@veclen@_@type_abbrev@_@acc_abbrev@,
-    sizeof(embedded_metric_@metric_name@_@veclen@_@type_abbrev@_@acc_abbrev@));
-}
-
-#endif
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
new file mode 100644
index 0000000000..b8c3d636e6
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
@@ -0,0 +1,25 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
+#include <cuvs/detail/jit_lto/ivf_flat/interleaved_scan_tags.hpp>
+#include "@embedded_header_file@"
+
+using namespace cuvs::neighbors::ivf_flat::detail;
+
+namespace {
+
+__attribute__((__constructor__)) void register_kernel()
+{
+  registerAlgorithm<tag_@type_abbrev@,
+    tag_acc_@acc_abbrev@>(
+    "@metric_name@_veclen_@veclen@",
+    embedded_fatbin,
+    sizeof(embedded_fatbin));
+}
+
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
new file mode 100644
index 0000000000..477ff6b8ec
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
@@ -0,0 +1,15 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <@header_file@>
+
+namespace cuvs::neighbors::ivf_flat::detail {
+
+// Instantiate the device function template
+template __device__ void compute_dist<@veclen@, @data_type@, @acc_type@>(@acc_type@&, @acc_type@, @acc_type@);
+
+}  // namespace cuvs::neighbors::ivf_flat::detail
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
new file mode 100644
index 0000000000..4701e9feb3
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
@@ -0,0 +1,48 @@
+{
+  "_metric": [
+    {
+      "metric_name": "euclidean",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_euclidean_dist.cuh"
+    },
+    {
+      "metric_name": "inner_prod",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_inner_product.cuh"
+    }
+  ],
+  "_data_type": [
+    {
+      "data_type": "float",
+      "acc_type": "float",
+      "veclen": ["1", "4"],
+      "type_abbrev": "f",
+      "acc_abbrev": "f"
+    },
+    {
+      "data_type": "__half",
+      "acc_type": "__half",
+      "veclen": ["1", "8"],
+      "type_abbrev": "h",
+      "acc_abbrev": "h"
+    },
+    {
+      "data_type": "uint8_t",
+      "acc_type": "uint32_t",
+      "veclen": ["1", "16"],
+      "type_abbrev": "uc",
+      "acc_abbrev": "ui"
+    },
+    {
+      "data_type": "int8_t",
+      "acc_type": "int32_t",
+      "veclen": ["1", "16"],
+      "type_abbrev": "sc",
+      "acc_abbrev": "i"
+    }
+  ],
+  "_index": [
+    {
+      "idx_type": "int64_t",
+      "idx_abbrev": "l"
+    }
+  ]
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda.cu.in
deleted file mode 100644
index abf156a133..0000000000
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda.cu.in
+++ /dev/null
@@ -1,32 +0,0 @@
-/*
- * SPDX-FileCopyrightText: Copyright (c) @year@, NVIDIA CORPORATION.
- * SPDX-License-Identifier: Apache-2.0
- */
-
-// This file is auto-generated. Do not edit manually.
-
-#ifdef BUILD_KERNEL
-
-#include <@header_file@>
-
-namespace cuvs::neighbors::ivf_flat::detail {
-
-// Instantiate the device function template
-template __device__ float post_process<float>(float);
-
-}  // namespace cuvs::neighbors::ivf_flat::detail
-
-#else
-
-#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
-#include "@post_lambda_name@.h"
-
-__attribute__((__constructor__)) static void register_@post_lambda_name@()
-{
-  registerAlgorithm(
-    "@post_lambda_name@",
-    embedded_@post_lambda_name@,
-    sizeof(embedded_@post_lambda_name@));
-}
-
-#endif
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
new file mode 100644
index 0000000000..345e641360
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
@@ -0,0 +1,21 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <cuvs/detail/jit_lto/RegisterKernelFragment.hpp>
+#include "@embedded_header_file@"
+
+namespace {
+
+__attribute__((__constructor__)) void register_kernel()
+{
+  registerAlgorithm(
+    "@post_lambda_name@",
+    embedded_fatbin,
+    sizeof(embedded_fatbin));
+}
+
+}
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
new file mode 100644
index 0000000000..87331c4b6a
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
@@ -0,0 +1,15 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// This file is auto-generated. Do not edit manually.
+
+#include <@header_file@>
+
+namespace cuvs::neighbors::ivf_flat::detail {
+
+// Instantiate the device function template
+template __device__ float post_process<float>(float);
+
+}  // namespace cuvs::neighbors::ivf_flat::detail
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_matrix.json b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_matrix.json
new file mode 100644
index 0000000000..86f915923a
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_matrix.json
@@ -0,0 +1,16 @@
+{
+  "_post_lambda": [
+    {
+      "post_lambda_name": "post_identity",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/post_identity.cuh"
+    },
+    {
+      "post_lambda_name": "post_sqrt",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/post_sqrt.cuh"
+    },
+    {
+      "post_lambda_name": "post_compose",
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/post_compose.cuh"
+    }
+  ]
+}

From 8f16a69c6e28da435265d4a613f9abe7983caef8 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 19 Dec 2025 18:14:47 +0000
Subject: [PATCH 02/20] Add tests and warnings to compute_matrix_product

---
 .pre-commit-config.yaml                       |   2 +-
 cpp/cmake/modules/compute_matrix_product.py   | 145 ++++++++--
 .../cuvs/tests/test_compute_matrix_product.py | 254 ++++++++++++++++++
 python/cuvs/cuvs/tests/test_doctests.py       |   1 +
 4 files changed, 373 insertions(+), 29 deletions(-)
 create mode 100644 python/cuvs/cuvs/tests/test_compute_matrix_product.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5e99298f18..54e0ba0769 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -25,7 +25,7 @@ repos:
                 args: [--fix]
               - id: ruff-format
       - repo: https://github.com/pre-commit/mirrors-mypy
-        rev: 'v0.971'
+        rev: 'v1.19.1'
         hooks:
               - id: mypy
                 additional_dependencies: [types-cachetools]
diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
index 296a03f68f..0c67fe67d6 100644
--- a/cpp/cmake/modules/compute_matrix_product.py
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -8,8 +8,11 @@
 # https://gitlab.kitware.com/cmake/cmake/-/merge_requests/11516, we may be able
 # to port the algorithm to CMake script and use it in other RAPIDS projects.
 
+import argparse
 import json
+import re
 import sys
+import warnings
 from itertools import chain
 from typing import TYPE_CHECKING
 
@@ -20,54 +23,140 @@
     Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"]
 
 
+class NoKeyError(ValueError):
+    pass
+
+
+class UnusedKeyWarning(UserWarning):
+    pass
+
+
+class UsedKeyWarning(UserWarning):
+    pass
+
+
+IDENTIFIER_RE: re.Pattern = re.compile("^(?P<underscores>_*)(?P<rest>.*)$")
+
+
 def iterate_matrix_product(
     matrix: "Matrix",
+    warn_unused=True,
+    warn_used=True,
 ) -> "Generator[dict[str, MatrixValue]]":
     def iterate_next_dimension(
-        queue: "Iterator[tuple[str | None, Matrix]]",
+        queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]",
         entry: "dict[str, MatrixValue]",
-    ) -> "Generator[dict[str, MatrixValue]]":
+    ) -> "Generator[tuple[dict[str, MatrixValue], bool]]":
         try:
-            key, matrix = next(queue)
+            path, key, matrix = next(queue)
         except StopIteration:
-            yield entry
+            yield (entry, False)
         else:
-            yield from iterate_impl(matrix, key, queue, entry)
+            used = False
+            for e, u in iterate_impl(path, key, matrix, queue, entry):
+                if u:
+                    used = True
+                yield e, u
+
+            try:
+                last = path[-1]
+            except IndexError:
+                pass
+            else:
+                if isinstance(last, str):
+                    match = IDENTIFIER_RE.search(last)
+                    assert match
+                    underscores = match.group("underscores")
+                    rest = match.group("rest")
+                    path_repr = "".join(
+                        f"[{json.dumps(i)}]" for i in path[:-1]
+                    )
+
+                    if warn_used and used and underscores:
+                        warnings.warn(
+                            f"Key {json.dumps(last)} at root{path_repr} "
+                            f"is used in a matrix product entry even though it "
+                            f"begins with {json.dumps(underscores)}. Consider "
+                            f"renaming it to {json.dumps(rest)} to indicate this.",
+                            category=UsedKeyWarning,
+                        )
+                    elif warn_unused and not used and not underscores:
+                        warnings.warn(
+                            f"Key {json.dumps(last)} at root{path_repr} "
+                            f"is never used in a matrix product entry and is used "
+                            f"only for grouping. Consider renaming it to "
+                            f"{json.dumps(f'_{last}')} to indicate this.",
+                            category=UnusedKeyWarning,
+                        )
 
     def iterate_impl(
-        matrix: "Matrix",
+        path: tuple[str | int, ...],
         key: str | None,
-        queue: "Iterator[tuple[str | None, Matrix]]",
+        matrix: "Matrix",
+        queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]",
         entry: "dict[str, MatrixValue]",
-    ) -> "Generator[dict[str, MatrixValue]]":
+    ) -> "Generator[tuple[dict[str, MatrixValue], bool]]":
         if isinstance(matrix, dict):
-            yield from iterate_next_dimension(
-                chain(sorted(matrix.items()), queue),
-                entry,
+            yield from (
+                (e, False)
+                for e, _ in iterate_next_dimension(
+                    chain(
+                        (
+                            ((*path, k), k, v)
+                            for (k, v) in sorted(matrix.items())
+                        ),
+                        queue,
+                    ),
+                    entry,
+                )
             )
         elif isinstance(matrix, list):
             queue_list = list(queue)
-            for v in matrix:
+            for i, v in enumerate(matrix):
                 yield from iterate_next_dimension(
-                    chain([(key, v)], queue_list), entry
+                    chain([((*path, i), key, v)], queue_list), entry
                 )
         else:
-            assert key is not None
-            yield from iterate_next_dimension(queue, {**entry, key: matrix})
+            if key is None:
+                raise NoKeyError
+            yield from (
+                (e, True)
+                for e, _ in iterate_next_dimension(
+                    queue, {**entry, key: matrix}
+                )
+            )
 
-    return iterate_impl(matrix, None, chain(), {})
+    yield from (
+        entry for entry, _ in iterate_impl((), None, matrix, chain(), {})
+    )
 
 
-try:
-    filename = sys.argv[1]
-except IndexError:
-    filename = "-"
-with sys.stdin if filename == "-" else open(filename) as f:
-    matrix: "Matrix" = json.load(f)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--warn-unused", action=argparse.BooleanOptionalAction, default=True
+    )
+    parser.add_argument(
+        "--warn-used", action=argparse.BooleanOptionalAction, default=True
+    )
+    parser.add_argument("filename", nargs="?", default="-")
+    namespace = parser.parse_args()
+    with (
+        sys.stdin
+        if namespace.filename == "-"
+        else open(namespace.filename) as f
+    ):
+        matrix: "Matrix" = json.load(f)
 
-json.dump(
-    list(iterate_matrix_product(matrix)),
-    sys.stdout,
-    indent=2,
-    sort_keys=True,
-)
+    json.dump(
+        list(
+            iterate_matrix_product(
+                matrix,
+                warn_unused=namespace.warn_unused,
+                warn_used=namespace.warn_used,
+            )
+        ),
+        sys.stdout,
+        indent=2,
+        sort_keys=True,
+    )
diff --git a/python/cuvs/cuvs/tests/test_compute_matrix_product.py b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
new file mode 100644
index 0000000000..523493a9ab
--- /dev/null
+++ b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
@@ -0,0 +1,254 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import pathlib
+import runpy
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator
+
+    MatrixValue = None | bool | int | float | str
+    Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"]
+
+try:
+    compute_matrix_product_script = runpy.run_path(
+        str(
+            pathlib.Path(__file__).parent
+            / "../../../../cpp/cmake/modules/compute_matrix_product.py"
+        )
+    )
+except FileNotFoundError:
+    pytest.skip(
+        "Could not find compute_matrix_product.py", allow_module_level=True
+    )
+else:
+    iterate_matrix_product: "Callable[[Matrix], Generator[dict[str, MatrixValue]]]" = compute_matrix_product_script[
+        "iterate_matrix_product"
+    ]
+    NoKeyError: type[ValueError] = compute_matrix_product_script["NoKeyError"]
+    UnusedKeyWarning: type[UserWarning] = compute_matrix_product_script[
+        "UnusedKeyWarning"
+    ]
+    UsedKeyWarning: type[UserWarning] = compute_matrix_product_script[
+        "UsedKeyWarning"
+    ]
+
+
+@pytest.mark.parametrize(
+    ["matrix", "warn_unused", "warn_used", "expected_product", "contexts"],
+    [
+        pytest.param(
+            {
+                "a": [1, 2],
+                "b": [3, 4],
+            },
+            True,
+            True,
+            [
+                {"a": 1, "b": 3},
+                {"a": 1, "b": 4},
+                {"a": 2, "b": 3},
+                {"a": 2, "b": 4},
+            ],
+            [],
+            id="basic",
+        ),
+        pytest.param(
+            {
+                "b": [3, 4],
+                "a": [1, 2],
+            },
+            True,
+            True,
+            [
+                {"a": 1, "b": 3},
+                {"a": 1, "b": 4},
+                {"a": 2, "b": 3},
+                {"a": 2, "b": 4},
+            ],
+            [],
+            id="sort",
+        ),
+        pytest.param(
+            {
+                "_a": [
+                    {"b": 1, "c": 2},
+                    {"b": 3, "c": 4},
+                ],
+                "_d": [
+                    {"e": 5, "f": 6},
+                    {"e": 7, "f": 8},
+                ],
+            },
+            True,
+            True,
+            [
+                {"b": 1, "c": 2, "e": 5, "f": 6},
+                {"b": 1, "c": 2, "e": 7, "f": 8},
+                {"b": 3, "c": 4, "e": 5, "f": 6},
+                {"b": 3, "c": 4, "e": 7, "f": 8},
+            ],
+            [],
+            id="subgroup",
+        ),
+        pytest.param(
+            {
+                "a": [[1, 2], [3, 4]],
+            },
+            True,
+            True,
+            [
+                {"a": 1},
+                {"a": 2},
+                {"a": 3},
+                {"a": 4},
+            ],
+            [],
+            id="nested-lists",
+        ),
+        pytest.param(
+            [
+                {"a": 1},
+                {"a": 2},
+            ],
+            True,
+            True,
+            [
+                {"a": 1},
+                {"a": 2},
+            ],
+            [],
+            id="list-root",
+        ),
+        pytest.param(
+            [1],
+            True,
+            True,
+            None,
+            [pytest.raises(NoKeyError)],
+            id="list-root-no-dict",
+        ),
+        pytest.param(
+            [[1]],
+            True,
+            True,
+            None,
+            [pytest.raises(NoKeyError)],
+            id="nested-list-root-no-dict",
+        ),
+        pytest.param(
+            1,
+            True,
+            True,
+            None,
+            [pytest.raises(NoKeyError)],
+            id="item-root",
+        ),
+        pytest.param(
+            {
+                "a": [
+                    {
+                        "b": {"c": 1},
+                    },
+                ],
+            },
+            True,
+            True,
+            [
+                {"c": 1},
+            ],
+            [
+                pytest.warns(
+                    UnusedKeyWarning,
+                    match=r'^Key "a" at root is never used in a matrix product '
+                    r"entry and is used only for grouping\. Consider renaming "
+                    r'it to "_a" to indicate this\.$',
+                ),
+                pytest.warns(
+                    UnusedKeyWarning,
+                    match=r'^Key "b" at root\["a"\]\[0\] is never used in a '
+                    r"matrix product entry and is used only for grouping\. "
+                    r'Consider renaming it to "_b" to indicate this\.$',
+                ),
+            ],
+            id="unused",
+        ),
+        pytest.param(
+            {
+                "a": [
+                    {
+                        "b": {"c": 1},
+                    },
+                ],
+            },
+            False,
+            True,
+            [
+                {"c": 1},
+            ],
+            [],
+            id="unused-no-warning",
+        ),
+        pytest.param(
+            {
+                "_a": [
+                    {
+                        "_b": {"_c": 1},
+                    },
+                ],
+                "__b": 1,
+            },
+            True,
+            True,
+            [
+                {"__b": 1, "_c": 1},
+            ],
+            [
+                pytest.warns(
+                    UsedKeyWarning,
+                    match=r'^Key "_c" at root\["_a"\]\[0\]\["_b"\] is used in a matrix product entry even though it begins with "_"\. Consider renaming it to "c" to indicate this\.$',
+                ),
+                pytest.warns(
+                    UsedKeyWarning,
+                    match=r'^Key "__b" at root is used in a matrix product entry even though it begins with "__"\. Consider renaming it to "b" to indicate this\.$',
+                ),
+            ],
+            id="used",
+        ),
+        pytest.param(
+            {
+                "_a": [
+                    {
+                        "_b": {"_c": 1},
+                    },
+                ],
+                "__b": 1,
+            },
+            True,
+            False,
+            [
+                {"__b": 1, "_c": 1},
+            ],
+            [],
+            id="used-no-warning",
+        ),
+    ],
+)
+def test_iterate_matrix_product(
+    matrix, warn_unused, warn_used, expected_product, contexts
+):
+    with contextlib.ExitStack() as stack:
+        for context in contexts:
+            stack.enter_context(context)
+        assert (
+            list(
+                iterate_matrix_product(
+                    matrix, warn_unused=warn_unused, warn_used=warn_used
+                )
+            )
+            == expected_product
+        )
diff --git a/python/cuvs/cuvs/tests/test_doctests.py b/python/cuvs/cuvs/tests/test_doctests.py
index 8804bb11f1..a73f648046 100644
--- a/python/cuvs/cuvs/tests/test_doctests.py
+++ b/python/cuvs/cuvs/tests/test_doctests.py
@@ -11,6 +11,7 @@
 import pytest
 
 import cuvs.cluster
+import cuvs.common
 import cuvs.distance
 import cuvs.neighbors
 import cuvs.preprocessing.quantize

From 87e6506a0202eb0c7b397d0c7660d1c85ed688d1 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 19 Dec 2025 18:21:57 +0000
Subject: [PATCH 03/20] Formatting

---
 python/cuvs/cuvs/tests/test_compute_matrix_product.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/python/cuvs/cuvs/tests/test_compute_matrix_product.py b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
index 523493a9ab..fe918fda09 100644
--- a/python/cuvs/cuvs/tests/test_compute_matrix_product.py
+++ b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
@@ -210,11 +210,15 @@
             [
                 pytest.warns(
                     UsedKeyWarning,
-                    match=r'^Key "_c" at root\["_a"\]\[0\]\["_b"\] is used in a matrix product entry even though it begins with "_"\. Consider renaming it to "c" to indicate this\.$',
+                    match=r'^Key "_c" at root\["_a"\]\[0\]\["_b"\] is used in '
+                    r'a matrix product entry even though it begins with "_"\. '
+                    r'Consider renaming it to "c" to indicate this\.$',
                 ),
                 pytest.warns(
                     UsedKeyWarning,
-                    match=r'^Key "__b" at root is used in a matrix product entry even though it begins with "__"\. Consider renaming it to "b" to indicate this\.$',
+                    match=r'^Key "__b" at root is used in a matrix product '
+                    r'entry even though it begins with "__"\. Consider '
+                    r'renaming it to "b" to indicate this\.$',
                 ),
             ],
             id="used",

From 75ea5e6f2b5d76c5502b39b4df2b3e3311607f9c Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 19 Dec 2025 19:37:51 +0000
Subject: [PATCH 04/20] pre-commit

---
 .pre-commit-config.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 54e0ba0769..29e2a11a45 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -30,7 +30,8 @@ repos:
               - id: mypy
                 additional_dependencies: [types-cachetools]
                 args: ["--config-file=pyproject.toml",
-                       "python/cuvs/cuvs"]
+                       "python/cuvs/cuvs",
+                       "cpp/cmake/modules/compute_matrix_product.py"]
                 pass_filenames: false
       - repo: https://github.com/PyCQA/pydocstyle
         rev: 6.1.1

From 5125722c7ebce875864f9062591b02ff2bdcae21 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 10:43:42 -0500
Subject: [PATCH 05/20] copyright

---
 .../neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in   | 2 +-
 cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in  | 2 +-
 .../ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in   | 2 +-
 .../ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in      | 2 +-
 .../neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in   | 2 +-
 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in  | 2 +-
 .../ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in        | 2 +-
 .../neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in | 2 +-
 8 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
index 9a1be3bf43..a5a7299b73 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_embedded.cpp.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
index c6f671f9f2..a4c2f18f53 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/filter_kernel.cu.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
index 7620880925..7e247543b0 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_embedded.cpp.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
index ed39e7eb25..b7ada4de5b 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/interleaved_scan_kernel.cu.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
index b8c3d636e6..b951476565 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_embedded.cpp.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
index 477ff6b8ec..a67956db58 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_kernel.cu.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
index 345e641360..a2e3f1ea03 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_embedded.cpp.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
index 87331c4b6a..363964dd42 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/post_lambda_kernel.cu.in
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 

From c7c63b2fb3761bbfbb9db32c529a6d314eae8583 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 10:48:23 -0500
Subject: [PATCH 06/20] Copyright

---
 cpp/cmake/modules/compute_matrix_product.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
index 0c67fe67d6..fb4393e987 100644
--- a/cpp/cmake/modules/compute_matrix_product.py
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 # This algorithm takes a JSON dictionary and computes a matrix product of all

From 0c1f6af6fb5ff89f4e86e7d622dd0a6af35807c2 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 10:54:51 -0500
Subject: [PATCH 07/20] More

---
 python/cuvs/cuvs/tests/test_compute_matrix_product.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/cuvs/cuvs/tests/test_compute_matrix_product.py b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
index fe918fda09..35007d5c30 100644
--- a/python/cuvs/cuvs/tests/test_compute_matrix_product.py
+++ b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 import contextlib

From 91ad66b9e7b1fc6aaf12bfaa4ac4c6475535e8fd Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 11:06:12 -0500
Subject: [PATCH 08/20] fix headers

---
 cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
index 4701e9feb3..0629e837b6 100644
--- a/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
+++ b/cpp/src/neighbors/ivf_flat/jit_lto_kernels/metric_matrix.json
@@ -2,11 +2,11 @@
   "_metric": [
     {
       "metric_name": "euclidean",
-      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_euclidean_dist.cuh"
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_euclidean.cuh"
     },
     {
       "metric_name": "inner_prod",
-      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_inner_product.cuh"
+      "header_file": "neighbors/ivf_flat/jit_lto_kernels/metric_inner_prod.cuh"
     }
   ],
   "_data_type": [

From 3b90ec93383f1e6ef861129edf53d9a34d322103 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 11:34:05 -0500
Subject: [PATCH 09/20] fPIC

---
 cpp/CMakeLists.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index b4ede2b42c..691ff579eb 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -427,6 +427,7 @@ if(NOT BUILD_CPU_ONLY)
       cuvs_jit_lto_kernels STATIC ${interleaved_scan_files} ${metric_files} ${filter_files}
                                   ${post_lambda_files}
     )
+    set_target_properties(cuvs_jit_lto_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON)
     target_include_directories(
       cuvs_jit_lto_kernels
       PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"

From 4f8746c1823e9d48a84498ba3b689883d60afdb7 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 12:21:12 -0500
Subject: [PATCH 10/20] r-string

---
 cpp/cmake/modules/compute_matrix_product.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
index fb4393e987..00b563b094 100644
--- a/cpp/cmake/modules/compute_matrix_product.py
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -35,7 +35,7 @@ class UsedKeyWarning(UserWarning):
     pass
 
 
-IDENTIFIER_RE: re.Pattern = re.compile("^(?P<underscores>_*)(?P<rest>.*)$")
+IDENTIFIER_RE: re.Pattern = re.compile(r"^(?P<underscores>_*)(?P<rest>.*)$")
 
 
 def iterate_matrix_product(

From a2f91c27284d8266ea0cbe8a516db24f88b530b4 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 15:07:49 -0500
Subject: [PATCH 11/20] execute_process()

---
 .../modules/generate_jit_lto_kernels.cmake    | 25 +++++++------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/cpp/cmake/modules/generate_jit_lto_kernels.cmake b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
index 5d30eda6b2..675a92677c 100644
--- a/cpp/cmake/modules/generate_jit_lto_kernels.cmake
+++ b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
@@ -15,27 +15,20 @@ function(compute_matrix_product output_var)
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
   if(_JIT_LTO_MATRIX_JSON_FILE)
-    set(commands #
-        COMMAND
-        "${Python3_EXECUTABLE}"
-        "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
-        "${_JIT_LTO_MATRIX_JSON_FILE}" #
+    execute_process(
+      COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
+              "${_JIT_LTO_MATRIX_JSON_FILE}" #
+      OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
     )
   else()
-    set(commands #
-        COMMAND
-        "${CMAKE_COMMAND}"
-        -E
-        echo
-        "${_JIT_LTO_MATRIX_JSON_STRING}" #
-        COMMAND
-        "${Python3_EXECUTABLE}"
-        "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
-        - #
+    execute_process(
+      COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}"
+      COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
+              -
+      OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
     )
   endif()
 
-  execute_process(${commands} OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY)
   set(${output_var}
       "${output}"
       PARENT_SCOPE

From 6aa8fbdff9831d4368f4faa61e551f8dcad653ad Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 15:49:37 -0500
Subject: [PATCH 12/20] Review feedback

---
 cpp/cmake/modules/compute_matrix_product.py   | 82 ++++++++++++++++++-
 .../modules/generate_jit_lto_kernels.cmake    |  1 -
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
index 00b563b094..a53b439fea 100644
--- a/cpp/cmake/modules/compute_matrix_product.py
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 # This algorithm takes a JSON dictionary and computes a matrix product of all
@@ -43,6 +43,86 @@ def iterate_matrix_product(
     warn_unused=True,
     warn_used=True,
 ) -> "Generator[dict[str, MatrixValue]]":
+    """Computes a matrix product of a JSON document
+
+    This algorithm computes the product of a matrix in a more sophisticated
+    way than can be done with itertools.product(). Multiple related values
+    can be grouped together, and a dimension can have sub-dimensions. Given
+    the following JSON document:
+
+    .. code-block:: json
+
+      {
+        "value": ["one", "two"],
+        "_group": [
+          {
+            "subgroup_value": "three"
+            "subgroup_subdim": ["four", "five"]
+          },
+          {
+            "subgroup_value": "six",
+            "subgroup_subdim": ["seven", "eight"]
+          }
+        ]
+      }
+
+    The following matrix product entries will be produced:
+
+    .. code-block:: json
+
+      {"subgroup_subdim": "four", "subgroup_value": "three", "value": "one"}
+      {"subgroup_subdim": "five", "subgroup_value": "three", "value": "one"}
+      {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "one"}
+      {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "one"}
+      {"subgroup_subdim": "four", "subgroup_value": "three", "value": "two"}
+      {"subgroup_subdim": "five", "subgroup_value": "three", "value": "two"}
+      {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "two"}
+      {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "two"}
+
+    Notice that the name ``_group`` does not appear in any of the matrix
+    product entries. This is because all leaf nodes beneath it are under
+    a different key that's closer to them in the hierarchy, which is used in
+    the final product. If a key is used only for grouping and does not appear
+    in the final product, it should be prefixed with an underscore (``_``) to
+    indicate that they are hidden. Likewise, keys that appear in the final
+    product should not be prefixed with an underscore. Failure to follow this
+    convention will not affect the proper functioning of the algorithm, but a
+    warning will be emitted unless the respective ``warn_unused`` and/or
+    ``warn_used`` parameters are set to ``False``.
+
+    Every leaf node in the input document must have at least one dictionary key
+    in its path, or else a ``NoKeyError`` will be thrown. For example, a
+    document consisting only of an array of strings is invalid.
+
+    Parameters
+    ----------
+    matrix : Matrix
+        JSON document on which to compute the product.
+    warn_unused : bool
+        Whether or not to warn if unused keys are not prefixed with underscores
+        (true by default).
+    warn_used : bool
+        Whether or not to warn if used keys are prefixed with underscores (true
+        by default).
+
+    Returns
+    -------
+    Generator[dict[str, MatrixValue]]
+        Iterator of matrix product entries.
+
+    Raises
+    ------
+    NoKeyError
+        If a leaf node in the document does not have any key leading to it.
+
+    Warns
+    -----
+    UnusedKeyWarning
+        If an unused key is not prefixed with an underscore.
+    UsedKeyWarning
+        If a used key is prefixed with an underscore.
+    """
+
     def iterate_next_dimension(
         queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]",
         entry: "dict[str, MatrixValue]",
diff --git a/cpp/cmake/modules/generate_jit_lto_kernels.cmake b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
index 675a92677c..36d5b28227 100644
--- a/cpp/cmake/modules/generate_jit_lto_kernels.cmake
+++ b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
@@ -146,7 +146,6 @@ function(generate_jit_lto_kernels source_list_var)
       MATRIX_JSON_ENTRY "${matrix_json_entry}"
       KERNEL_LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES}
     )
-    list(APPEND ${source_list_var} "${embedded_file}")
   endforeach()
 
   set(${source_list_var}

From 535d897db9d984c410dba7e9bc32a779c8ff30dc Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Tue, 17 Feb 2026 15:52:46 -0500
Subject: [PATCH 13/20] Copyright

---
 python/cuvs/cuvs/tests/test_compute_matrix_product.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/cuvs/cuvs/tests/test_compute_matrix_product.py b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
index 35007d5c30..a99a59a742 100644
--- a/python/cuvs/cuvs/tests/test_compute_matrix_product.py
+++ b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 import contextlib

From 3d0a7428f0204ab8f797c77b83c7ff4bb9a351c1 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Wed, 18 Feb 2026 10:48:12 -0500
Subject: [PATCH 14/20] C++/CUDA 20

---
 cpp/CMakeLists.txt | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 691ff579eb..79676404dd 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -365,12 +365,12 @@ if(NOT BUILD_CPU_ONLY)
       jit_lto_kernel_usage_requirements INTERFACE "$<$<COMPILE_LANGUAGE:CXX>:${CUVS_CXX_FLAGS}>"
                                                   "$<$<COMPILE_LANGUAGE:CUDA>:${CUVS_CUDA_FLAGS}>"
     )
-    target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_17)
+    target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_20 cxx_std_20)
     target_link_libraries(
       jit_lto_kernel_usage_requirements INTERFACE rmm::rmm raft::raft CCCL::CCCL
     )
 
-    set(cuda_architectures_old ${CMAKE_CUDA_ARCHITECTURES})
+    block()
     set(CMAKE_CUDA_ARCHITECTURES ${JIT_LTO_TARGET_ARCHITECTURE})
     generate_jit_lto_kernels(
       interleaved_scan_files
@@ -421,7 +421,7 @@ if(NOT BUILD_CPU_ONLY)
       OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/post_lambda"
       KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
     )
-    set(CMAKE_CUDA_ARCHITECTURES ${cuda_architectures_old})
+    endblock()
 
     add_library(
       cuvs_jit_lto_kernels STATIC ${interleaved_scan_files} ${metric_files} ${filter_files}

From a4ac30f2389f454047fa56094e27d3156c2f29e1 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Wed, 18 Feb 2026 10:57:08 -0500
Subject: [PATCH 15/20] Propagate

---
 cpp/CMakeLists.txt | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 79676404dd..cfdb75546c 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -365,12 +365,12 @@ if(NOT BUILD_CPU_ONLY)
       jit_lto_kernel_usage_requirements INTERFACE "$<$<COMPILE_LANGUAGE:CXX>:${CUVS_CXX_FLAGS}>"
                                                   "$<$<COMPILE_LANGUAGE:CUDA>:${CUVS_CUDA_FLAGS}>"
     )
-    target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_20 cxx_std_20)
+    target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_20)
     target_link_libraries(
       jit_lto_kernel_usage_requirements INTERFACE rmm::rmm raft::raft CCCL::CCCL
     )
 
-    block()
+    block(PROPAGATE interleaved_scan_files metric_files filter_files post_lambda_files)
     set(CMAKE_CUDA_ARCHITECTURES ${JIT_LTO_TARGET_ARCHITECTURE})
     generate_jit_lto_kernels(
       interleaved_scan_files
@@ -427,7 +427,9 @@ if(NOT BUILD_CPU_ONLY)
       cuvs_jit_lto_kernels STATIC ${interleaved_scan_files} ${metric_files} ${filter_files}
                                   ${post_lambda_files}
     )
-    set_target_properties(cuvs_jit_lto_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON)
+    set_target_properties(
+      cuvs_jit_lto_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 20
+    )
     target_include_directories(
       cuvs_jit_lto_kernels
       PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"

From 813d20e83fabe5194b517aa828d3065cd3ea7a79 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Wed, 18 Feb 2026 11:14:13 -0500
Subject: [PATCH 16/20] Add comment

---
 cpp/cmake/modules/generate_jit_lto_kernels.cmake | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/cpp/cmake/modules/generate_jit_lto_kernels.cmake b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
index 36d5b28227..1454bac97e 100644
--- a/cpp/cmake/modules/generate_jit_lto_kernels.cmake
+++ b/cpp/cmake/modules/generate_jit_lto_kernels.cmake
@@ -58,6 +58,9 @@ function(add_jit_lto_kernel kernel_target)
   cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})
 
   add_library(${kernel_target} OBJECT EXCLUDE_FROM_ALL "${_JIT_LTO_KERNEL_FILE}")
+  # Do not modify these properties, options, and libraries. Usage requirements (including CUDA
+  # version, etc.) should be propagated to the kernel targets via INTERFACE libraries passed in
+  # 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)
   set_target_properties(

From 8430f3dcccefa43096bca97c25edb4f040d94ce3 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 20 Feb 2026 16:34:44 -0500
Subject: [PATCH 17/20] Review feedback

---
 cpp/cmake/modules/compute_matrix_product.py   |  3 +-
 .../cuvs/tests/test_compute_matrix_product.py | 39 ++++++++-----------
 2 files changed, 19 insertions(+), 23 deletions(-)

diff --git a/cpp/cmake/modules/compute_matrix_product.py b/cpp/cmake/modules/compute_matrix_product.py
index a53b439fea..76262f6597 100644
--- a/cpp/cmake/modules/compute_matrix_product.py
+++ b/cpp/cmake/modules/compute_matrix_product.py
@@ -39,6 +39,7 @@ class UsedKeyWarning(UserWarning):
 
 
 def iterate_matrix_product(
+    *,
     matrix: "Matrix",
     warn_unused=True,
     warn_used=True,
@@ -231,7 +232,7 @@ def iterate_impl(
     json.dump(
         list(
             iterate_matrix_product(
-                matrix,
+                matrix=matrix,
                 warn_unused=namespace.warn_unused,
                 warn_used=namespace.warn_used,
             )
diff --git a/python/cuvs/cuvs/tests/test_compute_matrix_product.py b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
index a99a59a742..7ed54276ff 100644
--- a/python/cuvs/cuvs/tests/test_compute_matrix_product.py
+++ b/python/cuvs/cuvs/tests/test_compute_matrix_product.py
@@ -14,28 +14,23 @@
     MatrixValue = None | bool | int | float | str
     Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"]
 
-try:
-    compute_matrix_product_script = runpy.run_path(
-        str(
-            pathlib.Path(__file__).parent
-            / "../../../../cpp/cmake/modules/compute_matrix_product.py"
-        )
-    )
-except FileNotFoundError:
-    pytest.skip(
-        "Could not find compute_matrix_product.py", allow_module_level=True
+compute_matrix_product_script = runpy.run_path(
+    str(
+        pathlib.Path(__file__).parent
+        / "../../../../cpp/cmake/modules/compute_matrix_product.py"
     )
-else:
-    iterate_matrix_product: "Callable[[Matrix], Generator[dict[str, MatrixValue]]]" = compute_matrix_product_script[
-        "iterate_matrix_product"
-    ]
-    NoKeyError: type[ValueError] = compute_matrix_product_script["NoKeyError"]
-    UnusedKeyWarning: type[UserWarning] = compute_matrix_product_script[
-        "UnusedKeyWarning"
-    ]
-    UsedKeyWarning: type[UserWarning] = compute_matrix_product_script[
-        "UsedKeyWarning"
-    ]
+)
+
+iterate_matrix_product: "Callable[[Matrix], Generator[dict[str, MatrixValue]]]" = compute_matrix_product_script[
+    "iterate_matrix_product"
+]
+NoKeyError: type[ValueError] = compute_matrix_product_script["NoKeyError"]
+UnusedKeyWarning: type[UserWarning] = compute_matrix_product_script[
+    "UnusedKeyWarning"
+]
+UsedKeyWarning: type[UserWarning] = compute_matrix_product_script[
+    "UsedKeyWarning"
+]
 
 
 @pytest.mark.parametrize(
@@ -251,7 +246,7 @@ def test_iterate_matrix_product(
         assert (
             list(
                 iterate_matrix_product(
-                    matrix, warn_unused=warn_unused, warn_used=warn_used
+                    matrix=matrix, warn_unused=warn_unused, warn_used=warn_used
                 )
             )
             == expected_product

From 19dd0f4b2fd5edd3fd628d86619d87f328bd7305 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 20 Feb 2026 16:43:12 -0500
Subject: [PATCH 18/20] Keep JIT LTO files in cuvs_jit_lto_kernels target

---
 cpp/CMakeLists.txt | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index cfdb75546c..cc2454c2c6 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -345,7 +345,6 @@ if(NOT BUILD_CPU_ONLY)
 
   set(JIT_LTO_TARGET_ARCHITECTURE "")
   set(JIT_LTO_COMPILATION OFF)
-  set(JIT_LTO_FILES "")
   if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
     set(JIT_LTO_TARGET_ARCHITECTURE "75-real")
     set(JIT_LTO_COMPILATION ON)
@@ -424,8 +423,16 @@ if(NOT BUILD_CPU_ONLY)
     endblock()
 
     add_library(
-      cuvs_jit_lto_kernels STATIC ${interleaved_scan_files} ${metric_files} ${filter_files}
-                                  ${post_lambda_files}
+      cuvs_jit_lto_kernels STATIC
+      ${interleaved_scan_files}
+      ${metric_files}
+      ${filter_files}
+      ${post_lambda_files}
+      src/detail/jit_lto/AlgorithmLauncher.cu
+      src/detail/jit_lto/AlgorithmPlanner.cu
+      src/detail/jit_lto/FragmentDatabase.cu
+      src/detail/jit_lto/FragmentEntry.cu
+      src/detail/jit_lto/nvjitlink_checker.cpp
     )
     set_target_properties(
       cuvs_jit_lto_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 20
@@ -436,14 +443,6 @@ if(NOT BUILD_CPU_ONLY)
               "${CMAKE_CURRENT_SOURCE_DIR}/../c/include"
     )
     add_library(cuvs::cuvs_jit_lto_kernels ALIAS cuvs_jit_lto_kernels)
-
-    set(JIT_LTO_FILES
-        src/detail/jit_lto/AlgorithmLauncher.cu
-        src/detail/jit_lto/AlgorithmPlanner.cu
-        src/detail/jit_lto/FragmentDatabase.cu
-        src/detail/jit_lto/FragmentEntry.cu
-        src/detail/jit_lto/nvjitlink_checker.cpp
-    )
   endif()
 
   add_library(
@@ -659,7 +658,6 @@ if(NOT BUILD_CPU_ONLY)
     src/stats/silhouette_score.cu
     src/stats/trustworthiness_score.cu
     ${CUVS_MG_ALGOS}
-    $<$<BOOL:${JIT_LTO_COMPILATION}>:${JIT_LTO_FILES}>
   )
 
   set_target_properties(

From 6e8846b1ff86dec132d632ffb1e3a80fb8e1d665 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 20 Feb 2026 16:49:38 -0500
Subject: [PATCH 19/20] Broaden

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 29e2a11a45..4578f1db1b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -31,7 +31,7 @@ repos:
                 additional_dependencies: [types-cachetools]
                 args: ["--config-file=pyproject.toml",
                        "python/cuvs/cuvs",
-                       "cpp/cmake/modules/compute_matrix_product.py"]
+                       "cpp/cmake/modules"]
                 pass_filenames: false
       - repo: https://github.com/PyCQA/pydocstyle
         rev: 6.1.1

From f8a8d5ce6b565a835e780e16be1df92cb9d8efdf Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Fri, 20 Feb 2026 17:04:27 -0500
Subject: [PATCH 20/20] Link

---
 cpp/CMakeLists.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index cc2454c2c6..d90579812a 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -442,6 +442,7 @@ if(NOT BUILD_CPU_ONLY)
       PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"
               "${CMAKE_CURRENT_SOURCE_DIR}/../c/include"
     )
+    target_link_libraries(cuvs_jit_lto_kernels PRIVATE raft::raft)
     add_library(cuvs::cuvs_jit_lto_kernels ALIAS cuvs_jit_lto_kernels)
   endif()