GPU Acceleration Guide

Overview

pytcl.gpu accelerates tracking workloads by processing many tracks at once. Instead of running one filter per call, the batch functions take arrays with a leading track dimension – states of shape (n_tracks, state_dim), covariances of shape (n_tracks, state_dim, state_dim) – and advance every track in a single device operation.

Two backends are supported behind one API:

• CuPy – NVIDIA GPUs (CUDA), computes in float64

• MLX – Apple Silicon (unified memory), computes in float32

The backend is selected automatically. If neither is installed, everything falls back to NumPy on the CPU, so code written against pytcl.gpu runs anywhere.

Note

The batch functions accept plain NumPy arrays and move them to the device themselves. Use to_gpu() / to_cpu() when you want to control transfers explicitly, e.g. to keep intermediate results on the device across many steps.

Installation

NVIDIA GPU (CUDA):

pip install nrl-tracker[gpu]
# or directly, matching your CUDA version:
pip install cupy-cuda12x

Apple Silicon (MLX):

pip install nrl-tracker[gpu-apple]
# or directly:
pip install mlx

Check what you have:

from pytcl.gpu import get_backend, is_cupy_available, is_gpu_available, is_mlx_available

print("GPU available:", is_gpu_available())
print("Backend:      ", get_backend())
print("MLX:          ", is_mlx_available())
print("CuPy:         ", is_cupy_available())

Output on an Apple Silicon machine with MLX installed:

GPU available: True
Backend:       mlx
MLX:           True
CuPy:          False

Quick Start: Batch Linear Kalman Filter

Advance 1,000 constant-velocity tracks through one predict-update cycle. F, Q, H, and R may be shared across the batch (2-D, as here) or given per track as (n_tracks, dim, dim) stacks:

import numpy as np

from pytcl.gpu import batch_kf_predict, batch_kf_update, to_cpu, to_gpu

rng = np.random.default_rng(0)
n_tracks = 1000

# Constant-velocity model in 2D: state [x, vx, y, vy]
dt = 1.0
F = np.array(
    [[1, dt, 0, 0], [0, 1, 0, 0], [0, 0, 1, dt], [0, 0, 0, 1]], dtype=float
)
Q = np.eye(4) * 0.01
H = np.array([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=float)  # measure position
R = np.eye(2) * 0.1

x = rng.normal(size=(n_tracks, 4))
P = np.tile(np.eye(4), (n_tracks, 1, 1))
z = rng.normal(size=(n_tracks, 2))

# One transfer in ...
x_gpu, P_gpu = to_gpu(x), to_gpu(P)

# ... all tracks advance in parallel on the device ...
pred = batch_kf_predict(x_gpu, P_gpu, F, Q)
upd = batch_kf_update(pred.x, pred.P, z, H, R)

# ... one transfer out.
x_new = to_cpu(upd.x)
P_new = to_cpu(upd.P)

print("Updated states:     ", x_new.shape)
print("Updated covariances:", P_new.shape)
print("Innovations:        ", to_cpu(upd.y).shape)
print("Likelihoods:        ", to_cpu(upd.likelihood).shape)

Updated states:      (1000, 4)
Updated covariances: (1000, 4, 4)
Innovations:         (1000, 2)
Likelihoods:         (1000,)

The update result is a named tuple with x, P, y (innovations), S (innovation covariances), K (gains), and likelihood – one entry per track, ready for gating and association. There is also batch_kf_predict_update() for a fused step, and a stateful CuPyKalmanFilter class wrapping the same operations.

Nonlinear Filters: Batched Callbacks

The batch EKF and UKF take user callables, and those callables receive the whole batch as a single device array – they are called once per step, not once per track:

• f(x) and h(x) take (n_tracks, state_dim) and return (n_tracks, out_dim);

• F_jacobian(x) and H_jacobian(x) take (n_tracks, state_dim) and return (n_tracks, out_dim, state_dim).

Write callables against get_array_module() instead of NumPy directly, so the same code runs on MLX, CuPy, or the NumPy fallback. (Mixing a host NumPy array into a CuPy expression raises TypeError.)

import numpy as np

from pytcl.gpu import batch_ekf_predict, batch_ekf_update, get_array_module, to_cpu


def f(x):
    # Nearly constant velocity with mild drag on the velocity component.
    xp = get_array_module(x)
    return xp.stack([x[:, 0] + x[:, 1], 0.99 * x[:, 1]], axis=1)


def F_jac(x):
    # Constant Jacobian, broadcast over the batch: (n_tracks, 2, 2).
    xp = get_array_module(x)
    F = xp.array([[1.0, 1.0], [0.0, 0.99]])
    return xp.broadcast_to(F, (x.shape[0], 2, 2))


def h(x):
    # Range measurement: (n_tracks, 2) -> (n_tracks, 1).
    xp = get_array_module(x)
    return xp.sqrt(x[:, 0:1] ** 2 + 1.0)


def H_jac(x):
    # (n_tracks, 1, 2)
    xp = get_array_module(x)
    r = xp.sqrt(x[:, 0:1] ** 2 + 1.0)
    zeros = xp.zeros_like(r)
    return xp.stack([x[:, 0:1] / r, zeros], axis=2)


rng = np.random.default_rng(1)
n_tracks = 500
x = rng.normal(size=(n_tracks, 2))
P = np.tile(np.eye(2), (n_tracks, 1, 1))
Q = np.eye(2) * 0.01
R = np.array([[0.1]])
z = rng.normal(loc=1.5, size=(n_tracks, 1))

pred = batch_ekf_predict(x, P, f, F_jac, Q)
upd = batch_ekf_update(pred.x, pred.P, z, h, H_jac, R)

print("Predicted states:", to_cpu(pred.x).shape)
print("Updated states:  ", to_cpu(upd.x).shape)
print("Kalman gains:    ", to_cpu(upd.K).shape)

Predicted states: (500, 2)
Updated states:   (500, 2)
Kalman gains:     (500, 2, 1)

Pass None for a Jacobian argument to have it computed by finite differences on the device.

Unscented filter: batch_ukf_predict() and batch_ukf_update() use the same batched f/h contract (no Jacobians needed). One MLX-specific caveat: the Merwe sigma-point weights scale as 1 / alpha**2, and the conventional default alpha=1e-3 gives weights of order 1e6 – unresolvable in float32. The library emits a RuntimeWarning below alpha=1e-2; on MLX pass alpha of 0.1 or larger:

from pytcl.gpu import batch_ukf_predict, batch_ukf_update

pred = batch_ukf_predict(x, P, f, Q, alpha=0.5)
upd = batch_ukf_update(pred.x, pred.P, z, h, R, alpha=0.5)

Particle Filters

CuPyParticleFilter (the name is historical; it runs on either backend) keeps its particle set on the device across predict, update, and resample. The dynamics callable receives all particles as one (n_particles, state_dim) backend array; the likelihood callable receives the particles and one measurement and returns per-particle likelihoods.

import numpy as np

from pytcl.gpu import (
    CuPyParticleFilter,
    get_array_module,
    gpu_effective_sample_size,
    gpu_normalize_weights,
    gpu_resample_systematic,
)


def dynamics(particles):
    # Receives the whole particle set (n_particles, state_dim) as a
    # backend array; returns the propagated set with the same shape.
    return particles * 0.99


def likelihood(particles, measurement):
    # Backend-agnostic: get_array_module returns mlx.core, cupy, or numpy.
    xp = get_array_module(particles)
    diff = particles[:, 0] - measurement
    return xp.exp(-0.5 * diff**2)


np.random.seed(7)  # initialize() samples the prior with NumPy
pf = CuPyParticleFilter(n_particles=10000, state_dim=2)
pf.initialize(np.zeros(2), np.eye(2))
pf.predict(dynamics)
log_lik = pf.update(0.5, likelihood)

print("Estimate shape:", pf.get_estimate().shape)
print("ESS:           ", round(pf.get_ess(), 1))

# The helpers also work standalone on plain NumPy or device arrays:
weights = np.full(10000, 1.0 / 10000)
print("Standalone ESS:", round(gpu_effective_sample_size(weights), 1))
idx = gpu_resample_systematic(weights, seed=0)
print("Resample index:", idx.shape, idx.dtype)
w_norm, log_sum = gpu_normalize_weights(np.log(weights))
print("Normalized sum:", round(float(w_norm.sum()), 6))

Output on MLX:

Estimate shape: (2,)
ESS:            8402.5
Standalone ESS: 10000.0
Resample index: (10000,) mlx.core.int32
Normalized sum: 1.0

Resampling is automatic when the effective sample size drops below resample_threshold * n_particles; choose the scheme with resample_method ("systematic", "stratified", or "multinomial"). gpu_resample_stratified() and gpu_resample_multinomial() mirror the systematic helper. batch_particle_filter_update() updates many independent particle filters (shape (n_filters, n_particles, state_dim)) in one call.

Device Utilities

Transfers and introspection:

• to_gpu() / to_cpu() – move arrays to and from the active backend; both accept arrays that are already where they belong

• ensure_gpu_array() – like to_gpu but with a dtype guarantee

• get_array_module() – returns mlx.core, cupy, or numpy for a given array, for backend-agnostic callables

• sync_gpu() – block until queued device work completes; required for honest timing, since both backends are lazy or asynchronous

Memory:

from pytcl.gpu import clear_gpu_memory, get_gpu_memory_info, sync_gpu

info = get_gpu_memory_info()
print("Backend:", info["backend"])
print("Used bytes:", info["used"])

sync_gpu()  # block until queued device work completes (for timing)
clear_gpu_memory()  # release cached device memory

Backend: mlx
Used bytes: 0

On MLX the dictionary reports allocator state (used, peak, cache); free and total are -1 because unified memory has no separate device pool. On CuPy it reports the device’s free and total plus the CuPy memory pool’s usage. get_memory_pool() returns a MemoryPool manager wrapping the backend allocator, with get_stats(), set_limit(), and free_all().

Linear algebra helpers that run on the device and fall back transparently: gpu_cholesky(), gpu_cholesky_safe() (regularizing, returns a success flag), gpu_inv(), gpu_solve(), gpu_qr(), gpu_eigh(), and gpu_matrix_sqrt().

Measured Performance

The only benchmark we publish is one we have actually run. Conditions: Apple Silicon, MLX backend, batch linear Kalman predict+update versus a per-track CPU loop over the reference implementation, timed end-to-end including host-device transfers and result materialization, after warm-up (August 2026):

Batch size	Speedup vs per-track CPU loop
100 tracks	1.6x
1,000 tracks	13x
20,000 tracks	40x

The shape of that curve is the real lesson: the device does not make one filter step faster, it makes many filter steps simultaneous. At 100 tracks the fixed cost of dispatch and transfer eats most of the win; by 20,000 tracks it is negligible.

CuPy: correctness of the CuPy backend is validated against the CPU reference on real NVIDIA hardware – the recorded run covered all 85 CuPy-gated tests on a CUDA 13.0 / RTX 5090 host (see CHANGELOG.md’s v2.1.0 entry; an earlier pre-2.0.0 manual run on an RTX 5080/CUDA 12 host exercised the layer but left no artifact). We have not measured CuPy speedups, so this guide quotes none. Expect the same qualitative behavior – batch size pays for transfer overhead – and profile your own workload.

When the GPU helps:

• Hundreds to tens of thousands of tracks stepped together

• Particle filters with large particle counts

• Pipelines that keep data on the device across many steps

When it does not:

• A single track, or a handful – the CPU filters in pytcl.dynamic_estimation will be faster

• Per-step round-trips: converting to NumPy after every update discards the batching advantage

• Anything outside this module: pytcl.gpu accelerates batch Kalman, EKF, UKF, and particle filtering only. Assignment algorithms and coordinate conversions are CPU code paths and gain nothing from installing a GPU backend.

Precision on MLX

MLX computes in float32 and raises on float64 GPU operations, so the MLX backend converts inputs to float32 throughout. Consequences:

• Batch results match the CPU reference implementations to roughly float32 precision – measured about 1e-7 relative error for the linear and extended Kalman filters – rather than to machine epsilon.

• The UKF is the sensitive case: keep alpha at 0.1 or larger on MLX (see above). CuPy computes in float64 and has no such restriction.

• MLX linear-algebra kernels (inv, cholesky, solve, eigh) run on the CPU stream; the dispatch layer handles this transparently, and unified memory makes it a scheduling change rather than a copy.

Troubleshooting

“No GPU available” from to_gpu

No backend is installed (or you are on hardware without one). The batch functions themselves still work – they fall back to NumPy – but explicit to_gpu calls require a backend.

ImportError: libcublas.so.12 (or another lib…so.12)

Your system CUDA is 13.x (or missing entirely) and only cupy-cuda12x itself is installed. The [gpu] extra ships the CUDA 12 runtime libraries as pip wheels on Linux, so reinstalling with pip install nrl-tracker[gpu] resolves this; on Windows, install a system CUDA 12.8+ toolkit. NVRTC older than 12.8 also cannot compile kernels for Blackwell (RTX 50-series) GPUs – the pinned wheels cover that case too.

Slower than the CPU

Almost always one of: the batch is too small to amortize transfer overhead, or the loop transfers to NumPy every step. Keep results as backend arrays between steps and convert once at the end. When timing, call sync_gpu() before reading the clock – both backends queue work asynchronously, so un-synchronized timings measure dispatch, not compute.

Out of device memory

Process the track set in chunks along the batch dimension, and call clear_gpu_memory() between chunks if the allocator cache grows.

See Also

• Performance Optimization Guide - CPU optimization techniques

• Kalman Filter Tuning Guide - Filter tuning and diagnostics

• CuPy Documentation

• MLX Documentation