From a478ba3d8428c5fd394f7eb03a0c99712825898f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 00:01:31 +0200 Subject: [PATCH] perf(ml-alpha): block-per-batch attention pool bwd refactor (Phase B commit 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit attention_pool_bwd refactored from grid=(1,1,1) to grid=(B,1,1). The existing per-batch grad_ln_out writes were already uniquely indexed; only grad_Q needed scratch+reducer (1 scratch, 1 reducer launch). Adds 1 per-batch grad scratch buffer + 1 reduce_axis0 launch: attn_grad_q_scratch_d [B, HIDDEN_DIM] ~16 KB scratch at B=32 — trivially small. attn_pool bwd runs 1×/step (not in K-loop) so the absolute wall-time win here is tiny. With this commit every single-SM bwd kernel in the trainer has been refactored to block-per-batch + scratch+reducer. Phase B kernel work complete. Next: local + cluster A/B perf benchmark to verify acceptance gates 6, 7, 8 from the spec. All 9 perception_overfit smokes pass. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/attention_pool.cu | 143 +++++++++++----------- crates/ml-alpha/src/trainer/perception.rs | 34 ++++- 2 files changed, 102 insertions(+), 75 deletions(-) diff --git a/crates/ml-alpha/cuda/attention_pool.cu b/crates/ml-alpha/cuda/attention_pool.cu index b98e8e98f..65701480e 100644 --- a/crates/ml-alpha/cuda/attention_pool.cu +++ b/crates/ml-alpha/cuda/attention_pool.cu @@ -145,97 +145,100 @@ extern "C" __global__ void attention_pool_fwd( // iterates over (b, k, h) sequentially; with block_dim = HIDDEN_DIM, // thread h is sole writer of column h for ALL (b, k). +// Block-per-batch attn_pool bwd (Phase B commit 4). +// grid=(n_batch, 1, 1) block=(ATTN_BLOCK, 1, 1) +// +// Each block handles one batch's HIDDEN_DIM channels. grad_ln_out is +// per-batch indexed (already safe), so each block bi writes its +// [bi, :, :] slice via += onto whatever value grad_ln_out holds at +// kernel launch (the K-loop's contribution to LN_b output grad). +// +// grad_Q is a single [HIDDEN_DIM] shared across all batches → per-batch +// scratch [B, HIDDEN_DIM], reduced after the kernel returns. extern "C" __global__ void attention_pool_bwd( - const float* __restrict__ Q, // [HIDDEN_DIM] - const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] (= values) - const float* __restrict__ attn_weights, // [B, K] from fwd - const float* __restrict__ grad_context, // [B, HIDDEN_DIM] + const float* __restrict__ Q, // [HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] (= values) + const float* __restrict__ attn_weights, // [B, K] from fwd + const float* __restrict__ grad_context, // [B, HIDDEN_DIM] int n_batch, int k_seq, - float* __restrict__ grad_Q, // [HIDDEN_DIM] (+=, caller zero'd) - float* __restrict__ grad_ln_out // [B, K, HIDDEN_DIM] (overwrite) + float* __restrict__ grad_Q_scratch, // [B, HIDDEN_DIM] (+=) + float* __restrict__ grad_ln_out // [B, K, HIDDEN_DIM] (+= chained with K-loop) ) { + int bi = blockIdx.x; int tid = threadIdx.x; - if (tid >= ATTN_BLOCK) return; + if (bi >= n_batch || tid >= ATTN_BLOCK) return; extern __shared__ float smem[]; - float* s_dattn = smem; // [K] - float* s_dscores = smem + k_seq; // [K] - float* s_red = smem + 2 * k_seq; // [BLOCK] + float* s_dattn = smem; // [K] + float* s_dscores = smem + k_seq; // [K] + float* s_red = smem + 2 * k_seq; // [BLOCK] __shared__ float s_Q[ATTN_HIDDEN_DIM]; - __shared__ float s_attn[ATTN_MAX_K]; // K is small (<=512); ok as shared + __shared__ float s_attn[ATTN_MAX_K]; __shared__ float s_grad_ctx[ATTN_HIDDEN_DIM]; __shared__ float s_sum_attn_dattn; if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid]; __syncthreads(); - for (int bi = 0; bi < n_batch; ++bi) { - const float* ln_b = ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; - float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; - // Cache attn + grad_context for this batch. - for (int k = tid; k < k_seq; k += ATTN_BLOCK) { - s_attn[k] = attn_weights[(long long)bi * k_seq + k]; - } - if (tid < ATTN_HIDDEN_DIM) { - s_grad_ctx[tid] = grad_context[(long long)bi * ATTN_HIDDEN_DIM + tid]; - } - __syncthreads(); + const float* ln_b = ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; + float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM; + // Cache attn + grad_context for this block's batch. + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + s_attn[k] = attn_weights[(long long)bi * k_seq + k]; + } + if (tid < ATTN_HIDDEN_DIM) { + s_grad_ctx[tid] = grad_context[(long long)bi * ATTN_HIDDEN_DIM + tid]; + } + __syncthreads(); - // Pass 1: d_attn[k] = sum_h grad_context[h] * values[b, k, h]. - // K iterations; threads tile HIDDEN_DIM via tree-reduce. - for (int k = 0; k < k_seq; ++k) { - const float v = (tid < ATTN_HIDDEN_DIM) - ? s_grad_ctx[tid] * ln_b[k * ATTN_HIDDEN_DIM + tid] : 0.0f; - s_red[tid] = v; - __syncthreads(); - for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_dattn[k] = s_red[0]; - __syncthreads(); - } - - // Pass 2: sum_{kp} attn[kp] * d_attn[kp] (softmax Jacobian centring). - float my_sum = 0.0f; - for (int k = tid; k < k_seq; k += ATTN_BLOCK) { - my_sum += s_attn[k] * s_dattn[k]; - } - s_red[tid] = my_sum; + // Pass 1: d_attn[k] = sum_h grad_context[h] * values[b, k, h]. + for (int k = 0; k < k_seq; ++k) { + const float v = (tid < ATTN_HIDDEN_DIM) + ? s_grad_ctx[tid] * ln_b[k * ATTN_HIDDEN_DIM + tid] : 0.0f; + s_red[tid] = v; __syncthreads(); for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { if (tid < s) s_red[tid] += s_red[tid + s]; __syncthreads(); } - if (tid == 0) s_sum_attn_dattn = s_red[0]; - __syncthreads(); - - // Pass 3: d_scores[k] = attn[k] * (d_attn[k] - s_sum_attn_dattn). - for (int k = tid; k < k_seq; k += ATTN_BLOCK) { - s_dscores[k] = s_attn[k] * (s_dattn[k] - s_sum_attn_dattn); - } - __syncthreads(); - - // Pass 4: d_Q[h] += sum_k d_scores[k] * ln_out[b, k, h]. - // d_ln_out[b, k, h] += grad_context[h] * attn[k] + d_scores[k] * Q[h] - // Thread h owns column h. Loops over k. Both d_Q and d_ln_out - // are += so the caller can chain the attn-path gradient on top - // of the K-loop's contribution on the same `grad_ln_out` buffer - // (= grad_h_enriched_seq_d on the LN_b output). Caller MUST - // pre-zero grad_Q at step start; grad_ln_out should already - // hold the K-loop's contribution before this kernel runs. - if (tid < ATTN_HIDDEN_DIM) { - float dq_local = 0.0f; - for (int k = 0; k < k_seq; ++k) { - const float v = ln_b[k * ATTN_HIDDEN_DIM + tid]; - dq_local += s_dscores[k] * v; - grad_ln_b[k * ATTN_HIDDEN_DIM + tid] += - s_grad_ctx[tid] * s_attn[k] + s_dscores[k] * s_Q[tid]; - } - grad_Q[tid] += dq_local; - } + if (tid == 0) s_dattn[k] = s_red[0]; __syncthreads(); } + + // Pass 2: sum_{kp} attn[kp] * d_attn[kp] (softmax Jacobian centring). + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + my_sum += s_attn[k] * s_dattn[k]; + } + s_red[tid] = my_sum; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) s_sum_attn_dattn = s_red[0]; + __syncthreads(); + + // Pass 3: d_scores[k] = attn[k] * (d_attn[k] - s_sum_attn_dattn). + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + s_dscores[k] = s_attn[k] * (s_dattn[k] - s_sum_attn_dattn); + } + __syncthreads(); + + // Pass 4: grad_Q_scratch[bi, h] += sum_k d_scores[k] * ln_out[b, k, h]. + // d_ln_out[b, k, h] += grad_context[h] * attn[k] + d_scores[k] * Q[h]. + // Thread h owns column h. Loops over k. grad_ln_out += chains the + // attn-path gradient on top of the K-loop's contribution. + if (tid < ATTN_HIDDEN_DIM) { + float dq_local = 0.0f; + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * ATTN_HIDDEN_DIM + tid]; + dq_local += s_dscores[k] * v; + grad_ln_b[k * ATTN_HIDDEN_DIM + tid] += + s_grad_ctx[tid] * s_attn[k] + s_dscores[k] * s_Q[tid]; + } + grad_Q_scratch[(long long)bi * ATTN_HIDDEN_DIM + tid] += dq_local; + } } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index dd620bcff..78dce55e8 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -399,6 +399,8 @@ pub struct PerceptionTrainer { // 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] + // Attention pool per-batch grad scratch (Phase B commit 4). + attn_grad_q_scratch_d: CudaSlice, // [B, HIDDEN_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, @@ -798,6 +800,8 @@ impl PerceptionTrainer { let attn_weights_d = stream.alloc_zeros::(cfg.n_batch * cfg.seq_len)?; let grad_attn_q_d = stream.alloc_zeros::(HIDDEN_DIM)?; let opt_attn_q = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?; + // Phase B: attn pool per-batch grad scratch. + let attn_grad_q_scratch_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM)?; let k = cfg.seq_len; Ok(Self { @@ -871,6 +875,7 @@ impl PerceptionTrainer { grn_grad_b_skip_scratch_d, vsn_grad_w_scratch_d, vsn_grad_b_scratch_d, + attn_grad_q_scratch_d, reduce_axis0_fn, _reduce_module: reduce_module, loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, @@ -1464,9 +1469,9 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("zero cfc_grad_b_scratch: {e}"))?; self.stream.memset_zeros(&mut self.cfc_grad_tau_scratch_d) .map_err(|e| anyhow::anyhow!("zero cfc_grad_tau_scratch: {e}"))?; - // Phase 3: attention pool Q grad accumulator. - self.stream.memset_zeros(&mut self.grad_attn_q_d) - .map_err(|e| anyhow::anyhow!("zero grad_attn_q: {e}"))?; + // Phase B commit 4: attention pool per-batch grad_Q scratch. + self.stream.memset_zeros(&mut self.attn_grad_q_scratch_d) + .map_err(|e| anyhow::anyhow!("zero attn_grad_q_scratch: {e}"))?; // GRN per-batch grad scratch (Phase B commit 2): zero ONCE per // step; K-loop bwd accumulates into these, then reduce_axis0 // collapses → final grad buffers (OVERWRITE) after the K-loop. @@ -1824,12 +1829,14 @@ impl PerceptionTrainer { // clean overwrite for the Q grad. grad_h_enriched_seq_d already // holds the K-loop's contribution at this point — attn's // contribution adds on top. + // Phase B commit 4: block-per-batch attn bwd writes per-batch + // grad_Q scratch; reducer collapses → final grad_attn_q_d below. { let k_i32 = k_seq as i32; let n_batch_attn = b_sz as i32; let shared = (2 * k_seq + 128) * std::mem::size_of::(); let cfg_attn_bwd = LaunchConfig { - grid_dim: (1, 1, 1), + grid_dim: (b_sz as u32, 1, 1), block_dim: (128, 1, 1), // ATTN_BLOCK shared_mem_bytes: shared as u32, }; @@ -1840,10 +1847,27 @@ impl PerceptionTrainer { .arg(&self.attn_weights_d) .arg(&self.grad_h_carry_d) .arg(&n_batch_attn).arg(&k_i32) - .arg(&mut self.grad_attn_q_d) + .arg(&mut self.attn_grad_q_scratch_d) .arg(self.grad_h_enriched_seq_d.data_mut()); unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; } } + // Attn pool reducer: collapse [B, HIDDEN_DIM] → [HIDDEN_DIM]. + { + let n_batch_i = b_sz as i32; + let n_tail_i = HIDDEN_DIM as i32; + let cfg_red = LaunchConfig { + grid_dim: (HIDDEN_DIM as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); + launch + .arg(&self.attn_grad_q_scratch_d) + .arg(&n_batch_i) + .arg(&n_tail_i) + .arg(&mut self.grad_attn_q_d); + unsafe { launch.launch(cfg_red).context("reduce attn_grad_q")?; } + } // ── 7c. LayerNorm B backward (between m2 and CfC). Consumes: // x = m2.h_enriched_seq [B, K, H]