From 9607f33518c299ee2f01d54c2b57433c0b4ef556 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 23:57:32 +0200 Subject: [PATCH] perf(ml-alpha): block-per-row VSN bwd refactor (Phase B commit 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit variable_selection_bwd refactored from grid=(1,1,1) to grid=(B*K,1,1). VSN's n_rows = B*K positions (one row per (batch, K-position) pair); block-per-row matches the existing fwd kernel's layout. Adds 2 per-row grad scratch buffers + 2 reduce_axis0 launches: vsn_grad_w_scratch_d [B*K, FEATURE_DIM, FEATURE_DIM] vsn_grad_b_scratch_d [B*K, FEATURE_DIM] ~210 KB scratch at B=32, K=64. VSN bwd runs 1×/step (not K×) so the absolute wall-time win here is small versus commits 1+2. Done for pattern uniformity — every per-batch or per-row bwd in the trainer now uses scratch+reducer. All 9 perception_overfit smokes pass. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/variable_selection.cu | 139 ++++++++++++--------- crates/ml-alpha/src/trainer/perception.rs | 60 +++++++-- 2 files changed, 131 insertions(+), 68 deletions(-) diff --git a/crates/ml-alpha/cuda/variable_selection.cu b/crates/ml-alpha/cuda/variable_selection.cu index 3963b70d4..b86b69bfa 100644 --- a/crates/ml-alpha/cuda/variable_selection.cu +++ b/crates/ml-alpha/cuda/variable_selection.cu @@ -123,18 +123,35 @@ extern "C" __global__ void variable_selection_fwd( // ONE block per launch (n_rows internal loop), same pattern as the // 2-layer and GRN heads-bwd kernels. +// Block-per-row VSN bwd (Phase B commit 3). +// grid=(n_rows, 1, 1) block=(VSN_BLOCK, 1, 1) +// +// n_rows = B * K (one row per (batch, K-position) pair). Each block +// handles one row's FEATURE_DIM features. +// +// Per-row grad scratch tensors: +// grad_W_vsn_scratch [n_rows, FEATURE_DIM, FEATURE_DIM] +// grad_b_vsn_scratch [n_rows, FEATURE_DIM] +// Thread (row, tid) is sole writer to its slice — single launch per +// training step (variable_selection_bwd runs 1×/step, not K×). +// Caller zeroes scratch at step start; reduce_axis0 collapses n_rows +// → final grad after this kernel. +// +// grad_x is per-row indexed; block row is sole writer to its slice +// (overwrite). extern "C" __global__ void variable_selection_bwd( - const float* __restrict__ W_vsn, // [FEATURE_DIM, FEATURE_DIM] - const float* __restrict__ x, // [N_rows, FEATURE_DIM] - const float* __restrict__ gates, // [N_rows, FEATURE_DIM] saved by fwd - const float* __restrict__ grad_y, // [N_rows, FEATURE_DIM] + const float* __restrict__ W_vsn, // [FEATURE_DIM, FEATURE_DIM] + const float* __restrict__ x, // [N_rows, FEATURE_DIM] + const float* __restrict__ gates, // [N_rows, FEATURE_DIM] saved by fwd + const float* __restrict__ grad_y, // [N_rows, FEATURE_DIM] int n_rows, - float* __restrict__ grad_W_vsn, // [FEATURE_DIM, FEATURE_DIM] (+=) - float* __restrict__ grad_b_vsn, // [FEATURE_DIM] (+=) - float* __restrict__ grad_x // [N_rows, FEATURE_DIM] (overwrite) + float* __restrict__ grad_W_vsn_scratch, // [N_rows, FEATURE_DIM, FEATURE_DIM] (+=) + float* __restrict__ grad_b_vsn_scratch, // [N_rows, FEATURE_DIM] (+=) + float* __restrict__ grad_x // [N_rows, FEATURE_DIM] (overwrite) ) { + int row = blockIdx.x; int tid = threadIdx.x; - if (tid >= VSN_BLOCK) return; + if (row >= n_rows || tid >= VSN_BLOCK) return; __shared__ float s_dgates[VSN_FEATURE_DIM]; __shared__ float s_dlogit[VSN_FEATURE_DIM]; @@ -143,61 +160,61 @@ extern "C" __global__ void variable_selection_bwd( __shared__ float s_sumdot; __syncthreads(); - for (int row = 0; row < n_rows; ++row) { - // Pass 1: d_gates[i] = grad_y[i] * x[i]; cache gates[i] + x[i]. - float gates_i = 0.0f; - float x_i = 0.0f; - float dy_i = 0.0f; - if (tid < VSN_FEATURE_DIM) { - gates_i = gates[(long long)row * VSN_FEATURE_DIM + tid]; - x_i = x[(long long)row * VSN_FEATURE_DIM + tid]; - dy_i = grad_y[(long long)row * VSN_FEATURE_DIM + tid]; - s_gates[tid] = gates_i; - s_dgates[tid] = dy_i * x_i; - } - __syncthreads(); + // Per-row base offsets into scratch tensors. + const long long row_F = (long long)row * VSN_FEATURE_DIM; + const long long row_FF = row_F * VSN_FEATURE_DIM; - // Pass 2: dot_product = sum_j gates[j] * d_gates[j]. - s_red[tid] = (tid < VSN_FEATURE_DIM) ? s_gates[tid] * s_dgates[tid] : 0.0f; - __syncthreads(); - for (int s = VSN_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_sumdot = s_red[0]; - __syncthreads(); + // Pass 1: d_gates[i] = grad_y[i] * x[i]; cache gates[i] + x[i]. + float gates_i = 0.0f; + float x_i = 0.0f; + float dy_i = 0.0f; + if (tid < VSN_FEATURE_DIM) { + gates_i = gates[(long long)row * VSN_FEATURE_DIM + tid]; + x_i = x[(long long)row * VSN_FEATURE_DIM + tid]; + dy_i = grad_y[(long long)row * VSN_FEATURE_DIM + tid]; + s_gates[tid] = gates_i; + s_dgates[tid] = dy_i * x_i; + } + __syncthreads(); - // Pass 3: d_logit[i] = gates[i] * (d_gates[i] - sum_dot). - // grad_b_vsn[i] += d_logit[i]. (Thread tid sole writer.) - if (tid < VSN_FEATURE_DIM) { - const float dlogit_i = gates_i * (s_dgates[tid] - s_sumdot); - s_dlogit[tid] = dlogit_i; - grad_b_vsn[tid] += dlogit_i; - } - __syncthreads(); - - // Pass 4: grad_W_vsn[i=tid, j] += d_logit[tid] * x[row, j] - // Thread tid is sole writer of row tid. Loop over j. - if (tid < VSN_FEATURE_DIM) { - const float dl_t = s_dlogit[tid]; - #pragma unroll - for (int j = 0; j < VSN_FEATURE_DIM; ++j) { - const float xj = x[(long long)row * VSN_FEATURE_DIM + j]; - grad_W_vsn[tid * VSN_FEATURE_DIM + j] += dl_t * xj; - } - } - - // Pass 5: d_x_via_W[j=tid] = sum_i d_logit[i] * W_vsn[i, j=tid]. - // grad_x[row, j] = d_x_via_W[j] + grad_y[j] * gates[j]. - if (tid < VSN_FEATURE_DIM) { - float dx_via_W = 0.0f; - #pragma unroll - for (int i = 0; i < VSN_FEATURE_DIM; ++i) { - dx_via_W += s_dlogit[i] * W_vsn[i * VSN_FEATURE_DIM + tid]; - } - const float dx_via_gate = dy_i * gates_i; - grad_x[(long long)row * VSN_FEATURE_DIM + tid] = dx_via_W + dx_via_gate; - } + // Pass 2: dot_product = sum_j gates[j] * d_gates[j]. + s_red[tid] = (tid < VSN_FEATURE_DIM) ? s_gates[tid] * s_dgates[tid] : 0.0f; + __syncthreads(); + for (int s = VSN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; __syncthreads(); } + if (tid == 0) s_sumdot = s_red[0]; + __syncthreads(); + + // Pass 3: d_logit[i] = gates[i] * (d_gates[i] - sum_dot). + // grad_b_vsn_scratch[row, i] += d_logit[i]. (Thread tid sole writer per row.) + if (tid < VSN_FEATURE_DIM) { + const float dlogit_i = gates_i * (s_dgates[tid] - s_sumdot); + s_dlogit[tid] = dlogit_i; + grad_b_vsn_scratch[row_F + tid] += dlogit_i; + } + __syncthreads(); + + // Pass 4: grad_W_vsn_scratch[row, i=tid, j] += d_logit[tid] * x[row, j]. + if (tid < VSN_FEATURE_DIM) { + const float dl_t = s_dlogit[tid]; + #pragma unroll + for (int j = 0; j < VSN_FEATURE_DIM; ++j) { + const float xj = x[(long long)row * VSN_FEATURE_DIM + j]; + grad_W_vsn_scratch[row_FF + (long long)tid * VSN_FEATURE_DIM + j] + += dl_t * xj; + } + } + + // Pass 5: grad_x[row, j=tid] = sum_i d_logit[i] * W_vsn[i, j=tid] + grad_y[j] * gates[j]. + if (tid < VSN_FEATURE_DIM) { + float dx_via_W = 0.0f; + #pragma unroll + for (int i = 0; i < VSN_FEATURE_DIM; ++i) { + dx_via_W += s_dlogit[i] * W_vsn[i * VSN_FEATURE_DIM + tid]; + } + const float dx_via_gate = dy_i * gates_i; + grad_x[(long long)row * VSN_FEATURE_DIM + tid] = dx_via_W + dx_via_gate; + } } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 12d1e14d2..dd620bcff 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -396,6 +396,9 @@ pub struct PerceptionTrainer { grn_grad_b_main_scratch_d: CudaSlice, // [B, 5] grn_grad_w_skip_scratch_d: CudaSlice, // [B, 5, HIDDEN] grn_grad_b_skip_scratch_d: CudaSlice, // [B, 5] + // VSN per-row grad scratch (Phase B commit 3). n_rows = B * K. + vsn_grad_w_scratch_d: CudaSlice, // [B*K, FEATURE_DIM, FEATURE_DIM] + vsn_grad_b_scratch_d: CudaSlice, // [B*K, FEATURE_DIM] /// Cross-batch reducer kernel: `[B, N] → [N]` via block tree-reduce. /// Used for every per-batch grad scratch in the refactored bwd path. reduce_axis0_fn: CudaFunction, @@ -777,6 +780,11 @@ impl PerceptionTrainer { let opt_vsn_w = AdamW::new(dev, FEATURE_DIM * FEATURE_DIM, cfg.lr_cfc)?; let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?; opt_vsn_b.wd = 0.0; + // Phase B: VSN per-row grad scratch (n_rows = B * K). + let vsn_grad_w_scratch_d = stream.alloc_zeros::( + cfg.n_batch * cfg.seq_len * FEATURE_DIM * FEATURE_DIM)?; + let vsn_grad_b_scratch_d = stream.alloc_zeros::( + cfg.n_batch * cfg.seq_len * FEATURE_DIM)?; // ── Attention pool init (Phase 3) ── // Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K @@ -861,6 +869,8 @@ impl PerceptionTrainer { grn_grad_b_main_scratch_d, grn_grad_w_skip_scratch_d, grn_grad_b_skip_scratch_d, + vsn_grad_w_scratch_d, + vsn_grad_b_scratch_d, reduce_axis0_fn, _reduce_module: reduce_module, loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, @@ -1998,14 +2008,15 @@ impl PerceptionTrainer { // Param grads OVERWRITE (kernel uses += but the prior // memsets clear them at step start, so a single VSN bwd // per step writes the full accumulation cleanly). - self.stream.memset_zeros(&mut self.grad_vsn_w_d) - .map_err(|e| anyhow::anyhow!("zero grad_vsn_w: {e}"))?; - self.stream.memset_zeros(&mut self.grad_vsn_b_d) - .map_err(|e| anyhow::anyhow!("zero grad_vsn_b: {e}"))?; + // Phase B commit 3: VSN per-row grad scratch + reducer. + self.stream.memset_zeros(&mut self.vsn_grad_w_scratch_d) + .map_err(|e| anyhow::anyhow!("zero vsn_grad_w_scratch: {e}"))?; + self.stream.memset_zeros(&mut self.vsn_grad_b_scratch_d) + .map_err(|e| anyhow::anyhow!("zero vsn_grad_b_scratch: {e}"))?; { let n_rows_vsn: i32 = (b_sz * k_seq) as i32; let cfg_vsn_bwd = LaunchConfig { - grid_dim: (1, 1, 1), + grid_dim: (n_rows_vsn as u32, 1, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, }; @@ -2016,11 +2027,46 @@ impl PerceptionTrainer { .arg(&self.vsn_gates_d) .arg(self.mamba2_grads_buffers.d_x_from_in.cuda_data()) .arg(&n_rows_vsn) - .arg(&mut self.grad_vsn_w_d) - .arg(&mut self.grad_vsn_b_d) + .arg(&mut self.vsn_grad_w_scratch_d) + .arg(&mut self.vsn_grad_b_scratch_d) .arg(&mut self.vsn_grad_x_d); unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; } } + // VSN reducer: collapse n_rows (= B * K) → final grad buffers. + { + let n_rows_i = (b_sz * k_seq) as i32; + let reduce_block: u32 = 256; + let cfg_red_w = LaunchConfig { + grid_dim: ((FEATURE_DIM * FEATURE_DIM) as u32, 1, 1), + block_dim: (reduce_block, 1, 1), + shared_mem_bytes: 0, + }; + let n_tail_w = (FEATURE_DIM * FEATURE_DIM) as i32; + unsafe { + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.vsn_grad_w_scratch_d) + .arg(&n_rows_i) + .arg(&n_tail_w) + .arg(&mut self.grad_vsn_w_d); + launch.launch(cfg_red_w).context("reduce vsn_grad_w")?; + } + let cfg_red_b = LaunchConfig { + grid_dim: (FEATURE_DIM as u32, 1, 1), + block_dim: (reduce_block, 1, 1), + shared_mem_bytes: 0, + }; + let n_tail_b = FEATURE_DIM as i32; + unsafe { + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.vsn_grad_b_scratch_d) + .arg(&n_rows_i) + .arg(&n_tail_b) + .arg(&mut self.grad_vsn_b_d); + launch.launch(cfg_red_b).context("reduce vsn_grad_b")?; + } + } // ── 8c. (Phase B) Reduce CfC per-batch grad scratch → final grad buffers. // Block tree-reduce across bi; deterministic sum order.