feat(ml-alpha): add batched cfc + heads CUDA kernels (foundation for #8)
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user