diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index c652aea2a..9022ddc73 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -22,6 +22,7 @@ const KERNELS: &[&str] = &[ "variable_selection", // Phase 2D: TFT-style per-feature gating "attention_pool", // Phase 3: learned context summary at CfC k=0 "horizon_token_attention_pool", // v2-C: horizon-token K-prepend single-Q attention + "inverted_attention_pool", // v2-E: iTransformer-style cross-variate attention "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/inverted_attention_pool.cu b/crates/ml-alpha/cuda/inverted_attention_pool.cu new file mode 100644 index 000000000..139ddb956 --- /dev/null +++ b/crates/ml-alpha/cuda/inverted_attention_pool.cu @@ -0,0 +1,223 @@ +// inverted_attention_pool.cu — iTransformer-style inverted attention (v2 axis E). +// +// Standard attention treats time positions as tokens with HIDDEN_DIM +// as embedding. iTransformer inverts: each of HIDDEN_DIM features +// becomes a "variate token" with its K-trajectory as embedding. +// +// FORWARD MATH (per batch b; b indexing suppressed; H = HIDDEN_DIM): +// +// X_inv[h, k] = ln_out[k, h] # [H, K] +// scores[h, j] = (1 / sqrt(K)) · Σ_k X_inv[h, k] · X_inv[j, k] +// attn[h, j] = softmax_j(scores[h, j]) # [H, H] +// pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] # where +// # mean_k_X_inv[j] = (1/K) Σ_k X_inv[j, k] +// (which equals Σ_j attn[h, j] · (1/K) Σ_k X_inv[j, k] — +// i.e. mean-pool over K commutes with the per-j attention.) +// +// BACKWARD: gradients into X_inv from three independent chains: +// - value: d_pooled[h] · attn[h, j] · (1/K) added to X_inv[j, k] for each k +// - query: inv_scale · d_scores[h, j] · X_inv[j, k] (X_inv[h] as query) +// - key: inv_scale · d_scores[h, j] · X_inv[h, k] (X_inv[j] as key) +// +// Plus the softmax bwd: +// d_attn[h, j] = d_pooled[h] · mean_k_X_inv[j] +// d_scores[h, j] = attn[h, j] · (d_attn[h, j] − Σ_l attn[h, l] d_attn[h, l]) +// +// The bwd uses a DRAM scratch buffer `d_scores_scratch [B, H, H]` to +// pass per-(h, j) softmax-bwd values between the two phases (write all, +// __syncthreads, then read for value/query/key accumulation). +// +// PERFORMANCE: +// - Grid = (B). Block = (HIDDEN_DIM) = 128 threads = 4 warps. +// - Smem holds x_inv [H, K] = H·K floats. For K = 8..32 this is +// 4-16 KiB — well under the 48 KiB dynamic-shared limit. +// - No __syncthreads inside hot inner loops; only one block-wide +// barrier in bwd between the d_scores write and the accumulation. +// - All non-warp-uniform writes go through DRAM (per-batch slices), +// never via atomicAdd. + +#define IAP_HIDDEN_DIM 128 +#define IAP_BLOCK 128 +#define IAP_MAX_KSEQ 128 + +extern "C" __global__ void inverted_attention_pool_fwd( + const float* __restrict__ ln_out, // [B, K, H] + int n_batch, + int k_seq, + float inv_scale, // 1 / sqrt(K) + float* __restrict__ pooled_out, // [B, H] + float* __restrict__ attn_out // [B, H, H] +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= IAP_BLOCK) return; + + extern __shared__ float x_inv[]; // [H * K] + + if (tid < IAP_HIDDEN_DIM) { + const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; + for (int k = 0; k < k_seq; ++k) { + x_inv[tid * k_seq + k] = ln_b[k * IAP_HIDDEN_DIM + tid]; + } + } + __syncthreads(); + + if (tid < IAP_HIDDEN_DIM) { + const int h = tid; + const float* row_h = x_inv + (long long)h * k_seq; + + // Pass 1: scores[h, j] = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k] + // Track max while writing to attn_out (used as scratch). + float my_max = -INFINITY; + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + const float* row_j = x_inv + (long long)j * k_seq; + float dot = 0.0f; + #pragma unroll 8 + for (int k = 0; k < k_seq; ++k) { + dot += row_h[k] * row_j[k]; + } + const float s = dot * inv_scale; + attn_out[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)h * IAP_HIDDEN_DIM + j] = s; + if (s > my_max) my_max = s; + } + + // Softmax over j: read scores back from DRAM, exp+sum, divide. + float my_sum = 0.0f; + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + const long long idx = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)h * IAP_HIDDEN_DIM + j; + const float e = expf(attn_out[idx] - my_max); + attn_out[idx] = e; + my_sum += e; + } + const float inv_sum = 1.0f / my_sum; + + // pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] + float pooled_h = 0.0f; + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + const long long idx = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)h * IAP_HIDDEN_DIM + j; + const float a_hj = attn_out[idx] * inv_sum; + attn_out[idx] = a_hj; + float mean_k = 0.0f; + #pragma unroll 8 + for (int k = 0; k < k_seq; ++k) { + mean_k += x_inv[(long long)j * k_seq + k]; + } + mean_k *= (1.0f / (float)k_seq); + pooled_h += a_hj * mean_k; + } + pooled_out[(long long)b * IAP_HIDDEN_DIM + h] = pooled_h; + } +} + +extern "C" __global__ void inverted_attention_pool_bwd( + const float* __restrict__ ln_out, // [B, K, H] + const float* __restrict__ attn, // [B, H, H] + const float* __restrict__ grad_pooled, // [B, H] + int n_batch, + int k_seq, + float inv_scale, // 1 / sqrt(K) + float* __restrict__ d_scores_scratch, // [B, H, H] intermediate + float* __restrict__ grad_ln_out // [B, K, H] += per-batch +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= IAP_BLOCK) return; + + extern __shared__ float x_inv[]; + + if (tid < IAP_HIDDEN_DIM) { + const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; + for (int k = 0; k < k_seq; ++k) { + x_inv[tid * k_seq + k] = ln_b[k * IAP_HIDDEN_DIM + tid]; + } + } + __syncthreads(); + + const float dp_h = (tid < IAP_HIDDEN_DIM) + ? grad_pooled[(long long)b * IAP_HIDDEN_DIM + tid] : 0.0f; + + // Phase 1: compute d_scores[my_h, j] for all j; write to DRAM scratch. + // + // d_attn[my_h, j] = dp_h · mean_k_X_inv[j] + // dot = Σ_l attn[my_h, l] · d_attn[my_h, l] + // = dp_h · Σ_l attn[my_h, l] · mean_k_X_inv[l] + // = dp_h · pooled[my_h] ← already computed in fwd! + // But we don't have pooled here; recompute it + // cheaply since we have attn + mean_k_X_inv. + // d_scores[my_h, j] = attn[my_h, j] · (d_attn[my_h, j] − dot) + if (tid < IAP_HIDDEN_DIM) { + const int my_h = tid; + const long long attn_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)my_h * IAP_HIDDEN_DIM; + + // Compute dot = Σ_l attn[my_h, l] · mean_k_X_inv[l] in one pass. + float dot = 0.0f; + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + float mean_k = 0.0f; + #pragma unroll 8 + for (int k = 0; k < k_seq; ++k) { + mean_k += x_inv[(long long)j * k_seq + k]; + } + mean_k *= (1.0f / (float)k_seq); + dot += attn[attn_row_base + j] * mean_k; + } + const float dot_scaled = dp_h * dot; // = Σ_l attn · d_attn + + // Second pass: write d_scores[my_h, ·] to scratch. + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + float mean_k = 0.0f; + #pragma unroll 8 + for (int k = 0; k < k_seq; ++k) { + mean_k += x_inv[(long long)j * k_seq + k]; + } + mean_k *= (1.0f / (float)k_seq); + const float d_attn_my_h_j = dp_h * mean_k; + const float a_my_h_j = attn[attn_row_base + j]; + const float d_score = a_my_h_j * (d_attn_my_h_j - dot_scaled); + d_scores_scratch[attn_row_base + j] = d_score; + } + } + __syncthreads(); + + // Phase 2: accumulate d_ln_out[k, my_h] for each k. + // - value : V_h = (1/K) · Σ_h' dp[h'] · attn[h', my_h] (per-k constant) + // - query : q_k = inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k] + // - key : k_k = inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k] + if (tid < IAP_HIDDEN_DIM) { + const int my_h = tid; + + // V_h: read dp_h' from `dp_shared` (cache in smem-ish via DRAM + // re-read; cheap since H=128 reads). + float V_h = 0.0f; + for (int hp = 0; hp < IAP_HIDDEN_DIM; ++hp) { + const float a_hp_myh = attn[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)hp * IAP_HIDDEN_DIM + my_h]; + const float dp_hp = grad_pooled[(long long)b * IAP_HIDDEN_DIM + hp]; + V_h += dp_hp * a_hp_myh; + } + V_h *= (1.0f / (float)k_seq); + + float* d_ln_b = grad_ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM; + for (int k = 0; k < k_seq; ++k) { + float q_k = 0.0f; + float k_k = 0.0f; + for (int j = 0; j < IAP_HIDDEN_DIM; ++j) { + const float ds_myh_j = d_scores_scratch[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)my_h * IAP_HIDDEN_DIM + j]; + const float xj_k = x_inv[(long long)j * k_seq + k]; + q_k += ds_myh_j * xj_k; + } + for (int hp = 0; hp < IAP_HIDDEN_DIM; ++hp) { + const float ds_hp_myh = d_scores_scratch[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM + + (long long)hp * IAP_HIDDEN_DIM + my_h]; + const float xhp_k = x_inv[(long long)hp * k_seq + k]; + k_k += ds_hp_myh * xhp_k; + } + const float grad_h_k = V_h + inv_scale * (q_k + k_k); + d_ln_b[k * IAP_HIDDEN_DIM + my_h] += grad_h_k; + } + } +} diff --git a/crates/ml-alpha/src/inverted_attention_pool.rs b/crates/ml-alpha/src/inverted_attention_pool.rs new file mode 100644 index 000000000..2b0dddbd1 --- /dev/null +++ b/crates/ml-alpha/src/inverted_attention_pool.rs @@ -0,0 +1,111 @@ +//! Inverted attention pool host binding (v2 axis E, V3 commit). +//! +//! iTransformer-style cross-variate attention: each of HIDDEN_DIM +//! features becomes a "variate token" with its K-trajectory as +//! embedding. Single-block-per-batch, no learnable parameters in this +//! v2 (Q = K = V = X_inv self-attention with mean-K pool). +//! +//! See `cuda/inverted_attention_pool.cu` for math + the v2 spec §3.2. + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use std::sync::Arc; + +pub const IAP_HIDDEN_DIM: usize = 128; +pub const IAP_BLOCK: usize = 128; + +const CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/inverted_attention_pool.cubin")); + +pub struct InvertedAttentionPool { + _module: Arc, + fwd_fn: CudaFunction, + bwd_fn: CudaFunction, + stream: Arc, +} + +impl InvertedAttentionPool { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + let module = ctx + .load_cubin(CUBIN.to_vec()) + .context("load inverted_attention_pool cubin")?; + let fwd_fn = module + .load_function("inverted_attention_pool_fwd") + .context("load inverted_attention_pool_fwd")?; + let bwd_fn = module + .load_function("inverted_attention_pool_bwd") + .context("load inverted_attention_pool_bwd")?; + Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) + } + + /// Forward. `inv_scale` must equal `1 / sqrt(k_seq)` to match the + /// scaled-dot-product attention convention. + /// `ln_out` : [B, K, H] + /// `pooled_out` : [B, H] written + /// `attn_out` : [B, H, H] written (saved for backward) + pub fn forward( + &self, + ln_out: &CudaSlice, + n_batch: i32, + k_seq: i32, + inv_scale: f32, + pooled_out: &mut CudaSlice, + attn_out: &mut CudaSlice, + ) -> Result<()> { + // Smem: x_inv[H * K] + let smem_bytes = (IAP_HIDDEN_DIM * k_seq as usize * std::mem::size_of::()) as u32; + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (IAP_BLOCK as u32, 1, 1), + shared_mem_bytes: smem_bytes, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + unsafe { + launch + .arg(ln_out) + .arg(&n_batch).arg(&k_seq).arg(&inv_scale) + .arg(pooled_out).arg(attn_out) + .launch(cfg) + .context("inverted_attention_pool_fwd")?; + } + Ok(()) + } + + /// Backward. `grad_ln_out` accumulates += per-batch (no race). + /// `d_scores_scratch` is a per-step intermediate buffer of shape + /// `[B, H, H]` (must be at least `n_batch * H * H` floats). Trainer + /// allocates once and reuses across steps. + pub fn backward( + &self, + ln_out: &CudaSlice, + attn: &CudaSlice, + grad_pooled: &CudaSlice, + n_batch: i32, + k_seq: i32, + inv_scale: f32, + d_scores_scratch: &mut CudaSlice, + grad_ln_out: &mut CudaSlice, + ) -> Result<()> { + // Smem: x_inv[H * K] only (d_scores moved to DRAM scratch to fit + // under the 48 KiB dynamic-shared default on sm_86). + let smem_bytes = (IAP_HIDDEN_DIM * k_seq as usize * std::mem::size_of::()) as u32; + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (IAP_BLOCK as u32, 1, 1), + shared_mem_bytes: smem_bytes, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + unsafe { + launch + .arg(ln_out).arg(attn).arg(grad_pooled) + .arg(&n_batch).arg(&k_seq).arg(&inv_scale) + .arg(d_scores_scratch) + .arg(grad_ln_out) + .launch(cfg) + .context("inverted_attention_pool_bwd")?; + } + Ok(()) + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index cbca1f2b7..b5206705b 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 horizon_token_attention_pool; +pub mod inverted_attention_pool; pub mod isv; pub mod pinned; pub mod pinned_mem; diff --git a/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs b/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs new file mode 100644 index 000000000..3b4ed5ea3 --- /dev/null +++ b/crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs @@ -0,0 +1,127 @@ +//! Numgrad parity test for inverted_attention_pool (v2 axis E). + +#![cfg(feature = "cuda")] + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; +use ml_alpha::inverted_attention_pool::{InvertedAttentionPool, IAP_HIDDEN_DIM}; +use ml_alpha::pinned_mem::MappedF32Buffer; +use ml_core::device::MlDevice; +use std::sync::Arc; + +const B: usize = 2; +const K: usize = 8; +const H: usize = IAP_HIDDEN_DIM; + +fn rng(seed: u64) -> impl FnMut() -> f32 { + let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); + move || -> f32 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; + x * 2.0 - 1.0 + } +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n)?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), + )?; + } + Ok(dst) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), + )?; + } + stream.synchronize()?; + Ok(staging.read_all()) +} + +fn forward_loss( + pool: &InvertedAttentionPool, + stream: &Arc, + ln: &[f32], + grad_pooled: &[f32], +) -> Result { + let inv_scale = 1.0_f32 / (K as f32).sqrt(); + let ln_d = upload(stream, ln)?; + let mut pooled_d = stream.alloc_zeros::(B * H)?; + let mut attn_d = stream.alloc_zeros::(B * H * H)?; + pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?; + stream.synchronize()?; + let pooled = download(stream, &pooled_d)?; + Ok(pooled.iter().zip(grad_pooled).map(|(p, g)| p * g).sum()) +} + +#[test] +#[ignore = "requires CUDA"] +fn forward_then_backward_matches_central_difference() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream().context("stream")?.clone(); + let ctx_dev = dev.cuda_context().context("ctx")?; + let pool = InvertedAttentionPool::new(ctx_dev, stream.clone())?; + + let mut r = rng(20260518); + let ln: Vec = (0..B * K * H).map(|_| r() * 0.3).collect(); + let grad_pooled: Vec = (0..B * H).map(|_| r() * 0.1).collect(); + + let inv_scale = 1.0_f32 / (K as f32).sqrt(); + let ln_d = upload(&stream, &ln)?; + let mut pooled_d = stream.alloc_zeros::(B * H)?; + let mut attn_d = stream.alloc_zeros::(B * H * H)?; + pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?; + + let grad_pooled_d = upload(&stream, &grad_pooled)?; + let mut d_scores_scratch_d = stream.alloc_zeros::(B * H * H)?; + let mut grad_ln_d = stream.alloc_zeros::(B * K * H)?; + pool.backward( + &ln_d, &attn_d, &grad_pooled_d, + B as i32, K as i32, inv_scale, + &mut d_scores_scratch_d, + &mut grad_ln_d, + )?; + stream.synchronize()?; + + let grad_ln = download(&stream, &grad_ln_d)?; + + let eps = 1e-3_f32; + let mut r2 = rng(42); + let total = B * K * H; + let mut checks = 0; + for _ in 0..6 { + let idx = ((r2() + 1.0) * 0.5 * total as f32) as usize % total; + let mut ln_p = ln.clone(); + let mut ln_m = ln.clone(); + ln_p[idx] += eps; + ln_m[idx] -= eps; + let l_p = forward_loss(&pool, &stream, &ln_p, &grad_pooled)?; + let l_m = forward_loss(&pool, &stream, &ln_m, &grad_pooled)?; + let numgrad = (l_p - l_m) / (2.0 * eps); + let analytic = grad_ln[idx]; + let abs = (analytic - numgrad).abs(); + let rel = abs / numgrad.abs().max(1e-3); + assert!( + abs < 5e-3 || rel < 5e-2, + "ln_out[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}" + ); + checks += 1; + } + eprintln!("PASSED {checks} numgrad checks"); + Ok(()) +}