[JAX] FFI initialize stage captures XLA's compute stream, swallowing work other threads enqueue on it #3863

Describe the bug

wrapInStreamCapture (transformer_engine/jax/csrc/extensions/ffi.h) begins a capture on the stream XLA passes to the FFI "initialize" stage, which is XLA's per-device compute stream. Relaxed capture mode doesn't stop other threads from enqueueing on that stream while the capture is open. Any work another host thread puts on it during a program's first execution (an H2D transfer recording its compute-stream sync event, or a small jitted program run by a data-loading thread) goes into the graph that is then destroyed. The program then fails, crashes, or the other thread's work silently never runs.

Typical error:

XLA FFI call failed: .../util/vectorized_pointwise.h:367 in function VectorizedUnaryKernelLauncher: CUDA Error: operation failed due to a previous error during capture

We have also seen a segfault in cuEventRecord under CopyRawHostToDeviceAndReturnEvent → GetEventForComputeStreamSyncPoint, and hangs.

Steps/Code to reproduce bug

One GPU (B300), TE 2.20.2, jax 0.12. A background thread continuously device_puts batches and runs a small jitted dynamic_update_slice into a resident device buffer. Meanwhile the main thread first-executes a fresh jitted program using TE layernorm → activation → an NVFP4 GEMM (forward and backward), which goes through the NormFwd/Bwd, ActLu and RHTAmax initialize handlers. Each trial runs in its own process:

Proposed fix

Capture on a TE-owned stream, not XLA's. Module loading, which the initialize stage exists for (#2219), is per context, not per stream. The five handlers wrapped today only pass the stream to nvte_* launchers (the cuDNN norm path rebinds its handle every call), so nothing depends on which stream they ran on. With the patch, NVFP4 and MXFP8 forward and backward outputs are bitwise identical to stock, on both first and second execution.

+// Captures on a TE-owned stream, never the one XLA passes: work other threads enqueue on XLA's
+// stream while the capture is open would otherwise be swallowed into the discarded graph.
+inline cudaStream_t captureStreamForCurrentDevice() {
+  thread_local std::unordered_map<int, cudaStream_t> streams;
+  int device = 0;
+  NVTE_CHECK_CUDA(cudaGetDevice(&device));
+  auto [it, inserted] = streams.try_emplace(device, nullptr);
+  if (inserted) {
+    NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&it->second, cudaStreamNonBlocking));
+  }
+  return it->second;
+}
+
 template <typename... Args>
 Error_Type wrapInStreamCapture(std::function<Error_Type(cudaStream_t, Args...)> func,
-                               cudaStream_t stream, Args... args) {
+                               cudaStream_t /*stream*/, Args... args) {
+  cudaStream_t capture = captureStreamForCurrentDevice();
   cudaGraph_t graph{};
-  NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed));
-  Error_Type error = func(stream, std::forward<Args>(args)...);
-  NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph));
+  NVTE_CHECK_CUDA(cudaStreamBeginCapture(capture, cudaStreamCaptureModeRelaxed));
+  Error_Type error = func(capture, std::forward<Args>(args)...);
+  NVTE_CHECK_CUDA(cudaStreamEndCapture(capture, &graph));
   NVTE_CHECK_CUDA(cudaGraphDestroy(graph));

#2800 adds initialize stages to thirteen more FFIs, which would widen this window to most TE programs.

Environment

TE 2.20.2 built from the PyPI sdist against jax[cuda13] 0.12.0.1, CUDA 13, B300, single process.