From 59e236b4e8c1751396fc4c7b012f70c7c63273a7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 21:24:49 +0200 Subject: [PATCH] feat(ml-alpha): 2-layer heads backward kernel with ISV lambda (Phase 1.4) --- crates/ml-alpha/cuda/multi_horizon_heads.cu | 161 ++++++++++++++++++++ 1 file changed, 161 insertions(+) diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 344699e28..f8194dfdf 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -264,3 +264,164 @@ extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( probs[(long long)b_idx * N_HORIZONS_H + k] = 1.0f / (1.0f + expf(-z2)); } } + +// 2-layer heads backward — chain rule through sigmoid → linear → +// GELU → linear → trunk. ISV `lambda[5]` scales the trunk-grad +// contribution per horizon (per `pearl_adam_normalizes_loss_weights.md` +// the effective lever is the trunk gradient, not loss weight). +// +// Chain rule: +// d_z2[k] = grad_probs[k] * p * (1 - p) +// d_a1[k, m] = d_z2[k] * w2[k, m] +// d_z1[k, m] = d_a1[k, m] * gelu_prime(z1[k, m]) +// grad_w2[k, m] += d_z2[k] * a1[k, m] +// grad_b2[k] += d_z2[k] +// grad_w1[k, m, i] += d_z1[k, m] * h[i] +// grad_b1[k, m] += d_z1[k, m] +// grad_h[i] = sum_{k, m} lambda[k] * d_z1[k, m] * w1[k, m, i] + carry[i] +// +// Single-writer discipline: thread tid (< HIDDEN) is the sole writer +// of grad_h[bi, tid] and grad_w1[k, m, tid] (across all k, m). Both +// the per-block sample loop and the per-k loop are serial within one +// thread; no cross-block races since each block handles ONE sample. +// Same for grad_w2[k, m] / grad_b2[k] / grad_b1[k, m]: each is written +// by exactly one thread per block (m for grad_w2 / grad_b1 dims, m<5 +// for grad_b2). grad_w1 / grad_b1 / grad_w2 / grad_b2 accumulate via +// += across the n_batch samples (multiple blocks); single-writer per +// (block, output) holds within a block but blocks race on the +=. +// +// Resolution: launch ONE block per launch (not one per sample) and +// iterate over n_batch internally; this matches the existing +// single-layer kernel pattern. Block dim = HEAD_MID = 64; threads +// also tile over HIDDEN for grad_h / grad_w1 columns (HEAD_MID=64 < +// HIDDEN=128 so we serialise across 2 HIDDEN-column groups). +// +// Layout for the per-sample loop: +// - threads 0..63 each own one (m) hidden-unit dimension. +// - grad_b2[k]: thread m==k for k in 0..5 (within m<5 group). +// - grad_w2[k, m]: thread m writes column m for all k = serialised inside thread. +// - grad_b1[k, m]: thread m writes (k, m) for all k = serialised inside thread. +// - grad_w1[k, m, i] += d_z1[k, m] * h[bi, i]: the i dimension is +// HIDDEN=128; thread m needs to write to column i for ALL i. +// Each thread loops over i; threads contend on the same row (k, m) +// but write to DIFFERENT columns i — no race. +// - grad_h[bi, i]: thread responsible for column i. We tile i over +// the 64 threads in 2 strides (i = m, m+64). + +__device__ __forceinline__ float gelu_prime(float x) { + const float u = GELU_C0 * (x + GELU_C1 * x * x * x); + const float t = tanhf(u); + const float dut = GELU_C0 * (1.0f + 3.0f * GELU_C1 * x * x); + return 0.5f * (1.0f + t) + 0.5f * x * (1.0f - t * t) * dut; +} + +extern "C" __global__ void multi_horizon_heads_2layer_bwd_batched( + const float* __restrict__ w1, // [5, HEAD_MID, HIDDEN] + const float* __restrict__ w2, // [5, HEAD_MID] + const float* __restrict__ probs, // [n_batch, 5] + const float* __restrict__ grad_probs, // [n_batch, 5] + const float* __restrict__ z1, // [n_batch, 5, HEAD_MID] + const float* __restrict__ a1, // [n_batch, 5, HEAD_MID] + const float* __restrict__ h, // [n_batch, HIDDEN] + const float* __restrict__ grad_h_carry, // [n_batch, HIDDEN] (nullptr OK) + const float* __restrict__ lambda, // [5] + int n_batch, + float* __restrict__ grad_w1, // [5, HEAD_MID, HIDDEN] (+=) + float* __restrict__ grad_b1, // [5, HEAD_MID] (+=) + float* __restrict__ grad_w2, // [5, HEAD_MID] (+=) + float* __restrict__ grad_b2, // [5] (+=) + float* __restrict__ grad_h // [n_batch, HIDDEN] (overwrite + carry) +) { + int tid = threadIdx.x; + if (tid >= HEAD_MID_H) return; + + // Broadcast lambda[5] via shared mem (5 floats). + __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 mem for d_z1[k, m] and d_z2[k] per sample. + __shared__ float s_d_z1[N_HORIZONS_H * HEAD_MID_H]; + __shared__ float s_d_z2[N_HORIZONS_H]; + __syncthreads(); + + for (int bi = 0; bi < n_batch; ++bi) { + // Pass 1: compute d_z2[k] for k=0..4 (only first 5 threads). + if (tid < N_HORIZONS_H) { + const int k = tid; + const float p = probs[(long long)bi * N_HORIZONS_H + k]; + const float dp = grad_probs[(long long)bi * N_HORIZONS_H + k]; + s_d_z2[k] = dp * p * (1.0f - p); + } + __syncthreads(); + + // Pass 2: per-(k, m) compute d_a1 and d_z1, stash in shared. + // Also accumulate grad_w2[k, m] += d_z2[k] * a1[bi, k, m] and + // grad_b1[k, m] += d_z1[k, m]. Thread m is sole writer. + const int m = tid; + #pragma unroll + for (int k = 0; k < N_HORIZONS_H; ++k) { + const float d_z2 = s_d_z2[k]; + const float w2_km = w2[k * HEAD_MID_H + m]; + const float d_a1 = d_z2 * w2_km; + const float z1_km = z1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m]; + const float a1_km = a1[((long long)bi * N_HORIZONS_H + k) * HEAD_MID_H + m]; + const float d_z1_km = d_a1 * gelu_prime(z1_km); + s_d_z1[k * HEAD_MID_H + m] = d_z1_km; + + // grad_w2[k, m] += d_z2[k] * a1[k, m] + grad_w2[k * HEAD_MID_H + m] += d_z2 * a1_km; + // grad_b1[k, m] += d_z1[k, m] + grad_b1[k * HEAD_MID_H + m] += d_z1_km; + } + // grad_b2[k] (only first 5 threads). + if (tid < N_HORIZONS_H) { + grad_b2[tid] += s_d_z2[tid]; + } + __syncthreads(); + + // Pass 3: grad_w1[k, m, i] += d_z1[k, m] * h[bi, i] + // and grad_h[bi, i] = sum_{k, m} lambda[k] * d_z1[k, m] * w1[k, m, i] + carry[i] + // + // Thread m owns row (k, m) of grad_w1 for ALL k — loops over + // HIDDEN cols i serially. Thread m also computes grad_h[bi, i] + // for i = m and i = m + HEAD_MID (since HIDDEN = 2 * HEAD_MID). + // + // Two passes for grad_h tiling: i = tid (0..63), then i = tid + 64 (64..127). + // Thread for column i is the sole writer. + const float* h_row = h + (long long)bi * HIDDEN_H; + + // grad_w1: thread m writes (k, m, i) for all (k, i). Loop over i. + for (int i = 0; i < HIDDEN_H; ++i) { + const float h_bi = h_row[i]; + #pragma unroll + for (int k = 0; k < N_HORIZONS_H; ++k) { + grad_w1[((long long)(k * HEAD_MID_H) + m) * HIDDEN_H + i] += + s_d_z1[k * HEAD_MID_H + m] * h_bi; + } + } + + // grad_h[bi, i] — tile i over 2 strides of HEAD_MID. + // Thread tid handles columns 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; + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < N_HORIZONS_H; ++k) { + #pragma unroll + for (int mm = 0; mm < HEAD_MID_H; ++mm) { + acc += s_lambda[k] + * s_d_z1[k * HEAD_MID_H + mm] + * w1[((long long)(k * HEAD_MID_H) + mm) * HIDDEN_H + i]; + } + } + const float carry = (grad_h_carry != nullptr) + ? grad_h_carry[(long long)bi * HIDDEN_H + i] : 0.0f; + grad_h[(long long)bi * HIDDEN_H + i] = acc + carry; + } + __syncthreads(); + } +}