FEAT: Cache numeric Python UDF kernels

.github/workflows/pr.yaml

@@ -352,11 +352,13 @@ jobs:
       packages: read
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main
+    # TODO: Return to @main after rapidsai/shared-workflows#629 merges.
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52
     if: fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_builds || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_common || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_cudf
     with:
       build_type: pull-request
       script: "ci/test_python_cudf.sh"
+      numba-cache-key-prefix: cudf-numba-udf-v1
       # https://github.com/NVIDIA/cudf/issues/23498
       matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200"))
   conda-python-other-tests:
@@ -549,7 +551,8 @@ jobs:
       packages: read
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main
+    # TODO: Return to @main after rapidsai/shared-workflows#629 merges.
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52
     if: fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf_streaming || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf
     with:
       build_type: pull-request
@@ -559,6 +562,7 @@ jobs:
         RUN_CUDF_TESTS=${{ fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf }}
         RUN_CUDF_STREAMING_TESTS=${{ fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf_streaming || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf }}
         ci/test_wheel_cudf.sh
+      numba-cache-key-prefix: cudf-numba-udf-v1
       # https://github.com/NVIDIA/cudf/issues/23498
       matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200"))
   wheel-tests-cudf-polars:

.github/workflows/test.yaml

@@ -116,13 +116,15 @@ jobs:
       packages: read
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main
+    # TODO: Return to @main after rapidsai/shared-workflows#629 merges.
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       script: "ci/test_python_cudf.sh"
+      numba-cache-key-prefix: cudf-numba-udf-v1
       # https://github.com/NVIDIA/cudf/issues/23498
       matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200"))
   conda-python-other-tests:
@@ -200,13 +202,15 @@ jobs:
       packages: read
       pull-requests: read
     secrets: inherit # zizmor: ignore[secrets-inherit]
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main
+    # TODO: Return to @main after rapidsai/shared-workflows#629 merges.
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       script: ci/test_wheel_cudf.sh
+      numba-cache-key-prefix: cudf-numba-udf-v1
   wheel-tests-dask-cudf:
     permissions:
       actions: read

python/cudf/cudf/core/udf/groupby_utils.py

@@ -197,6 +197,11 @@ def f(group):
     def kernel_type(self):
         return "groupby_apply"
 
+    @property
+    def _requires_linked_udf_shim(self):
+        # GroupBy UDF reductions call the device functions in UDF_SHIM_FILE.
+        return True
+
     def _get_frame_type(self):
         return _get_frame_groupby_type(
             np.dtype(list(_all_dtypes_from_frame(self.frame).items())),

python/cudf/cudf/core/udf/scalar_function.py

@@ -1,6 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import warnings
+from contextlib import nullcontext
 from functools import cache
 
 from numba import cuda
@@ -9,6 +11,7 @@
 from cudf.core.dtype.validators import is_dtype_obj_string
 from cudf.core.udf.api import Masked, pack_return
 from cudf.core.udf.masked_typing import MaskedType
+from cudf.core.udf.nrt_utils import CaptureNRTUsage, nrt_enabled
 from cudf.core.udf.strings_typing import string_view
 from cudf.core.udf.templates import (
     masked_input_initializer_template,
@@ -17,10 +20,74 @@
 )
 from cudf.core.udf.udf_kernel_base import ApplyKernelBase
 from cudf.core.udf.utils import (
+    DEPRECATED_SM_REGEX,
     _mask_get,
 )
 
 
+def _compile_unmasked_series_apply_kernel(f_, sig, nrt):
+    """Compile a source-backed, disk-cacheable Series.apply kernel."""
+    ctx = nrt_enabled() if nrt else nullcontext()
+    with ctx:
+        with warnings.catch_warnings():
+            warnings.simplefilter("default")
+            warnings.filterwarnings(
+                "ignore",
+                message=DEPRECATED_SM_REGEX,
+                category=UserWarning,
+                module=r"^numba\.cuda(\.|$)",
+            )
+
+            @cuda.jit(
+                sig,
+                cache=True,
+            )
+            def _kernel(retval, size, input_col_0, offset_0):
+                i = cuda.grid(1)
+                ret_data_arr, ret_mask_arr = retval
+
+                if i < size:
+                    masked_0 = Masked(input_col_0[i + offset_0], True)
+                    ret_masked = pack_return(f_(masked_0))
+                    ret_data_arr[i] = ret_masked.value
+                    ret_mask_arr[i] = ret_masked.valid
+
+    return _kernel
+
+
+def _compile_masked_series_apply_kernel(f_, sig, nrt):
+    """Compile a source-backed, disk-cacheable Series.apply kernel."""
+    ctx = nrt_enabled() if nrt else nullcontext()
+    with ctx:
+        with warnings.catch_warnings():
+            warnings.simplefilter("default")
+            warnings.filterwarnings(
+                "ignore",
+                message=DEPRECATED_SM_REGEX,
+                category=UserWarning,
+                module=r"^numba\.cuda(\.|$)",
+            )
+
+            @cuda.jit(
+                sig,
+                cache=True,
+            )
+            def _kernel(retval, size, input_col_0, offset_0):
+                i = cuda.grid(1)
+                ret_data_arr, ret_mask_arr = retval
+
+                if i < size:
+                    d_0, m_0 = input_col_0
+                    masked_0 = Masked(
+                        d_0[i + offset_0], _mask_get(m_0, i + offset_0)
+                    )
+                    ret_masked = pack_return(f_(masked_0))
+                    ret_data_arr[i] = ret_masked.value
+                    ret_mask_arr[i] = ret_masked.valid
+
+    return _kernel
+
+
 class SeriesApplyKernel(ApplyKernelBase):
     """
     Class representing a kernel that computes the result of
@@ -58,6 +125,34 @@ def _get_kernel_string(self):
             extra_args=extra_args, masked_initializer=masked_initializer
         )
 
+    def compile_kernel(self):
+        # The static no-argument kernels below have a real source location, so
+        # Numba can persist their compiled specializations. Other signatures
+        # keep the existing generated-kernel path for this prototype.
+        if self.args or is_dtype_obj_string(self.frame.dtype):
+            return super().compile_kernel()
+
+        capture_nrt_usage = CaptureNRTUsage()
+        with capture_nrt_usage:
+            return_type = self._get_udf_return_type()
+
+        # String allocation uses symbols from UDF_SHIM_FILE, which Numba's
+        # on-disk cache cannot serialize. Keep the existing linked path for
+        # those specializations.
+        if capture_nrt_usage.use_nrt:
+            return super().compile_kernel()
+
+        self.sig = self._construct_signature(return_type)
+        kernel_factory = (
+            _compile_masked_series_apply_kernel
+            if self.frame._column.mask is not None
+            else _compile_unmasked_series_apply_kernel
+        )
+        kernel = kernel_factory(
+            self.device_func, self.sig, capture_nrt_usage.use_nrt
+        )
+        return kernel, return_type
+
     @cache
     def _get_kernel_string_exec_context(self):
         # This is the global execution context that will be used

python/cudf/cudf/core/udf/udf_kernel_base.py

@@ -1,9 +1,12 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 import warnings
 from abc import ABC, abstractmethod
 from contextlib import nullcontext
+from hashlib import sha256
+from pickle import dumps
+from textwrap import indent
 
 import numpy as np
 from numba import cuda, typeof
@@ -12,6 +15,7 @@
 from numba.types import CPointer, Poison, Tuple, boolean, int64, void
 
 from cudf.api.types import is_scalar
+from cudf.core.dtype.validators import is_dtype_obj_string
 from cudf.core.udf.masked_typing import MaskedType
 from cudf.core.udf.nrt_utils import CaptureNRTUsage, nrt_enabled
 from cudf.core.udf.strings_typing import str_view_arg_handler
@@ -75,6 +79,11 @@ def _get_kernel_string_exec_context(self):
         string.
         """
 
+    @property
+    def _requires_linked_udf_shim(self):
+        """Whether this API has device-library dependencies beyond NRT."""
+        return False
+
     @staticmethod
     def _format_arg_list(prefix, count):
         """Build a comma-separated parameter list like 'prefix_0, prefix_1, ...'."""
@@ -156,12 +165,65 @@ def compile_kernel(self):
 
         return kernel, return_type
 
+    def _disk_cache_identity(self):
+        """Return a stable identity for globals used by a generated kernel."""
+        key = _generate_cache_key(
+            self.frame, self.func, self.args, suffix=self.kernel_type
+        )
+        return int.from_bytes(sha256(dumps(key)).digest()[:8], "big")
+
+    def _can_use_disk_cache(self, nrt):
+        return (
+            not self._requires_linked_udf_shim
+            and not nrt
+            and not any(
+                is_dtype_obj_string(col.dtype) for col in self.frame._columns
+            )
+        )
+
     def compile_kernel_string(self, kernel_string, nrt=False):
         global_exec_context = self._get_kernel_string_exec_context()
-        global_exec_context["f_"] = self.device_func
+        use_disk_cache = self._can_use_disk_cache(nrt)
+        if use_disk_cache:
+            # The closure carries the UDF and generated-kernel globals into
+            # Numba's cache key. Compiling with this module's filename gives
+            # the dispatcher a persistent source locator.
+            kernel_lines = indent(
+                kernel_string, "    ", lambda _line: True
+            ).splitlines()
+            kernel_def_idx = next(
+                i
+                for i, line in enumerate(kernel_lines)
+                if line.startswith("    def _kernel")
+            )
+            kernel_lines[kernel_def_idx + 1 : kernel_def_idx + 1] = [
+                "        if cache_identity < 0:",
+                "            return",
+            ]
+            factory_string = "\n".join(
+                [
+                    "def _make_kernel(f_, cache_identity):",
+                    *kernel_lines,
+                    "    return _kernel",
+                ]
+            )
+            exec(
+                compile(factory_string, __file__, "exec"), global_exec_context
+            )
+            _kernel = global_exec_context["_make_kernel"](
+                self.device_func, self._disk_cache_identity()
+            )
+        else:
+            global_exec_context["f_"] = self.device_func
+            exec(kernel_string, global_exec_context)
+            _kernel = global_exec_context["_kernel"]
 
-        exec(kernel_string, global_exec_context)
-        _kernel = global_exec_context["_kernel"]
+        jit_kwargs = {}
+        if use_disk_cache:
+            jit_kwargs["cache"] = True
+        else:
+            jit_kwargs["link"] = [UDF_SHIM_FILE]
+            jit_kwargs["extensions"] = [str_view_arg_handler]
         ctx = nrt_enabled() if nrt else nullcontext()
         with ctx:
             with warnings.catch_warnings():
@@ -174,8 +236,7 @@ def compile_kernel_string(self, kernel_string, nrt=False):
                 )
                 kernel = cuda.jit(
                     self.sig,
-                    link=[UDF_SHIM_FILE],
-                    extensions=[str_view_arg_handler],
+                    **jit_kwargs,
                 )(_kernel)
         return kernel