Select devices and use several GPUs

Choose the CUDA device for each RayD backend, replicate a Torch scene across GPUs, and run one process per GPU.

RayD runs every query on one NVIDIA GPU. This guide shows how each backend picks that GPU and the three ways to use more than one.

Pick the device#

The Dr.Jit backend follows Dr.Jit's current-thread CUDA device. Select it before you create meshes, scenes, or arrays:

import rayd.drjit as rd
 
print(rd.device_count())    # visible CUDA devices
rd.set_device(0)            # initialize_optix=True also creates the OptiX context
print(rd.current_device())

Do not reuse scenes or OptiX resources across a device switch in the same process.

Replicate one Torch scene across GPUs#

Scene(devices=[0, 1]) builds a full copy of the scene on each device. Each query splits its batch of rays across the copies and returns results and gradients on the first device. Per-ray results are bitwise identical to a single-device run; accumulated grids match up to float32 summation order.

scene = rt.Scene(devices=[0, 1])
scene.add_mesh(rt.Mesh(vertices, faces))
scene.build()
hit = scene.intersect(ray)            # ray tensors on cuda:0; result on cuda:0

MultiDeviceOptions tunes the split:

FieldDefaultEffect
weightsequalshare of the batch per device, as a ratio
operation_weightsNoneweights for one operation family ("intersect") or exact shape ("trace_reflections:4")
require_peer_accessTruereject devices without CUDA peer access
require_homogeneous_devicesTruereject a mix of GPU models
warm_upTruepre-link OptiX pipelines on every device during build()
chunk_rays, tape_memory_budget_bytes, offloadNonerun the batch in chunks to bound memory
pipeline, pipeline_chunks_per_deviceTrue, 4overlap copies and compute
min_rays_per_device, min_lanes_per_device262144below this size, run on the first device only
opts = rt.MultiDeviceOptions(weights=[0.6, 0.4])
scene = rt.Scene(devices=[0, 1], options=opts)

To measure the split instead of guessing it, call scene.calibrate_devices() after build(). It times a probe (intersect by default, trace_reflections with max_bounces set) on each device and stores the fastest split for that operation. Read it with scene.device_weights_for("intersect"). A scene with fewer than two devices raises RuntimeError.

Splitting pays off for compute-heavy work and large accumulations. Cheap queries with wide results can be slower on two GPUs, because the rays and results cross the interconnect.

Note

add_instance is not available on a multi-device scene. Diffraction accumulation splits its Monte Carlo sample window, not a ray batch.

One process per GPU#

The Dr.Jit backend has no in-process multi-device mode. Run one process per GPU and pin each one with CUDA_VISIBLE_DEVICES. The same recipe works for Torch and across nodes, for example with torchrun.

  • Give every process its own OPTIX_CACHE_PATH so that concurrent processes do not share one OptiX disk cache.
  • To split a Monte Carlo accumulation such as accum_dfr_direct, give each rank a disjoint window of the sample space with lane_offset and lane_count, then sum the grids (for example with all_reduce). The sum equals the single-run grid up to summation order.

The RayD repository has runnable recipes under examples/torch/distributed.

Related#