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

feat(ml-alpha): wire attention pool into PerceptionTrainer (Phase 3.2)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 22:38:05 +02:00
parent f76437c0c7
commit 7a558b88b7
1 changed files with 158 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<CudaModule>,
// ── 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<f32>,
/// Per-sample pooled context `[B, HIDDEN_DIM]` — fed as h_old at k=0.
attn_context_d: CudaSlice<f32>,
/// Saved post-softmax attention weights `[B, K]` for bwd.
attn_weights_d: CudaSlice<f32>,
grad_attn_q_d: CudaSlice<f32>,
pub opt_attn_q: AdamW,
attn_fwd_fn: CudaFunction,
attn_bwd_fn: CudaFunction,
_attn_module: Arc<CudaModule>,
/// 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<f32> = (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::<f32>(cfg.n_batch * HIDDEN_DIM)?;
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)?;
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::<f32>(N_HORIZONS)?,
lambda_d: stream.alloc_zeros::<f32>(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::<f32>();
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::<f32>();
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::<f32>();
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::<f32>();
let kb_nh_mid_bytes = b_sz * N_HORIZONS * HEAD_MID_DIM * std::mem::size_of::<f32>();
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
};