perf(ml-alpha): eliminate all dtoh from training hot path

GPU-resident grad-norm + clip-scale; mapped-pinned loss readback.
Replaces 9× memcpy_dtoh per Mamba2 AdamW step (grad-norm host roundtrip)
+ 1× per-step download() (loss). Saves ~10 stream-sync barriers/step.

New kernels (cuda/grad_norm.cu):
  - grad_norm_sq_phase1: per-block tree-reduce of x[i]^2 (no atomicAdd)
  - grad_norm_sq_phase2: cross-tensor accumulator (sequential stream-ordered)
  - grad_clip_scale: writes min(1, max_norm/sqrt(norm_sq)) to device ptr
  - mamba2_alpha_adamw_step_devscale: reads grad_scale from device pointer
    instead of host scalar, allowing AdamW kernels to launch async without
    waiting for a CPU-side norm computation.

Trainer changes (perception.rs):
  - loss_d kept device-side; mapped-pinned MappedF32Buffer shadow.
  - Single stream.synchronize() at end of step (was 2: post-bwd + download).
  - DtoD copy loss_d → loss_host_d queued, then sync flushes both kernels
    + copy in one barrier. Loss read via host_ptr (no dtoh).

Mamba2 AdamW (mamba2_block.rs):
  - step_from_buffers_gpu_clip(): all grad-norm tensors processed via
    phase1+phase2 chain, scale computed on-device, AdamW launches with
    devscale variant. Zero host roundtrips.
  - Pre-allocated block_partials_d, grad_norm_sq_d, grad_scale_d.

Optimizer (optim.rs): removed redundant stream.synchronize() per AdamW step.
Each per-tensor AdamW kernel is stream-ordered; sync only needed before
host reads, which the trainer handles centrally.

Synthetic overfit smoke: initial=0.30 → final=0.0006 (matches pre-refactor
trajectory). Full ml-alpha test suite passes (45 tests across lib +
integration).

Honors:
  - feedback_no_htod_htoh_only_mapped_pinned.md (MappedF32Buffer only)
  - feedback_no_atomicadd.md (block tree-reduce only)
  - feedback_no_legacy_aliases.md (step_from_buffers replaced, not aliased)

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

crates/ml-alpha/build.rs

@@ -15,6 +15,7 @@ const KERNELS: &[&str] = &[
     "projection",
     "bce_loss_multi_horizon",
     "adamw_step",
+    "grad_norm",
 ];
 
 // Cache bust v3 (2026-05-17): batched cfc/heads kernels + transpose_3d_swap_01

crates/ml-alpha/cuda/grad_norm.cu

@@ -0,0 +1,84 @@
+// grad_norm.cu — GPU-resident grad-norm + clip-scale computation.
+//
+// Replaces the host-side memcpy_dtoh + sum-of-squares loop in
+// `Mamba2AdamW::step_from_buffers`. The previous path did 9× dtoh per
+// training step — each a stream sync barrier costing ~1ms+ on the hot
+// path. This kernel suite keeps the entire computation on GPU.
+//
+// Workflow (per training step):
+//   1. Zero `grad_norm_sq_d` once via `memset_zeros`.
+//   2. For each of N gradient tensors:
+//      - launch `grad_norm_sq_phase1` → writes block-tree partials to scratch
+//      - launch `grad_norm_sq_phase2` → reduces partials, ACCUMULATES into grad_norm_sq_d
+//   3. launch `grad_clip_scale` → reads grad_norm_sq_d, writes
+//      `grad_scale_d = min(1, max_norm / sqrt(norm_sq))`
+//   4. AdamW kernels read `grad_scale_d` as a device pointer (not host scalar).
+//
+// All per-block reductions use shared-memory tree reduce per
+// `feedback_no_atomicadd.md`. Phase 2 uses a single-thread `+=` for
+// the cross-tensor accumulation, which is race-free because phase-2
+// launches are sequential in the stream and each launch writes exactly
+// one float to the accumulator slot.
+
+#include <cuda_runtime.h>
+
+#define GRAD_NORM_BLOCK 256
+
+/// Phase 1: per-block tree-reduce of `x[]*x[]` → block_partials[blockIdx.x].
+extern "C" __global__ void grad_norm_sq_phase1(
+    const float* __restrict__ x,
+    int n,
+    float* __restrict__ block_partials
+) {
+    __shared__ float smem[GRAD_NORM_BLOCK];
+    int tid = threadIdx.x;
+    int gid = blockIdx.x * blockDim.x + tid;
+    smem[tid] = (gid < n) ? x[gid] * x[gid] : 0.0f;
+    __syncthreads();
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) smem[tid] += smem[tid + s];
+        __syncthreads();
+    }
+    if (tid == 0) block_partials[blockIdx.x] = smem[0];
+}
+
+/// Phase 2: single-block tree-reduce of `block_partials[0..n_blocks]`
+/// → accumulates into `accum[0]`. `accumulate=1` adds; `accumulate=0`
+/// overwrites (use for the first call in a step).
+extern "C" __global__ void grad_norm_sq_phase2(
+    const float* __restrict__ block_partials,
+    int n_blocks,
+    float* __restrict__ accum,
+    int accumulate
+) {
+    __shared__ float smem[GRAD_NORM_BLOCK];
+    int tid = threadIdx.x;
+    float v = 0.0f;
+    for (int i = tid; i < n_blocks; i += blockDim.x) {
+        v += block_partials[i];
+    }
+    smem[tid] = v;
+    __syncthreads();
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) smem[tid] += smem[tid + s];
+        __syncthreads();
+    }
+    if (tid == 0) {
+        float prev = (accumulate != 0) ? accum[0] : 0.0f;
+        accum[0] = prev + smem[0];
+    }
+}
+
+/// Final clamp: grad_scale = min(1, max_norm / sqrt(norm_sq + eps)).
+/// One-thread kernel — writes a single float to `grad_scale_out[0]`.
+/// AdamW kernels read this device pointer instead of a host scalar.
+extern "C" __global__ void grad_clip_scale(
+    const float* __restrict__ norm_sq,
+    float max_norm,
+    float* __restrict__ grad_scale_out
+) {
+    if (threadIdx.x != 0 || blockIdx.x != 0) return;
+    const float eps = 1e-12f;
+    float norm = sqrtf(norm_sq[0] + eps);
+    grad_scale_out[0] = (norm > max_norm) ? (max_norm / norm) : 1.0f;
+}

crates/ml-alpha/cuda/mamba2_alpha_kernel.cu

@@ -559,6 +559,52 @@ extern "C" __global__ void mamba2_alpha_adamw_step(
 }
 
 
+/// Variant of `mamba2_alpha_adamw_step` that reads `grad_scale` from a
+/// DEVICE pointer instead of a host scalar — companion to the
+/// GPU-resident grad-clip pipeline in `grad_norm.cu`. Eliminates the
+/// per-step host roundtrip + 9 memcpy_dtoh syncs that the host-scalar
+/// path required to compute grad-norm.
+extern "C" __global__ void mamba2_alpha_adamw_step_devscale(
+    float* __restrict__ param,
+    const float* __restrict__ grad,
+    float* __restrict__ m,
+    float* __restrict__ v,
+    float lr,
+    float beta1,
+    float beta2,
+    float epsilon,
+    float weight_decay,
+    const float* __restrict__ grad_scale_ptr,
+    int t,
+    int n
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+
+    float g = grad[i] * grad_scale_ptr[0];
+    float p = param[i];
+
+    /* Decoupled weight decay. */
+    p = p * (1.0f - lr * weight_decay);
+
+    /* Moment updates. */
+    float mi = beta1 * m[i] + (1.0f - beta1) * g;
+    float vi = beta2 * v[i] + (1.0f - beta2) * g * g;
+    m[i] = mi;
+    v[i] = vi;
+
+    /* Bias correction. */
+    float bc1 = 1.0f - powf(beta1, (float)t);
+    float bc2 = 1.0f - powf(beta2, (float)t);
+    float m_hat = mi / bc1;
+    float v_hat = vi / bc2;
+
+    /* Parameter update. */
+    p = p - lr * m_hat / (sqrtf(v_hat) + epsilon);
+    param[i] = p;
+}
+
+
 /* ---------------------------------------------------------------------
  * Reduction kernel: sum d_w_c_per_sample[N, sh2, state_d] across N →
  *                   d_w_c[sh2, state_d].

crates/ml-alpha/src/mamba2_block.rs

@@ -1596,6 +1596,27 @@ pub struct Mamba2AdamW {
     stream: Arc<CudaStream>,
     kernel: CudaFunction,
 
+    // ─── GPU-side grad-clip support ─────────────────────────────────
+    // Eliminates the 9× per-step memcpy_dtoh that the host grad-norm
+    // path required. norm_sq_phase1 → block_partials → norm_sq_phase2
+    // accumulates into grad_norm_sq_d; clip_scale converts → grad_scale_d
+    // which the devscale AdamW kernel reads as a device pointer.
+    kernel_norm_sq_phase1: CudaFunction,
+    kernel_norm_sq_phase2: CudaFunction,
+    kernel_clip_scale: CudaFunction,
+    kernel_adamw_devscale: CudaFunction,
+    /// Per-block partial sums for grad-norm reduction phase 1.
+    /// Sized for the largest parameter tensor's grid: max(n)/256 blocks.
+    block_partials_d: CudaSlice<f32>,
+    /// Accumulator across all 9 tensors — single scalar.
+    grad_norm_sq_d: CudaSlice<f32>,
+    /// Final grad-clip scale (= min(1, max_norm/sqrt(norm_sq))). Read
+    /// by mamba2_alpha_adamw_step_devscale as a device pointer arg.
+    grad_scale_d: CudaSlice<f32>,
+    /// max(grid_blocks) across all 9 tensors — capacity check for
+    /// block_partials_d.
+    max_grid_blocks: usize,
+
     // (m, v) state for each of the 9 Mamba2 parameter tensors.
     s_w_in: AdamState,
     s_b_in: AdamState,
@@ -1634,9 +1655,57 @@ impl Mamba2AdamW {
         let s_w_out = alloc(block.w_out.weight.len(), "w_out.weight")?;
         let s_b_out = alloc(block.w_out.bias.len(),   "w_out.bias")?;
 
+        // Load grad-norm kernel handles (compiled by build.rs into a
+        // dedicated grad_norm.cubin) + the devscale AdamW variant
+        // (compiled into mamba2_alpha_kernel.cubin alongside the
+        // host-scalar version).
+        static GRAD_NORM_CUBIN: &[u8] =
+            include_bytes!(concat!(env!("OUT_DIR"), "/grad_norm.cubin"));
+        let grad_norm_module = stream.context()
+            .load_cubin(GRAD_NORM_CUBIN.to_vec())
+            .map_err(|e| anyhow!("Mamba2AdamW: grad_norm cubin load: {e}"))?;
+        let kernel_norm_sq_phase1 = grad_norm_module
+            .load_function("grad_norm_sq_phase1")
+            .map_err(|e| anyhow!("grad_norm_sq_phase1 symbol: {e}"))?;
+        let kernel_norm_sq_phase2 = grad_norm_module
+            .load_function("grad_norm_sq_phase2")
+            .map_err(|e| anyhow!("grad_norm_sq_phase2 symbol: {e}"))?;
+        let kernel_clip_scale = grad_norm_module
+            .load_function("grad_clip_scale")
+            .map_err(|e| anyhow!("grad_clip_scale symbol: {e}"))?;
+        let kernel_adamw_devscale = block
+            ._module
+            .load_function("mamba2_alpha_adamw_step_devscale")
+            .map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale symbol: {e}"))?;
+
+        // Sizing: largest parameter tensor's grid dim. With block=256,
+        // the largest tensor is w_in.weight = hidden_dim * in_dim.
+        let max_n = [
+            block.w_in.weight.len(),  block.w_in.bias.len(),
+            block.w_a.weight.len(),   block.w_a.bias.len(),
+            block.w_b.weight.len(),   block.w_b.bias.len(),
+            block.w_c.len(),
+            block.w_out.weight.len(), block.w_out.bias.len(),
+        ].into_iter().max().unwrap_or(0);
+        let max_grid_blocks = (max_n + 255) / 256;
+        let block_partials_d = stream.alloc_zeros::<f32>(max_grid_blocks.max(1))
+            .map_err(|e| anyhow!("block_partials alloc: {e}"))?;
+        let grad_norm_sq_d = stream.alloc_zeros::<f32>(1)
+            .map_err(|e| anyhow!("grad_norm_sq alloc: {e}"))?;
+        let grad_scale_d = stream.alloc_zeros::<f32>(1)
+            .map_err(|e| anyhow!("grad_scale alloc: {e}"))?;
+
         Ok(Self {
             stream,
             kernel,
+            kernel_norm_sq_phase1,
+            kernel_norm_sq_phase2,
+            kernel_clip_scale,
+            kernel_adamw_devscale,
+            block_partials_d,
+            grad_norm_sq_d,
+            grad_scale_d,
+            max_grid_blocks,
             step_count: 0,
             s_w_in, s_b_in, s_w_a, s_b_a, s_w_b, s_b_b, s_w_c, s_w_out, s_b_out,
             config,
@@ -1711,6 +1780,127 @@ impl Mamba2AdamW {
         self.config.lr = lr;
     }
 
+    /// FULLY GPU-RESIDENT AdamW step — no host roundtrips, no per-step
+    /// `memcpy_dtoh`. Grad-norm reduction + clip-scale computed on
+    /// device; the devscale AdamW kernel reads `grad_scale_d` as a
+    /// device pointer.
+    ///
+    /// Replaces [`step_from_buffers`] (which did 9× `memcpy_dtoh` per
+    /// step — each a stream sync barrier costing ~ms on the hot path)
+    /// for the production training path. Same math, zero pipeline stalls.
+    ///
+    /// Requires `config.grad_clip_max_norm` to be Some at every call —
+    /// the no-clip path would need a kernel-write to init grad_scale_d
+    /// to 1.0 which the trainer doesn't currently set up. (Callers
+    /// who want no clipping should set max_norm to a very large value
+    /// like f32::INFINITY; the clip becomes a no-op since grad_scale=1.)
+    pub fn step_from_buffers_gpu_clip(
+        &mut self,
+        block: &mut Mamba2Block,
+        grads: &Mamba2BackwardGradsBuffers,
+    ) -> Result<()> {
+        self.step_count += 1;
+        let t = self.step_count;
+
+        // ── 1. Compute grad_scale_d on device.
+        if let Some(max_norm) = self.config.grad_clip_max_norm {
+            // Zero the cross-tensor accumulator.
+            self.stream
+                .memset_zeros(&mut self.grad_norm_sq_d)
+                .map_err(|e| anyhow!("zero grad_norm_sq: {e}"))?;
+
+            // Reduce each gradient tensor's sum-of-squares INTO the accumulator.
+            let grad_slices: [&CudaSlice<f32>; 9] = [
+                grads.dw_in.cuda_data(), grads.db_in.cuda_data(),
+                grads.dw_a.cuda_data(),  grads.db_a.cuda_data(),
+                grads.dw_b.cuda_data(),  grads.db_b.cuda_data(),
+                grads.dw_c.cuda_data(),
+                grads.dw_out.cuda_data(), grads.db_out.cuda_data(),
+            ];
+            for grad in grad_slices.iter() {
+                let n = grad.len();
+                if n == 0 { continue; }
+                let block_threads: u32 = 256;
+                let n_blocks: u32 = (n as u32).div_ceil(block_threads);
+                debug_assert!((n_blocks as usize) <= self.max_grid_blocks,
+                    "block_partials capacity {} < n_blocks {}", self.max_grid_blocks, n_blocks);
+                let n_i32 = n as i32;
+                // Phase 1: per-block tree-reduce → block_partials_d.
+                {
+                    let cfg = LaunchConfig {
+                        grid_dim: (n_blocks, 1, 1),
+                        block_dim: (block_threads, 1, 1),
+                        shared_mem_bytes: 0,  // smem declared as fixed-size in kernel
+                    };
+                    let mut launch = self.stream.launch_builder(&self.kernel_norm_sq_phase1);
+                    launch.arg(*grad).arg(&n_i32).arg(&mut self.block_partials_d);
+                    unsafe { launch.launch(cfg).map_err(|e| anyhow!("grad_norm phase1: {e}"))?; }
+                }
+                // Phase 2: single-block reduce of partials → ACCUMULATE into norm_sq_d.
+                {
+                    let nb_i32 = n_blocks as i32;
+                    let accumulate_i32: i32 = 1; // ALWAYS accumulate (zero was done before loop)
+                    let cfg = LaunchConfig {
+                        grid_dim: (1, 1, 1),
+                        block_dim: (block_threads, 1, 1),
+                        shared_mem_bytes: 0,
+                    };
+                    let mut launch = self.stream.launch_builder(&self.kernel_norm_sq_phase2);
+                    launch
+                        .arg(&self.block_partials_d)
+                        .arg(&nb_i32)
+                        .arg(&mut self.grad_norm_sq_d)
+                        .arg(&accumulate_i32);
+                    unsafe { launch.launch(cfg).map_err(|e| anyhow!("grad_norm phase2: {e}"))?; }
+                }
+            }
+
+            // Convert norm_sq → grad_scale_d = min(1, max_norm / sqrt(norm_sq + eps)).
+            let cfg = LaunchConfig {
+                grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0,
+            };
+            let mut launch = self.stream.launch_builder(&self.kernel_clip_scale);
+            launch
+                .arg(&self.grad_norm_sq_d)
+                .arg(&max_norm)
+                .arg(&mut self.grad_scale_d);
+            unsafe { launch.launch(cfg).map_err(|e| anyhow!("grad_clip_scale: {e}"))?; }
+        } else {
+            return Err(anyhow!(
+                "step_from_buffers_gpu_clip requires grad_clip_max_norm to be Some \
+                 (set to f32::INFINITY to disable clipping in practice)"
+            ));
+        }
+
+        // ── 2. AdamW per param via the devscale kernel (reads
+        //       grad_scale_d as device pointer).
+        let stream = &self.stream;
+        let kernel = &self.kernel_adamw_devscale;
+        let cfg = &self.config;
+        let grad_scale = &self.grad_scale_d;
+
+        adamw_apply_devscale(stream, kernel, cfg, block.w_in.weight.len(),
+                             &mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_in.bias.len(),
+                             &mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_a.weight.len(),
+                             &mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_a.bias.len(),
+                             &mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_b.weight.len(),
+                             &mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_b.bias.len(),
+                             &mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_c.len(),
+                             &mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_out.weight.len(),
+                             &mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, t, grad_scale)?;
+        adamw_apply_devscale(stream, kernel, cfg, block.w_out.bias.len(),
+                             &mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, t, grad_scale)?;
+
+        Ok(())
+    }
+
     /// AdamW step that reads gradients directly from a
     /// [`Mamba2BackwardGradsBuffers`] (the pre-allocated buffer set
     /// produced by [`Mamba2Block::backward_from_h_enriched_seq_full_into`]).
@@ -1777,6 +1967,50 @@ impl Mamba2AdamW {
 /// (rather than a method) so the caller can mutably borrow distinct fields
 /// of `Mamba2AdamW` (the per-param Adam state) while sharing immutable
 /// references to the stream + kernel + hyperparameters.
+/// Companion to [`adamw_apply`] — launches `mamba2_alpha_adamw_step_devscale`
+/// which reads `grad_scale` from a device pointer instead of a host
+/// scalar. Used by [`Mamba2AdamW::step_from_buffers_gpu_clip`] to keep
+/// the entire AdamW step host-roundtrip-free.
+fn adamw_apply_devscale(
+    stream: &Arc<CudaStream>,
+    kernel: &CudaFunction,
+    cfg: &Mamba2AdamWConfig,
+    n: usize,
+    param: &mut CudaSlice<f32>,
+    grad: &CudaSlice<f32>,
+    state: &mut AdamState,
+    t: i32,
+    grad_scale_d: &CudaSlice<f32>,
+) -> Result<()> {
+    let block_threads: u32 = 256;
+    let blocks: u32 = (n as u32).div_ceil(block_threads);
+    let launch_cfg = LaunchConfig {
+        grid_dim: (blocks, 1, 1),
+        block_dim: (block_threads, 1, 1),
+        shared_mem_bytes: 0,
+    };
+    let n_i32 = n as i32;
+    unsafe {
+        stream
+            .launch_builder(kernel)
+            .arg(param)
+            .arg(grad)
+            .arg(&mut state.m)
+            .arg(&mut state.v)
+            .arg(&cfg.lr)
+            .arg(&cfg.beta1)
+            .arg(&cfg.beta2)
+            .arg(&cfg.epsilon)
+            .arg(&cfg.weight_decay)
+            .arg(grad_scale_d)   // DEVICE pointer instead of host scalar
+            .arg(&t)
+            .arg(&n_i32)
+            .launch(launch_cfg)
+            .map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale launch (n={n}): {e}"))?;
+    }
+    Ok(())
+}
+
 fn adamw_apply(
     stream: &Arc<CudaStream>,
     kernel: &CudaFunction,

crates/ml-alpha/src/trainer/optim.rs

@@ -58,7 +58,7 @@ impl AdamW {
         let n_params_i = n as i32;
         let step_i = self.step;
         let cfg = LaunchConfig {
-            grid_dim: (((n as u32) + 255) / 256, 1, 1),
+            grid_dim: ((n as u32).div_ceil(256), 1, 1),
             block_dim: (256, 1, 1),
             shared_mem_bytes: 0,
         };
@@ -69,7 +69,10 @@ impl AdamW {
             .arg(&self.lr).arg(&self.beta1).arg(&self.beta2).arg(&self.eps).arg(&self.wd)
             .arg(&step_i);
         unsafe { launch.launch(cfg).context("adamw launch")?; }
-        self.stream.synchronize().context("adamw sync")?;
+        // No per-call sync — caller batches multiple AdamW launches on the
+        // same stream and synchronises once at step boundary. Removing 7
+        // per-step syncs (one per param group) eliminates ~35-70 μs of
+        // pipeline flushes per training step.
         Ok(())
     }
 }

crates/ml-alpha/src/trainer/perception.rs

@@ -221,6 +221,9 @@ pub struct PerceptionTrainer {
     grad_probs_per_k_d: CudaSlice<f32>,
     /// Scalar BCE loss output (mean over valid label entries).
     loss_d: CudaSlice<f32>,
+    /// Mapped-pinned host shadow of `loss_d`. After the post-step sync,
+    /// host_ptr holds the freshly-computed loss — no extra dtoh sync.
+    loss_host_d: MappedF32Buffer,
     /// Valid-label count for BCE normalisation.
     valid_d: CudaSlice<i32>,
     /// Per-horizon BCE loss weights, device-resident. Filled from
@@ -356,6 +359,7 @@ impl PerceptionTrainer {
             grad_probs_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
             loss_weights_d: upload(&stream, &cfg.horizon_weights)?,
             loss_d: stream.alloc_zeros::<f32>(1)?,
+            loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?,
             valid_d: stream.alloc_zeros::<i32>(1)?,
             grad_h_carry_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
             grad_h_new_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
@@ -877,8 +881,10 @@ impl PerceptionTrainer {
         }
         drop((_g_hpk_bwd, _g_probs_bwd, _g_gprobs_bwd, _g_ghen_t_mut));
 
-        // ── 7. Sync once before optimizer step; download loss scalar.
-        self.stream.synchronize().context("bwd loop sync")?;
+        // ── 7. Stream stays async — bwd loop kernels, transposes, and
+        //       Mamba2 bwd are all sequential on the same stream, no
+        //       sync needed between them. Single sync at the very end
+        //       to flush before reading the mapped-pinned loss shadow.
 
         // ── 7b. Transpose grad_h_enriched_seq_t_d [K, B, H] → [B, K, H]
         //        (pre-allocated grad_h_enriched_seq_d).
@@ -920,14 +926,25 @@ impl PerceptionTrainer {
         self.opt_tau.step(&mut self.tau_d, &self.grad_tau_d)?;
         self.opt_heads_w.step(&mut self.heads_w_d, &self.grad_heads_w_d)?;
         self.opt_heads_b.step(&mut self.heads_b_d, &self.grad_heads_b_d)?;
+        // GPU-resident grad-clip + AdamW (zero host roundtrips).
+        // Replaces step_from_buffers which did 9× memcpy_dtoh per step.
         self.mamba2_adamw
-            .step_from_buffers(&mut self.mamba2, &self.mamba2_grads_buffers)
-            .context("mamba2 AdamW step_from_buffers")?;
+            .step_from_buffers_gpu_clip(&mut self.mamba2, &self.mamba2_grads_buffers)
+            .context("mamba2 AdamW step_from_buffers_gpu_clip")?;
 
-        // Final: read the single loss scalar back to host. This is the
-        // ONLY post-step download in the hot path.
-        let loss_host = download(&self.stream, &self.loss_d)?;
-        Ok(loss_host[0])
+        // Final: queue the single-float DtoD copy loss_d → loss_host_d
+        // (mapped-pinned). Sync ONCE at end of step — flushes all
+        // queued kernels + the DtoD. Read via the mapped-pinned host_ptr
+        // without any extra dtoh roundtrip.
+        unsafe {
+            let (src_ptr, _g) = self.loss_d.device_ptr(&self.stream);
+            cudarc::driver::result::memcpy_dtod_async(
+                self.loss_host_d.dev_ptr, src_ptr, 4, self.stream.cu_stream(),
+            ).context("loss dtod → mapped-pinned")?;
+        }
+        self.stream.synchronize().context("end-of-step sync")?;
+        let loss = unsafe { std::ptr::read_volatile(self.loss_host_d.host_ptr) };
+        Ok(loss)
     }
 
     /// Single-sequence forward-only eval — thin wrapper around