From 7a558b88b7110724c66917e1d9958733f6960f96 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 22:38:05 +0200 Subject: [PATCH] feat(ml-alpha): wire attention pool into PerceptionTrainer (Phase 3.2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces CfC's zero-initialised h_old at k=0 with the attention-pooled context vector — a learned content-addressable summary over all K LN_b output positions. The K-loop's recurrent semantics (h_old at k+1 = h_new at k) are preserved; only the INITIAL state at k=0 changes from zero to the pooled context. Forward chain change: ... → m2 → LN_b → ln_out_d [B, K, HIDDEN_DIM] → attention_pool_fwd(Q, ln_out_d) → attn_context_d [B, HIDDEN_DIM] (used as h_old@k=0) → attn_weights_d [B, K] (saved for bwd) → K-loop CfC (h_old@k=0 = attn_context, not zero) Backward chain change: K-loop bwd ends with grad_h_carry_d holding the gradient that would have flowed into the initial h_old = grad on attn_context. attention_pool_bwd consumes: Q, ln_out_d, attn_weights_d (forward state) grad_h_carry_d = grad_context Writes (BOTH `+=`): grad_attn_q_d (accumulates Q gradient — pre-zeroed at step start) grad_h_enriched_seq_d (ADDS attn-path contribution onto LN_b output gradient — chains with K-loop contribution) LN_b bwd then consumes the now-summed grad_h_enriched_seq_d. Trainer state additions (8 fields): attn_q_d, attn_context_d, attn_weights_d, grad_attn_q_d, attn_fwd_fn, attn_bwd_fn, _attn_module, opt_attn_q Q is tiny (HIDDEN_DIM=128 floats); initialised near zero so initial attention ≈ uniform 1/K (context ≈ mean of LN_b output). Model learns content addressing from a near-uniform starting point. Eval path mirrors training: attn_pool_fwd runs after LN_b fwd, attn_context_d feeds the eval K-loop at k=0. Trainer now manages 22 AdamW: CfC×4 + GRN heads×10 + LN×2 + LN_a×2 + VSN×2 + Attn Q×1 + Mamba2×2 grouped. Synthetic overfit smoke: stride=1 0.29 → 0.0000 in 50 steps (faster than pre-attn 0.30), stride=4 0.30 → 0.0000. All 8 perception_overfit tests PASS. Demonstrates the full Phase 1+2+3 stack (VSN → m1 → LN_a → m2 → LN_b → attn pool → CfC + GRN heads) is wired forward + backward end-to-end with every gradient flowing through every learned param. Phase 1+2+3 capacity scale-up complete. The cumulative architectural lift over the 3-fix-stack baseline (496q7 mean_auc=0.716 / h6000=0.704) will be measured by deploying this stack head-to-head against bsml6. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/perception.rs | 164 +++++++++++++++++++++- 1 file changed, 158 insertions(+), 6 deletions(-) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 138c3fce3..0a038986c 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -61,6 +61,7 @@ const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_mult const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin")); const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin")); const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin")); +const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin")); #[derive(Clone, Debug)] pub struct PerceptionTrainerConfig { @@ -357,6 +358,21 @@ pub struct PerceptionTrainer { vsn_fwd_fn: CudaFunction, vsn_bwd_fn: CudaFunction, _vsn_module: Arc, + + // ── Attention pool (Phase 3) ── + // Single-head attention pool over LN_b output. Replaces CfC's + // zero-initialised h_old at k=0 with a learned content-addressable + // summary over all K positions. Single learned param: Q [HIDDEN_DIM]. + pub attn_q_d: CudaSlice, + /// Per-sample pooled context `[B, HIDDEN_DIM]` — fed as h_old at k=0. + attn_context_d: CudaSlice, + /// Saved post-softmax attention weights `[B, K]` for bwd. + attn_weights_d: CudaSlice, + grad_attn_q_d: CudaSlice, + pub opt_attn_q: AdamW, + attn_fwd_fn: CudaFunction, + attn_bwd_fn: CudaFunction, + _attn_module: Arc, /// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised /// (sentinel); the horizon_ema_and_lambda kernel detects the /// sentinel on step 1 and replaces directly per @@ -482,6 +498,15 @@ impl PerceptionTrainer { let vsn_bwd_fn = vsn_module .load_function("variable_selection_bwd") .context("variable_selection_bwd symbol")?; + let attn_module = ctx + .load_cubin(ATTENTION_POOL_CUBIN.to_vec()) + .context("attention_pool cubin")?; + let attn_fwd_fn = attn_module + .load_function("attention_pool_fwd") + .context("attention_pool_fwd symbol")?; + let attn_bwd_fn = attn_module + .load_function("attention_pool_bwd") + .context("attention_pool_bwd symbol")?; let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched")?; let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?; let step_batched_fn = step_module.load_function("cfc_step_batched")?; @@ -683,6 +708,19 @@ impl PerceptionTrainer { let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?; opt_vsn_b.wd = 0.0; + // ── Attention pool init (Phase 3) ── + // Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K + // → context ≈ mean of LN_b output. Model learns content + // addressing from a near-uniform starting point. + let attn_q_scale = (1.0_f32 / HIDDEN_DIM as f32).sqrt(); + let attn_q_init: Vec = (0..HIDDEN_DIM) + .map(|_| r.gen_range(-attn_q_scale..attn_q_scale)).collect(); + let attn_q_d = upload(&stream, &attn_q_init)?; + let attn_context_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM)?; + 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)?; + let k = cfg.seq_len; Ok(Self { cfg: cfg.clone(), @@ -729,6 +767,15 @@ impl PerceptionTrainer { vsn_fwd_fn, vsn_bwd_fn, _vsn_module: vsn_module, + // Attention pool (Phase 3). + attn_q_d, + attn_context_d, + attn_weights_d, + grad_attn_q_d, + opt_attn_q, + attn_fwd_fn, + attn_bwd_fn, + _attn_module: attn_module, loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, lambda_d: stream.alloc_zeros::(N_HORIZONS)?, horizon_lambda_fn, @@ -1272,6 +1319,31 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; } } + // ── 2d. Attention pool forward (Phase 3) — produces + // attn_context_d [B, HIDDEN_DIM] = learned content-summary + // over all K LN_b output positions. Replaces CfC's + // zero-initialised h_old at k=0 with this context vector. + // Shared mem: K floats (scores) + BLOCK floats (reduce) + + // HIDDEN_DIM floats (context) = (K + 128 + 128) * 4 bytes. + { + let k_i32 = k_seq as i32; + let n_batch_attn = b_sz as i32; + let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::(); + let cfg_attn = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (128, 1, 1), // ATTN_BLOCK + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_fwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .arg(&n_batch_attn).arg(&k_i32) + .arg(&mut self.attn_context_d) + .arg(&mut self.attn_weights_d); + unsafe { launch.launch(cfg_attn).context("attention_pool_fwd")?; } + } + // ── 3. Labels DtoD: staging (filled in step_batched before // captured region) → device. let total_labels = k_seq * b_sz * N_HORIZONS; @@ -1292,6 +1364,9 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("zero grad_b: {e}"))?; self.stream.memset_zeros(&mut self.grad_tau_d) .map_err(|e| anyhow::anyhow!("zero grad_tau: {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}"))?; // GRN heads: 10 grad accumulators. self.stream.memset_zeros(&mut self.grad_heads_w1_d) .map_err(|e| anyhow::anyhow!("zero grad_heads_w1: {e}"))?; @@ -1386,10 +1461,18 @@ impl PerceptionTrainer { // [K, B, H] layout. Pointer-offset trick (single mut // borrow + raw u64 arithmetic for slot pointers) keeps // the launches stream-ordered with no syncs. - let zero_h_ptr = { + // Phase 3: attention pool produces attn_context_d which replaces + // the zero initial h_old at k=0. zero_h_ptr is no longer used in + // the fwd K-loop but is retained for the bwd k=0 case (cfc bwd + // needs an h_old slot for the input gradient to read from). + let _zero_h_ptr_unused = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; + let attn_context_ptr = { + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); + p + }; let henr_t_base = { let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream); p @@ -1406,7 +1489,7 @@ impl PerceptionTrainer { for k in 0..k_seq { let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64; let h_old_k_ptr = if k == 0 { - zero_h_ptr + attn_context_ptr } else { h_per_k_base + ((k - 1) * kb_hid_bytes) as u64 }; @@ -1498,10 +1581,17 @@ impl PerceptionTrainer { self.stream.memset_zeros(&mut self.grad_h_carry_d) .map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?; - let zero_h_ptr_bwd = { + // Phase 3: at k=0 the bwd kernel reads `h_old` = attn_context_d + // (mirroring the forward pass). zero_h_ptr_bwd retained as a + // legacy fallback / unused alias. + let _zero_h_ptr_bwd_unused = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; + let attn_context_ptr_bwd = { + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); + p + }; let henr_t_base_bwd = { let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream); p @@ -1527,7 +1617,7 @@ impl PerceptionTrainer { let gate_k_ptr = gate_base_bwd + (k * kb_nh_bytes) as u64; let main_k_ptr = main_base_bwd + (k * kb_nh_bytes) as u64; let h_old_k_ptr = if k == 0 { - zero_h_ptr_bwd + attn_context_ptr_bwd } else { h_per_k_base_bwd + ((k - 1) * kb_hid_bytes) as u64 }; @@ -1606,6 +1696,41 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; } } + // ── 7c-pre. Attention pool backward (Phase 3). Consumes: + // Q = self.attn_q_d [HIDDEN_DIM] + // ln_out (values)= self.ln_out_d [B, K, HIDDEN_DIM] + // attn_weights = self.attn_weights_d (saved by fwd) [B, K] + // grad_context = self.grad_h_carry_d (= grad on initial + // h_old at k=0, which IS attn_context) [B, HIDDEN_DIM] + // Writes (BOTH ARE +=): + // grad_attn_q_d += attn-path contribution to Q + // grad_h_enriched_seq_d (LN_b output grad) += attn-path + // contribution to ln_out + // The pre-zero of grad_attn_q_d at step start makes the += a + // 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. + { + 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), + block_dim: (128, 1, 1), // ATTN_BLOCK + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_bwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .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(self.grad_h_enriched_seq_d.data_mut()); + unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; } + } + // ── 7c. LayerNorm B backward (between m2 and CfC). Consumes: // x = m2.h_enriched_seq [B, K, H] // gain = self.ln_gain_d [H] @@ -1815,6 +1940,7 @@ impl PerceptionTrainer { self.opt_ln_a_bias.step(&mut self.ln_a_bias_d, &self.grad_ln_a_bias_d)?; self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?; self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?; + self.opt_attn_q.step(&mut self.attn_q_d, &self.grad_attn_q_d)?; // GPU-resident grad-clip + AdamW (zero host roundtrips). // Replaces step_from_buffers which did 9× memcpy_dtoh per step. self.mamba2_adamw @@ -2064,6 +2190,27 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; } } + // Attention pool fwd — same as training, populates attn_context_d + // for use as k=0 h_old in the eval K-loop below. + { + let k_i32 = k_seq as i32; + let n_batch_attn = b_sz as i32; + let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::(); + let cfg_attn = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: shared as u32, + }; + let mut launch = self.stream.launch_builder(&self.attn_fwd_fn); + launch + .arg(&self.attn_q_d) + .arg(&self.ln_out_d) + .arg(&n_batch_attn).arg(&k_i32) + .arg(&mut self.attn_context_d) + .arg(&mut self.attn_weights_d); + unsafe { launch.launch(cfg_attn).context("eval attention_pool_fwd")?; } + } + // Upload labels [K, B, N_HORIZONS]. let total_labels = k_seq * b_sz * N_HORIZONS; { @@ -2105,10 +2252,15 @@ impl PerceptionTrainer { let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::(); let kb_nh_mid_bytes = b_sz * N_HORIZONS * HEAD_MID_DIM * std::mem::size_of::(); let n_batch_i = b_sz as i32; - let zero_h_ptr = { + let _zero_h_ptr_unused_eval = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; + // Phase 3: eval uses attn_context_d as h_old at k=0, mirroring training. + let attn_context_ptr = { + let (p, _g) = self.attn_context_d.device_ptr(&self.stream); + p + }; let henr_t_base = { let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream); p @@ -2124,7 +2276,7 @@ impl PerceptionTrainer { for k in 0..k_seq { let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64; let h_old_k_ptr = if k == 0 { - zero_h_ptr + attn_context_ptr } else { h_per_k_base + ((k - 1) * kb_hid_bytes) as u64 };