RAFT (pylibraft) Reference

Review: 2026-10-01. Code below is illustrative unless explicitly described as CPU-tested. GPU execution, performance, GDS, and multi-GPU behavior require validation on target hardware.

RAFT (Reusable Accelerated Functions and Tools) is a RAPIDS library of GPU-accelerated building blocks for machine learning and information retrieval. It provides low-level primitives — sparse eigensolvers, device memory management, random graph generation, and multi-GPU communication — that higher-level libraries like cuML and cuGraph are built on. Use pylibraft directly when you need these primitives without the overhead of a full ML framework.

Full documentation: https://docs.nvidia.com/raft/26.08/ Note: Vector search and clustering algorithms have been migrated to cuVS. Use cuVS for nearest neighbor search, not RAFT.

Table of Contents

1. Installation and Setup
2. Core Concepts
3. Device Memory Management
4. Sparse Eigenvalue Problems
5. Random Graph Generation
6. Multi-Node Multi-GPU with raft-dask
7. Interoperability
8. Performance Tips
9. Common Pitfalls

Installation and Setup

Use uv add in standalone examples; follow the user's existing project package manager when one is already configured.

# pylibraft (core library)
uv add --extra-index-url=https://pypi.nvidia.com "pylibraft-cu12==26.8.*"   # For CUDA 12.x
uv add --extra-index-url=https://pypi.nvidia.com "pylibraft-cu13==26.8.*"   # For CUDA 13.x

# raft-dask (multi-node multi-GPU support, optional)
uv add --extra-index-url=https://pypi.nvidia.com "raft-dask-cu12==26.8.*"   # For CUDA 12.x
uv add --extra-index-url=https://pypi.nvidia.com "raft-dask-cu13==26.8.*"   # For CUDA 13.x

pylibraft and raft-dask wheels (including the companion libraft wheel) are also published directly to PyPI, so the extra index is optional.

Verify:

import pylibraft
from pylibraft.common import DeviceResources
handle = DeviceResources()
handle.sync()
print("pylibraft is working")

Core Concepts

DeviceResources (CUDA Resource Handle)

DeviceResources manages expensive CUDA resources (streams, stream pools, library handles for cuBLAS/cuSOLVER). Create one and reuse it across multiple RAFT calls to avoid repeated allocation overhead.

from pylibraft.common import DeviceResources, Stream

# Default stream
handle = DeviceResources()

# Custom stream
stream = Stream()
handle = DeviceResources(stream)

# With a CuPy stream
import cupy
cupy_stream = cupy.cuda.Stream()
handle = DeviceResources(stream=cupy_stream.ptr)

# Always sync before reading results
handle.sync()

RAFT functions are asynchronous by default — they return immediately and work continues on the GPU. You must call handle.sync() before accessing output data on the CPU. If you don't pass a handle, RAFT allocates temporary resources internally and synchronizes before returning (convenient but slower for repeated calls).

Stream

A thin wrapper around cudaStream_t for ordering GPU operations:

from pylibraft.common import Stream

stream = Stream()
stream.sync()                  # Synchronize all work on this stream
ptr = stream.get_ptr()         # Get the raw cudaStream_t pointer (uintptr_t)

Device Memory Management

device_ndarray

device_ndarray is RAFT's lightweight GPU array type. It implements __cuda_array_interface__, making it interoperable with CuPy, Numba, PyTorch, and other GPU libraries.

from pylibraft.common import device_ndarray
import numpy as np

# Allocate empty GPU array
gpu_arr = device_ndarray.empty((1000, 50), dtype=np.float32)

# From a NumPy array (copies data to GPU)
cpu_data = np.random.rand(1000, 50).astype(np.float32)
gpu_arr = device_ndarray(cpu_data)

# Back to NumPy (copies data to CPU)
result = gpu_arr.copy_to_host()

# Properties
print(gpu_arr.shape)          # (1000, 50)
print(gpu_arr.dtype)          # float32
print(gpu_arr.c_contiguous)   # True (row-major)
print(gpu_arr.f_contiguous)   # False

Configuring Output Types

You can configure all RAFT compute APIs to return CuPy arrays or PyTorch tensors instead of device_ndarray:

import pylibraft.config

pylibraft.config.set_output_as("cupy")    # All APIs return cupy arrays
pylibraft.config.set_output_as("torch")   # All APIs return torch tensors

# Custom conversion
pylibraft.config.set_output_as(lambda arr: arr.copy_to_host())  # Return numpy

Sparse Eigenvalue Problems

eigsh — Sparse Symmetric Eigenvalue Decomposition

GPU-accelerated Lanczos method for finding eigenvalues/eigenvectors of large sparse symmetric matrices. A narrower interface than SciPy: the 26.08 implementation accepts real float32/float64 CSR values and int32/int64 indices. It does not implement SciPy generalized/shift-invert parameters such as M or sigma.

import cupy as cp
import cupyx.scipy.sparse as sp
from pylibraft.sparse.linalg import eigsh
from pylibraft.common import DeviceResources

# Create a sparse symmetric matrix (CSR format)
n = 10000
density = 0.01
A = sp.random(n, n, density=density, dtype=cp.float32, format='csr')
A = A + A.T  # Make symmetric

# Find 6 largest eigenvalues
handle = DeviceResources()
eigenvalues, eigenvectors = eigsh(A, k=6, which='LM', handle=handle)
handle.sync()

print(f"Eigenvalues shape: {eigenvalues.shape}")      # (6,)
print(f"Eigenvectors shape: {eigenvectors.shape}")     # (10000, 6)

Parameters:

• A — Sparse symmetric CSR matrix (cupyx.scipy.sparse.csr_matrix)
• k — Number of eigenvalues to compute (default: 6). Must be 1 <= k < n
• which — Which eigenvalues:
• 'LM': largest in magnitude (default)
• 'LA': largest algebraic
• 'SA': smallest algebraic
• 'SM': smallest in magnitude
• v0 — Starting vector (optional, random if None)
• ncv — Number of Lanczos vectors. Must be k + 1 < ncv < n
• maxiter — Maximum iterations
• tol — Convergence tolerance (0 = machine precision)
• seed — Random seed for reproducibility
• handle — Optional DeviceResources handle

When to use: Spectral methods (spectral clustering, graph partitioning, PageRank-like computations), dimensionality reduction on sparse data, physics simulations with large sparse Hamiltonians, structural analysis (vibration modes).

Random Graph Generation

rmat — R-MAT Graph Generation

Generates random graphs using the Recursive Matrix (R-MAT) model, commonly used for benchmarking graph algorithms with realistic structure (power-law degree distribution, community structure).

import cupy as cp
from pylibraft.random import rmat
from pylibraft.common import DeviceResources

n_edges = 100000
r_scale = 16          # log2 of source node count (2^16 = 65536 nodes)
c_scale = 16          # log2 of destination node count
theta_len = max(r_scale, c_scale) * 4

# Output: edge list as (src, dst) pairs
out = cp.empty((n_edges, 2), dtype=cp.int32)
# Probability distribution at each R-MAT level
theta = cp.tile(cp.asarray([0.57, 0.19, 0.19, 0.05], dtype=cp.float32),
                max(r_scale, c_scale))  # Each four-probability level sums to one

handle = DeviceResources()
rmat(out, theta, r_scale, c_scale, seed=42, handle=handle)
handle.sync()

print(f"Generated {n_edges} edges")
print(f"Edge list shape: {out.shape}")       # (100000, 2)
print(f"Sample edges:\n{out[:5].get()}")     # First 5 edges on CPU

When to use: Benchmarking graph algorithms, generating synthetic social/web graphs, testing graph processing pipelines at scale.

Multi-Node Multi-GPU with raft-dask

raft-dask provides a Comms class for managing NCCL and UCX communication across workers in a Dask cluster. This is the foundation for distributed GPU computing in RAPIDS.

from dask_cuda import LocalCUDACluster
from dask.distributed import Client
from raft_dask.common import Comms, local_handle

# Set up a local multi-GPU Dask cluster
cluster = LocalCUDACluster()
client = Client(cluster)

def run_on_gpu(sessionId):
    handle = local_handle(sessionId)
    # Use handle with RAFT or cuML algorithms
    return "done"

# Initialize multi-GPU communication
comms = Comms(client=client)
comms.init()

# Submit work to each GPU worker
futures = [
    client.submit(run_on_gpu, comms.sessionId, workers=[w], pure=False)
    for w in comms.worker_addresses
]

# Wait for results
from dask.distributed import wait
wait(futures, timeout=60)
results = client.gather(futures)  # Propagate worker failures, not just completion

# Clean up
comms.destroy()
client.close()
cluster.close()

Comms parameters:

• comms_p2p (bool) — Enable UCX peer-to-peer communication (default: False). Enable for algorithms that need direct GPU-to-GPU transfers.
• client — Dask distributed client
• verbose (bool) — Enable verbose logging
• streams_per_handle (int) — Number of CUDA streams per handle

Interoperability

RAFT's device_ndarray implements __cuda_array_interface__, enabling zero-copy sharing with other GPU libraries:

import cupy as cp
import torch
import numpy as np
from pylibraft.common import device_ndarray

# pylibraft -> CuPy (zero-copy)
raft_arr = device_ndarray(np.random.rand(100).astype(np.float32))
cupy_arr = cp.asarray(raft_arr)

# pylibraft -> PyTorch (zero-copy)
torch_tensor = torch.as_tensor(raft_arr, device='cuda')

# CuPy -> pylibraft (pass directly — RAFT APIs accept __cuda_array_interface__)
cupy_data = cp.random.rand(100, 50, dtype=cp.float32)
# Can pass cupy_data directly to pylibraft functions like eigsh()

# pylibraft -> NumPy (copy)
numpy_arr = raft_arr.copy_to_host()

RAFT functions accept any object implementing __cuda_array_interface__ as input — you don't need to convert to device_ndarray first. This means CuPy arrays, Numba device arrays, PyTorch CUDA tensors, and cuDF columns all work directly.

Performance Tips

1. Reuse DeviceResources. Creating a DeviceResources allocates CUDA library handles (cuBLAS, cuSOLVER). Create once, pass to all calls.

1. Batch your syncs. RAFT calls are asynchronous. Queue multiple operations before calling handle.sync() rather than syncing after each one.

1. Choose precision by residual and conditioning checks. Throughput ratios are hardware-specific; float32 is not automatically adequate for spectral problems.

1. Pre-allocate outputs. Many RAFT functions accept an out parameter. Pre-allocating avoids repeated GPU memory allocation.

1. Keep data on GPU. RAFT interoperates with CuPy, cuDF, and cuML via __cuda_array_interface__. Pass GPU arrays directly between libraries instead of round-tripping through CPU.

Common Pitfalls

• Forgetting to sync. RAFT operations are asynchronous. Reading results without calling handle.sync() gives undefined/stale data. If you omit the handle parameter, RAFT syncs internally (safe but slower).

• Using RAFT for vector search. Vector search (k-NN, IVFPQ, CAGRA, etc.) has been migrated to cuVS. RAFT no longer maintains these algorithms.

• Wrong sparse format. eigsh() requires cupyx.scipy.sparse.csr_matrix. Other sparse formats (COO, CSC) must be converted first.

• Matrix contract. The 26.08 Python implementation supports real symmetric float32/float64 CSR matrices. Despite broad Hermitian wording in the docstring, complex values raise an unsupported-dtype error. Validate symmetry, residuals and orthogonality; compare eigenspaces for repeated eigenvalues rather than vector signs.

• dtype mismatch. RAFT functions are picky about dtypes. Use float32 or float64 explicitly — don't rely on implicit conversion.