[BUG] cuvs-java device matrices are capped at a quarter of GPU memory #2866

Describe the bug

Since #2035, cuvsRMMAlloc allocates from the RAFT workspace resource instead of the current device memory resource. cuvs-java allocates device matrices (CuVSMatrix.deviceBuilder) through cuvsRMMAlloc, so a device matrix can't be larger than the workspace limit, which defaults to a quarter of total device memory. On a 96 GB GPU that is 6,224,112 rows of 1,024 float32. cuvsResourcesSetWorkspacePool, which is documented as configuring an uncapped pool, keeps the same limit.

I hit this building Vamana graphs through cuvs-java. An 8M × 1,024 dataset couldn't be copied to the GPU, so my code built from host memory instead, and the build ran over 2 hours where 4M took about 5 minutes from a device copy.

Steps/Code to reproduce bug

The script below calls the C API through ctypes on one otherwise idle GPU.

import ctypes

import cupy

lib = ctypes.CDLL("libcuvs_c.so")
lib.cuvsResourcesCreate.argtypes = [ctypes.POINTER(ctypes.c_size_t)]
lib.cuvsRMMAlloc.argtypes = [ctypes.c_size_t, ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
lib.cuvsRMMFree.argtypes = [ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t]
lib.cuvsResourcesSetWorkspacePool.argtypes = [ctypes.c_size_t, ctypes.c_size_t]
lib.cuvsGetLastErrorText.restype = ctypes.c_char_p
CUVS_SUCCESS = 1

free, total = cupy.cuda.runtime.memGetInfo()
row = 1024 * 4  # one 1,024-dimensional float32 vector
quarter_rows = total // 4 // row
print(f"total device memory {total} bytes, a quarter holds {quarter_rows} rows")


def alloc(rows, pool=None):
    res = ctypes.c_size_t()
    assert lib.cuvsResourcesCreate(ctypes.byref(res)) == CUVS_SUCCESS
    if pool is not None:
        assert lib.cuvsResourcesSetWorkspacePool(res, pool) == CUVS_SUCCESS
    ptr = ctypes.c_void_p()
    ok = lib.cuvsRMMAlloc(res, ctypes.byref(ptr), rows * row) == CUVS_SUCCESS
    print(f"{rows} rows{' after SetWorkspacePool' if pool else ''}:",
          "OK" if ok else lib.cuvsGetLastErrorText().decode().rsplit(": ", 1)[-1])
    if ok:
        lib.cuvsRMMFree(res, ptr, rows * row)


alloc(quarter_rows)
alloc(quarter_rows + 1)
alloc(8_000_000)
alloc(8_000_000, pool=1 << 30)
x = cupy.empty((8_000_000, 1024), dtype=cupy.float32)
print("8000000 rows with CuPy: OK")

Output

total device memory 101975851008 bytes, a quarter holds 6224112 rows
6224112 rows: OK
6224113 rows: Exceeded memory limit 25493962752; Allocated bytes 0; Requested bytes 25493966848
8000000 rows: Exceeded memory limit 25493962752; Allocated bytes 0; Requested bytes 32768000000
8000000 rows after SetWorkspacePool: Exceeded memory limit 25493962752; Allocated bytes 0; Requested bytes 32768000000
8000000 rows with CuPy: OK

In cuvs-java, CuVSMatrix.deviceBuilder(resources, 8_000_000, 1024, CuVSMatrix.DataType.FLOAT).build() fails the same way.

Expected behavior

A device matrix that holds a dataset should be limited by device memory, as it was before #2035. RAFT's comment on default_allocation_limit describes the workspace limit as "merely an order for algorithm internals to plan the batching accordingly", so long-lived allocations like datasets shouldn't go through it.

Environment details (please complete the following information):

• Environment location: Bare-metal
• Method of cuVS install: conda, libcuvs 26.12.00a12 from rapidsai-nightly (commit c740818), with cuvs-java 26.12.0 built from source. cuvsRMMAlloc is unchanged on main as of 32776b9.
• GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition (96 GB), driver 595.71.05

Additional context

#2035 changed cuvsRMMAlloc and cuvsRMMFree from rmm::mr::get_current_device_resource_ref() to raft::resource::get_workspace_resource_ref(). cuvs-java uses cuvsRMMAlloc both for short-lived buffers, such as filter bitsets, and for long-lived device matrices. The workspace pool suits the first and caps the second.

A possible fix is to keep the workspace pool for scratch buffers and give device matrices their own allocation and free path through the current device memory resource. Unpooled cudaMalloc and cudaFree synchronize the device, so the cuvs-java docs should also recommend enableRMMPooledMemory or enableRMMAsyncMemory for multithreaded applications. I'm happy to open a PR with this.

#682 reports the same workspace limit inside balanced k-means.