feat(ml-alpha): 2-layer GELU MLP heads — forward kernel (Phase 1.3)

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));
+    }
+}