feat(ml-alpha): attention pool forward+backward CUDA kernels (Phase 3.1)
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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() {
|
||||
|
||||
241
crates/ml-alpha/cuda/attention_pool.cu
Normal file
241
crates/ml-alpha/cuda/attention_pool.cu
Normal file
@@ -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();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user