gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(ml-alpha): block-per-row VSN bwd refactor (Phase B commit 3)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 23:57:32 +02:00
parent 5c2c3b65a8
commit 9607f33518
2 changed files with 131 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
crates/ml-alpha/cuda/variable_selection.cu
View File

@@ -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;
}
}

60
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -396,6 +396,9 @@ pub struct PerceptionTrainer {
grn_grad_b_main_scratch_d: CudaSlice<f32>, // [B, 5]
grn_grad_w_skip_scratch_d: CudaSlice<f32>, // [B, 5, HIDDEN]
grn_grad_b_skip_scratch_d: CudaSlice<f32>, // [B, 5]
// VSN per-row grad scratch (Phase B commit 3). n_rows = B * K.
vsn_grad_w_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM, FEATURE_DIM]
vsn_grad_b_scratch_d: CudaSlice<f32>, // [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::<f32>(
cfg.n_batch * cfg.seq_len * FEATURE_DIM * FEATURE_DIM)?;
let vsn_grad_b_scratch_d = stream.alloc_zeros::<f32>(
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::<f32>(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.