cuda: infrastructure for whole-step CUDA graph capture

- Thread-local launch stream (xserv_cuda::stream): every kernel
  wrapper, cublasSetStream, and NCCL collective now launches on
  current_stream_raw() — the legacy null stream by default (behavior
  unchanged), or the capture stream installed via push_stream during
  graph capture. Capture is impossible on the legacy stream.
- Allocator retain mode: blocks freed inside a retain window are
  quarantined (RetainedBlocks) instead of pooled, so an instantiated
  graph keeps exclusive ownership of every intermediate buffer it
  references across replays.
- Capture mode GLOBAL -> THREAD_LOCAL: concurrent TP rank threads
  must not poison each other's captures with their own cudaMallocs.
- embedding_device_ids / rope_inplace_device_pos: variants reading
  token ids / positions from persistent device buffers, replacing the
  per-call host upload that a captured region cannot contain.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

crates/xserv-distributed/src/lib.rs

@@ -14,10 +14,13 @@ use xserv_cuda::GpuBuffer;
 
 pub use ffi::NcclUniqueId as UniqueId;
 
-/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
-/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
-/// after the producing matmul and before the consuming kernel — no extra sync.
-const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();
+/// NCCL is issued on the thread's current launch stream (legacy null stream
+/// by default, the capture stream during CUDA graph capture). The model's
+/// kernels run on the same stream, so AllReduce stays correctly ordered after
+/// the producing matmul and before the consuming kernel — no extra sync.
+fn launch_stream() -> xserv_cuda::ffi::CudaStream {
+    xserv_cuda::stream::current_stream_raw()
+}
 
 /// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
 /// to all ranks out-of-band (e.g. via a shared variable across threads).
@@ -80,7 +83,7 @@ impl TpContext {
                     ffi::NCCL_BF16,
                     ffi::NCCL_SUM,
                     self.comm,
-                    NULL_STREAM,
+                    launch_stream(),
                 )
             },
             "ncclAllReduce",
@@ -135,7 +138,7 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) },
+            unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
             "ncclSend",
         );
     }
@@ -146,7 +149,7 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) },
+            unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
             "ncclRecv",
         );
     }