perf: GPU AdamW + grad-norm

Eliminate the per-step GPU↔host roundtrip of every parameter/gradient.

- optim.cu: adamw_step (m/v on device, in-place param update), sumsq_accum
  (block-reduced global grad sum-of-squares), scale_inplace.
- GpuAdamW: device m/v state per param; step launches the kernel reading
  each param's .grad() and rewriting the param buffer in place — no host
  roundtrip. Host AdamW kept as the torch-parity reference.
- clip_grad_norm_gpu: device sum-of-squares reduction (only the scalar norm
  comes back), in-place rescale of grads by pre_scale·clip_factor.
- train_loop: use GpuAdamW + clip_grad_norm_gpu.
- test: GPU AdamW vs host reference parity (max abs err < 1e-6).

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

crates/xtrain-optim/src/lib.rs

@@ -113,7 +113,7 @@ impl AdamW {
 mod gpu {
     use super::AdamW;
     use xtrain_autodiff::tape::Var;
-    use xtrain_tensor::{Device, Tensor};
+    use xtrain_tensor::{DType, Device, Tensor};
 
     impl AdamW {
         /// Apply one AdamW step to every parameter `Var`, using `lr` for this step
@@ -125,6 +125,10 @@ mod gpu {
         ///
         /// Does NOT zero grads — the caller does that (matching the GD-step
         /// template in the T5 overfit test).
+        ///
+        /// This is the host-roundtrip reference path; training uses
+        /// [`GpuAdamW`] (kernel, m/v on device). Both are checked against the
+        /// torch parity in tests.
         pub fn step(&mut self, lr: f32, params: &[Var]) {
             let device = params[0].value().device();
             let shapes: Vec<Vec<usize>> =
@@ -151,4 +155,93 @@ mod gpu {
             }
         }
     }
+
+    /// GPU AdamW (Phase T7): the optimizer state (m/v moments) lives on the device
+    /// as one tensor pair per parameter, and the update runs as a CUDA kernel that
+    /// reads each param's `.grad()` and rewrites the param buffer in place — no
+    /// per-step GPU↔host roundtrip of params/grads. Same math as
+    /// [`AdamW::step_host`] (the parity reference).
+    pub struct GpuAdamW {
+        beta1: f32,
+        beta2: f32,
+        eps: f32,
+        weight_decay: f32,
+        t: u64,
+        /// Per-parameter (m, v) device buffers, sized lazily on first step.
+        state: Vec<(Tensor, Tensor)>,
+    }
+
+    impl GpuAdamW {
+        /// PyTorch-default betas/eps; you set lr (per-step) + weight decay.
+        pub fn new(weight_decay: f32) -> Self {
+            Self {
+                beta1: 0.9,
+                beta2: 0.999,
+                eps: 1e-8,
+                weight_decay,
+                t: 0,
+                state: Vec::new(),
+            }
+        }
+
+        pub fn step_count(&self) -> u64 {
+            self.t
+        }
+
+        /// One in-place AdamW step over every parameter `Var` at learning rate
+        /// `lr`. Updates the param value buffer and the device m/v state via the
+        /// `adamw_step_f32` kernel. Params are mutated in place, so the leaf `Var`
+        /// identities stay stable across steps (no `set_value`). Does NOT zero
+        /// grads — the caller does. A param without a grad is skipped this step.
+        pub fn step(&mut self, lr: f32, params: &[Var]) {
+            let device = params[0].value().device();
+            if self.state.is_empty() {
+                self.state = params
+                    .iter()
+                    .map(|p| {
+                        let shape = p.value().shape().to_vec();
+                        (
+                            Tensor::zeros(&shape, DType::F32, device),
+                            Tensor::zeros(&shape, DType::F32, device),
+                        )
+                    })
+                    .collect();
+            }
+            assert_eq!(self.state.len(), params.len(), "param count changed");
+
+            self.t += 1;
+            let bc1 = 1.0 - self.beta1.powi(self.t as i32);
+            let bc2 = 1.0 - self.beta2.powi(self.t as i32);
+
+            for (p, (m, v)) in params.iter().zip(&self.state) {
+                let g = match p.grad() {
+                    Some(g) => g,
+                    None => continue,
+                };
+                let pv = p.value();
+                let n = pv.numel() as i32;
+                unsafe {
+                    xtrain_cuda::ffi::launch_adamw_step_f32(
+                        pv.data_ptr() as *mut f32,
+                        g.data_ptr() as *const f32,
+                        m.data_ptr() as *mut f32,
+                        v.data_ptr() as *mut f32,
+                        lr,
+                        self.beta1,
+                        self.beta2,
+                        self.eps,
+                        self.weight_decay,
+                        bc1,
+                        bc2,
+                        n,
+                        std::ptr::null_mut(),
+                    );
+                }
+            }
+            xtrain_cuda::device::synchronize().expect("adamw step sync failed");
+        }
+    }
 }
+
+#[cfg(not(no_cuda))]
+pub use gpu::GpuAdamW;