feat(ml-alpha): per-horizon attention pool kernel + numgrad parity (C21)
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
];
|
||||
|
||||
|
||||
243
crates/ml-alpha/cuda/per_horizon_attention_pool.cu
Normal file
243
crates/ml-alpha/cuda/per_horizon_attention_pool.cu
Normal file
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
128
crates/ml-alpha/src/per_horizon_attention_pool.rs
Normal file
128
crates/ml-alpha/src/per_horizon_attention_pool.rs
Normal file
@@ -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<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl PerHorizonAttentionPool {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
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<f32>,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
context_h: &mut CudaSlice<f32>,
|
||||
attn_h: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let smem_bytes = ((k_seq as usize + PHA_BLOCK + PHA_HIDDEN_DIM) * std::mem::size_of::<f32>())
|
||||
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<f32>,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
attn_h: &CudaSlice<f32>,
|
||||
grad_context_h: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
grad_qh_scratch: &mut CudaSlice<f32>,
|
||||
grad_ln_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let smem_bytes =
|
||||
((2 * k_seq as usize + PHA_BLOCK) * std::mem::size_of::<f32>()) 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(())
|
||||
}
|
||||
}
|
||||
152
crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs
Normal file
152
crates/ml-alpha/tests/per_horizon_attention_pool_numgrad.rs
Normal file
@@ -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<MlDevice> {
|
||||
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<cudarc::driver::CudaStream>, host: &[f32]) -> CudaSlice<f32> {
|
||||
let mut buf = stream.alloc_zeros::<f32>(host.len()).expect("alloc");
|
||||
stream.memcpy_htod(host, &mut buf).expect("htod");
|
||||
buf
|
||||
}
|
||||
|
||||
fn download(stream: &std::sync::Arc<cudarc::driver::CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
|
||||
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<cudarc::driver::CudaStream>,
|
||||
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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM).unwrap();
|
||||
let mut attn = stream.alloc_zeros::<f32>(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<f32> = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM)
|
||||
.map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let ln_host: Vec<f32> = (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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut grad_ln_d = stream.alloc_zeros::<f32>(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<usize> = (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<usize> = (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(())
|
||||
}
|
||||
Reference in New Issue
Block a user