JIT compile `interleaved_scan_kernel` for CUDA 13 by divyegala · Pull Request #1405 · rapidsai/cuvs

divyegala · 2025-10-02T18:34:46Z

This PR introduces infrastructure to generate LTO-IR representations of kernels at compile time. At runtime, this LTO-IR representations are JITted to SASS using the architecture of the native GPU with LTO (Link Time Optimizations) of device functions called within the kernel.

The changes in this PR can be divided in 4 parts:

CMake code to generate fatbins to JIT at runtime
Architecture to handle JITing and caching of kernels
Changes to the build system to bring in nvJitLink as a CUDA 13 dependency
Using all of the above to JIT interleaved_scan_kernel, where we use a Python script to generate the files that are using for first compiling and then for embedding fatbins

Binary sizes of interleaved_scan_*.cu TUs in CUDA 13:

TU name	Size
ivf_flat_interleaved_scan_half_int64_t.cu.o	5.658 MB
vf_flat_interleaved_scan_half_int64_t_bitset.cu.o	5.998 MB
ivf_flat_interleaved_scan_uint8_t_int64_t.cu.o	6.402 MB
ivf_flat_interleaved_scan_int8_t_int64_t.cu.o	6.338 MB
ivf_flat_interleaved_scan_float_int64_t_bitset.cu.o	6.273 MB
ivf_flat_interleaved_scan_float_int64_t.cu.o	5.872 MB
ivf_flat_interleaved_scan_uint8_t_int64_t_bitset.cu.o	6.851 MB
ivf_flat_interleaved_scan_int8_t_int64_t_bitset.cu.o	6.901 MB

Total size of aforementioned TUs: 50.293 MB
Total size of libcuvs.so on 27/10/2025 in main 468.79 MB
Total size of libcuvs.so in this PR: 444.70 MB

copy-pr-bot · 2025-10-02T18:34:50Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

cpp/cmake/modules/embed_fatbins.cmake

cpp/cmake/modules/generate_interleaved_scan_kernels.cmake

divyegala · 2025-10-09T02:25:03Z

Benchmark:

@tfeher as requested, I ran the JITed kernels on a small batch size with a warmup time of 4s. The results are practically the same

…uvs into jit-lto-ivf-flat-interleaved

cpp/cmake/modules/generate_jit_lto_kernels.cmake

…uvs into jit-lto-ivf-flat-interleaved

…rleaved

dantegd

Mainly one bug in the norm computation, and a couple other minor comments/questions

cpp/src/detail/jit_lto/AlgorithmPlanner.cu

cpp/src/neighbors/ivf_flat/jit_lto_kernels/ivf_flat_interleaved_scan_kernel.cuh

cpp/CMakeLists.txt

…rleaved

divyegala · 2026-02-14T04:31:48Z

/merge

…wheels against mix of CTK versions (#1862) The changes from #1405 introduced linking against nvJitLink. nvJitLink has versioned symbols that are added in each new CTK release, and some of those are exposed in `libcuvs.so`. `libcuvs` wheels are built against the latest CTK supported in RAPIDS (CUDA 13.1.1 as of this writing), so when those wheels are used in environments with older nvJitLink, runtime errors like this can happen: > libcugraph.so: undefined symbol: __nvJitLinkGetErrorLog_13_1, version libnvJitLink.so.13 For more details, see rapidsai/cugraph#5443 This tries to fix that. Contributes to rapidsai/build-planning#257 * builds CUDA 13 wheels with the 13.0 CTK * ensures CUDA 13 wheels ship with a runtime dependency of `nvidia-nvjitlink>={whatever-minor-version-they-were-built-against}` Contributes to rapidsai/build-planning#256 * updates wheel tests to cover a range of CTK versions (we previously, accidentally, were only testing the latest 12.x and 13.x) Other changes * ensures conda packages also take on floors of `libnvjitlink>={whatever-minor-version-they-were-built-against}` ## Notes for Reviewers ### How I tested this This uses wheels from similar PRs from RAPIDS dependencies, at build and test time: * rapidsai/raft#2971 * rapidsai/rmm#2270 * rapidsai/ucxx#604 ### Other Options 1. avoiding those versioned symbols with a build-time shim (#1855 does this, but hasn't been successful yet) 2. statically linking libnvJitLink (hasn't been successful yet) Authors: - James Lamb (https://github.com/jameslamb) Approvers: - Gil Forsyth (https://github.com/gforsyth) URL: #1862

…wheels against mix of CTK versions (rapidsai#1862) The changes from rapidsai#1405 introduced linking against nvJitLink. nvJitLink has versioned symbols that are added in each new CTK release, and some of those are exposed in `libcuvs.so`. `libcuvs` wheels are built against the latest CTK supported in RAPIDS (CUDA 13.1.1 as of this writing), so when those wheels are used in environments with older nvJitLink, runtime errors like this can happen: > libcugraph.so: undefined symbol: __nvJitLinkGetErrorLog_13_1, version libnvJitLink.so.13 For more details, see rapidsai/cugraph#5443 This tries to fix that. Contributes to rapidsai/build-planning#257 * builds CUDA 13 wheels with the 13.0 CTK * ensures CUDA 13 wheels ship with a runtime dependency of `nvidia-nvjitlink>={whatever-minor-version-they-were-built-against}` Contributes to rapidsai/build-planning#256 * updates wheel tests to cover a range of CTK versions (we previously, accidentally, were only testing the latest 12.x and 13.x) Other changes * ensures conda packages also take on floors of `libnvjitlink>={whatever-minor-version-they-were-built-against}` This uses wheels from similar PRs from RAPIDS dependencies, at build and test time: * rapidsai/raft#2971 * rapidsai/rmm#2270 * rapidsai/ucxx#604 1. avoiding those versioned symbols with a build-time shim (rapidsai#1855 does this, but hasn't been successful yet) 2. statically linking libnvJitLink (hasn't been successful yet) Authors: - James Lamb (https://github.com/jameslamb) Approvers: - Gil Forsyth (https://github.com/gforsyth) URL: rapidsai#1862

jit lto interleaved scan

a024f61

github-project-automation bot added this to Unstructured Data Processing Oct 2, 2025

github-project-automation bot moved this to Todo in Unstructured Data Processing Oct 2, 2025

divyegala added 6 commits October 2, 2025 18:40

fix dependencies.yaml

45da4aa

generate files at build time, use tags to avoid compilation of types

a7c8621

passing tests

eb2d74b

update gitignore

d2318e8

separate out distance function from main kernel

5e6afcd

fix deps

6eee4da

KyleFromNVIDIA requested changes Oct 6, 2025

View reviewed changes

cpp/cmake/modules/embed_fatbins.cmake Outdated Show resolved Hide resolved

cpp/cmake/modules/generate_interleaved_scan_kernels.cmake Outdated Show resolved Hide resolved

cpp/cmake/modules/generate_interleaved_scan_kernels.cmake Outdated Show resolved Hide resolved

add filters as jit device functions, rework caching logic

1de8f28

divyegala changed the title ~~jit lto interleaved scan~~ JIT compile interleaved_scan_kernel and LTO device functions Oct 7, 2025

divyegala changed the title ~~JIT compile interleaved_scan_kernel and LTO device functions~~ JIT compile interleaved_scan_kernel with LTO of device functions Oct 7, 2025

divyegala added 2 commits October 7, 2025 03:13

lto post lambda, cleanup files, generate cmake in build dir

84c6020

don't read hardcoded kernels, use generator properly

22680c8

divyegala marked this pull request as ready for review October 8, 2025 20:31

divyegala requested review from a team as code owners October 8, 2025 20:31

divyegala requested a review from KyleFromNVIDIA October 8, 2025 20:31

divyegala added feature request New feature or request non-breaking Introduces a non-breaking change labels Oct 8, 2025

divyegala added 3 commits October 8, 2025 20:35

random cmake changes carried over from 25.10

37f1163

cmake format

0ae5383

remove dep on kernel list

fe56aec

attempt to solve overlinking problem

40c8fd6

divyegala requested a review from a team as a code owner October 9, 2025 02:55

reorder if-else in compiler check

e87a8c7

Merge branch 'jit-lto-ivf-flat-interleaved' of github.com:divyegala/c…

30a8a9f

…uvs into jit-lto-ivf-flat-interleaved

divyegala requested review from dantegd and lowener February 4, 2026 21:15

lowener reviewed Feb 4, 2026

View reviewed changes

cpp/cmake/modules/generate_jit_lto_kernels.cmake Outdated Show resolved Hide resolved

divyegala added 5 commits February 4, 2026 20:55

Merge branch 'main' into jit-lto-ivf-flat-interleaved

72ddb36

add to docs and log about jit

4bd2102

Merge branch 'jit-lto-ivf-flat-interleaved' of github.com:divyegala/c…

fb722f0

…uvs into jit-lto-ivf-flat-interleaved

Merge remote-tracking branch 'origin/main' into jit-lto-ivf-flat-inte…

3523b96

…rleaved

address review

ba758a2

divyegala requested a review from lowener February 10, 2026 20:23

divyegala added 3 commits February 10, 2026 23:30

rename inner_product to inner_prod

42b78ae

Merge remote-tracking branch 'origin/main' into jit-lto-ivf-flat-inte…

2e3a471

…rleaved

include header and form better log

f6377fa

dantegd requested changes Feb 11, 2026

View reviewed changes

cpp/src/detail/jit_lto/AlgorithmPlanner.cu Outdated Show resolved Hide resolved

cpp/src/neighbors/ivf_flat/jit_lto_kernels/ivf_flat_interleaved_scan_kernel.cuh Outdated Show resolved Hide resolved

cpp/CMakeLists.txt Show resolved Hide resolved

divyegala added 4 commits February 11, 2026 06:05

address review and move

533b770

Merge remote-tracking branch 'origin/main' into jit-lto-ivf-flat-inte…

af23585

…rleaved

one more fix

78c59d9

correct path

7f8802b

dantegd approved these changes Feb 12, 2026

View reviewed changes

lowener approved these changes Feb 13, 2026

View reviewed changes

robertmaynard approved these changes Feb 13, 2026

View reviewed changes

merge upstream

27acbb6

rapids-bot bot merged commit da6136e into rapidsai:main Feb 14, 2026
149 of 151 checks passed

github-project-automation bot moved this from In Progress to Done in Unstructured Data Processing Feb 14, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

JIT compile `interleaved_scan_kernel` for CUDA 13#1405

JIT compile `interleaved_scan_kernel` for CUDA 13#1405
rapids-bot[bot] merged 105 commits intorapidsai:mainfrom
divyegala:jit-lto-ivf-flat-interleaved

divyegala commented Oct 2, 2025 •

edited by KyleFromNVIDIA

Loading

Uh oh!

copy-pr-bot bot commented Oct 2, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyegala commented Oct 9, 2025 •

edited

Loading

Uh oh!

Uh oh!

dantegd left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyegala commented Feb 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Conversation

divyegala commented Oct 2, 2025 • edited by KyleFromNVIDIA Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot bot commented Oct 2, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyegala commented Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

dantegd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyegala commented Feb 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

divyegala commented Oct 2, 2025 •

edited by KyleFromNVIDIA

Loading

divyegala commented Oct 9, 2025 •

edited

Loading