cuda: device caching allocator (pool GpuBuffer alloc)

Every tape op allocates its output via Tensor::zeros -> GpuBuffer::alloc ->
cudaMalloc, a synchronous process-serialized driver call. Under the single-
process thread-per-GPU DDP model the rank threads' hundreds of per-step allocs
serialize through the driver (KI-5 root cause); it costs single-GPU too.

Add a per-device, size-classed caching pool: GpuBuffer::alloc serves from a
free-list (request rounded up to a size class so repeating training shapes
reuse buffers), only cudaMalloc on a miss; Drop returns the buffer to the pool
instead of cudaFree. Thread-safe via a global registry keyed by device id with
each device's free-list behind its own Mutex (registry lock held only to clone
out the per-device Arc<Mutex<_>>, so rank threads don't contend across devices).
The buffer records its alloc-time device so Drop returns to the right pool.

Transparent: physical capacity may be rounded up, but len()/memset/copy bounds
all use the requested length, so the rounded tail is never read and numerics are
unchanged. zeros() still memsets (reused buffers hold stale bytes).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-cuda/src/memory.rs

@@ -1,18 +1,37 @@
 use crate::error::{self, Result};
 use crate::ffi;
+use crate::pool;
 
-/// RAII wrapper around a GPU memory allocation. Dropping frees the memory.
+/// RAII wrapper around a GPU memory allocation. Dropping returns the buffer to
+/// the per-device caching pool (see [`crate::pool`]) for reuse instead of
+/// calling `cudaFree`.
+///
+/// `len` is the logical (requested) length used for all copy/memset bounds and
+/// exposed via [`GpuBuffer::len`]; `cap` is the physical size class the pool
+/// rounded up to (>= `len`), used only to bucket the buffer for reuse. The
+/// extra `cap - len` bytes are never exposed to callers, so pooling is
+/// numerically transparent. `device` records which device pool to return to.
 pub struct GpuBuffer {
     ptr: *mut u8,
     len: usize,
+    cap: usize,
+    device: i32,
 }
 
 impl GpuBuffer {
+    /// Allocate at least `len` bytes on the calling thread's current device,
+    /// reusing a pooled buffer when one of the matching size class is free.
+    /// The contents are **uninitialized** (a reused buffer holds stale bytes);
+    /// callers that need zeros must memset (see [`crate::Storage::zeros`]).
     pub fn alloc(len: usize) -> Result<Self> {
         assert!(len > 0, "cannot allocate 0 bytes on GPU");
-        let mut ptr = std::ptr::null_mut();
-        error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
-        Ok(Self { ptr, len })
+        let a = pool::acquire(len)?;
+        Ok(Self {
+            ptr: a.ptr,
+            len,
+            cap: a.cap,
+            device: a.device,
+        })
     }
 
     pub fn len(&self) -> usize {
@@ -56,9 +75,10 @@ impl GpuBuffer {
 
 impl Drop for GpuBuffer {
     fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            unsafe { ffi::cudaFree(self.ptr) };
-        }
+        // Return to the device pool for reuse (no cudaFree). The pool retains
+        // the raw pointer for the process lifetime; on process exit the OS
+        // reclaims the device context, so this is not a leak.
+        pool::release(self.ptr, self.device, self.cap);
     }
 }