diff --git a/crates/ml-alpha/cuda/cfc_step.cu b/crates/ml-alpha/cuda/cfc_step.cu index 89acf7de5..09c2cb062 100644 --- a/crates/ml-alpha/cuda/cfc_step.cu +++ b/crates/ml-alpha/cuda/cfc_step.cu @@ -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; } } diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 469864b33..c58849bfd 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -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) { diff --git a/crates/ml-alpha/cuda/reduce_axis0.cu b/crates/ml-alpha/cuda/reduce_axis0.cu index 451857e3a..73f4b44e2 100644 --- a/crates/ml-alpha/cuda/reduce_axis0.cu +++ b/crates/ml-alpha/cuda/reduce_axis0.cu @@ -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]; } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 64b3094c9..0b36eb2f2 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -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::()) 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::()) 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::()) 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, out: &mut CudaSlice, 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::()) 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),