feat(ml-alpha): inverted_attention_pool kernel + numgrad (v2 E) [V3]
iTransformer-style cross-variate attention pass: each of HIDDEN_DIM
features becomes a "variate token" with its K-trajectory as embedding.
FORWARD:
X_inv[h, k] = ln_out[k, h] # transpose
scores[h, j] = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k]
attn[h, j] = softmax_j(scores)
pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] # mean-K commutes out
BACKWARD: three independent chains into ln_out:
- value path: (1/K) · attn[h', my_h] · d_pooled[h'] (per-k constant)
- query path: inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k]
- key path: inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k]
Softmax bwd: d_scores[h, j] = attn · (d_attn - Σ_l attn · d_attn)
IMPLEMENTATION NOTES:
- First attempt cached attn [H, H] = 64 KB in shared mem → tripped
the 48 KB dynamic-shared limit on sm_86 (CUDA_ERROR_INVALID_VALUE).
- Fixed by moving d_scores to a DRAM scratch buffer; shared mem
holds only X_inv [H, K] (≤ 16 KB at K = 32). One block-wide barrier
between the d_scores write and the value/query/key accumulation.
- All per-batch slice writes; no atomicAdd, no cross-block race.
- Pooled computation uses the mean-K commute (Σ_k attn · X_inv =
attn · mean_k_X_inv), saving an entire H×K accumulation pass.
LOCAL VERIFICATION (RTX 3050 sm_86):
forward_then_backward_matches_central_difference PASSES 6 numgrad
checks on ln_out positions within 5e-2 rel / 5e-3 abs.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
];
|
||||
|
||||
|
||||
223
crates/ml-alpha/cuda/inverted_attention_pool.cu
Normal file
223
crates/ml-alpha/cuda/inverted_attention_pool.cu
Normal file
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
111
crates/ml-alpha/src/inverted_attention_pool.rs
Normal file
111
crates/ml-alpha/src/inverted_attention_pool.rs
Normal file
@@ -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<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl InvertedAttentionPool {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
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<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
inv_scale: f32,
|
||||
pooled_out: &mut CudaSlice<f32>,
|
||||
attn_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// Smem: x_inv[H * K]
|
||||
let smem_bytes = (IAP_HIDDEN_DIM * k_seq as usize * std::mem::size_of::<f32>()) 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<f32>,
|
||||
attn: &CudaSlice<f32>,
|
||||
grad_pooled: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
inv_scale: f32,
|
||||
d_scores_scratch: &mut CudaSlice<f32>,
|
||||
grad_ln_out: &mut CudaSlice<f32>,
|
||||
) -> 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::<f32>()) 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(())
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
127
crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs
Normal file
127
crates/ml-alpha/tests/inverted_attention_pool_numgrad.rs
Normal file
@@ -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<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
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::<f32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
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<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
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::<f32>();
|
||||
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<CudaStream>,
|
||||
ln: &[f32],
|
||||
grad_pooled: &[f32],
|
||||
) -> Result<f32> {
|
||||
let inv_scale = 1.0_f32 / (K as f32).sqrt();
|
||||
let ln_d = upload(stream, ln)?;
|
||||
let mut pooled_d = stream.alloc_zeros::<f32>(B * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(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<f32> = (0..B * K * H).map(|_| r() * 0.3).collect();
|
||||
let grad_pooled: Vec<f32> = (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::<f32>(B * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(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::<f32>(B * H * H)?;
|
||||
let mut grad_ln_d = stream.alloc_zeros::<f32>(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(())
|
||||
}
|
||||
Reference in New Issue
Block a user