From 829ddfa62c249fb7ef4324e53fe164e987c51a3e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 10:20:27 +0200 Subject: [PATCH] feat(ml-alpha): add batched cfc + heads CUDA kernels (foundation for #8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds 4 new kernel symbols alongside the existing single-sample ones — zero changes to current call sites, so the in-flight qf5mj baseline is unaffected. The next commit wires these into PerceptionTrainer's K loop and exposes --batch-size in the CLI. cfc_step_batched processes [n_batch, n_in/n_hid] tensors cfc_step_backward_batched same; shared mem holds sd_pre[B, n_hid] + sdecay[n_hid] (~16 KiB at B=32, well under L40S 48 KiB block limit). Param grads (grad_b/grad_w_in/grad_w_rec/ grad_tau) accumulated via += — thread i is sole writer to its row across all samples, so no atomicAdd and no per- batch scratch buffer. multi_horizon_heads_batched [n_batch, 5] sigmoid outputs from [n_batch, 128] hidden inputs. multi_horizon_heads_backward_batched shared mem holds sd_z[B, 5]. grad_w / grad_b += across batch (thread tid sole writer). grad_h carries the optional per-sample grad_h_carry (cfc recurrence chain). Design notes: - Threading: one block of n_hid threads. Each thread loops over b ∈ 0..B internally. This avoids cross-block races on grad_* buffers and keeps the existing "no atomicAdd" discipline. Cost: less raw parallelism than grid-batching, but the bottleneck is Mamba2 (already batch-parallel via its own kernel grid). - Per-thread accumulators: grad_b / grad_tau land in registers, flushed once at end. grad_w_in / grad_w_rec written += per-b (thread sole writer to its row, safe). - All B samples processed in stream order inside one kernel launch — saves K * (B-1) launches per sequence vs serialising B independent calls. 77 ml-alpha tests pass (kernels not yet exercised — wiring is the next commit). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/cfc_step.cu | 148 ++++++++++++++++++++ crates/ml-alpha/cuda/multi_horizon_heads.cu | 84 +++++++++++ 2 files changed, 232 insertions(+) diff --git a/crates/ml-alpha/cuda/cfc_step.cu b/crates/ml-alpha/cuda/cfc_step.cu index dbb07207b..9be0c80e6 100644 --- a/crates/ml-alpha/cuda/cfc_step.cu +++ b/crates/ml-alpha/cuda/cfc_step.cu @@ -141,3 +141,151 @@ extern "C" __global__ void cfc_step_backward( } } } + + +// ─── Batched variants ──────────────────────────────────────────────── +// +// Same math as the single-sample kernels above, but process B sequences +// in parallel within one launch. Layout: x[B, n_in], h_old[B, n_hid], +// h_new[B, n_hid] (row-major, B-major). Thread i handles hidden unit i +// across ALL B samples — per-thread loop over b ∈ 0..B. +// +// Param-grad writes (grad_b/grad_w_in/grad_w_rec/grad_tau) use `+=` into +// SHARED-across-batch buffers; thread i is the sole writer to its row of +// each grad matrix, so no atomic / no race regardless of B. Sample- +// scoped outputs (h_new, grad_h_old, grad_x) write into per-b slots. +// +// Shared memory budget for the backward kernel: `sd_pre[B, n_hid] + +// sdecay[n_hid]` floats = (B+1) * n_hid * 4 bytes. At B=32, n_hid=128: +// 33 * 128 * 4 = 16.5 KiB per block — well under the 48 KiB L40S limit. + +extern "C" __global__ void cfc_step_batched( + const float* __restrict__ w_in, // [n_hid, n_in], row-major + const float* __restrict__ w_rec, // [n_hid, n_hid], row-major + const float* __restrict__ b, // [n_hid] + const float* __restrict__ tau, // [n_hid] + const float* __restrict__ x, // [n_batch, n_in] + const float* __restrict__ h_old, // [n_batch, n_hid] + float dt_s, + int n_in, + int n_hid, + int n_batch, + float* __restrict__ h_new // [n_batch, n_hid] +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_hid) return; + + const float bias_i = b[i]; + const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f)); + const float one_minus_decay = 1.0f - decay; + + for (int bi = 0; bi < n_batch; ++bi) { + const float* x_b = x + (long long)bi * n_in; + const float* h_old_b = h_old + (long long)bi * n_hid; + float pre = bias_i; + for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k]; + for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k]; + h_new[(long long)bi * n_hid + i] = h_old_b[i] * decay + one_minus_decay * tanhf(pre); + } +} + + +extern "C" __global__ void cfc_step_backward_batched( + 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_batch, n_in] + const float* __restrict__ h_old, // [n_batch, n_hid] + const float* __restrict__ grad_h_new, // [n_batch, n_hid] + float dt_s, + int n_in, + int n_hid, + int n_batch, + float* __restrict__ grad_w_in, // [n_hid, n_in] accum += + float* __restrict__ grad_w_rec, // [n_hid, n_hid] accum += + float* __restrict__ grad_b, // [n_hid] accum += + float* __restrict__ grad_tau, // [n_hid] accum += + float* __restrict__ grad_h_old, // [n_batch, n_hid] overwrite + float* __restrict__ grad_x // [n_batch, n_in] overwrite (nullptr OK) +) { + // Shared mem: sd_pre[B, n_hid] then sdecay[n_hid]. + extern __shared__ float smem[]; + float* sd_pre = smem; // [n_batch * n_hid] + float* sdecay = sd_pre + (long long)n_batch * n_hid; // [n_hid] + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_hid) return; + + const float bias_i = b[i]; + const float tau_eps = 1e-6f; + const float tau_eff = fmaxf(tau[i], tau_eps); + const float decay = expf(-dt_s / tau_eff); + sdecay[i] = decay; + + // Per-thread (private) accumulators across batch — avoids any + // cross-thread race since thread i is sole writer to its row. + float grad_b_acc = 0.0f; + float grad_tau_acc = 0.0f; + // grad_w_in_acc / grad_w_rec_acc would be n_in / n_hid floats per + // thread — too big for registers. Instead we accumulate them via + // direct `+=` into global memory (thread i is sole writer to its + // rows; no race even across the per-b loop). + + // Pass 1: per-b compute d_pre, store in shared, accumulate + // per-row param grads (grad_w_in, grad_w_rec rows). + for (int bi = 0; bi < n_batch; ++bi) { + const float* x_b = x + (long long)bi * n_in; + const float* h_old_b = h_old + (long long)bi * n_hid; + const float dh_b = grad_h_new[(long long)bi * n_hid + i]; + + float pre = bias_i; + for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x_b[k]; + for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old_b[k]; + const float s = tanhf(pre); + const float d_pre_b = dh_b * (1.0f - decay) * (1.0f - s * s); + sd_pre[(long long)bi * n_hid + i] = d_pre_b; + + grad_b_acc += d_pre_b; + // grad_w_in / grad_w_rec rows: thread i sole writer → safe +=. + for (int k = 0; k < n_in; ++k) { + grad_w_in[i * n_in + k] += d_pre_b * x_b[k]; + } + for (int k = 0; k < n_hid; ++k) { + grad_w_rec[i * n_hid + k] += d_pre_b * h_old_b[k]; + } + if (tau[i] > tau_eps) { + const float d_decay_b = dh_b * (h_old_b[i] - s); + grad_tau_acc += d_decay_b * decay * dt_s / (tau_eff * tau_eff); + } + } + grad_b[i] += grad_b_acc; + grad_tau[i] += grad_tau_acc; + __syncthreads(); + + // Pass 2: per-b emit grad_h_old[bi, i] = sum_j sd_pre[bi, j] * W_rec[j, i] + dh_bi * decay. + for (int bi = 0; bi < n_batch; ++bi) { + float gh = grad_h_new[(long long)bi * n_hid + i] * decay; + for (int j = 0; j < n_hid; ++j) { + gh += sd_pre[(long long)bi * n_hid + j] * w_rec[j * n_hid + i]; + } + grad_h_old[(long long)bi * n_hid + i] = gh; + } + + // Pass 3: per-b emit grad_x[bi, k] via thread-0 reduction (n_in is + // small, this is amortised over the K-loop work). Mirrors the + // single-sample kernel's pattern. + if (grad_x != nullptr && i == 0) { + for (int bi = 0; bi < n_batch; ++bi) { + const float* sd_pre_b = sd_pre + (long long)bi * n_hid; + float* grad_x_b = grad_x + (long long)bi * n_in; + for (int k = 0; k < n_in; ++k) { + float gx = 0.0f; + for (int j = 0; j < n_hid; ++j) { + gx += sd_pre_b[j] * w_in[j * n_in + k]; + } + grad_x_b[k] = gx; + } + } + } +} diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 15174f36a..7ab21731e 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -80,3 +80,87 @@ extern "C" __global__ void multi_horizon_heads_backward( grad_h[tid] = acc + carry; } } + + +// ─── Batched variants ──────────────────────────────────────────────── +// Process N_BATCH samples per launch. Layout: h[B, 128], probs[B, 5]. +// Thread tid handles either head k (tid < 5) or hidden unit i (tid < +// 128). Block dim = 128, one block per launch — each thread loops over +// b internally. Param grads (grad_w, grad_b) shared across batch and +// accumulated via `+=` (no race: thread tid is sole writer to its row). + +extern "C" __global__ void multi_horizon_heads_batched( + const float* __restrict__ w, // [5, 128] + const float* __restrict__ b, // [5] + const float* __restrict__ h, // [n_batch, 128] + int n_batch, + float* __restrict__ probs // [n_batch, 5] +) { + int k = threadIdx.x; + if (k >= 5) return; + const float bias_k = b[k]; + for (int bi = 0; bi < n_batch; ++bi) { + const float* h_b = h + (long long)bi * 128; + float z = bias_k; + for (int i = 0; i < 128; ++i) z += w[k * 128 + i] * h_b[i]; + probs[(long long)bi * 5 + k] = 1.0f / (1.0f + expf(-z)); + } +} + + +extern "C" __global__ void multi_horizon_heads_backward_batched( + const float* __restrict__ w, // [5, 128] + const float* __restrict__ probs, // [n_batch, 5] + const float* __restrict__ h, // [n_batch, 128] + const float* __restrict__ grad_probs, // [n_batch, 5] + const float* __restrict__ grad_h_carry, // [n_batch, 128] (nullptr OK) + int n_batch, + float* __restrict__ grad_w, // [5, 128] accum += + float* __restrict__ grad_b, // [5] accum += + float* __restrict__ grad_h // [n_batch, 128] overwrite + carry +) { + // Shared mem: d_z[B, 5]. At B=32: 32 * 5 = 160 floats = 640 bytes. + extern __shared__ float sd_z[]; // size = n_batch * 5 + + int tid = threadIdx.x; + + // Pass 1: per-(bi, k) compute d_z and stash in shared. + if (tid < 5) { + for (int bi = 0; bi < n_batch; ++bi) { + const float p = probs[(long long)bi * 5 + tid]; + const float dp = grad_probs[(long long)bi * 5 + tid]; + sd_z[(long long)bi * 5 + tid] = dp * p * (1.0f - p); + } + } + __syncthreads(); + + // grad_b[k] += sum_b d_z[b, k]. Thread k (k < 5) is sole writer. + if (tid < 5) { + float acc = 0.0f; + for (int bi = 0; bi < n_batch; ++bi) { + acc += sd_z[(long long)bi * 5 + tid]; + } + grad_b[tid] += acc; + } + + // grad_w[k, i] += sum_b d_z[b, k] * h[b, i]. Thread i (i < 128) is + // sole writer to its column of grad_w. (Below: parallelise over i + // across the 5 head rows.) + if (tid < 128) { + // grad_h: sum_k d_z[b, k] * W[k, i] for each b, plus optional carry. + for (int bi = 0; bi < n_batch; ++bi) { + const float h_bi = h[(long long)bi * 128 + tid]; + // Accumulate W-row grads via += into global (thread tid sole writer of column tid). + for (int k = 0; k < 5; ++k) { + grad_w[k * 128 + tid] += sd_z[(long long)bi * 5 + k] * h_bi; + } + float acc = 0.0f; + for (int k = 0; k < 5; ++k) { + acc += sd_z[(long long)bi * 5 + k] * w[k * 128 + tid]; + } + const float carry = (grad_h_carry != nullptr) + ? grad_h_carry[(long long)bi * 128 + tid] : 0.0f; + grad_h[(long long)bi * 128 + tid] = acc + carry; + } + } +}