From 11eb74766aae0e18244d8dcb3d3efdf4c5759705 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 15:56:20 +0200 Subject: [PATCH] feat(per-horizon-cfc): add cfc_step_per_branch.cu (fused fwd+bwd) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/cfc_step_per_branch.cu | 230 +++++++++++++++++ crates/ml-alpha/tests/cfc_step_per_branch.rs | 247 +++++++++++++++++++ 3 files changed, 478 insertions(+) create mode 100644 crates/ml-alpha/cuda/cfc_step_per_branch.cu create mode 100644 crates/ml-alpha/tests/cfc_step_per_branch.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 3fde4d31a..1a056efa3 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -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. diff --git a/crates/ml-alpha/cuda/cfc_step_per_branch.cu b/crates/ml-alpha/cuda/cfc_step_per_branch.cu new file mode 100644 index 000000000..151e9bcb4 --- /dev/null +++ b/crates/ml-alpha/cuda/cfc_step_per_branch.cu @@ -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; +} diff --git a/crates/ml-alpha/tests/cfc_step_per_branch.rs b/crates/ml-alpha/tests/cfc_step_per_branch.rs new file mode 100644 index 000000000..4fb9b4d56 --- /dev/null +++ b/crates/ml-alpha/tests/cfc_step_per_branch.rs @@ -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, + fn_name: &str, +) -> Result { + 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( + stream: &std::sync::Arc, + host: &[T], +) -> Result> { + let mut d = stream.alloc_zeros::(host.len())?; + stream.memcpy_htod(host, &mut d)?; + Ok(d) +} + +fn download( + stream: &std::sync::Arc, + d: &CudaSlice, +) -> Result> { + 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 = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect(); + let h_old: Vec = (0..b_sz * HIDDEN_DIM) + .map(|i| ((i as f32) * 0.003).sin()) + .collect(); + let w_in: Vec = (0..HIDDEN_DIM * HIDDEN_DIM) + .map(|i| ((i % 17) as f32) * 0.01) + .collect(); + let w_rec: Vec = (0..HIDDEN_DIM * HIDDEN_DIM) + .map(|i| ((i % 13) as f32) * 0.005) + .collect(); + let b_vec: Vec = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect(); + let tau_all: Vec = (0..HIDDEN_DIM) + .map(|c| ((c / 25 + 1) as f32) * 10.0) + .collect(); + let bucket_channel_offset: Vec = vec![0, 25, 50, 75, 100, 128]; + let bucket_dim_k: Vec = 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::(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::()) 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 = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect(); + let h_old: Vec = (0..b_sz * HIDDEN_DIM) + .map(|i| ((i as f32) * 0.003).sin()) + .collect(); + let w_in: Vec = (0..HIDDEN_DIM * HIDDEN_DIM) + .map(|i| ((i % 17) as f32) * 0.01) + .collect(); + let w_rec: Vec = (0..HIDDEN_DIM * HIDDEN_DIM) + .map(|i| ((i % 13) as f32) * 0.005) + .collect(); + let b_vec: Vec = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect(); + let tau_all: Vec = (0..HIDDEN_DIM) + .map(|c| ((c / 25 + 1) as f32) * 10.0) + .collect(); + let grad_h_new: Vec = (0..b_sz * HIDDEN_DIM) + .map(|i| ((i as f32) * 0.007).cos() * 0.1) + .collect(); + let bucket_channel_offset: Vec = vec![0, 25, 50, 75, 100, 128]; + let bucket_dim_k: Vec = 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::(b_sz * HIDDEN_DIM * HIDDEN_DIM)?; + let mut grad_w_rec_d = + stream.alloc_zeros::(b_sz * HIDDEN_DIM * HIDDEN_DIM)?; + let mut grad_b_d = stream.alloc_zeros::(b_sz * HIDDEN_DIM)?; + let mut grad_tau_d = stream.alloc_zeros::(b_sz * HIDDEN_DIM)?; + let mut grad_h_old_d = stream.alloc_zeros::(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::()) 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(()) +}