From e0a497da9e38e39bf76a2699496f297ad53b4493 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 21:23:46 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=202-layer=20GELU=20MLP=20heads?= =?UTF-8?q?=20=E2=80=94=20forward=20kernel=20(Phase=201.3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/ml-alpha/cuda/multi_horizon_heads.cu | 78 +++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 83f54fe54..344699e28 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -186,3 +186,81 @@ extern "C" __global__ void multi_horizon_heads_backward_batched( } } } + +// ───────────────────────────────────────────────────────────────────── +// Phase 1: 2-layer MLP heads — per-horizon [HIDDEN=128 → HEAD_MID=64] +// GELU hidden → [HEAD_MID=64 → 1] sigmoid output. +// +// Forward layout: +// w1 [5, HEAD_MID, HIDDEN] per-horizon first-layer weights +// b1 [5, HEAD_MID] per-horizon first-layer biases +// w2 [5, HEAD_MID] per-horizon output weights (scalar per horizon) +// b2 [5] per-horizon output bias +// h [n_batch, HIDDEN] trunk output (LayerNorm-normalised) +// probs [n_batch, 5] sigmoid output +// z1_out [n_batch, 5, HEAD_MID] saved pre-GELU for backward +// a1_out [n_batch, 5, HEAD_MID] saved post-GELU for backward +// +// GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) +// +// Launch: grid=(n_batch, 1, 1), block=(HEAD_MID, 1, 1). +// Each block handles ONE sample and loops over the 5 horizons. + +#define N_HORIZONS_H 5 +#define HIDDEN_H 128 +#define HEAD_MID_H 64 +#define GELU_C0 0.7978845608f // sqrt(2/pi) +#define GELU_C1 0.044715f + +__device__ __forceinline__ float gelu_act(float x) { + const float u = GELU_C0 * (x + GELU_C1 * x * x * x); + return 0.5f * x * (1.0f + tanhf(u)); +} + +extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched( + const float* __restrict__ w1, // [5, HEAD_MID, HIDDEN] + const float* __restrict__ b1, // [5, HEAD_MID] + const float* __restrict__ w2, // [5, HEAD_MID] + const float* __restrict__ b2, // [5] + const float* __restrict__ h, // [n_batch, HIDDEN] + int n_batch, + float* __restrict__ probs, // [n_batch, 5] + float* __restrict__ z1_out, // [n_batch, 5, HEAD_MID] + float* __restrict__ a1_out // [n_batch, 5, HEAD_MID] +) { + int b_idx = blockIdx.x; + int m = threadIdx.x; + if (b_idx >= n_batch || m >= HEAD_MID_H) return; + + const float* h_row = h + (long long)b_idx * HIDDEN_H; + + // Shared per-horizon post-GELU activations [5 * 64], indexed as [k*64 + m]. + __shared__ float s_a1[N_HORIZONS_H * HEAD_MID_H]; + + // Pass 1: per-horizon hidden unit. Thread `m` computes the m-th + // hidden unit for each of the 5 horizons sequentially. + #pragma unroll + 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]; + } + const float a = gelu_act(z); + s_a1[k * HEAD_MID_H + m] = a; + z1_out[((long long)b_idx * N_HORIZONS_H + k) * HEAD_MID_H + m] = z; + a1_out[((long long)b_idx * N_HORIZONS_H + k) * HEAD_MID_H + m] = a; + } + __syncthreads(); + + // Pass 2: per-horizon output logit. Threads m=0..4 each handle one + // horizon by computing the dot product over a1[k, :] and W2[k, :]. + if (m < N_HORIZONS_H) { + const int k = m; + float z2 = b2[k]; + #pragma unroll + for (int u = 0; u < HEAD_MID_H; ++u) { + z2 += w2[k * HEAD_MID_H + u] * s_a1[k * HEAD_MID_H + u]; + } + probs[(long long)b_idx * N_HORIZONS_H + k] = 1.0f / (1.0f + expf(-z2)); + } +}