From f76437c0c71ee8345cbfec892da02d9dcbdf4e80 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 22:32:00 +0200 Subject: [PATCH] feat(ml-alpha): attention pool forward+backward CUDA kernels (Phase 3.1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-head attention pool over Mamba2 K-positions, designed to replace the CfC's zero-initialised `h_old` at k=0 with a learned content- addressable summary over all K LN_b output positions. Forward math: scores[k] = Q · keys[b, k, :] # [K] attn[k] = softmax_k(scores) # [K] context[h] = sum_k attn[k] * values[b, k, h] # [HIDDEN_DIM] For our attention pool, keys == values == LN_b output [B, K, HIDDEN_DIM]. Single learned param: Q [HIDDEN_DIM]. Tiny (128 params). Forward layout: grid = (B, 1, 1), block = HIDDEN_DIM=128 threads. Three passes: (1) K dot-products with tree-reduce over HIDDEN_DIM, (2) softmax over K with max-subtract+sum, (3) weighted sum into context. Backward chain rule: d_attn[k] = sum_h grad_context[h] * values[b, k, h] d_scores[k] = attn[k] * (d_attn[k] - sum_kp attn[kp] * d_attn[kp]) d_Q[h] += sum_{b, k} d_scores[k] * values[b, k, h] d_values[b, k, h] += grad_context[h] * attn[k] + d_scores[k] * Q[h] Both `d_Q` and `d_values` use += semantics: - d_Q: accumulates across batch (single block, internal n_batch loop). - d_values: writes ADD onto whatever grad_ln_out already holds, so the trainer can chain it on top of the K-loop's contribution to the LN_b output gradient (no separate add-kernel needed). Single-writer (no atomicAdd): one block per launch, thread h owns column h of grad_ln_out for ALL (b, k). Internal n_batch loop matches the GRN / VSN bwd pattern. build.rs: - "attention_pool" added to KERNELS - Cache bust → v10 Wiring into PerceptionTrainer (Phase 3.2) is the follow-up commit: add attn_q_d learned param + per-batch context + attn_weights buffers, run attn_pool_fwd between LN_b fwd and the K-loop, use attn_context as the K-loop's k=0 h_old (instead of zero_h_d), and chain attn_pool_bwd after the K-loop reverse pass. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 7 +- crates/ml-alpha/cuda/attention_pool.cu | 241 +++++++++++++++++++++++++ 2 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 crates/ml-alpha/cuda/attention_pool.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index b2add2f98..53074cad7 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -19,11 +19,12 @@ const KERNELS: &[&str] = &[ "horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda) "layer_norm", // Phase 1: trunk pre-CfC normalisation "variable_selection", // Phase 2D: TFT-style per-feature gating + "attention_pool", // Phase 3: learned context summary at CfC k=0 ]; -// Cache bust v9 (2026-05-17): Phase 2D TFT VSN — new -// variable_selection.cu kernel pair (fwd + bwd) for per-feature -// softmax-normalised gating at the trunk entry. Force fresh nvcc +// Cache bust v10 (2026-05-17): Phase 3 attention pool — new +// attention_pool.cu kernel pair (fwd + bwd) for the learned context +// summary replacing CfC's zero-initialised h_old. Force fresh nvcc // compile so the new cubin lands in OUT_DIR. fn main() { diff --git a/crates/ml-alpha/cuda/attention_pool.cu b/crates/ml-alpha/cuda/attention_pool.cu new file mode 100644 index 000000000..b98e8e98f --- /dev/null +++ b/crates/ml-alpha/cuda/attention_pool.cu @@ -0,0 +1,241 @@ +// attention_pool.cu — Single-head attention pool over Mamba2 K-positions. +// Phase 3 (2026-05-17). +// +// Replaces the CfC's zero-initialised `h_old` at k=0 with a learned +// attention-pooled summary over all K LN_b output positions. The K-loop +// CfC keeps its recurrent semantics (h_old at k+1 = h_new at k); only +// the initial state changes from zero to a content-addressable lookup. +// +// Forward math (per sample b): +// scores[k] = Q · keys[b, k, :] # [K] — dot product over HIDDEN_DIM +// attn[k] = softmax_k(scores) # [K] +// context[h] = sum_k attn[k] * values[b, k, h] # [HIDDEN_DIM] +// +// For this attention pool, keys == values == LN_b output [B, K, HIDDEN_DIM]. +// Single learned param: Q [HIDDEN_DIM]. +// +// Saved for backward: +// attn_weights [B, K] (post-softmax) +// +// Block layout (forward): +// grid_dim = (B, 1, 1), block_dim = (HIDDEN_DIM=128, 1, 1). +// Thread tid handles dimension h of HIDDEN_DIM. +// +// Per `feedback_no_atomicadd.md`: block tree-reduce only, no atomics. + +#define ATTN_HIDDEN_DIM 128 +#define ATTN_BLOCK 128 +#define ATTN_MAX_K 512 // safety cap on K; smoke uses K=16-64. + +extern "C" __global__ void attention_pool_fwd( + const float* __restrict__ Q, // [HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] + int n_batch, + int k_seq, + float* __restrict__ context, // [B, HIDDEN_DIM] + float* __restrict__ attn_weights // [B, K] — saved post-softmax +) { + int b_idx = blockIdx.x; + int tid = threadIdx.x; + if (b_idx >= n_batch || tid >= ATTN_BLOCK) return; + + extern __shared__ float smem[]; + float* s_scores = smem; // [K] + float* s_red = smem + k_seq; // [BLOCK] — reduce scratch + float* s_context = smem + k_seq + ATTN_BLOCK; // [HIDDEN_DIM] — final output buffer + + // Cache Q in shared mem (broadcast). + __shared__ float s_Q[ATTN_HIDDEN_DIM]; + if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid]; + __syncthreads(); + + // Pass 1: scores[k] = Q · ln_out[b, k, :]. + // Each k is handled sequentially by ALL threads via tree-reduce + // over HIDDEN_DIM. K outer loop, threads tile HIDDEN_DIM inner. + // For HIDDEN_DIM=128 and ATTN_BLOCK=128 we have one-to-one. + const float* ln_b = ln_out + (long long)b_idx * k_seq * ATTN_HIDDEN_DIM; + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * ATTN_HIDDEN_DIM + tid] * s_Q[tid]; + 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_scores[k] = s_red[0]; + __syncthreads(); + } + + // Pass 2: softmax over K. Numerically stable: max-subtract + exp + sum. + // Block-wide max over s_scores[0..k_seq). + float my_max = -INFINITY; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float v = s_scores[k]; + if (v > my_max) my_max = v; + } + s_red[tid] = my_max; + __syncthreads(); + for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) { + const float a = s_red[tid]; + const float b = s_red[tid + s]; + s_red[tid] = (a > b) ? a : b; + } + __syncthreads(); + } + __shared__ float s_max; + if (tid == 0) s_max = s_red[0]; + __syncthreads(); + + // exp(score - max) + sum. + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float e = expf(s_scores[k] - s_max); + s_scores[k] = e; // reuse as exp_shifted + my_sum += e; + } + 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(); + } + __shared__ float s_sum; + if (tid == 0) s_sum = s_red[0]; + __syncthreads(); + + // Pass 3: attn[k] = exp/sum; save; context[h] = sum_k attn[k] * ln_out[b, k, h]. + // Save attn weights AND build the context output. + for (int k = tid; k < k_seq; k += ATTN_BLOCK) { + const float a = s_scores[k] / s_sum; + s_scores[k] = a; // overwrite again — now holds true attn weights + attn_weights[(long long)b_idx * k_seq + k] = a; + } + __syncthreads(); + + // Accumulate context[h=tid] = sum_k attn[k] * ln_out[b, k, h=tid]. + if (tid < ATTN_HIDDEN_DIM) { + float c = 0.0f; + for (int k = 0; k < k_seq; ++k) { + c += s_scores[k] * ln_b[k * ATTN_HIDDEN_DIM + tid]; + } + s_context[tid] = c; + context[(long long)b_idx * ATTN_HIDDEN_DIM + tid] = c; + } +} + +// Attention pool backward — chain rule: +// +// d_attn[k] = sum_h grad_context[h] * values[b, k, h] +// + (this is from context = sum_k attn[k] * values[b, k, :]) +// +// d_values[b, k, h] += grad_context[h] * attn[k] +// + d_scores[k] * Q[h] +// +// d_scores via softmax Jacobian: +// d_scores[k] = attn[k] * (d_attn[k] - sum_kp attn[kp] * d_attn[kp]) +// +// d_Q[h] += sum_{b, k} d_scores[k] * values[b, k, h] +// (here values == ln_out) +// +// Single-writer discipline: +// - d_Q is [HIDDEN_DIM] shared across all (b, k). ONE block per +// launch, internal batch loop, tile over HIDDEN_DIM. += into d_Q. +// - d_values[b, k, h] is [B, K, HIDDEN_DIM] — one block per launch +// iterates over (b, k, h) sequentially; with block_dim = HIDDEN_DIM, +// thread h is sole writer of column h for ALL (b, k). + +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] + 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) +) { + int tid = threadIdx.x; + if (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] + + __shared__ float s_Q[ATTN_HIDDEN_DIM]; + __shared__ float s_attn[ATTN_MAX_K]; // K is small (<=512); ok as shared + __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(); + + // 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; + __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; + } + __syncthreads(); + } +}