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

feat(ml-alpha): backward kernels — heads + cfc_step (K=1 BPTT)

Browse Source 
multi_horizon_heads_backward: sigmoid + linear chain rule. One block,
HIDDEN_DIM=128 threads. Computes grad_w, grad_b, grad_h_in. No
atomicAdd; per-thread accumulation only.

cfc_step_backward: truncated K=1 BPTT through one CfC time step.
Forward pre/decay/tanh recomputed inside the kernel; emits grad_w_in,
grad_w_rec, grad_b, grad_h_old. tau is held frozen (structural
log-uniform init per Hasani 2022; backprop through tau deferred to
Phase A v2 if the gate needs it). Uses dynamic shared memory for the
d_pre relay between threads (size = 2 * n_hid * 4 bytes).

Tests (4/4 on sm_86) validate via on-GPU finite-difference:
  - heads grad_h vs forward(h±eps) → matches at eps=1e-3, rel<=1%
  - heads grad_b vs forward(b±eps) → matches at eps=1e-3, rel<=1%
  - cfc grad_b vs forward(b±eps) → matches at eps=1e-3, rel<=5%
  - cfc grad_h_old vs forward(h_old±eps) → matches at eps=1e-3, rel<=5%

CPU is not the reference (per feedback_no_cpu_test_fallbacks.md). The
kernel is the truth; numerical perturbation validates the analytic
gradient against the kernel's own forward.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-16 22:20:54 +02:00
parent a6f4772617
commit 753284f2ae
5 changed files with 383 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
crates/ml-alpha/cuda/cfc_step.cu
View File

@@ -34,3 +34,77 @@ extern "C" __global__ void cfc_step(
const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
h_new[i] = h_old[i] * decay + (1.0f - decay) * tanhf(pre);
}
// Backward through ONE cfc_step (truncated BPTT, K=1).
//
// Forward (for reference):
// pre[i] = b[i] + sum_k W_in[i,k] x[k] + sum_k W_rec[i,k] h_old[k]
// decay[i] = exp(-dt / max(tau[i], 1e-6))
// s[i] = tanh(pre[i])
// h_new[i] = h_old[i] * decay + (1 - decay) * s
//
// Given grad_h_new[i] = dL/dh_new[i], compute:
// d_pre[i] = grad_h_new[i] * (1 - decay) * (1 - s^2)
// grad_b[i] += d_pre[i]
// grad_W_in[i,k] += d_pre[i] * x[k]
// grad_W_rec[i,k] += d_pre[i] * h_old[k]
// grad_h_old[i] += grad_h_new[i] * decay + sum_j d_pre[j] * W_rec[j, i]
//
// tau is held frozen (structurally log-uniform initialised per spec
// Section 2); no grad_tau path in this version.
//
// One thread per hidden unit (128 threads in one block). Block tree-
// reduce not needed — each thread owns its own grad_W rows and grad_h
// component; grad_h_old accumulates via shared `d_pre` reads.
extern "C" __global__ void cfc_step_backward(
const float* __restrict__ w_in, // [n_hid, n_in]
const float* __restrict__ w_rec, // [n_hid, n_hid]
const float* __restrict__ b,
const float* __restrict__ tau,
const float* __restrict__ x, // [n_in]
const float* __restrict__ h_old, // [n_hid]
const float* __restrict__ grad_h_new, // [n_hid]
float dt_s,
int n_in,
int n_hid,
float* __restrict__ grad_w_in, // [n_hid, n_in]
float* __restrict__ grad_w_rec, // [n_hid, n_hid]
float* __restrict__ grad_b, // [n_hid]
float* __restrict__ grad_h_old // [n_hid]
) {
extern __shared__ float smem[];
float* sd_pre = smem; // [n_hid]
float* sdecay = sd_pre + n_hid; // [n_hid]
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_hid) return;
// Recompute forward pre + tanh to derive d_pre.
float pre = b[i];
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x[k];
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old[k];
const float s = tanhf(pre);
const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
const float dh = grad_h_new[i];
const float d_pre = dh * (1.0f - decay) * (1.0f - s * s);
sd_pre[i] = d_pre;
sdecay[i] = decay;
__syncthreads();
grad_b[i] = d_pre;
for (int k = 0; k < n_in; ++k) {
grad_w_in[i * n_in + k] = d_pre * x[k];
}
for (int k = 0; k < n_hid; ++k) {
grad_w_rec[i * n_hid + k] = d_pre * h_old[k];
}
// grad_h_old[i] = dh * decay + sum_j d_pre[j] * W_rec[j, i]
float gh = dh * decay;
for (int j = 0; j < n_hid; ++j) {
gh += sd_pre[j] * w_rec[j * n_hid + i];
}
grad_h_old[i] = gh;
}

52
crates/ml-alpha/cuda/multi_horizon_heads.cu
View File

@@ -19,3 +19,55 @@ extern "C" __global__ void multi_horizon_heads(
}
probs[k] = 1.0f / (1.0f + expf(-z));
}
// Backward through multi_horizon_heads.
//
// Given:
// grad_probs[5] = dL / dp
// probs[5] = forward output (sigmoid)
// h[128] = forward input
// Computes:
// grad_w[5, 128] = dL / dW = d_z[k] * h[i]
// grad_b[5] = dL / db = d_z[k]
// grad_h[128] = dL / dh = sum_k d_z[k] * W[k, i]
// where d_z[k] = grad_probs[k] * p[k] * (1 - p[k]).
//
// Block tree-reduce for the grad_h sum (no atomicAdd). One thread per
// hidden unit (128 threads); thread 0 also writes grad_b.
extern "C" __global__ void multi_horizon_heads_backward(
const float* __restrict__ w, // [5, 128]
const float* __restrict__ probs, // [5]
const float* __restrict__ h, // [128]
const float* __restrict__ grad_probs, // [5]
float* __restrict__ grad_w, // [5, 128]
float* __restrict__ grad_b, // [5]
float* __restrict__ grad_h // [128]
) {
__shared__ float d_z[5];
int tid = threadIdx.x;
if (tid < 5) {
const float p = probs[tid];
d_z[tid] = grad_probs[tid] * p * (1.0f - p);
}
__syncthreads();
if (tid < 5) {
grad_b[tid] = d_z[tid];
}
// grad_w[k, i] = d_z[k] * h[i]
// grid stride: each thread covers one i for k in 0..5.
if (tid < 128) {
const float h_i = h[tid];
for (int k = 0; k < 5; ++k) {
grad_w[k * 128 + tid] = d_z[k] * h_i;
}
// grad_h[i] = sum_k d_z[k] * W[k, i]
float acc = 0.0f;
for (int k = 0; k < 5; ++k) {
acc += d_z[k] * w[k * 128 + tid];
}
grad_h[tid] = acc;
}
}

60
crates/ml-alpha/src/cfc/step.rs
View File

@@ -20,6 +20,66 @@ pub struct CfcWeights {
pub n_hid: usize,
}
/// Backward through ONE cfc_step (truncated BPTT, K=1).
///
/// Returns (grad_w_in, grad_w_rec, grad_b, grad_h_old).
pub fn cfc_step_backward_gpu(
dev: &MlDevice,
w: &CfcWeights,
x: &[f32],
h_old: &[f32],
grad_h_new: &[f32],
dt_s: f32,
) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>)> {
assert_eq!(x.len(), w.n_in);
assert_eq!(h_old.len(), w.n_hid);
assert_eq!(grad_h_new.len(), w.n_hid);
let stream: &Arc<CudaStream> = dev.cuda_stream().context("cfc-bwd stream")?;
let ctx = dev.cuda_context().context("cfc-bwd ctx")?;
let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load cfc cubin")?;
let func = module.load_function("cfc_step_backward").context("cfc_bwd fn")?;
let w_in_d = upload(stream, &w.w_in)?;
let w_rec_d = upload(stream, &w.w_rec)?;
let b_d = upload(stream, &w.b)?;
let tau_d = upload(stream, &w.tau)?;
let x_d = upload(stream, x)?;
let h_old_d = upload(stream, h_old)?;
let grad_h_new_d = upload(stream, grad_h_new)?;
let mut grad_w_in_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_in).context("grad_w_in alloc")?;
let mut grad_w_rec_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_hid).context("grad_w_rec alloc")?;
let mut grad_b_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_b alloc")?;
let mut grad_h_old_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_h_old alloc")?;
let n_in_i = w.n_in as i32;
let n_hid_i = w.n_hid as i32;
let block_dim = 128.min(w.n_hid as u32);
let grid_dim = ((w.n_hid as u32) + block_dim - 1) / block_dim;
let shared_mem = (2 * w.n_hid * std::mem::size_of::<f32>()) as u32;
let cfg = LaunchConfig {
grid_dim: (grid_dim, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: shared_mem,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&w_in_d).arg(&w_rec_d).arg(&b_d).arg(&tau_d)
.arg(&x_d).arg(&h_old_d).arg(&grad_h_new_d)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
.arg(&mut grad_w_in_d).arg(&mut grad_w_rec_d)
.arg(&mut grad_b_d).arg(&mut grad_h_old_d);
unsafe { launch.launch(cfg).context("cfc_bwd launch")?; }
Ok((
download(stream, &grad_w_in_d)?,
download(stream, &grad_w_rec_d)?,
download(stream, &grad_b_d)?,
download(stream, &grad_h_old_d)?,
))
}
pub fn cfc_step_gpu(
dev: &MlDevice,
w: &CfcWeights,

41
crates/ml-alpha/src/heads.rs
View File

@@ -39,6 +39,47 @@ pub struct MultiHorizonHeads;
pub struct Projection;
/// Backward through `multi_horizon_heads`. Returns (grad_w, grad_b, grad_h).
pub fn multi_horizon_heads_backward_gpu(
dev: &MlDevice,
w: &HeadsWeights,
probs: &[f32; N_HORIZONS],
h: &[f32],
grad_probs: &[f32; N_HORIZONS],
) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
assert_eq!(h.len(), HIDDEN_DIM);
let stream: &Arc<CudaStream> = dev.cuda_stream().context("heads-bwd stream")?;
let ctx = dev.cuda_context().context("heads-bwd ctx")?;
let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
let func = module.load_function("multi_horizon_heads_backward").context("heads_bwd fn")?;
let w_d = upload(stream, &w.w)?;
let probs_d = upload(stream, probs.as_slice())?;
let h_d = upload(stream, h)?;
let grad_p_d = upload(stream, grad_probs.as_slice())?;
let mut grad_w_d = stream.alloc_zeros::<f32>(N_HORIZONS * HIDDEN_DIM).context("grad_w alloc")?;
let mut grad_b_d = stream.alloc_zeros::<f32>(N_HORIZONS).context("grad_b alloc")?;
let mut grad_h_d = stream.alloc_zeros::<f32>(HIDDEN_DIM).context("grad_h alloc")?;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&w_d).arg(&probs_d).arg(&h_d).arg(&grad_p_d)
.arg(&mut grad_w_d).arg(&mut grad_b_d).arg(&mut grad_h_d);
unsafe { launch.launch(cfg).context("heads_bwd launch")?; }
Ok((
download(stream, &grad_w_d)?,
download(stream, &grad_b_d)?,
download(stream, &grad_h_d)?,
))
}
pub fn projection_gpu(
dev: &MlDevice,
w: &ProjectionWeights,

156
crates/ml-alpha/tests/backward_finite_diff.rs Normal file
View File

@@ -0,0 +1,156 @@
//! Finite-difference validation of cfc_step_backward + multi_horizon_heads_backward.
//!
//! Per `feedback_no_cpu_test_fallbacks.md`: the kernel is the reference.
//! We perturb a single input, re-run the GPU forward, and compare the
//! resulting finite-difference gradient to the analytic backward.
use approx::assert_relative_eq;
use ml_alpha::cfc::step::{cfc_step_backward_gpu, cfc_step_gpu, CfcWeights};
use ml_alpha::heads::{
multi_horizon_heads_backward_gpu, multi_horizon_heads_gpu, HeadsWeights, HIDDEN_DIM, N_HORIZONS,
};
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
fn test_device() -> MlDevice {
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}
fn rand_heads(seed: u64) -> HeadsWeights {
let mut r = ChaCha8Rng::seed_from_u64(seed);
HeadsWeights {
w: (0..N_HORIZONS * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(),
b: (0..N_HORIZONS).map(|_| r.gen_range(-0.01..0.01)).collect(),
}
}
fn rand_cfc(seed: u64, n_in: usize, n_hid: usize) -> CfcWeights {
let mut r = ChaCha8Rng::seed_from_u64(seed);
CfcWeights {
w_in: (0..n_hid * n_in).map(|_| r.gen_range(-0.1..0.1)).collect(),
w_rec: (0..n_hid * n_hid).map(|_| r.gen_range(-0.05..0.05)).collect(),
b: (0..n_hid).map(|_| r.gen_range(-0.01..0.01)).collect(),
tau: (0..n_hid).map(|_| r.gen_range(0.05..2.0)).collect(),
n_in,
n_hid,
}
}
#[test]
fn heads_backward_grad_h_matches_finite_diff() {
let dev = test_device();
let w = rand_heads(0x1111);
let mut r = ChaCha8Rng::seed_from_u64(0x2222);
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
let probs = multi_horizon_heads_gpu(&dev, &w, &h).unwrap();
// Surrogate loss L = sum_k probs[k] -> dL/dp = 1 for all k.
let grad_probs = [1.0f32; N_HORIZONS];
let (_, _, grad_h_analytic) =
multi_horizon_heads_backward_gpu(&dev, &w, &probs, &h, &grad_probs).unwrap();
let eps = 1e-3_f32;
for &i in &[0usize, 32, 64, 96, 127] {
let mut h_up = h.clone();
h_up[i] += eps;
let mut h_dn = h.clone();
h_dn[i] -= eps;
let p_up = multi_horizon_heads_gpu(&dev, &w, &h_up).unwrap();
let p_dn = multi_horizon_heads_gpu(&dev, &w, &h_dn).unwrap();
let l_up: f32 = p_up.iter().sum();
let l_dn: f32 = p_dn.iter().sum();
let fd = (l_up - l_dn) / (2.0 * eps);
assert_relative_eq!(grad_h_analytic[i], fd, epsilon = 1e-3, max_relative = 1e-2);
}
}
#[test]
fn heads_backward_grad_b_matches_finite_diff() {
let dev = test_device();
let w = rand_heads(0x3333);
let mut r = ChaCha8Rng::seed_from_u64(0x4444);
let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
let probs = multi_horizon_heads_gpu(&dev, &w, &h).unwrap();
let grad_probs = [1.0f32; N_HORIZONS];
let (_, grad_b_analytic, _) =
multi_horizon_heads_backward_gpu(&dev, &w, &probs, &h, &grad_probs).unwrap();
let eps = 1e-3_f32;
for k in 0..N_HORIZONS {
let mut w_up = w.clone();
w_up.b[k] += eps;
let mut w_dn = w.clone();
w_dn.b[k] -= eps;
let p_up = multi_horizon_heads_gpu(&dev, &w_up, &h).unwrap();
let p_dn = multi_horizon_heads_gpu(&dev, &w_dn, &h).unwrap();
let l_up: f32 = p_up.iter().sum();
let l_dn: f32 = p_dn.iter().sum();
let fd = (l_up - l_dn) / (2.0 * eps);
assert_relative_eq!(grad_b_analytic[k], fd, epsilon = 1e-3, max_relative = 1e-2);
let _ = w; // suppress
let _ = w_up;
let _ = w_dn;
}
let _ = w;
}
#[test]
fn cfc_backward_grad_b_matches_finite_diff() {
let dev = test_device();
let n_in = 8;
let n_hid = 16;
let w = rand_cfc(0x5555, n_in, n_hid);
let mut r = ChaCha8Rng::seed_from_u64(0x6666);
let x: Vec<f32> = (0..n_in).map(|_| r.gen_range(-1.0..1.0)).collect();
let h_old: Vec<f32> = (0..n_hid).map(|_| r.gen_range(-1.0..1.0)).collect();
let dt_s = 0.02_f32;
// Surrogate loss L = sum_i h_new[i].
let grad_h_new = vec![1.0f32; n_hid];
let (_, _, grad_b_analytic, _) =
cfc_step_backward_gpu(&dev, &w, &x, &h_old, &grad_h_new, dt_s).unwrap();
let eps = 1e-3_f32;
for &i in &[0usize, 4, 8, 15] {
let mut w_up = w.clone();
w_up.b[i] += eps;
let mut w_dn = w.clone();
w_dn.b[i] -= eps;
let h_up = cfc_step_gpu(&dev, &w_up, &x, &h_old, dt_s).unwrap();
let h_dn = cfc_step_gpu(&dev, &w_dn, &x, &h_old, dt_s).unwrap();
let l_up: f32 = h_up.iter().sum();
let l_dn: f32 = h_dn.iter().sum();
let fd = (l_up - l_dn) / (2.0 * eps);
assert_relative_eq!(grad_b_analytic[i], fd, epsilon = 5e-3, max_relative = 5e-2);
}
}
#[test]
fn cfc_backward_grad_h_old_matches_finite_diff() {
let dev = test_device();
let n_in = 8;
let n_hid = 16;
let w = rand_cfc(0x7777, n_in, n_hid);
let mut r = ChaCha8Rng::seed_from_u64(0x8888);
let x: Vec<f32> = (0..n_in).map(|_| r.gen_range(-1.0..1.0)).collect();
let h_old: Vec<f32> = (0..n_hid).map(|_| r.gen_range(-1.0..1.0)).collect();
let dt_s = 0.02_f32;
let grad_h_new = vec![1.0f32; n_hid];
let (_, _, _, grad_h_old_analytic) =
cfc_step_backward_gpu(&dev, &w, &x, &h_old, &grad_h_new, dt_s).unwrap();
let eps = 1e-3_f32;
for &i in &[0usize, 4, 8, 15] {
let mut h_up = h_old.clone();
h_up[i] += eps;
let mut h_dn = h_old.clone();
h_dn[i] -= eps;
let h_new_up = cfc_step_gpu(&dev, &w, &x, &h_up, dt_s).unwrap();
let h_new_dn = cfc_step_gpu(&dev, &w, &x, &h_dn, dt_s).unwrap();
let l_up: f32 = h_new_up.iter().sum();
let l_dn: f32 = h_new_dn.iter().sum();
let fd = (l_up - l_dn) / (2.0 * eps);
assert_relative_eq!(grad_h_old_analytic[i], fd, epsilon = 5e-3, max_relative = 5e-2);
}
}