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

feat(ml-alpha): 2-layer heads backward kernel with ISV lambda (Phase 1.4)

Browse Source
This commit is contained in:
jgrusewski
2026-05-17 21:24:49 +02:00
parent e0a497da9e
commit 59e236b4e8
1 changed files with 161 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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();
}
}