feat(per-horizon-cfc): add cfc_step_per_branch.cu (fused fwd+bwd)
Per spec §5.4 point 1. Single fused kernel: grid=(B, N_HORIZONS, 1), block=(MAX_BUCKET_DIM=28, 1, 1) uniform predicate. Shared x_local + h_old_local cooperative-staged into shared memory once per block per pearl_cooperative_staging_eliminates_redundant_reads. No atomicAdd. No host branches. Forward GPU oracle: matches naive per-channel CfC math (tol 1e-4) on sm_86. Backward smoke: kernel loads + writes non-NaN gradients. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -25,6 +25,7 @@ const KERNELS: &[&str] = &[
|
||||
"smoothness_lambda_controller", // CRT.train: ISV-driven λ controller anchored on h30 jitter
|
||||
"gpu_log_ring", // GPU diagnostic log ring — tick kernel + log_record helper
|
||||
"bucket_transition_kernels", // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, tau_reorder, heads_compact
|
||||
"cfc_step_per_branch", // Per-horizon CfC Phase 2: fused per-(batch, branch) fwd + bwd over [25,25,25,25,28] buckets
|
||||
];
|
||||
|
||||
// Cache bust v15 (2026-05-21): gpu_log_helpers.cuh extracted + smoothness_lambda_controller emits records.
|
||||
|
||||
230
crates/ml-alpha/cuda/cfc_step_per_branch.cu
Normal file
230
crates/ml-alpha/cuda/cfc_step_per_branch.cu
Normal file
@@ -0,0 +1,230 @@
|
||||
// cfc_step_per_branch.cu — fused per-branch CfC step (forward + backward).
|
||||
//
|
||||
// Per spec §5.4 point 1 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md):
|
||||
// Single fused kernel covering all 5 branches × n_batch in one launch.
|
||||
// Grid: (B, N_HORIZONS, 1), block: (MAX_BUCKET_DIM=28, 1, 1) uniform predicate.
|
||||
//
|
||||
// Uniform predicate (`threadIdx.x >= bucket_dim_k[branch]` early-return) handles
|
||||
// uneven bucket sizes [25, 25, 25, 25, 28] without warp divergence — the
|
||||
// comparison is against a per-block constant, so all threads in a warp take
|
||||
// the same branch.
|
||||
//
|
||||
// Per pearl_cooperative_staging_eliminates_redundant_reads: the per-batch
|
||||
// `x` and `h_old` rows are SHARED across all output channels in a block
|
||||
// (every thread reads HIDDEN_DIM elements of both in the K-loop). Stage
|
||||
// them into shared memory once at block entry. W_in / W_rec rows are
|
||||
// per-channel (different per thread c), so they remain direct DRAM reads.
|
||||
//
|
||||
// Per feedback_no_atomicadd.md: no atomicAdd anywhere. Each thread writes
|
||||
// to its own (batch, c) slot in forward, and to its own per-batch grad
|
||||
// slice in backward. Cross-batch reduction is the caller's responsibility
|
||||
// via the existing `reduce_axis0` infrastructure.
|
||||
//
|
||||
// Per pearl_no_host_branches_in_captured_graph: no host branching; all
|
||||
// control flow uses device-resident metadata read from device tensors.
|
||||
//
|
||||
// Per feedback_nvidia_grade_perf_for_kernels.md: warp-uniform branches,
|
||||
// no atomicAdd, coalesced loads via cooperative staging.
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define N_HORIZONS 5
|
||||
#define MAX_BUCKET_DIM 28
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// cfc_step_per_branch_fwd: forward pass.
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, N_HORIZONS, 1)
|
||||
// block = (MAX_BUCKET_DIM = 28, 1, 1)
|
||||
// shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = 1024 bytes
|
||||
//
|
||||
// Per (batch, branch, channel-in-bucket) thread computes:
|
||||
// c = bucket_channel_offset[branch] + threadIdx.x
|
||||
// pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k]
|
||||
// decay = exp(-dt / max(tau_all[c], 1e-6))
|
||||
// h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre)
|
||||
//
|
||||
// x_local[HIDDEN_DIM] and h_old_local[HIDDEN_DIM] are cooperative-staged
|
||||
// in shared memory; without them every thread re-reads the full row
|
||||
// (HIDDEN_DIM × bucket_dim threads → up to 28 × redundant reads per row).
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void cfc_step_per_branch_fwd(
|
||||
const float* __restrict__ w_in, // [HIDDEN_DIM × HIDDEN_DIM]
|
||||
const float* __restrict__ w_rec, // [HIDDEN_DIM × HIDDEN_DIM]
|
||||
const float* __restrict__ b, // [HIDDEN_DIM]
|
||||
const float* __restrict__ tau_all, // [HIDDEN_DIM] (bucket-grouped)
|
||||
const float* __restrict__ x, // [B × HIDDEN_DIM]
|
||||
const float* __restrict__ h_old, // [B × HIDDEN_DIM]
|
||||
float dt,
|
||||
int B,
|
||||
const unsigned int* __restrict__ bucket_channel_offset, // [N_HORIZONS + 1]
|
||||
const unsigned int* __restrict__ bucket_dim_k, // [N_HORIZONS]
|
||||
float* __restrict__ h_new // [B × HIDDEN_DIM]
|
||||
) {
|
||||
extern __shared__ float smem[];
|
||||
float* x_local = smem; // [HIDDEN_DIM]
|
||||
float* h_old_local = smem + HIDDEN_DIM; // [HIDDEN_DIM]
|
||||
|
||||
int batch = blockIdx.x;
|
||||
int branch = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
if (batch >= B) return;
|
||||
|
||||
// Cooperative staging of per-batch x[batch, *] and h_old[batch, *] rows.
|
||||
// All threads (including those that will early-return on the per-branch
|
||||
// predicate) participate in the staging — the row is shared across the
|
||||
// block's output channels, so we need every thread to help load.
|
||||
//
|
||||
// Each thread loads HIDDEN_DIM / blockDim.x = 128 / 28 ≈ 5 elements
|
||||
// (rounded up via the stride loop).
|
||||
for (int i = tid; i < HIDDEN_DIM; i += blockDim.x) {
|
||||
x_local[i] = x[batch * HIDDEN_DIM + i];
|
||||
h_old_local[i] = h_old[batch * HIDDEN_DIM + i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Uniform per-branch predicate: idle threads in shorter buckets early-return.
|
||||
// The comparison is against a per-block constant (bucket_dim_k[branch]),
|
||||
// so all threads in a warp take the same path → no warp divergence.
|
||||
unsigned int bucket_dim = bucket_dim_k[branch];
|
||||
if ((unsigned int)tid >= bucket_dim) return;
|
||||
|
||||
unsigned int bucket_start = bucket_channel_offset[branch];
|
||||
int c = (int)bucket_start + tid;
|
||||
// Defensive: c must be < HIDDEN_DIM by construction of bucket_channel_offset.
|
||||
// (bucket_channel_offset[5] == HIDDEN_DIM == 128.)
|
||||
if (c >= HIDDEN_DIM) return;
|
||||
|
||||
// pre = b[c] + Σ_k W_in[c, k] * x_local[k] + Σ_k W_rec[c, k] * h_old_local[k]
|
||||
float pre = b[c];
|
||||
for (int k = 0; k < HIDDEN_DIM; ++k) {
|
||||
pre += w_in[c * HIDDEN_DIM + k] * x_local[k];
|
||||
}
|
||||
for (int k = 0; k < HIDDEN_DIM; ++k) {
|
||||
pre += w_rec[c * HIDDEN_DIM + k] * h_old_local[k];
|
||||
}
|
||||
|
||||
float decay = expf(-dt / fmaxf(tau_all[c], 1e-6f));
|
||||
h_new[batch * HIDDEN_DIM + c] =
|
||||
h_old_local[c] * decay + (1.0f - decay) * tanhf(pre);
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// cfc_step_per_branch_bwd: backward pass.
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, N_HORIZONS, 1)
|
||||
// block = (MAX_BUCKET_DIM = 28, 1, 1)
|
||||
// shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = 1024 bytes
|
||||
//
|
||||
// Per-batch grad slices are written by this kernel; cross-batch
|
||||
// reduction is the caller's responsibility (existing `reduce_axis0`
|
||||
// pattern; see crates/ml-alpha/cuda/cfc_step.cu Phase B comment block).
|
||||
//
|
||||
// grad shapes:
|
||||
// grad_w_in : [B × HIDDEN_DIM × HIDDEN_DIM] per-batch scratch (overwrite)
|
||||
// grad_w_rec : [B × HIDDEN_DIM × HIDDEN_DIM] per-batch scratch (overwrite)
|
||||
// grad_b : [B × HIDDEN_DIM] per-batch scratch (overwrite)
|
||||
// grad_tau_all: [B × HIDDEN_DIM] per-batch scratch (overwrite)
|
||||
// grad_h_old : [B × HIDDEN_DIM] per-batch (overwrite)
|
||||
//
|
||||
// Note: this kernel uses OVERWRITE semantics (the caller zeros scratch
|
||||
// per training step). The single-step CfC backward in the existing
|
||||
// `cfc_step_backward_batched` uses += for K-loop accumulation; for
|
||||
// per-branch usage in Phase 2 the K-loop's per-position accumulation is
|
||||
// the caller's concern (matching existing scratch + reduce pattern).
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void cfc_step_per_branch_bwd(
|
||||
const float* __restrict__ w_in,
|
||||
const float* __restrict__ w_rec,
|
||||
const float* __restrict__ b,
|
||||
const float* __restrict__ tau_all,
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ h_old,
|
||||
const float* __restrict__ grad_h_new,
|
||||
float dt,
|
||||
int B,
|
||||
const unsigned int* __restrict__ bucket_channel_offset,
|
||||
const unsigned int* __restrict__ bucket_dim_k,
|
||||
float* __restrict__ grad_w_in, // [B × HIDDEN_DIM × HIDDEN_DIM]
|
||||
float* __restrict__ grad_w_rec, // [B × HIDDEN_DIM × HIDDEN_DIM]
|
||||
float* __restrict__ grad_b, // [B × HIDDEN_DIM]
|
||||
float* __restrict__ grad_tau_all, // [B × HIDDEN_DIM]
|
||||
float* __restrict__ grad_h_old // [B × HIDDEN_DIM]
|
||||
) {
|
||||
extern __shared__ float smem[];
|
||||
float* x_local = smem; // [HIDDEN_DIM]
|
||||
float* h_old_local = smem + HIDDEN_DIM; // [HIDDEN_DIM]
|
||||
|
||||
int batch = blockIdx.x;
|
||||
int branch = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
if (batch >= B) return;
|
||||
|
||||
// Cooperative staging — same pattern as forward.
|
||||
for (int i = tid; i < HIDDEN_DIM; i += blockDim.x) {
|
||||
x_local[i] = x[batch * HIDDEN_DIM + i];
|
||||
h_old_local[i] = h_old[batch * HIDDEN_DIM + i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
unsigned int bucket_dim = bucket_dim_k[branch];
|
||||
if ((unsigned int)tid >= bucket_dim) return;
|
||||
|
||||
unsigned int bucket_start = bucket_channel_offset[branch];
|
||||
int c = (int)bucket_start + tid;
|
||||
if (c >= HIDDEN_DIM) return;
|
||||
|
||||
// Recompute forward pre + tanh to derive d_pre.
|
||||
float pre = b[c];
|
||||
for (int k = 0; k < HIDDEN_DIM; ++k) {
|
||||
pre += w_in[c * HIDDEN_DIM + k] * x_local[k];
|
||||
}
|
||||
for (int k = 0; k < HIDDEN_DIM; ++k) {
|
||||
pre += w_rec[c * HIDDEN_DIM + k] * h_old_local[k];
|
||||
}
|
||||
const float tau_eps = 1e-6f;
|
||||
const float tau_raw = tau_all[c];
|
||||
const float tau_eff = fmaxf(tau_raw, tau_eps);
|
||||
const float decay = expf(-dt / tau_eff);
|
||||
const float s = tanhf(pre);
|
||||
const float dh = grad_h_new[batch * HIDDEN_DIM + c];
|
||||
const float d_pre = dh * (1.0f - decay) * (1.0f - s * s);
|
||||
const float d_decay = dh * (h_old_local[c] - s);
|
||||
|
||||
// Per-batch grad scratch writes — each thread (batch, branch, c) is the
|
||||
// sole writer of grad_*[batch, c, ...]. No race; no atomicAdd.
|
||||
grad_b[batch * HIDDEN_DIM + c] = d_pre;
|
||||
|
||||
// grad_tau receives 0 when tau is at the clamp floor (strict d(max)/d(tau) = 0).
|
||||
const float gate = (tau_raw > tau_eps) ? 1.0f : 0.0f;
|
||||
grad_tau_all[batch * HIDDEN_DIM + c] =
|
||||
gate * d_decay * decay * dt / (tau_eff * tau_eff);
|
||||
|
||||
// grad_w_in[batch, c, k] = d_pre * x_local[k]
|
||||
// grad_w_rec[batch, c, k] = d_pre * h_old_local[k]
|
||||
//
|
||||
// Layout note: thread serves the (batch, c) role, sweeping k in the inner
|
||||
// loop. Writes along k are coalesced within a thread but strided across
|
||||
// threads (each thread writes its own [c, *] row in a per-(batch) slab).
|
||||
// Per pearl_coalesce_via_thread_role_swap: this is single-thread-writes-
|
||||
// single-row, the warp-coalescing concern is on cross-thread writes to
|
||||
// the SAME row, which we don't do here. No swap needed.
|
||||
for (int k = 0; k < HIDDEN_DIM; ++k) {
|
||||
const long long off =
|
||||
(long long)batch * HIDDEN_DIM * HIDDEN_DIM
|
||||
+ (long long)c * HIDDEN_DIM + k;
|
||||
grad_w_in[off] = d_pre * x_local[k];
|
||||
grad_w_rec[off] = d_pre * h_old_local[k];
|
||||
}
|
||||
|
||||
// grad_h_old[batch, c] receives the direct decay contribution. The cross-
|
||||
// channel term Σ_j d_pre[j] * W_rec[j, c] requires a reduction across
|
||||
// threads in the block — the caller may run a second-pass reduction
|
||||
// kernel when grad_h_old needs to feed backward through a prior CfC
|
||||
// step (currently CfC is K=1 in per-branch Phase 2, so this direct
|
||||
// contribution is the only one).
|
||||
grad_h_old[batch * HIDDEN_DIM + c] = dh * decay;
|
||||
}
|
||||
247
crates/ml-alpha/tests/cfc_step_per_branch.rs
Normal file
247
crates/ml-alpha/tests/cfc_step_per_branch.rs
Normal file
@@ -0,0 +1,247 @@
|
||||
//! GPU oracle tests for cfc_step_per_branch.cu fused per-branch CfC step.
|
||||
//!
|
||||
//! Run with: SQLX_OFFLINE=true cargo test -p ml-alpha --test cfc_step_per_branch -- --ignored --nocapture
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
const KERNEL_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step_per_branch.cubin"));
|
||||
const HIDDEN_DIM: usize = 128;
|
||||
const N_HORIZONS: usize = 5;
|
||||
const MAX_BUCKET_DIM: usize = 28;
|
||||
|
||||
fn load_kernel(
|
||||
stream: &std::sync::Arc<CudaStream>,
|
||||
fn_name: &str,
|
||||
) -> Result<cudarc::driver::CudaFunction> {
|
||||
let ctx = stream.context();
|
||||
let module = ctx
|
||||
.load_cubin(KERNEL_CUBIN.to_vec())
|
||||
.context("load cfc_step_per_branch cubin")?;
|
||||
module
|
||||
.load_function(fn_name)
|
||||
.with_context(|| format!("load {fn_name}"))
|
||||
}
|
||||
|
||||
fn upload<T: cudarc::driver::DeviceRepr + cudarc::driver::ValidAsZeroBits + Copy>(
|
||||
stream: &std::sync::Arc<CudaStream>,
|
||||
host: &[T],
|
||||
) -> Result<CudaSlice<T>> {
|
||||
let mut d = stream.alloc_zeros::<T>(host.len())?;
|
||||
stream.memcpy_htod(host, &mut d)?;
|
||||
Ok(d)
|
||||
}
|
||||
|
||||
fn download<T: cudarc::driver::DeviceRepr + Default + Copy>(
|
||||
stream: &std::sync::Arc<CudaStream>,
|
||||
d: &CudaSlice<T>,
|
||||
) -> Result<Vec<T>> {
|
||||
let mut h = vec![T::default(); d.len()];
|
||||
stream.memcpy_dtoh(d, &mut h)?;
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn fwd_matches_naive_per_branch() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let func = load_kernel(&stream, "cfc_step_per_branch_fwd")?;
|
||||
|
||||
let b_sz: usize = 2;
|
||||
// Inputs: x [B, HIDDEN_DIM], h_old [B, HIDDEN_DIM]
|
||||
// W_in, W_rec shared; b [HIDDEN_DIM]; tau_all [HIDDEN_DIM] (one per channel, bucket-grouped)
|
||||
let x: Vec<f32> = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect();
|
||||
let h_old: Vec<f32> = (0..b_sz * HIDDEN_DIM)
|
||||
.map(|i| ((i as f32) * 0.003).sin())
|
||||
.collect();
|
||||
let w_in: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
|
||||
.map(|i| ((i % 17) as f32) * 0.01)
|
||||
.collect();
|
||||
let w_rec: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
|
||||
.map(|i| ((i % 13) as f32) * 0.005)
|
||||
.collect();
|
||||
let b_vec: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect();
|
||||
let tau_all: Vec<f32> = (0..HIDDEN_DIM)
|
||||
.map(|c| ((c / 25 + 1) as f32) * 10.0)
|
||||
.collect();
|
||||
let bucket_channel_offset: Vec<u32> = vec![0, 25, 50, 75, 100, 128];
|
||||
let bucket_dim_k: Vec<u32> = vec![25, 25, 25, 25, 28];
|
||||
|
||||
// Upload
|
||||
let x_d = upload(&stream, &x)?;
|
||||
let h_old_d = upload(&stream, &h_old)?;
|
||||
let w_in_d = upload(&stream, &w_in)?;
|
||||
let w_rec_d = upload(&stream, &w_rec)?;
|
||||
let b_d = upload(&stream, &b_vec)?;
|
||||
let tau_d = upload(&stream, &tau_all)?;
|
||||
let bco_d = upload(&stream, &bucket_channel_offset)?;
|
||||
let bdk_d = upload(&stream, &bucket_dim_k)?;
|
||||
let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
|
||||
|
||||
let dt: f32 = 1.0;
|
||||
let b_i32: i32 = b_sz as i32;
|
||||
|
||||
// shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = x_local + h_old_local
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, N_HORIZONS as u32, 1),
|
||||
block_dim: (MAX_BUCKET_DIM as u32, 1, 1),
|
||||
shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
|
||||
};
|
||||
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(&dt)
|
||||
.arg(&b_i32)
|
||||
.arg(&bco_d)
|
||||
.arg(&bdk_d)
|
||||
.arg(&mut h_new_d);
|
||||
unsafe {
|
||||
launch.launch(cfg)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
|
||||
let h_new = download(&stream, &h_new_d)?;
|
||||
|
||||
// Naive reference: same math per-channel
|
||||
for batch in 0..b_sz {
|
||||
for c in 0..HIDDEN_DIM {
|
||||
let mut pre = b_vec[c];
|
||||
for k in 0..HIDDEN_DIM {
|
||||
pre += w_in[c * HIDDEN_DIM + k] * x[batch * HIDDEN_DIM + k];
|
||||
}
|
||||
for k in 0..HIDDEN_DIM {
|
||||
pre += w_rec[c * HIDDEN_DIM + k] * h_old[batch * HIDDEN_DIM + k];
|
||||
}
|
||||
let decay = (-dt / tau_all[c].max(1e-6)).exp();
|
||||
let expected =
|
||||
h_old[batch * HIDDEN_DIM + c] * decay + (1.0 - decay) * pre.tanh();
|
||||
let got = h_new[batch * HIDDEN_DIM + c];
|
||||
assert!(
|
||||
(got - expected).abs() < 1e-4,
|
||||
"h_new[{},{}]={} expected {}",
|
||||
batch,
|
||||
c,
|
||||
got,
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn bwd_finite_diff_matches() -> Result<()> {
|
||||
// Backward smoke for Task 4: kernel loads cleanly + writes non-NaN gradients.
|
||||
// Full numeric finite-diff validation deferred to Task 14 local validation.
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream()?.clone();
|
||||
let func_bwd = load_kernel(&stream, "cfc_step_per_branch_bwd")?;
|
||||
|
||||
let b_sz: usize = 2;
|
||||
let x: Vec<f32> = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect();
|
||||
let h_old: Vec<f32> = (0..b_sz * HIDDEN_DIM)
|
||||
.map(|i| ((i as f32) * 0.003).sin())
|
||||
.collect();
|
||||
let w_in: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
|
||||
.map(|i| ((i % 17) as f32) * 0.01)
|
||||
.collect();
|
||||
let w_rec: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
|
||||
.map(|i| ((i % 13) as f32) * 0.005)
|
||||
.collect();
|
||||
let b_vec: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect();
|
||||
let tau_all: Vec<f32> = (0..HIDDEN_DIM)
|
||||
.map(|c| ((c / 25 + 1) as f32) * 10.0)
|
||||
.collect();
|
||||
let grad_h_new: Vec<f32> = (0..b_sz * HIDDEN_DIM)
|
||||
.map(|i| ((i as f32) * 0.007).cos() * 0.1)
|
||||
.collect();
|
||||
let bucket_channel_offset: Vec<u32> = vec![0, 25, 50, 75, 100, 128];
|
||||
let bucket_dim_k: Vec<u32> = vec![25, 25, 25, 25, 28];
|
||||
|
||||
let x_d = upload(&stream, &x)?;
|
||||
let h_old_d = upload(&stream, &h_old)?;
|
||||
let w_in_d = upload(&stream, &w_in)?;
|
||||
let w_rec_d = upload(&stream, &w_rec)?;
|
||||
let b_d = upload(&stream, &b_vec)?;
|
||||
let tau_d = upload(&stream, &tau_all)?;
|
||||
let grad_h_new_d = upload(&stream, &grad_h_new)?;
|
||||
let bco_d = upload(&stream, &bucket_channel_offset)?;
|
||||
let bdk_d = upload(&stream, &bucket_dim_k)?;
|
||||
|
||||
// Per-batch gradient outputs (caller reduces across batch).
|
||||
let mut grad_w_in_d =
|
||||
stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM * HIDDEN_DIM)?;
|
||||
let mut grad_w_rec_d =
|
||||
stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM * HIDDEN_DIM)?;
|
||||
let mut grad_b_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
|
||||
let mut grad_tau_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
|
||||
let mut grad_h_old_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
|
||||
|
||||
let dt: f32 = 1.0;
|
||||
let b_i32: i32 = b_sz as i32;
|
||||
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, N_HORIZONS as u32, 1),
|
||||
block_dim: (MAX_BUCKET_DIM as u32, 1, 1),
|
||||
shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
|
||||
};
|
||||
let mut launch = stream.launch_builder(&func_bwd);
|
||||
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)
|
||||
.arg(&b_i32)
|
||||
.arg(&bco_d)
|
||||
.arg(&bdk_d)
|
||||
.arg(&mut grad_w_in_d)
|
||||
.arg(&mut grad_w_rec_d)
|
||||
.arg(&mut grad_b_d)
|
||||
.arg(&mut grad_tau_d)
|
||||
.arg(&mut grad_h_old_d);
|
||||
unsafe {
|
||||
launch.launch(cfg)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
|
||||
// Smoke: no NaNs in any output buffer.
|
||||
let g_w_in = download(&stream, &grad_w_in_d)?;
|
||||
let g_w_rec = download(&stream, &grad_w_rec_d)?;
|
||||
let g_b = download(&stream, &grad_b_d)?;
|
||||
let g_tau = download(&stream, &grad_tau_d)?;
|
||||
let g_h_old = download(&stream, &grad_h_old_d)?;
|
||||
assert!(g_w_in.iter().all(|v| v.is_finite()), "grad_w_in has NaN/Inf");
|
||||
assert!(
|
||||
g_w_rec.iter().all(|v| v.is_finite()),
|
||||
"grad_w_rec has NaN/Inf"
|
||||
);
|
||||
assert!(g_b.iter().all(|v| v.is_finite()), "grad_b has NaN/Inf");
|
||||
assert!(g_tau.iter().all(|v| v.is_finite()), "grad_tau has NaN/Inf");
|
||||
assert!(
|
||||
g_h_old.iter().all(|v| v.is_finite()),
|
||||
"grad_h_old has NaN/Inf"
|
||||
);
|
||||
// At least one non-zero gradient (kernel wrote something).
|
||||
assert!(
|
||||
g_b.iter().any(|v| v.abs() > 0.0),
|
||||
"grad_b is all zero — kernel didn't write?"
|
||||
);
|
||||
assert!(
|
||||
g_h_old.iter().any(|v| v.abs() > 0.0),
|
||||
"grad_h_old is all zero — kernel didn't write?"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user