From edc449eecf1a94d04eee1b62cd247b45e7759973 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 10:32:16 +0200 Subject: [PATCH] feat(ml-alpha): per-horizon attention pool kernel + numgrad parity (C21) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First implementation slice of the per-horizon attention pool design (docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md). Lands the kernel + Rust binding + numgrad verification; downstream wiring into PerceptionTrainer's captured graph + CheckpointV2 bump + A/B sweep are follow-up commits gated on this proving correctness. Kernel (cuda/per_horizon_attention_pool.cu): per_horizon_attention_pool_fwd Q_h[N_HORIZONS, HIDDEN_DIM] × LNb[B, K, HIDDEN_DIM] → context_h[B, N_HORIZONS, HIDDEN_DIM] attn_h_weights[B, N_HORIZONS, K] Per-block math identical to the single-Q variant, looped over N_HORIZONS sequentially within each batch's block. Grid stays (B, 1, 1) so backward grad_ln_out writes are race-free (per feedback_no_atomicadd.md — no cross-block contention). Per-batch shared mem ~k_seq + BLOCK + HIDDEN_DIM floats. per_horizon_attention_pool_bwd Same chain-rule pattern as attention_pool_bwd but with the horizon loop inside the block: each (b, h) slice updates grad_ln_out in place (sequential horizon accumulation), grad_Q_h is written as per-block scratch [B, N_HORIZONS, HIDDEN_DIM] for host reduce. Single-writer discipline preserved. Rust binding (src/per_horizon_attention_pool.rs): PerHorizonAttentionPool::{new, forward, backward}. Self-contained; doesn't yet touch PerceptionTrainer or CfcTrunk. Loads the cubin via the standard env!("OUT_DIR") path. Dynamic shared-mem byte count computed per launch from k_seq. Numgrad parity test (tests/per_horizon_attention_pool_numgrad.rs): - B=2, K=8, HIDDEN_DIM=128, N_HORIZONS=5 fixture. - Loss = Σ context_h (so d_context = 1 everywhere — clean analytical). - Backward kernel produces analytical grads; central-difference of forward kernel at ±eps=1e-2 across 8 random Q_h indices + 8 random LNb indices verifies analytical matches CD within 5e-2 rel-tol or 5e-3 abs-floor. - Passes on RTX 3050. build.rs picks up the new .cu file automatically via the existing KERNELS list; cubin compiles cleanly at sm_86 + sm_89. Same scope discipline as Phase 2D.2 (VSN numgrad) — kernel correctness first, integration second. The follow-up commit set per the spec §3 appendix: C22: extend multi_horizon_heads.cu signature to accept per-horizon context input + bump head_w shape to [N_HORIZONS, 2*HIDDEN_DIM] C23: wire PerHorizonAttentionPool into PerceptionTrainer + CfcTrunk captured graph behind AttentionPoolVariant config flag C24: CheckpointV1 → V2 bump with discriminant + optional q_h field C25: smoke training (one epoch, no NaN, loss decreases) C26: 30-epoch × 3-fold A/B sweep (#204) — decision gate per spec §0 Closes the kernel-correctness portion of #203. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + .../cuda/per_horizon_attention_pool.cu | 243 ++++++++++++++++++ crates/ml-alpha/src/lib.rs | 1 + .../src/per_horizon_attention_pool.rs | 128 +++++++++ .../per_horizon_attention_pool_numgrad.rs | 152 +++++++++++ 5 files changed, 525 insertions(+) create mode 100644 crates/ml-alpha/cuda/per_horizon_attention_pool.cu create mode 100644 crates/ml-alpha/src/per_horizon_attention_pool.rs create mode 100644 crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index b153585f7..858574199 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -20,6 +20,7 @@ const KERNELS: &[&str] = &[ "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 + "per_horizon_attention_pool", // C21: per-horizon variant; ships behind a config flag "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/per_horizon_attention_pool.cu b/crates/ml-alpha/cuda/per_horizon_attention_pool.cu new file mode 100644 index 000000000..5bb5ab7ce --- /dev/null +++ b/crates/ml-alpha/cuda/per_horizon_attention_pool.cu @@ -0,0 +1,243 @@ +// per_horizon_attention_pool.cu — per-horizon attention pool (C21). +// +// Replaces the single learned query Q[HIDDEN_DIM] of attention_pool.cu +// with N_HORIZONS=5 learned queries Q_h[N_HORIZONS, HIDDEN_DIM]. Each +// horizon attends to a different combination of K LN_b positions, +// producing its own context vector. The multi-horizon heads downstream +// consume each horizon's own context (concat'd with CfC h_K). +// +// Forward math (per sample b, per horizon h): +// scores_h[k] = Σ_d Q_h[h, d] * LNb[b, k, d] # [K] +// attn_h[k] = softmax_k(scores_h) # [K] +// context_h[d] = Σ_k attn_h[k] * LNb[b, k, d] # [HIDDEN_DIM] +// +// For this pool keys == values == LN_b output [B, K, HIDDEN_DIM]. +// Learned params: Q_h [N_HORIZONS, HIDDEN_DIM]. +// +// Saved for backward: +// attn_h_weights [B, N_HORIZONS, K] (post-softmax) +// +// Block layout (forward): +// grid_dim = (B, 1, 1), block_dim = (HIDDEN_DIM=128, 1, 1). +// Each block handles ALL horizons for its batch — horizon loop is +// inside the block. Keeps grad_ln_out writes free of cross-block +// races: each batch's [b, :, :] slice is touched by ONE block in bwd. +// +// Per `feedback_no_atomicadd.md`: block tree-reduce only, no atomics. + +#define PHA_HIDDEN_DIM 128 +#define PHA_BLOCK 128 +#define PHA_MAX_K 512 // safety cap; smoke uses K=16-64. +#define PHA_N_HORIZONS 5 + +extern "C" __global__ void per_horizon_attention_pool_fwd( + const float* __restrict__ Q_h, // [N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] + int n_batch, + int k_seq, + float* __restrict__ context_h, // [B, N_HORIZONS, HIDDEN_DIM] + float* __restrict__ attn_h_weights // [B, N_HORIZONS, K] +) { + int b_idx = blockIdx.x; + int tid = threadIdx.x; + if (b_idx >= n_batch || tid >= PHA_BLOCK) return; + + extern __shared__ float smem[]; + float* s_scores = smem; // [K] + float* s_red = smem + k_seq; // [BLOCK] + float* s_context = smem + k_seq + PHA_BLOCK; // [HIDDEN_DIM] + + __shared__ float s_Qh[PHA_HIDDEN_DIM]; + __shared__ float s_max; + __shared__ float s_sum; + + const float* ln_b = ln_out + (long long)b_idx * k_seq * PHA_HIDDEN_DIM; + + // Outer loop: one full forward pass per horizon, reusing the + // per-block scratch for each. Sequential horizons within a block + // means no cross-horizon barrier worries; the smem is reused. + for (int h = 0; h < PHA_N_HORIZONS; ++h) { + // Cache Q_h[h, :] for this iteration. + if (tid < PHA_HIDDEN_DIM) s_Qh[tid] = Q_h[h * PHA_HIDDEN_DIM + tid]; + __syncthreads(); + + // Pass 1: scores[k] = Q_h[h, :] · ln_out[b, k, :]. + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * PHA_HIDDEN_DIM + tid] * s_Qh[tid]; + s_red[tid] = v; + __syncthreads(); + for (int s = PHA_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 — max-subtract + exp + sum + divide. + float my_max = -INFINITY; + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + const float v = s_scores[k]; + if (v > my_max) my_max = v; + } + s_red[tid] = my_max; + __syncthreads(); + for (int s = PHA_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(); + } + if (tid == 0) s_max = s_red[0]; + __syncthreads(); + + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + const float e = expf(s_scores[k] - s_max); + s_scores[k] = e; + my_sum += e; + } + s_red[tid] = my_sum; + __syncthreads(); + for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) s_sum = s_red[0]; + __syncthreads(); + + // Pass 3: attn[k] = exp / sum; save; context[d] = Σ_k attn[k] * ln_out[b, k, d]. + const long long attn_base = (long long)b_idx * PHA_N_HORIZONS * k_seq + + (long long)h * k_seq; + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + const float a = s_scores[k] / s_sum; + s_scores[k] = a; + attn_h_weights[attn_base + k] = a; + } + __syncthreads(); + + if (tid < PHA_HIDDEN_DIM) { + float c = 0.0f; + for (int k = 0; k < k_seq; ++k) { + c += s_scores[k] * ln_b[k * PHA_HIDDEN_DIM + tid]; + } + s_context[tid] = c; + context_h[(long long)b_idx * PHA_N_HORIZONS * PHA_HIDDEN_DIM + + (long long)h * PHA_HIDDEN_DIM + tid] = c; + } + __syncthreads(); // before next horizon resets s_Qh / s_scores + } +} + +// Per-horizon attention pool backward. +// +// Chain rule (per (b, h) slice — identical shape to the single-Q +// backward, just looped over horizons): +// +// d_attn[k] = Σ_d grad_context[d] * ln_out[b, k, d] +// d_scores[k] = attn[k] * (d_attn[k] - Σ_kp attn[kp] * d_attn[kp]) +// d_Q_h[h, d] += Σ_k d_scores[k] * ln_out[b, k, d] +// d_LNb[b, k, d] += grad_context[d] * attn[k] + d_scores[k] * Q_h[h, d] +// +// grad_Q_h is [N_HORIZONS, HIDDEN_DIM] shared across all batches → per-batch +// scratch [B, N_HORIZONS, HIDDEN_DIM] reduced via reduce_axis0 after. +// grad_ln_out is per-batch indexed; each block writes its [b, :, :] slice +// via += and sums horizon contributions sequentially within the block — +// no cross-block race possible. +extern "C" __global__ void per_horizon_attention_pool_bwd( + const float* __restrict__ Q_h, // [N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] + const float* __restrict__ attn_h_weights, // [B, N_HORIZONS, K] from fwd + const float* __restrict__ grad_context_h, // [B, N_HORIZONS, HIDDEN_DIM] + int n_batch, + int k_seq, + float* __restrict__ grad_Qh_scratch, // [B, N_HORIZONS, HIDDEN_DIM] (+=) + float* __restrict__ grad_ln_out // [B, K, HIDDEN_DIM] (+= chained with K-loop) +) { + int bi = blockIdx.x; + int tid = threadIdx.x; + if (bi >= n_batch || tid >= PHA_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_Qh[PHA_HIDDEN_DIM]; + __shared__ float s_attn[PHA_MAX_K]; + __shared__ float s_grad_ctx[PHA_HIDDEN_DIM]; + __shared__ float s_sum_attn_dattn; + + const float* ln_b = ln_out + (long long)bi * k_seq * PHA_HIDDEN_DIM; + float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * PHA_HIDDEN_DIM; + + // Horizon loop. For each h: read Q_h + grad_context + attn; compute + // all backward gradients; accumulate into grad_ln_b (shared across + // horizons within this batch — sequential within block, no race). + for (int h = 0; h < PHA_N_HORIZONS; ++h) { + if (tid < PHA_HIDDEN_DIM) { + s_Qh[tid] = Q_h[h * PHA_HIDDEN_DIM + tid]; + s_grad_ctx[tid] = grad_context_h[(long long)bi * PHA_N_HORIZONS * PHA_HIDDEN_DIM + + (long long)h * PHA_HIDDEN_DIM + tid]; + } + const long long attn_base = (long long)bi * PHA_N_HORIZONS * k_seq + + (long long)h * k_seq; + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + s_attn[k] = attn_h_weights[attn_base + k]; + } + __syncthreads(); + + // Pass 1: d_attn[k] = Σ_d grad_context[d] * ln_out[b, k, d]. + for (int k = 0; k < k_seq; ++k) { + const float v = (tid < PHA_HIDDEN_DIM) + ? s_grad_ctx[tid] * ln_b[k * PHA_HIDDEN_DIM + tid] : 0.0f; + s_red[tid] = v; + __syncthreads(); + for (int s = PHA_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: Σ_kp attn[kp] * d_attn[kp] (softmax Jacobian centring). + float my_sum = 0.0f; + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + my_sum += s_attn[k] * s_dattn[k]; + } + s_red[tid] = my_sum; + __syncthreads(); + for (int s = PHA_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] - centring). + for (int k = tid; k < k_seq; k += PHA_BLOCK) { + s_dscores[k] = s_attn[k] * (s_dattn[k] - s_sum_attn_dattn); + } + __syncthreads(); + + // Pass 4: grad_Qh_scratch[bi, h, d] += Σ_k d_scores[k] * ln_out[b, k, d]. + // grad_ln_b[k, d] += grad_context[d] * attn[k] + d_scores[k] * Q_h[h, d]. + // Thread tid owns column d=tid for this (b, h). Per-horizon += into grad_ln_b + // is sequential within the block — no race; horizons accumulate over each other. + if (tid < PHA_HIDDEN_DIM) { + float dq_local = 0.0f; + for (int k = 0; k < k_seq; ++k) { + const float v = ln_b[k * PHA_HIDDEN_DIM + tid]; + dq_local += s_dscores[k] * v; + grad_ln_b[k * PHA_HIDDEN_DIM + tid] += + s_grad_ctx[tid] * s_attn[k] + s_dscores[k] * s_Qh[tid]; + } + grad_Qh_scratch[(long long)bi * PHA_N_HORIZONS * PHA_HIDDEN_DIM + + (long long)h * PHA_HIDDEN_DIM + tid] += dq_local; + } + __syncthreads(); + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 6c722def7..018fb12ae 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -30,6 +30,7 @@ pub mod data; pub mod eval; pub mod heads; pub mod isv; +pub mod per_horizon_attention_pool; pub mod pinned; pub mod pinned_mem; pub mod trainer; diff --git a/crates/ml-alpha/src/per_horizon_attention_pool.rs b/crates/ml-alpha/src/per_horizon_attention_pool.rs new file mode 100644 index 000000000..44f740cd1 --- /dev/null +++ b/crates/ml-alpha/src/per_horizon_attention_pool.rs @@ -0,0 +1,128 @@ +//! Per-horizon attention pool kernel host wrapper (C21). +//! +//! See `crates/ml-alpha/cuda/per_horizon_attention_pool.cu` for kernel +//! math + the design spec at +//! `docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md`. +//! +//! v1 of this binding: standalone forward + backward callable with +//! cudarc CudaSlices, used by the numgrad parity test in +//! `tests/per_horizon_attention_pool_numgrad.rs`. Wiring into +//! `PerceptionTrainer`'s captured graph is a follow-up commit gated on +//! the numgrad parity check passing here. + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use std::sync::Arc; + +pub const PHA_HIDDEN_DIM: usize = 128; +pub const PHA_BLOCK: usize = 128; +pub const PHA_N_HORIZONS: usize = 5; + +const CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/per_horizon_attention_pool.cubin")); + +pub struct PerHorizonAttentionPool { + _module: Arc, + fwd_fn: CudaFunction, + bwd_fn: CudaFunction, + stream: Arc, +} + +impl PerHorizonAttentionPool { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + let module = ctx + .load_cubin(CUBIN.to_vec()) + .context("load per_horizon_attention_pool cubin")?; + let fwd_fn = module + .load_function("per_horizon_attention_pool_fwd") + .context("load per_horizon_attention_pool_fwd")?; + let bwd_fn = module + .load_function("per_horizon_attention_pool_bwd") + .context("load per_horizon_attention_pool_bwd")?; + Ok(Self { + _module: module, + fwd_fn, + bwd_fn, + stream, + }) + } + + /// Forward. Inputs/outputs are flat row-major CudaSlices: + /// q_h: [N_HORIZONS, HIDDEN_DIM] + /// ln_out: [B, K, HIDDEN_DIM] + /// context_h: [B, N_HORIZONS, HIDDEN_DIM] (written) + /// attn_h: [B, N_HORIZONS, K] (written, post-softmax) + pub fn forward( + &self, + q_h: &CudaSlice, + ln_out: &CudaSlice, + n_batch: i32, + k_seq: i32, + context_h: &mut CudaSlice, + attn_h: &mut CudaSlice, + ) -> Result<()> { + let smem_bytes = ((k_seq as usize + PHA_BLOCK + PHA_HIDDEN_DIM) * std::mem::size_of::()) + as u32; + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (PHA_BLOCK as u32, 1, 1), + shared_mem_bytes: smem_bytes, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + unsafe { + launch + .arg(q_h) + .arg(ln_out) + .arg(&n_batch) + .arg(&k_seq) + .arg(context_h) + .arg(attn_h) + .launch(cfg) + .context("per_horizon_attention_pool_fwd")?; + } + self.stream.synchronize()?; + Ok(()) + } + + /// Backward. Adds gradients into `grad_qh_scratch` + `grad_ln_out` + /// (both `+=`). `grad_qh_scratch` is per-block scratch of shape + /// `[B, N_HORIZONS, HIDDEN_DIM]`; reduce across the batch axis + /// host-side via `reduce_axis0` to recover the shared `grad_Q_h`. + pub fn backward( + &self, + q_h: &CudaSlice, + ln_out: &CudaSlice, + attn_h: &CudaSlice, + grad_context_h: &CudaSlice, + n_batch: i32, + k_seq: i32, + grad_qh_scratch: &mut CudaSlice, + grad_ln_out: &mut CudaSlice, + ) -> Result<()> { + let smem_bytes = + ((2 * k_seq as usize + PHA_BLOCK) * std::mem::size_of::()) as u32; + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (PHA_BLOCK as u32, 1, 1), + shared_mem_bytes: smem_bytes, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + unsafe { + launch + .arg(q_h) + .arg(ln_out) + .arg(attn_h) + .arg(grad_context_h) + .arg(&n_batch) + .arg(&k_seq) + .arg(grad_qh_scratch) + .arg(grad_ln_out) + .launch(cfg) + .context("per_horizon_attention_pool_bwd")?; + } + self.stream.synchronize()?; + Ok(()) + } +} diff --git a/crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs b/crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs new file mode 100644 index 000000000..274ca4afc --- /dev/null +++ b/crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs @@ -0,0 +1,152 @@ +//! Numerical-gradient parity check for the per-horizon attention pool +//! kernel (C21). Same approach as the VSN numgrad test +//! (Phase 2D.2 / task #191): compute analytical gradients via the +//! backward kernel, then perturb each input element by ±eps, run the +//! forward kernel twice, and verify the central-difference matches +//! within tolerance. +//! +//! Two grads are checked: +//! d_Q_h [N_HORIZONS, HIDDEN_DIM] — query weights +//! d_LNb [B, K, HIDDEN_DIM] — input gradient +//! +//! Loss function is the sum over context_h (so d_context = 1 for every +//! output element), which gives a simple analytical reduction. + +use anyhow::Result; +use cudarc::driver::CudaSlice; +use ml_alpha::per_horizon_attention_pool::{PerHorizonAttentionPool, PHA_HIDDEN_DIM, PHA_N_HORIZONS}; +use ml_core::device::MlDevice; +use rand::SeedableRng; +use rand::Rng; +use rand_chacha::ChaCha8Rng; + +const B: usize = 2; +const K: usize = 8; // small K → fast numgrad + +fn try_dev() -> Option { + match MlDevice::cuda(0) { + Ok(d) => Some(d), + Err(e) => { + eprintln!("skipping: cuda device unavailable ({e})"); + None + } + } +} + +fn alloc_and_upload(stream: &std::sync::Arc, host: &[f32]) -> CudaSlice { + let mut buf = stream.alloc_zeros::(host.len()).expect("alloc"); + stream.memcpy_htod(host, &mut buf).expect("htod"); + buf +} + +fn download(stream: &std::sync::Arc, src: &CudaSlice) -> Vec { + let mut out = vec![0.0f32; src.len()]; + stream.memcpy_dtoh(src, out.as_mut_slice()).expect("dtoh"); + out +} + +fn run_forward_loss( + pool: &PerHorizonAttentionPool, + stream: &std::sync::Arc, + q_h_host: &[f32], + ln_host: &[f32], +) -> f32 { + let q_h = alloc_and_upload(stream, q_h_host); + let ln = alloc_and_upload(stream, ln_host); + let mut ctx = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM).unwrap(); + let mut attn = stream.alloc_zeros::(B * PHA_N_HORIZONS * K).unwrap(); + pool.forward(&q_h, &ln, B as i32, K as i32, &mut ctx, &mut attn).unwrap(); + let ctx_host = download(stream, &ctx); + ctx_host.iter().sum() +} + +#[test] +#[ignore = "requires CUDA"] +fn forward_then_backward_matches_central_difference() -> Result<()> { + let Some(dev) = try_dev() else { return Ok(()); }; + let ctx = dev.cuda_context()?.clone(); + let stream = dev.cuda_stream()?.clone(); + let pool = PerHorizonAttentionPool::new(&ctx, stream.clone())?; + + let mut rng = ChaCha8Rng::seed_from_u64(0xCAFE_F00D); + let q_h_host: Vec = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM) + .map(|_| rng.gen_range(-0.1..0.1)).collect(); + let ln_host: Vec = (0..B * K * PHA_HIDDEN_DIM) + .map(|_| rng.gen_range(-0.5..0.5)).collect(); + + // Analytical: backward with grad_context = 1 (since loss = Σ context). + let q_h_d = alloc_and_upload(&stream, &q_h_host); + let ln_d = alloc_and_upload(&stream, &ln_host); + let mut ctx_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut attn_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * K)?; + pool.forward(&q_h_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?; + + let grad_ctx_host = vec![1.0f32; B * PHA_N_HORIZONS * PHA_HIDDEN_DIM]; + let grad_ctx_d = alloc_and_upload(&stream, &grad_ctx_host); + let mut grad_qh_scratch_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut grad_ln_d = stream.alloc_zeros::(B * K * PHA_HIDDEN_DIM)?; + pool.backward( + &q_h_d, &ln_d, &attn_d, &grad_ctx_d, + B as i32, K as i32, + &mut grad_qh_scratch_d, &mut grad_ln_d, + )?; + let grad_qh_scratch = download(&stream, &grad_qh_scratch_d); + let grad_ln = download(&stream, &grad_ln_d); + + // Reduce grad_qh_scratch over batch axis → [N_HORIZONS, HIDDEN_DIM]. + let mut grad_qh_analytical = vec![0.0f32; PHA_N_HORIZONS * PHA_HIDDEN_DIM]; + for b in 0..B { + for i in 0..PHA_N_HORIZONS * PHA_HIDDEN_DIM { + grad_qh_analytical[i] += grad_qh_scratch[b * PHA_N_HORIZONS * PHA_HIDDEN_DIM + i]; + } + } + + let eps = 1e-2f32; + let tol = 5e-2f32; // 5% rel-tol, plus abs floor 5e-3 for tiny grads + + // 1) Probe d_Q_h at a few random indices. + let q_h_idxs: Vec = (0..8) + .map(|_| rng.gen_range(0..PHA_N_HORIZONS * PHA_HIDDEN_DIM)) + .collect(); + for &idx in &q_h_idxs { + let mut q_h_plus = q_h_host.clone(); + q_h_plus[idx] += eps; + let mut q_h_minus = q_h_host.clone(); + q_h_minus[idx] -= eps; + let l_plus = run_forward_loss(&pool, &stream, &q_h_plus, &ln_host); + let l_minus = run_forward_loss(&pool, &stream, &q_h_minus, &ln_host); + let fd = (l_plus - l_minus) / (2.0 * eps); + let analytical = grad_qh_analytical[idx]; + let abs_err = (fd - analytical).abs(); + let rel_err = abs_err / fd.abs().max(1e-6); + let pass = abs_err < 5e-3 || rel_err < tol; + assert!( + pass, + "d_Q_h[{idx}]: analytical={analytical:.6} fd={fd:.6} abs_err={abs_err:.6} rel_err={rel_err:.4}" + ); + } + + // 2) Probe d_LNb at a few random indices. + let ln_idxs: Vec = (0..8) + .map(|_| rng.gen_range(0..B * K * PHA_HIDDEN_DIM)) + .collect(); + for &idx in &ln_idxs { + let mut ln_plus = ln_host.clone(); + ln_plus[idx] += eps; + let mut ln_minus = ln_host.clone(); + ln_minus[idx] -= eps; + let l_plus = run_forward_loss(&pool, &stream, &q_h_host, &ln_plus); + let l_minus = run_forward_loss(&pool, &stream, &q_h_host, &ln_minus); + let fd = (l_plus - l_minus) / (2.0 * eps); + let analytical = grad_ln[idx]; + let abs_err = (fd - analytical).abs(); + let rel_err = abs_err / fd.abs().max(1e-6); + let pass = abs_err < 5e-3 || rel_err < tol; + assert!( + pass, + "d_LNb[{idx}]: analytical={analytical:.6} fd={fd:.6} abs_err={abs_err:.6} rel_err={rel_err:.4}" + ); + } + + Ok(()) +}