FFI and CUDA stream capture errors

[tilmantroester]

When running CUDA code through FFI, I sometimes get these errors:

2025-02-27 20:52:23.535817: E external/xla/xla/stream_executor/cuda/cuda_command_buffer.cc:715] CUDA error: Failed to destroy CUDA graph: CUDA_ERROR_INVALID_VALUE: invalid argument
E0227 20:52:23.535843 3010756 pjrt_stream_executor_client.cc:3026] Execution of replica 0 failed: INTERNAL: CUDA error: Failed to end stream capture: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error during capture

Frustratingly, this only happens within a larger program and not for smaller test cases with just the FFI call.

Setting XLA_FLAGSG=--xla_gpu_enable_command_buffer= to disable CUDA graphs (as found here) seems to prevent this error.

My understanding of XLA and CUDA is limited, so I am a bit stuck on how to address this.

• What causes the CUDA graph capture in JAX?
• Which CUDA operations are problematic? About FFI failure: Failed to destroy GPU Graph #25141 (comment) mentions device to host memory copy and stream synchronization. In my case the offending calls seems to be cudaMallocManaged, which gives a code 900 error: operation not permitted when stream is capturing. Are there other ways to allocate memory within the FFI call but outside CUDA kernels?