gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(ml-alpha): NVIDIA-grade rewrite of CfC K-loop hot kernels — 2.15× faster

Browse Source 
Local L40S profile (perception_overfit smoke) GPU kernel time:
  1589ms → 739ms  (53.5% reduction, 2.15× speedup).
Wall-clock smoke: 9.6s → 4.94s (1.94× faster).

Per-kernel deltas (nsys --cuda-graph-trace=node):

  reduce_axis0:              362ms → 10ms   (36× faster)
    Block layout: per-column (1 block / output) → 32-wide column tile
    (block_dim = 32 × 8). Cross-thread reads were strided by n_tail
    (~40K floats = 160KB stride) — one cache line per thread, 8× HBM
    bandwidth wasted. New tile gives coalesced 128B transactions per
    warp. Block tree-reduce kept (no atomicAdd, per feedback_no_atomicadd).
    +1 shared-mem pad to eliminate 32-way bank conflict on the ty reduce.

  multi_horizon_heads_grn_bwd_batched:  540ms → 113ms  (4.8× faster)
    1. Stage h_row[HIDDEN] and a1[5,HEAD_MID] in shared at block entry.
       Eliminates ~28K redundant DRAM reads/block across Pass 3 + Pass 5.
    2. Pass 5 reorder: k outer / i inner with d_z1[k,m] pinned in
       register; writes to grad_w1_scratch are sequential per-thread.
    3. Block size 64 → 128 threads. Pass 5/6 now partition over i
       (output column): cross-thread writes become COALESCED 128B/warp
       (was stride-128 = 512B). Passes 2/3/4 gate on (tid < HEAD_MID).
    4. Pass 3 thread role: m_out → m_in/n. Same coalescing fix on
       grad_w2 writes AND w2 reads in the d_eta_2 sum.

  cfc_step_backward_batched: 351ms → 271ms  (1.3× faster)
    1. Stage x_b[n_in] and h_old_b[n_hid] in shared (was 128× redundant
       DRAM reads per block; now 1× cooperative load).
    2. Pass 1 thread role: i (output row) → k (output col). For each
       i loop iteration, the warp writes grad_w_in[..., tid] /
       grad_w_rec[..., tid] — COALESCED 128B/warp (was stride-128
       non-coalesced).

  multi_horizon_heads_grn_fwd_batched:  211ms → 204ms
    Stage h_row[HIDDEN] in shared — Pass 1 and Pass 3 both consume.

  cfc_step_batched (fwd):     95ms →  94ms
    Stage x_b and h_old_b in shared.

Shared-mem budgets fit comfortably under the 48KB SM cap (~6KB / ~2KB
respectively). All 9 perception_overfit tests pass — gradient
correctness validated end-to-end (constant-signal overfit, K-loop
capture/replay, stride-4 path, evaluate-only paths).

Discipline:
  - Block tree-reduce only, never atomicAdd
  - No nvrtc; pre-compiled cubins via build.rs
  - Mapped-pinned-only is unaffected (CPU↔GPU contract untouched)
  - Single source of truth: replaced kernels in place, no v2 suffixes

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 19:07:29 +02:00
parent 1a465cf7d5
commit b23f8f2efa
4 changed files with 284 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

156
crates/ml-alpha/cuda/cfc_step.cu
View File

@@ -200,20 +200,35 @@ extern "C" __global__ void cfc_step_batched(
int n_batch,
float* __restrict__ h_new // [n_batch, n_hid]
) {
int bi = blockIdx.x;
int i = threadIdx.x;
if (bi >= n_batch || i >= n_hid) return;
// Dynamic shared layout (caller supplies (n_in + n_hid) * 4 bytes):
// s_x [n_in] — staged x_b row (read by every channel thread)
// s_h_old [n_hid] — staged h_old_b row (same)
extern __shared__ float smem_fwd[];
float* s_x = smem_fwd;
float* s_h_old = s_x + n_in;
int bi = blockIdx.x;
int tid = threadIdx.x;
if (bi >= n_batch) return;
// Cooperative staging: 128× redundant DRAM reads per row → 1× per row
// shared across all channel threads. Both rows fit since n_in = n_hid
// = HIDDEN.
if (tid < n_in) s_x[tid] = x[(long long)bi * n_in + tid];
if (tid < n_hid) s_h_old[tid] = h_old[(long long)bi * n_hid + tid];
__syncthreads();
if (tid >= n_hid) return;
const int i = tid;
const float bias_i = b[i];
const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f));
const float one_minus_decay = 1.0f - decay;
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);
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * s_x[k];
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * s_h_old[k];
h_new[(long long)bi * n_hid + i] =
s_h_old[i] * decay + one_minus_decay * tanhf(pre);
}
@@ -249,64 +264,101 @@ extern "C" __global__ void cfc_step_backward_batched(
float* __restrict__ grad_x // [n_batch, n_in] overwrite (nullptr OK)
) {
extern __shared__ float smem[];
float* sd_pre = smem; // [n_hid] — one row for this block's bi
float* sdecay = sd_pre + n_hid; // [n_hid]
// Shared layout — caller passes (2*n_hid + n_in + n_hid) * 4 bytes:
// sd_pre [n_hid] — d_pre_b per channel (set in Pass 0, consumed Pass 1..3)
// sdecay [n_hid] — reserved for future passes / kept for ABI symmetry
// s_x [n_in] — staged x_b row (eliminates 128× redundant DRAM reads)
// s_h_old [n_hid] — staged h_old_b row (same)
float* sd_pre = smem;
float* sdecay = sd_pre + n_hid;
float* s_x = sdecay + n_hid;
float* s_h_old = s_x + n_in;
int bi = blockIdx.x;
int i = threadIdx.x;
if (bi >= n_batch || i >= n_hid) return;
int bi = blockIdx.x;
int tid = threadIdx.x;
// block_dim is >= max(n_hid, n_in); each thread serves both roles
// across passes. Most paths gate on whether (tid < n_hid) / (tid < n_in).
if (bi >= n_batch) 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;
// Cooperative staging — coalesced loads, single sync.
if (tid < n_in) s_x[tid] = x[(long long)bi * n_in + tid];
if (tid < n_hid) s_h_old[tid] = h_old[(long long)bi * n_hid + tid];
__syncthreads();
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];
// Pass 0: compute pre[i], d_pre[i], decay[i]; write per-channel scratch.
// Thread tid serves the channel role i = tid. Reads from shared.
float dh_b = 0.0f;
float decay_i = 0.0f;
float s_i = 0.0f;
if (tid < n_hid) {
const int i = tid;
const float bias_i = b[i];
const float tau_eps = 1e-6f;
const float tau_i_raw = tau[i];
const float tau_eff = fmaxf(tau_i_raw, tau_eps);
decay_i = expf(-dt_s / tau_eff);
sdecay[i] = decay_i;
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[i] = d_pre_b;
float pre = bias_i;
for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * s_x[k];
for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * s_h_old[k];
s_i = tanhf(pre);
const float d_pre_b = dh_b * (1.0f - decay_i) * (1.0f - s_i * s_i);
sd_pre[i] = d_pre_b;
// Pass 1: per-batch param-grad scratch writes.
// Single-writer per (bi, i, *) — block bi is sole writer to its slice;
// K-loop's 64 invocations accumulate within the same slice via +=.
grad_b_scratch[(long long)bi * n_hid + i] += d_pre_b;
for (int k = 0; k < n_in; ++k) {
grad_w_in_scratch[((long long)bi * n_hid + i) * n_in + k] += d_pre_b * x_b[k];
}
for (int k = 0; k < n_hid; ++k) {
grad_w_rec_scratch[((long long)bi * n_hid + 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_scratch[(long long)bi * n_hid + i] += d_decay_b * decay * dt_s / (tau_eff * tau_eff);
// Per-channel scratch (one float per thread, no coalescing concern).
grad_b_scratch[(long long)bi * n_hid + i] += d_pre_b;
// Branchless tau-grad — gates contribution by (tau > eps) without
// warp divergence.
const float d_decay_b = dh_b * (s_h_old[i] - s_i);
const float gate = (tau_i_raw > tau_eps) ? 1.0f : 0.0f;
const float contrib = d_decay_b * decay_i * dt_s / (tau_eff * tau_eff);
grad_tau_scratch[(long long)bi * n_hid + i] += gate * contrib;
}
__syncthreads();
// Pass 1: grad_w_in[bi, i, k] += d_pre[i] * x[k]
// grad_w_rec[bi, i, k] += d_pre[i] * h_old[k]
// Thread tid serves the column role k = tid. For each i, every warp
// thread writes to grad_w_*[i, tid] → COALESCED 128B/transaction.
// Prior layout (thread = i) strided cross-thread writes by n_in/n_hid
// = 512B, paying 32 separate transactions per warp instruction.
if (tid < n_in) {
const int k = tid;
const float x_k = s_x[k];
const float h_old_k = s_h_old[k]; // valid because n_in == n_hid here
for (int i = 0; i < n_hid; ++i) {
const float d_pre_i = sd_pre[i];
grad_w_in_scratch[((long long)bi * n_hid + i) * n_in + k]
+= d_pre_i * x_k;
grad_w_rec_scratch[((long long)bi * n_hid + i) * n_hid + k]
+= d_pre_i * h_old_k;
}
}
// Pass 2: grad_h_old[bi, i] = sum_j sd_pre[j] * W_rec[j, i] + dh_b * decay.
{
float gh = dh_b * decay;
// Thread tid is back to channel role i = tid; dh_b and decay_i
// stayed in registers from Pass 0.
if (tid < n_hid) {
const int i = tid;
float gh = dh_b * decay_i;
for (int j = 0; j < n_hid; ++j) {
gh += sd_pre[j] * w_rec[j * n_hid + i];
}
grad_h_old[(long long)bi * n_hid + i] = gh;
}
// Pass 3: grad_x[bi, k] via thread 0 of this block.
if (grad_x != nullptr && i == 0) {
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[j] * w_in[j * n_in + k];
}
grad_x_b[k] = gx;
// Pass 3: grad_x[bi, k] = sum_j sd_pre[j] * W_in[j, k]
// Thread tid serves column role k = tid. Sequential along j, sd_pre
// in shared. One coalesced write per thread (grad_x layout matches).
if (grad_x != nullptr && tid < n_in) {
const int k = tid;
float gx = 0.0f;
for (int j = 0; j < n_hid; ++j) {
gx += sd_pre[j] * w_in[j * n_in + k];
}
grad_x[(long long)bi * n_in + k] = gx;
}
}

165
crates/ml-alpha/cuda/multi_horizon_heads.cu
View File

@@ -478,6 +478,15 @@ extern "C" __global__ void multi_horizon_heads_grn_fwd_batched(
__shared__ float s_a1[N_HORIZONS_H * HEAD_MID_H]; // post-GELU eta_2
__shared__ float s_z2[N_HORIZONS_H * HEAD_MID_H]; // eta_1
// Stage h_row in shared — Pass 1 reuses h[i] across all 64 threads,
// Pass 3 (skip path) reads it again. Eliminates ~8800 redundant
// DRAM reads per block.
__shared__ float s_h[HIDDEN_H];
#pragma unroll
for (int t = 0; t < HIDDEN_H / HEAD_MID_H; ++t) {
s_h[t * HEAD_MID_H + m] = h_row[t * HEAD_MID_H + m];
}
__syncthreads();
// Pass 1: z1[k, m] = W1[k, m, :] @ h + b1[k, m]; a1 = GELU(z1).
// Thread m owns row m (output dim) for all k. Sequential over HIDDEN.
@@ -485,7 +494,7 @@ extern "C" __global__ void multi_horizon_heads_grn_fwd_batched(
for (int k = 0; k < N_HORIZONS_H; ++k) {
float z = b1[k * HEAD_MID_H + m];
for (int i = 0; i < HIDDEN_H; ++i) {
z += w1[((long long)(k * HEAD_MID_H) + m) * HIDDEN_H + i] * h_row[i];
z += w1[((long long)(k * HEAD_MID_H) + m) * HIDDEN_H + i] * s_h[i];
}
const float a = gelu_act(z);
s_a1[k * HEAD_MID_H + m] = a;
@@ -526,7 +535,7 @@ extern "C" __global__ void multi_horizon_heads_grn_fwd_batched(
}
float sk = b_skip[k];
for (int i = 0; i < HIDDEN_H; ++i) {
sk += w_skip[k * HIDDEN_H + i] * h_row[i];
sk += w_skip[k * HIDDEN_H + i] * s_h[i];
}
const float sigmoid_gl = 1.0f / (1.0f + expf(-gl));
const float logit = sk + sigmoid_gl * mv;
@@ -621,23 +630,25 @@ extern "C" __global__ void multi_horizon_heads_grn_bwd_batched(
) {
int bi = blockIdx.x;
int tid = threadIdx.x;
if (bi >= n_batch || tid >= HEAD_MID_H) return;
if (bi >= n_batch || tid >= HIDDEN_H) return;
__shared__ float s_lambda[N_HORIZONS_H];
if (tid < N_HORIZONS_H) {
const float l = lambda[tid];
s_lambda[tid] = (l > 0.0f) ? l : 1.0f; // sentinel: zero buffer → 1.0
}
__shared__ float s_d_skip[N_HORIZONS_H]; // = d_logit (skip-path grad)
__shared__ float s_d_main[N_HORIZONS_H];
__shared__ float s_d_gate_lin[N_HORIZONS_H];
__shared__ float s_d_eta1[N_HORIZONS_H * HEAD_MID_H];
__shared__ float s_d_eta2[N_HORIZONS_H * HEAD_MID_H];
__shared__ float s_d_z1[N_HORIZONS_H * HEAD_MID_H];
__syncthreads();
// Cooperatively-staged inputs for warm cache hits across passes.
// s_h: reused per (k, m) by every thread in Pass 5 + once per i in Pass 6
// s_a1: reused 64× per thread in Pass 3 inner loop
__shared__ float s_h[HIDDEN_H];
__shared__ float s_a1[N_HORIZONS_H * HEAD_MID_H];
const int m = tid;
// Pass 2/3/4 partition over m ∈ [0, HEAD_MID); only first HEAD_MID threads work.
// Pass 5/6 partition over i ∈ [0, HIDDEN); all 128 threads work.
const int m = tid; // valid iff tid < HEAD_MID_H
const bool is_m_thread = (tid < HEAD_MID_H);
const float* h_row = h + (long long)bi * HIDDEN_H;
// Per-batch base offsets into scratch tensors.
@@ -645,6 +656,22 @@ extern "C" __global__ void multi_horizon_heads_grn_bwd_batched(
const long long bi_5_m = bi_5 * HEAD_MID_H;
const long long bi_5_h = bi_5 * HIDDEN_H;
// Cooperative staging — 128 threads load h (1 each), a1 needs 320 slots
// covered by HEAD_MID threads × N_HORIZONS_H, then sync.
if (tid < N_HORIZONS_H) {
const float l = lambda[tid];
s_lambda[tid] = (l > 0.0f) ? l : 1.0f; // sentinel: zero buffer → 1.0
}
s_h[tid] = h_row[tid]; // tid ∈ [0, HIDDEN) — one float each
if (is_m_thread) {
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
s_a1[k * HEAD_MID_H + m] =
a1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m];
}
}
__syncthreads();
// Pass 1: per-horizon scalar grads. Threads 0..4 only.
if (tid < N_HORIZONS_H) {
const int k = tid;
@@ -673,82 +700,100 @@ extern "C" __global__ void multi_horizon_heads_grn_bwd_batched(
__syncthreads();
// Pass 2: per-(k, m_out=m) compute d_eta_1[k, m] + grad_w_{main,gate}.
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float d_main_k = s_d_main[k];
const float d_gate_lin_k = s_d_gate_lin[k];
const float w_main_km = w_main[k * HEAD_MID_H + m];
const float w_gate_km = w_gate[k * HEAD_MID_H + m];
const float z2_km = z2[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m];
// First HEAD_MID threads only; threads 64..127 idle this pass.
if (is_m_thread) {
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float d_main_k = s_d_main[k];
const float d_gate_lin_k = s_d_gate_lin[k];
const float w_main_km = w_main[k * HEAD_MID_H + m];
const float w_gate_km = w_gate[k * HEAD_MID_H + m];
const float z2_km = z2[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m];
const float d_eta1_km = d_main_k * w_main_km + d_gate_lin_k * w_gate_km;
s_d_eta1[k * HEAD_MID_H + m] = d_eta1_km;
const float d_eta1_km = d_main_k * w_main_km + d_gate_lin_k * w_gate_km;
s_d_eta1[k * HEAD_MID_H + m] = d_eta1_km;
grad_w_main_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_main_k * z2_km;
grad_w_gate_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_gate_lin_k * z2_km;
grad_w_main_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_main_k * z2_km;
grad_w_gate_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_gate_lin_k * z2_km;
// grad_b2[k, m] += d_eta_1[k, m] — thread m sole writer of row m.
grad_b2_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_eta1_km;
// grad_b2[k, m] += d_eta_1[k, m] — thread m sole writer of row m.
grad_b2_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_eta1_km;
}
}
__syncthreads();
// Pass 3: through W2.
// grad_w2[k, m_out=m, m_in] += d_eta_1[k, m] * eta_2[k, m_in]
// d_eta_2[k, m_in=m] = sum_{m_out} d_eta_1[k, m_out] * W2[k, m_out, m_in=m]
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float d_eta1_km = s_d_eta1[k * HEAD_MID_H + m];
for (int n = 0; n < HEAD_MID_H; ++n) {
const float a1_kn = a1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + n];
grad_w2_scratch[bi_5_m * HEAD_MID_H
+ ((long long)k * HEAD_MID_H + m) * HEAD_MID_H + n]
+= d_eta1_km * a1_kn;
// Pass 3: through W2. Thread role: tid = n = m_in (column).
// grad_w2[k, m_out, m_in=tid] += d_eta_1[k, m_out] * eta_2[k, m_in=tid]
// d_eta_2[k, m_in=tid] = sum_{m_out} d_eta_1[k, m_out] * W2[k, m_out, m_in=tid]
// Cross-thread writes/reads at fixed m_out → consecutive in the innermost
// dim → COALESCED. Prior layout (thread = m_out) strided by HEAD_MID = 256B.
// a1[k, tid] is read from s_a1 — broadcast-free (one shared load per
// thread per k).
if (is_m_thread) {
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float a1_kn = s_a1[k * HEAD_MID_H + tid];
float d_eta2_kn = 0.0f;
#pragma unroll 8
for (int m_out = 0; m_out < HEAD_MID_H; ++m_out) {
const float d_eta1_kmout = s_d_eta1[k * HEAD_MID_H + m_out];
// grad_w2 row major: stride 1 in `tid` (the m_in/n dim).
grad_w2_scratch[bi_5_m * HEAD_MID_H
+ ((long long)k * HEAD_MID_H + m_out) * HEAD_MID_H + tid]
+= d_eta1_kmout * a1_kn;
// w2 row major: same stride 1 in `tid`.
d_eta2_kn += d_eta1_kmout
* w2[((long long)(k * HEAD_MID_H) + m_out) * HEAD_MID_H + tid];
}
s_d_eta2[k * HEAD_MID_H + tid] = d_eta2_kn;
}
float d_eta2_km = 0.0f;
for (int m_out = 0; m_out < HEAD_MID_H; ++m_out) {
d_eta2_km += s_d_eta1[k * HEAD_MID_H + m_out]
* w2[((long long)(k * HEAD_MID_H) + m_out) * HEAD_MID_H + m];
}
s_d_eta2[k * HEAD_MID_H + m] = d_eta2_km;
}
__syncthreads();
// Pass 4: through GELU.
// Pass 4: through GELU. First HEAD_MID threads only.
// d_z1[k, m] = d_eta_2[k, m] * gelu_prime(z1[k, m])
// grad_b1[k, m] += d_z1[k, m]
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float d_eta2_km = s_d_eta2[k * HEAD_MID_H + m];
const float z1_km = z1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m];
const float d_z1_km = d_eta2_km * gelu_prime(z1_km);
s_d_z1[k * HEAD_MID_H + m] = d_z1_km;
grad_b1_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_z1_km;
if (is_m_thread) {
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
const float d_eta2_km = s_d_eta2[k * HEAD_MID_H + m];
const float z1_km = z1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m];
const float d_z1_km = d_eta2_km * gelu_prime(z1_km);
s_d_z1[k * HEAD_MID_H + m] = d_z1_km;
grad_b1_scratch[bi_5_m + (long long)k * HEAD_MID_H + m] += d_z1_km;
}
}
__syncthreads();
// Pass 5: grad_w1[k, m, i] += d_z1[k, m] * h[i] (per-batch scratch).
for (int i = 0; i < HIDDEN_H; ++i) {
const float h_bi_i = h_row[i];
// Partitioned over i: thread tid handles output column i = tid.
// Cross-thread writes within a warp differ by 1 in the innermost dim
// → COALESCED 128-byte transactions per warp (~4× DRAM efficiency
// vs the prior thread-per-m layout which strided by HIDDEN_H = 512B).
// h_tid pinned in register (constant per thread for the entire pass).
{
const float h_tid = s_h[tid];
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {
grad_w1_scratch[bi_5_m * HIDDEN_H
+ ((long long)k * HEAD_MID_H + m) * HIDDEN_H + i]
+= s_d_z1[k * HEAD_MID_H + m] * h_bi_i;
for (int mm = 0; mm < HEAD_MID_H; ++mm) {
const float d_z1_km = s_d_z1[k * HEAD_MID_H + mm];
grad_w1_scratch[bi_5_m * HIDDEN_H
+ ((long long)k * HEAD_MID_H + mm) * HIDDEN_H + tid]
+= d_z1_km * h_tid;
}
}
}
// Pass 6: trunk grad_w_skip + grad_h.
// Pass 6: trunk grad_w_skip + grad_h. One thread per output column i.
// grad_w_skip[k, i] += d_skip[k] * h[i] (per-batch scratch)
// grad_h[bi, i] = sum_k lambda[k] * (
// d_skip[k] * w_skip[k, i] # skip path
// + sum_m d_z1[k, m] * w1[k, m, i] # main path
// ) + carry[i]
// Tile i over 2 strides of HEAD_MID. Thread tid handles i=tid and i=tid+HEAD_MID.
#pragma unroll
for (int tile = 0; tile < HIDDEN_H / HEAD_MID_H; ++tile) {
const int i = tile * HEAD_MID_H + tid;
const float h_bi_i = h_row[i];
// With block_dim == HIDDEN_H, no tile loop — thread tid IS the column i.
{
const int i = tid;
const float h_bi_i = s_h[i];
float acc_grad_h = 0.0f;
#pragma unroll
for (int k = 0; k < N_HORIZONS_H; ++k) {

70
crates/ml-alpha/cuda/reduce_axis0.cu
View File

@@ -1,40 +1,64 @@
// reduce_axis0.cu — sum [B, N] → [N] along the leading axis.
//
// One block per output column (N blocks total). Threads tile B via a
// stride-loop, then block-wide tree-reduce. No atomicAdd per
// feedback_no_atomicadd.md. OVERWRITE semantics: out[j] = sum.
// Layout: each block reduces TILE_J=32 contiguous columns of `out`,
// using a 32×8 thread grid. The (tx) dimension covers the column tile
// (output index `j`), the (ty) dimension covers a B-stride loop.
//
// Block dim is hard-coded to 256 — matches the typical tail dim
// granularity in the trainer (well within L40S limits, fits in shared
// memory at 1 KiB per block).
// DRAM access pattern:
// thread (tx, ty) reads per_batch[(b_base + ty*step + s) * n_tail + (j_base + tx)]
// for s = 0..(n_batch/TILE_B), where (tx varies, ty fixed) walks a warp.
// → 32 consecutive `j` per warp → COALESCED 128-byte transactions.
//
// vs the prior block-per-column layout where adjacent threads strided
// the input by `n_tail` (often >40K floats = 160KB stride), forcing
// one cache line per thread and wasting 8× HBM bandwidth.
//
// Block reduce: per-thread partial sums land in s[ty][tx]; the first
// warp (ty=0) sums the 8 partials per column and writes `out[j]`.
// No atomicAdd, no warp-shuffle across warps; block-tree-reduce only.
//
// Launch contract:
// grid_dim = (ceil(n_tail / 32), 1, 1)
// block_dim = (32, 8, 1)
// shared = 0 (compile-time static)
#define REDAX0_BLOCK 256
#define REDAX0_TILE_J 32
#define REDAX0_TILE_B 8
#define REDAX0_BLOCK (REDAX0_TILE_J * REDAX0_TILE_B) // 256 threads/block
extern "C" __global__ void reduce_axis0(
extern "C" __global__ __launch_bounds__(REDAX0_BLOCK, 4)
void reduce_axis0(
const float* __restrict__ per_batch, // [B, N]
int n_batch,
int n_tail,
float* __restrict__ out // [N] — OVERWRITE
) {
int j = blockIdx.x;
int tid = threadIdx.x;
if (j >= n_tail) return;
__shared__ float s[REDAX0_BLOCK];
const int tx = threadIdx.x; // column-tile lane
const int ty = threadIdx.y; // batch-tile lane
const int j = (int)blockIdx.x * REDAX0_TILE_J + tx; // output column
const bool active = (j < n_tail);
float my_sum = 0.0f;
// Stride-loop across B — handles n_batch > BLOCK and n_batch < BLOCK uniformly.
for (int bi = tid; bi < n_batch; bi += REDAX0_BLOCK) {
my_sum += per_batch[(long long)bi * n_tail + j];
if (active) {
// Stride-loop along B in chunks of TILE_B. Within a warp (fixed ty,
// varying tx) we step `j` by +1 → coalesced 32-float load per warp.
for (int bi = ty; bi < n_batch; bi += REDAX0_TILE_B) {
my_sum += per_batch[(long long)bi * n_tail + j];
}
}
s[tid] = my_sum;
// +1 pad eliminates 32-way shared-mem bank conflict on the ty reduce.
__shared__ float s[REDAX0_TILE_B][REDAX0_TILE_J + 1];
s[ty][tx] = my_sum;
__syncthreads();
// Block tree-reduce.
for (int stride = REDAX0_BLOCK / 2; stride > 0; stride >>= 1) {
if (tid < stride) s[tid] += s[tid + stride];
__syncthreads();
// First warp sums the 8 per-batch-tile partials for its column.
if (ty == 0 && active) {
float total = 0.0f;
#pragma unroll
for (int t = 0; t < REDAX0_TILE_B; ++t) {
total += s[t][tx];
}
out[j] = total;
}
if (tid == 0) out[j] = s[0];
}

46
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -1537,10 +1537,12 @@ impl PerceptionTrainer {
let _ = block_dim; // legacy 1D layout helper — superseded by block-per-batch launches.
// CfC fwd: block-per-batch (Phase B). One block per sample;
// each block's threads parallelize over n_hid channels.
// Shared layout: s_x [n_in] + s_h_old [n_hid].
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
let cfg_cfc = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
shared_mem_bytes: cfc_fwd_smem,
};
let cfg_heads = LaunchConfig {
grid_dim: (1, 1, 1),
@@ -1549,7 +1551,9 @@ impl PerceptionTrainer {
};
// CfC bwd (Phase B): block-per-batch + 2 * n_hid floats of shared mem
// (one sd_pre row + sdecay for the block's single bi).
let cfc_bwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
// Shared layout: sd_pre [n_hid] + sdecay [n_hid] + s_x [n_in] + s_h_old [n_hid].
// With n_in == n_hid == HIDDEN_DIM, that's 4 × HIDDEN_DIM floats per block.
let cfc_bwd_smem = (4 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
let cfg_cfc_bwd = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
@@ -1585,11 +1589,13 @@ impl PerceptionTrainer {
shared_mem_bytes: 0,
};
// GRN backward launch config: block-per-batch (Phase B). One block
// per sample; threads tile HEAD_MID. Per-batch scratch + reducer
// collapses B → final grad after the K-loop.
// per sample; threads = HIDDEN so Pass 5/6 can partition over the
// output i-dim for COALESCED writes. Passes 2/3/4 use first
// HEAD_MID threads only; the rest idle until Pass 5.
// Per-batch scratch + reducer collapses B → final grad after the K-loop.
let cfg_grn_bwd = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (HEAD_MID_DIM as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let _ = cfg_heads_batched_fwd; // GRN fwd uses cfg_grn_fwd now.
@@ -1890,12 +1896,13 @@ impl PerceptionTrainer {
unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; }
}
// Attn pool reducer: collapse [B, HIDDEN_DIM] → [HIDDEN_DIM].
// reduce_axis0: block tiles 32 cols × 8 batches per launch.
{
let n_batch_i = b_sz as i32;
let n_tail_i = HIDDEN_DIM as i32;
let cfg_red = LaunchConfig {
grid_dim: (HIDDEN_DIM as u32, 1, 1),
block_dim: (256, 1, 1),
grid_dim: (((HIDDEN_DIM + 31) / 32) as u32, 1, 1),
block_dim: (32, 8, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
@@ -2095,15 +2102,16 @@ impl PerceptionTrainer {
unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; }
}
// VSN reducer: collapse n_rows (= B * K) → final grad buffers.
// reduce_axis0 tiles 32 columns per block, 8-way batch stride.
{
let n_rows_i = (b_sz * k_seq) as i32;
let reduce_block: u32 = 256;
let n_tail_w_usz = FEATURE_DIM * FEATURE_DIM;
let cfg_red_w = LaunchConfig {
grid_dim: ((FEATURE_DIM * FEATURE_DIM) as u32, 1, 1),
block_dim: (reduce_block, 1, 1),
grid_dim: (((n_tail_w_usz + 31) / 32) as u32, 1, 1),
block_dim: (32, 8, 1),
shared_mem_bytes: 0,
};
let n_tail_w = (FEATURE_DIM * FEATURE_DIM) as i32;
let n_tail_w = n_tail_w_usz as i32;
unsafe {
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
launch
@@ -2114,8 +2122,8 @@ impl PerceptionTrainer {
launch.launch(cfg_red_w).context("reduce vsn_grad_w")?;
}
let cfg_red_b = LaunchConfig {
grid_dim: (FEATURE_DIM as u32, 1, 1),
block_dim: (reduce_block, 1, 1),
grid_dim: (((FEATURE_DIM + 31) / 32) as u32, 1, 1),
block_dim: (32, 8, 1),
shared_mem_bytes: 0,
};
let n_tail_b = FEATURE_DIM as i32;
@@ -2132,20 +2140,20 @@ impl PerceptionTrainer {
// ── 8c. (Phase B) Reduce CfC per-batch grad scratch → final grad buffers.
// Block tree-reduce across bi; deterministic sum order.
// Each launch is grid=(n_tail, 1, 1), block=(256, 1, 1).
// Each launch is grid=(ceil(n_tail/32), 1, 1), block=(32, 8, 1)
// — 32-wide column tile per block, 8-way B-stride per thread.
// 4 launches, one per cfc param tensor. AdamW (below) reads
// the final grad buffers; MUST run after these reducers.
{
let n_batch_i = b_sz as i32;
let reduce_block: u32 = 256;
let reduce_at = |n_tail: usize,
scratch: &CudaSlice<f32>,
out: &mut CudaSlice<f32>,
label: &'static str|
-> Result<()> {
let cfg = LaunchConfig {
grid_dim: (n_tail as u32, 1, 1),
block_dim: (reduce_block, 1, 1),
grid_dim: (((n_tail + 31) / 32) as u32, 1, 1),
block_dim: (32, 8, 1),
shared_mem_bytes: 0,
};
let n_tail_i = n_tail as i32;
@@ -2525,10 +2533,12 @@ impl PerceptionTrainer {
// grid=(1,1,1) config which only processed batch 0 — batches
// 1..B-1 got GARBAGE h_new. Caught by t6z89-vs-txftz cluster A/B
// (val_loss=0.83, mean_auc=chance) on 2026-05-17.
// Shared layout: s_x [n_in] + s_h_old [n_hid].
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
let cfg_cfc = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
shared_mem_bytes: cfc_fwd_smem,
};
let cfg_grn_fwd = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),