[cuda.compute]: cached algorithm wrappers and operator cache keys keep the caller's device arrays alive, leading to OOM in long sweeps #12187

Summary

Two caching paths in cuda.compute retain user device arrays after the user has dropped every reference on their side:

1. Cached algorithm wrappers retain the arrays of the last call. Algorithm factories such as make_unique_by_key, make_reduce_into, etc. are decorated with @cache_with_registered_key_functions, which caches the returned algorithm object keyed by the argument types. On every call, set_cccl_iterator_state rebinds the cached wrapper's iterators to the new pointers and passes the user's array as the owner of that state (cccl_it.bind_pointer_state(ptr, input_it)). The cached Pointer / Iterator objects therefore hold the arrays of the most recent call for a given (algorithm, type signature) until the next call with the same signature, or for the lifetime of the process.

2. Operator cache keys retain closure-captured arrays. A Python operator is keyed in the wrapper cache, the build-results cache, and the _jit._compile_op_impl memo through CachableFunction, which stores the function itself in _func. If the operator captures a device array in its closure (the stateful-operator path), the function keeps the array alive, and the cache keys keep the function alive. This is bounded to one array per distinct function body, since equal CachableFunctions share an entry, but the retained array is never released.

How it showed up (1)

Running the python/cuda_cccl/benchmarks/compute suite in profile mode (one pass over every axis value) on a 48 GB GPU. Each state allocates fresh inputs/outputs, uses them once, and returns. After ~20 states with distinct type combinations the pool held ~46 GB and later states failed:

Fail: Unexpected error: OutOfMemoryError: Out of memory allocating 2,147,483,648 bytes (allocated so far: 45,902,695,424 bytes).

Affected sweeps: select/unique_by_key (60 of 336 states), segmented_sort/keys (65 of 342), merge_sort/pairs (30 of 192), histogram/even (13), radix_sort/pairs (11), select/if (6). The matching C++ nvbench binaries run the same states without issue.

Instrumenting one benchmark confirmed the cause. At the start of each state, after cp.get_default_memory_pool().free_all_blocks() and gc.collect():

state 1: used=0.00G  live cupy arrays=0
state 2: used=4.50G  live cupy arrays=5   referrers: Iterator, Pointer
state 3: used=9.50G  live cupy arrays=10  referrers: Iterator, Pointer
state 6: used=25.0G  live cupy arrays=25  referrers: Iterator, Pointer

Five arrays per state (two inputs, two outputs, num_selected), each referenced only through a cuda.compute Iterator / Pointer, i.e. through the cached wrapper. Calling the private _caching._clear_wrapper_caches() before each state brings the live array count to zero and removes every OOM, so the cached wrappers are the only thing holding them.

• Reproducer: python python/cuda_cccl/benchmarks/compute/select/unique_by_key.py -d 0 --profile -a "Elements{io}[pow2]=28"

Reproducer (2)

import gc, weakref
import cupy as cp
import numpy as np
from cuda.compute import unary_transform

def run_once():
    scale = cp.asarray([3], dtype=np.int32)  # captured by the op below

    def op(x):
        return x * scale[0]

    d_in = cp.arange(16, dtype=np.int32)
    d_out = cp.empty_like(d_in)
    unary_transform(d_in=d_in, d_out=d_out, op=op, num_items=16)
    cp.cuda.Device().synchronize()
    return weakref.ref(scale)

ref = run_once()
gc.collect()
print("captured array alive after caller dropped it:", ref() is not None)

Prints True. Inspecting the referrers shows exactly one CachableFunction whose _func closure holds the array, and that object is itself held by the two cache-key tuples (wrapper cache, build-results cache) and the _compile_op_impl memo dict. This reproduces with (1) already fixed, so it is an independent retention path.

Why it matters

For a user this is a silent leak: any array passed on the first call of each (algorithm, type signature) stays resident until the process exits, and any array captured by an operator closure stays resident for as long as that operator's cache entries do. Code that iterates over many dtypes or many distinct output types, or a long-lived service, accumulates device memory it cannot release. clear_all_caches() works around it but also drops the compiled build results, which is not what a user wants.