From b6fb720acd9521ce3420d3beda70eca181a1fac4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 13:57:02 +0200 Subject: [PATCH] perf(ml-alpha): eliminate all dtoh from training hot path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/grad_norm.cu | 84 +++++++ crates/ml-alpha/cuda/mamba2_alpha_kernel.cu | 46 ++++ crates/ml-alpha/src/mamba2_block.rs | 234 ++++++++++++++++++++ crates/ml-alpha/src/trainer/optim.rs | 7 +- crates/ml-alpha/src/trainer/perception.rs | 33 ++- 6 files changed, 395 insertions(+), 10 deletions(-) create mode 100644 crates/ml-alpha/cuda/grad_norm.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 3ca33dd23..b58fe184c 100644 --- a/crates/ml-alpha/build.rs +++ b/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 diff --git a/crates/ml-alpha/cuda/grad_norm.cu b/crates/ml-alpha/cuda/grad_norm.cu new file mode 100644 index 000000000..2be0d0ad8 --- /dev/null +++ b/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 + +#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; +} diff --git a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu index a3ec29fd9..e9ac42c58 100644 --- a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu +++ b/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]. diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index bb7566409..3259f8334 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -1596,6 +1596,27 @@ pub struct Mamba2AdamW { stream: Arc, 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, + /// Accumulator across all 9 tensors — single scalar. + grad_norm_sq_d: CudaSlice, + /// 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, + /// 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::(max_grid_blocks.max(1)) + .map_err(|e| anyhow!("block_partials alloc: {e}"))?; + let grad_norm_sq_d = stream.alloc_zeros::(1) + .map_err(|e| anyhow!("grad_norm_sq alloc: {e}"))?; + let grad_scale_d = stream.alloc_zeros::(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; 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, + kernel: &CudaFunction, + cfg: &Mamba2AdamWConfig, + n: usize, + param: &mut CudaSlice, + grad: &CudaSlice, + state: &mut AdamState, + t: i32, + grad_scale_d: &CudaSlice, +) -> 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, kernel: &CudaFunction, diff --git a/crates/ml-alpha/src/trainer/optim.rs b/crates/ml-alpha/src/trainer/optim.rs index 0ef5c5a83..e6d35a900 100644 --- a/crates/ml-alpha/src/trainer/optim.rs +++ b/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(()) } } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index c391df2ff..b89b55754 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -221,6 +221,9 @@ pub struct PerceptionTrainer { grad_probs_per_k_d: CudaSlice, /// Scalar BCE loss output (mean over valid label entries). loss_d: CudaSlice, + /// 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, /// Per-horizon BCE loss weights, device-resident. Filled from @@ -356,6 +359,7 @@ impl PerceptionTrainer { grad_probs_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * N_HORIZONS)?, loss_weights_d: upload(&stream, &cfg.horizon_weights)?, loss_d: stream.alloc_zeros::(1)?, + loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?, valid_d: stream.alloc_zeros::(1)?, grad_h_carry_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, grad_h_new_d: stream.alloc_zeros::(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