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

perf(ml-alpha): block-per-batch attention pool bwd refactor (Phase B commit 4)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 00:01:31 +02:00
parent 9607f33518
commit a478ba3d84
2 changed files with 102 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
crates/ml-alpha/cuda/attention_pool.cu
View File

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

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

@@ -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<f32>, // [B*K, FEATURE_DIM, FEATURE_DIM]
vsn_grad_b_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM]
// Attention pool per-batch grad scratch (Phase B commit 4).
attn_grad_q_scratch_d: CudaSlice<f32>, // [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::<f32>(cfg.n_batch * cfg.seq_len)?;
let grad_attn_q_d = stream.alloc_zeros::<f32>(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::<f32>(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::<f32>(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::<f32>();
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]