// multi_horizon_heads.cu
//
// 128-dim hidden -> N_HORIZONS sigmoid logits, one per horizon
// (post-2026-05-22 rebase: 3 horizons {10, 100, 1000} snapshots forward).
// Single-block kernel, N_HORIZONS threads (one per head). Each thread
// computes its own dot product against the 128-dim hidden vector + bias,
// then sigmoid.

// File-wide layout constants (must agree with
// `crates/ml-alpha/src/heads.rs::{N_HORIZONS, HIDDEN_DIM, HEAD_MID_DIM}`).
#define N_HORIZONS_H 3
#define HIDDEN_H     128
#define HEAD_MID_H   64
#define GELU_C0      0.7978845608f   // sqrt(2/pi)
#define GELU_C1      0.044715f

extern "C" __global__ void multi_horizon_heads(
    const float* __restrict__ w,   // [N_HORIZONS, 128]
    const float* __restrict__ b,   // [N_HORIZONS]
    const float* __restrict__ h,   // [128]
    float* __restrict__ probs      // [N_HORIZONS]
) {
    int k = threadIdx.x;
    if (k >= N_HORIZONS_H) return;
    float z = b[k];
    for (int i = 0; i < 128; ++i) {
        z += w[k * 128 + i] * h[i];
    }
    probs[k] = 1.0f / (1.0f + expf(-z));
}

// Backward through multi_horizon_heads.
//
// Given:
//   grad_probs[N_HORIZONS] = dL / dp
//   probs[N_HORIZONS]      = forward output (sigmoid)
//   h[128]                 = forward input
// Computes:
//   grad_w[N_HORIZONS, 128] = dL / dW   = d_z[k] * h[i]
//   grad_b[N_HORIZONS]      = dL / db   = d_z[k]
//   grad_h[128]             = dL / dh   = sum_k d_z[k] * W[k, i]
// where d_z[k] = grad_probs[k] * p[k] * (1 - p[k]).
//
// Block tree-reduce for the grad_h sum (no atomicAdd). One thread per
// hidden unit (128 threads); thread 0 also writes grad_b.

extern "C" __global__ void multi_horizon_heads_backward(
    const float* __restrict__ w,             // [N_HORIZONS, 128]
    const float* __restrict__ probs,         // [N_HORIZONS]
    const float* __restrict__ h,             // [128]
    const float* __restrict__ grad_probs,    // [N_HORIZONS]
    const float* __restrict__ grad_h_carry,  // [128] — optional carry (nullptr OK)
    float* __restrict__ grad_w,              // [N_HORIZONS, 128]  (+=)
    float* __restrict__ grad_b,              // [N_HORIZONS]       (+=)
    float* __restrict__ grad_h               // [128]              overwrite-with-carry
) {
    __shared__ float d_z[N_HORIZONS_H];

    int tid = threadIdx.x;
    if (tid < N_HORIZONS_H) {
        const float p = probs[tid];
        d_z[tid] = grad_probs[tid] * p * (1.0f - p);
    }
    __syncthreads();

    // Parameter-grad writes use `+=`: per-position supervision invokes
    // this kernel K times per training step and needs the gradients
    // accumulated. Callers MUST pre-zero grad_w / grad_b at step start.
    if (tid < N_HORIZONS_H) {
        grad_b[tid] += d_z[tid];
    }
    // grad_w[k, i] += d_z[k] * h[i]
    // grid stride: each thread covers one i for k in 0..N_HORIZONS.
    if (tid < 128) {
        const float h_i = h[tid];
        for (int k = 0; k < N_HORIZONS_H; ++k) {
            grad_w[k * 128 + tid] += d_z[k] * h_i;
        }
        // grad_h[i] = sum_k d_z[k] * W[k, i] + carry[i]
        // The optional carry is the grad_h_old emitted by the cfc_step_bwd
        // of position k+1 — when the CfC is RECURRENT across positions,
        // h_new_k flows into h_old_{k+1}, so grad_h_new at position k must
        // sum the heads-side gradient with the carry from the upstream step.
        float acc = 0.0f;
        for (int k = 0; k < N_HORIZONS_H; ++k) {
            acc += d_z[k] * w[k * 128 + tid];
        }
        const float carry = (grad_h_carry != nullptr) ? grad_h_carry[tid] : 0.0f;
        grad_h[tid] = acc + carry;
    }
}


// ─── Batched variants ────────────────────────────────────────────────
// Process N_BATCH samples per launch. Layout: h[B, 128], probs[B, N_HORIZONS].
// Thread tid handles either head k (tid < N_HORIZONS) or hidden unit i
// (tid < 128). Block dim = 128, one block per launch — each thread loops
// over b internally. Param grads (grad_w, grad_b) shared across batch and
// accumulated via `+=` (no race: thread tid is sole writer to its row).

extern "C" __global__ void multi_horizon_heads_batched(
    const float* __restrict__ w,         // [N_HORIZONS, 128]
    const float* __restrict__ b,         // [N_HORIZONS]
    const float* __restrict__ h,         // [n_batch, 128]
    int n_batch,
    float* __restrict__ probs            // [n_batch, N_HORIZONS]
) {
    int k = threadIdx.x;
    if (k >= N_HORIZONS_H) return;
    const float bias_k = b[k];
    for (int bi = 0; bi < n_batch; ++bi) {
        const float* h_b = h + (long long)bi * 128;
        float z = bias_k;
        for (int i = 0; i < 128; ++i) z += w[k * 128 + i] * h_b[i];
        probs[(long long)bi * N_HORIZONS_H + k] = 1.0f / (1.0f + expf(-z));
    }
}


extern "C" __global__ void multi_horizon_heads_backward_batched(
    const float* __restrict__ w,              // [N_HORIZONS, 128]
    const float* __restrict__ probs,          // [n_batch, N_HORIZONS]
    const float* __restrict__ h,              // [n_batch, 128]
    const float* __restrict__ grad_probs,     // [n_batch, N_HORIZONS]
    const float* __restrict__ grad_h_carry,   // [n_batch, 128]  (nullptr OK)
    const float* __restrict__ lambda,         // [N_HORIZONS] — per-horizon trunk-grad scaler
    int n_batch,
    float* __restrict__ grad_w,               // [N_HORIZONS, 128]   accum +=
    float* __restrict__ grad_b,               // [N_HORIZONS]        accum +=
    float* __restrict__ grad_h                // [n_batch, 128]  overwrite + carry
) {
    // Shared mem: d_z[B, N_HORIZONS]. At B=32: 32 * 3 = 96 floats = 384 bytes.
    extern __shared__ float sd_z[];           // size = n_batch * N_HORIZONS_H

    int tid = threadIdx.x;

    // Pass 1: per-(bi, k) compute d_z and stash in shared.
    if (tid < N_HORIZONS_H) {
        for (int bi = 0; bi < n_batch; ++bi) {
            const float p = probs[(long long)bi * N_HORIZONS_H + tid];
            const float dp = grad_probs[(long long)bi * N_HORIZONS_H + tid];
            sd_z[(long long)bi * N_HORIZONS_H + tid] = dp * p * (1.0f - p);
        }
    }
    __syncthreads();

    // Per-horizon trunk-gradient scaler, broadcast via shared mem so
    // every thread reads the same N_HORIZONS floats without repeated global loads.
    __shared__ float s_lambda[N_HORIZONS_H];
    if (tid < N_HORIZONS_H) {
        // Sentinel: a fully zero lambda buffer (initial state, before
        // the EMA kernel has run even once) is treated as "scaler = 1"
        // so the kernel reduces to its pre-ISV behavior. After step 1
        // every entry is clamped to [0.5, 2.0] by horizon_lambda.cu
        // and the > 0 check is always true.
        const float l = lambda[tid];
        s_lambda[tid] = (l > 0.0f) ? l : 1.0f;
    }
    __syncthreads();

    // grad_b[k] += sum_b d_z[b, k]. Thread k (k < N_HORIZONS) is sole writer.
    // NOTE: lambda does NOT scale the head's own bias gradient — only
    // the trunk gradient. We want the per-horizon head to keep learning
    // its own bias normally; lambda only biases how much the horizon's
    // error signal LEAKS into the shared trunk.
    if (tid < N_HORIZONS_H) {
        float acc = 0.0f;
        for (int bi = 0; bi < n_batch; ++bi) {
            acc += sd_z[(long long)bi * N_HORIZONS_H + tid];
        }
        grad_b[tid] += acc;
    }

    // grad_w[k, i] += sum_b d_z[b, k] * h[b, i]. Same rationale as
    // grad_b: lambda does NOT scale the head's own weight gradient.
    if (tid < 128) {
        // grad_h: sum_k LAMBDA[k] * d_z[b, k] * W[k, i] for each b,
        // plus optional carry. This is the only path where lambda is
        // applied — per `pearl_adam_normalizes_loss_weights.md` the
        // effective lever for horizon-aware learning is the trunk
        // gradient, not the head loss weight.
        for (int bi = 0; bi < n_batch; ++bi) {
            const float h_bi = h[(long long)bi * 128 + tid];
            // Accumulate W-row grads via += into global (thread tid sole writer of column tid).
            for (int k = 0; k < N_HORIZONS_H; ++k) {
                grad_w[k * 128 + tid] += sd_z[(long long)bi * N_HORIZONS_H + k] * h_bi;
            }
            float acc = 0.0f;
            for (int k = 0; k < N_HORIZONS_H; ++k) {
                acc += s_lambda[k] * sd_z[(long long)bi * N_HORIZONS_H + k] * w[k * 128 + tid];
            }
            const float carry = (grad_h_carry != nullptr)
                ? grad_h_carry[(long long)bi * 128 + tid] : 0.0f;
            grad_h[(long long)bi * 128 + tid] = acc + carry;
        }
    }
}

// ─────────────────────────────────────────────────────────────────────
// Phase 1: 2-layer MLP heads — per-horizon [HIDDEN=128 → HEAD_MID=64]
// GELU hidden → [HEAD_MID=64 → 1] sigmoid output.
//
// Forward layout:
//   w1     [N_HORIZONS, HEAD_MID, HIDDEN]  per-horizon first-layer weights
//   b1     [N_HORIZONS, HEAD_MID]          per-horizon first-layer biases
//   w2     [N_HORIZONS, HEAD_MID]          per-horizon output weights (scalar per horizon)
//   b2     [N_HORIZONS]                    per-horizon output bias
//   h      [n_batch, HIDDEN]               trunk output (LayerNorm-normalised)
//   probs  [n_batch, N_HORIZONS]           sigmoid output
//   z1_out [n_batch, N_HORIZONS, HEAD_MID] saved pre-GELU for backward
//   a1_out [n_batch, N_HORIZONS, 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 N_HORIZONS horizons.

__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,        // [N_HORIZONS, HEAD_MID, HIDDEN]
    const float* __restrict__ b1,        // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ w2,        // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ b2,        // [N_HORIZONS]
    const float* __restrict__ h,         // [n_batch, HIDDEN]
    int n_batch,
    float* __restrict__ probs,           // [n_batch, N_HORIZONS]
    float* __restrict__ z1_out,          // [n_batch, N_HORIZONS, HEAD_MID]
    float* __restrict__ a1_out           // [n_batch, N_HORIZONS, 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 [N_HORIZONS_H * 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 N_HORIZONS 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..(N_HORIZONS-1) 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));
    }
}

// 2-layer heads backward — chain rule through sigmoid → linear →
// GELU → linear → trunk. ISV `lambda[N_HORIZONS]` 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<N_HORIZONS
// 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..N_HORIZONS (within m<N_HORIZONS 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,            // [N_HORIZONS, HEAD_MID, HIDDEN]
    const float* __restrict__ w2,            // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ probs,         // [n_batch, N_HORIZONS]
    const float* __restrict__ grad_probs,    // [n_batch, N_HORIZONS]
    const float* __restrict__ z1,            // [n_batch, N_HORIZONS, HEAD_MID]
    const float* __restrict__ a1,            // [n_batch, N_HORIZONS, HEAD_MID]
    const float* __restrict__ h,             // [n_batch, HIDDEN]
    const float* __restrict__ grad_h_carry,  // [n_batch, HIDDEN] (nullptr OK)
    const float* __restrict__ lambda,        // [N_HORIZONS]
    int n_batch,
    float* __restrict__ grad_w1,             // [N_HORIZONS, HEAD_MID, HIDDEN] (+=)
    float* __restrict__ grad_b1,             // [N_HORIZONS, HEAD_MID] (+=)
    float* __restrict__ grad_w2,             // [N_HORIZONS, HEAD_MID] (+=)
    float* __restrict__ grad_b2,             // [N_HORIZONS] (+=)
    float* __restrict__ grad_h               // [n_batch, HIDDEN] (overwrite + carry)
) {
    int tid = threadIdx.x;
    if (tid >= HEAD_MID_H) return;

    // Broadcast lambda[N_HORIZONS] via shared mem (N_HORIZONS 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..(N_HORIZONS-1) (only first N_HORIZONS 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 N_HORIZONS 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();
    }
}

// ─────────────────────────────────────────────────────────────────────
//  TFT Gated Residual Network (GRN) heads — Phase 1.7 (2026-05-17)
// ─────────────────────────────────────────────────────────────────────
//
// Per-horizon GRN structure (Lim et al. 2021 §3.3, adapted to scalar
// output per horizon):
//
//   eta_2[k, m_out]  = GELU(W1[k, m_out, i] @ h[i] + b1[k, m_out])   # [HIDDEN] → [HEAD_MID]
//   eta_1[k, m_out]  = W2[k, m_out, m_in] @ eta_2[k, m_in] + b2[...]  # [HEAD_MID] → [HEAD_MID]
//   gate_lin[k]      = W_gate[k, m] @ eta_1[k, m] + b_gate[k]         # [HEAD_MID] → scalar
//   main[k]          = W_main[k, m] @ eta_1[k, m] + b_main[k]
//   skip[k]          = W_skip[k, i] @ h[i] + b_skip[k]                # [HIDDEN] → scalar
//   logit[k]         = skip[k] + sigmoid(gate_lin[k]) * main[k]
//   p[k]             = sigmoid(logit[k])
//
// Naming convention follows the existing 2-layer kernel: w1 = first
// linear (HIDDEN → HEAD_MID, GELU-activated), w2 = second linear
// (HEAD_MID → HEAD_MID, no activation; produces eta_1).
//
// Thread layout (block_dim = HEAD_MID = 64):
//   m = threadIdx.x  -- owns "output_row" dim of (k, m_out=m) for w1/w2/eta_1/eta_2
//   for skip / gate / main scalars: thread tid<N_HORIZONS handles horizon k=tid

extern "C" __global__ void multi_horizon_heads_grn_fwd_batched(
    const float* __restrict__ w1,         // [N_HORIZONS, HEAD_MID, HIDDEN]
    const float* __restrict__ b1,         // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ w2,         // [N_HORIZONS, HEAD_MID_out, HEAD_MID_in]
    const float* __restrict__ b2,         // [N_HORIZONS, HEAD_MID_out]
    const float* __restrict__ w_gate,     // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ b_gate,     // [N_HORIZONS]
    const float* __restrict__ w_main,     // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ b_main,     // [N_HORIZONS]
    const float* __restrict__ w_skip,     // [N_HORIZONS, HIDDEN]
    const float* __restrict__ b_skip,     // [N_HORIZONS]
    const float* __restrict__ h,          // [B, HIDDEN]
    int n_batch,
    float* __restrict__ probs,            // [B, N_HORIZONS]
    float* __restrict__ z1_out,           // [B, N_HORIZONS, HEAD_MID] — pre-GELU z1 (= W1 @ h + b1)
    float* __restrict__ a1_out,           // [B, N_HORIZONS, HEAD_MID] — post-GELU eta_2
    float* __restrict__ z2_out,           // [B, N_HORIZONS, HEAD_MID] — eta_1 (= W2 @ eta_2 + b2)
    float* __restrict__ gate_logit_out,   // [B, N_HORIZONS] — pre-sigmoid gate scalar
    float* __restrict__ main_out,         // [B, N_HORIZONS]
    float* __restrict__ logit_out         // [B, N_HORIZONS] — pre-final-sigmoid logit (skip + sigma(gate)*main)
) {
    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__ 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.
    #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] * s_h[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: z2[k, m] = W2[k, m_out=m, m_in=n] @ a1[k, n] + b2[k, m].
    //   Pure linear — no activation. Thread m owns output row m for all k.
    #pragma unroll
    for (int k = 0; k < N_HORIZONS_H; ++k) {
        float z = b2[k * HEAD_MID_H + m];
        for (int n = 0; n < HEAD_MID_H; ++n) {
            z += w2[((long long)(k * HEAD_MID_H) + m) * HEAD_MID_H + n]
               * s_a1[k * HEAD_MID_H + n];
        }
        s_z2[k * HEAD_MID_H + m] = z;
        z2_out[((long long)b_idx * N_HORIZONS_H + k) * HEAD_MID_H + m] = z;
    }
    __syncthreads();

    // Pass 3: per-horizon scalars. Threads 0..(N_HORIZONS-1) each handle one horizon.
    //   gate_lin[k] = W_gate[k, :] @ z2[k, :] + b_gate[k]
    //   main[k]     = W_main[k, :] @ z2[k, :] + b_main[k]
    //   skip[k]     = W_skip[k, :] @ h + b_skip[k]
    //   logit[k]    = skip[k] + sigmoid(gate_lin[k]) * main[k]
    //   p[k]        = sigmoid(logit[k])
    if (m < N_HORIZONS_H) {
        const int k = m;
        float gl = b_gate[k];
        float mv = b_main[k];
        #pragma unroll
        for (int n = 0; n < HEAD_MID_H; ++n) {
            const float z2_kn = s_z2[k * HEAD_MID_H + n];
            gl += w_gate[k * HEAD_MID_H + n] * z2_kn;
            mv += w_main[k * HEAD_MID_H + n] * z2_kn;
        }
        float sk = b_skip[k];
        for (int i = 0; i < HIDDEN_H; ++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;
        const float p = 1.0f / (1.0f + expf(-logit));

        gate_logit_out[(long long)b_idx * N_HORIZONS_H + k] = gl;
        main_out[(long long)b_idx * N_HORIZONS_H + k] = mv;
        logit_out[(long long)b_idx * N_HORIZONS_H + k] = logit;
        probs[(long long)b_idx * N_HORIZONS_H + k] = p;
    }
}

// GRN backward — chain rule:
//
//   d_logit  = grad_probs * p * (1 - p)
//   d_skip   = d_logit
//   d_main   = d_logit * sigma(gate_lin)
//   d_gate_a = d_logit * main
//   d_gate_lin = d_gate_a * sigma(gate_lin) * (1 - sigma(gate_lin))
//
// Per-horizon scalar grads:
//   grad_w_skip[k, i] += d_skip[k] * h[i]
//   grad_b_skip[k]    += d_skip[k]
//   grad_w_main[k, m] += d_main[k] * eta_1[k, m]
//   grad_b_main[k]    += d_main[k]
//   grad_w_gate[k, m] += d_gate_lin[k] * eta_1[k, m]
//   grad_b_gate[k]    += d_gate_lin[k]
//
// Through eta_1:
//   d_eta_1[k, m] = d_main * w_main[k, m] + d_gate_lin * w_gate[k, m]
//   grad_w2[k, m_out, m_in] += d_eta_1[k, m_out] * eta_2[k, m_in]
//   grad_b2[k, m_out]        += d_eta_1[k, m_out]
//   d_eta_2[k, m_in] = sum_{m_out} d_eta_1[k, m_out] * w2[k, m_out, m_in]
//
// Through GELU + W1:
//   d_z1[k, m] = d_eta_2[k, m] * gelu_prime(z1[k, m])
//   grad_w1[k, m, i] += d_z1[k, m] * h[i]
//   grad_b1[k, m]    += d_z1[k, m]
//
// Trunk gradient (lambda-scaled, sums two paths):
//   grad_h[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 through W2→W1)
//               ) + carry[i]
//
// Single-writer discipline:
//   Thread m (= m_out for w2, m for w1) owns row m of grad_w2 (writes
//   cols m_in for all m_in) and row m of grad_w1 (writes cols i for
//   all i). Thread m also owns col m of d_eta_2 (sums over m_out).
//   Thread tid (with tid < N_HORIZONS) owns horizon-scalar grads.
//   Trunk grad_h tiles i over 2 strides of HEAD_MID.

// Block-per-batch GRN bwd (Phase B).
//   grid=(n_batch, 1, 1)  block=(HEAD_MID, 1, 1)
//
// Param-grad scratch tensors hold per-batch slices [B, ...]. Thread
// (bi, m) is sole writer to its slice for every K-iteration; the
// K-loop's 64 invocations accumulate via += within (bi, ..., m).
// Caller zeroes scratch once per training step; reduce_axis0 collapses
// B → final grad after the K-loop.
//
// grad_h is per-batch indexed; block bi is sole writer to its
// [bi, :] slice (overwrite + carry).
extern "C" __global__ void multi_horizon_heads_grn_bwd_batched(
    const float* __restrict__ w1,                  // [N_HORIZONS, HEAD_MID, HIDDEN]
    const float* __restrict__ w2,                  // [N_HORIZONS, HEAD_MID_out, HEAD_MID_in]
    const float* __restrict__ w_gate,              // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ w_main,              // [N_HORIZONS, HEAD_MID]
    const float* __restrict__ w_skip,              // [N_HORIZONS, HIDDEN]
    const float* __restrict__ probs,               // [B, N_HORIZONS]
    const float* __restrict__ grad_probs,          // [B, N_HORIZONS]
    const float* __restrict__ z1,                  // [B, N_HORIZONS, HEAD_MID]
    const float* __restrict__ a1,                  // [B, N_HORIZONS, HEAD_MID]
    const float* __restrict__ z2,                  // [B, N_HORIZONS, HEAD_MID]
    const float* __restrict__ gate_logit,          // [B, N_HORIZONS]
    const float* __restrict__ main_val,            // [B, N_HORIZONS]
    const float* __restrict__ h,                   // [B, HIDDEN]
    const float* __restrict__ grad_h_carry,        // [B, HIDDEN] (nullptr OK)
    const float* __restrict__ lambda,              // [N_HORIZONS]
    int n_batch,
    float* __restrict__ grad_w1_scratch,           // [B, N_HORIZONS, HEAD_MID, HIDDEN] (+=)
    float* __restrict__ grad_b1_scratch,           // [B, N_HORIZONS, HEAD_MID]         (+=)
    float* __restrict__ grad_w2_scratch,           // [B, N_HORIZONS, HEAD_MID, HEAD_MID] (+=)
    float* __restrict__ grad_b2_scratch,           // [B, N_HORIZONS, HEAD_MID]         (+=)
    float* __restrict__ grad_w_gate_scratch,       // [B, N_HORIZONS, HEAD_MID]         (+=)
    float* __restrict__ grad_b_gate_scratch,       // [B, N_HORIZONS]                   (+=)
    float* __restrict__ grad_w_main_scratch,       // [B, N_HORIZONS, HEAD_MID]         (+=)
    float* __restrict__ grad_b_main_scratch,       // [B, N_HORIZONS]                   (+=)
    float* __restrict__ grad_w_skip_scratch,       // [B, N_HORIZONS, HIDDEN]           (+=)
    float* __restrict__ grad_b_skip_scratch,       // [B, N_HORIZONS]                   (+=)
    float* __restrict__ grad_h                     // [B, HIDDEN] (overwrite + carry)
) {
    int bi = blockIdx.x;
    int tid = threadIdx.x;
    if (bi >= n_batch || tid >= HIDDEN_H) return;

    __shared__ float s_lambda[N_HORIZONS_H];
    __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];
    // 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];

    // 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.
    const long long bi_5    = (long long)bi * N_HORIZONS_H;
    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..(N_HORIZONS-1) only.
    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];
        const float d_logit = dp * p * (1.0f - p);

        const float gl = gate_logit[(long long)bi * N_HORIZONS_H + k];
        const float sigmoid_gl = 1.0f / (1.0f + expf(-gl));
        const float mv = main_val[(long long)bi * N_HORIZONS_H + k];

        const float d_skip = d_logit;
        const float d_main = d_logit * sigmoid_gl;
        const float d_gate_a = d_logit * mv;
        const float d_gate_lin = d_gate_a * sigmoid_gl * (1.0f - sigmoid_gl);

        s_d_skip[k] = d_skip;
        s_d_main[k] = d_main;
        s_d_gate_lin[k] = d_gate_lin;

        // Per-batch scratch writes for horizon-scalar biases.
        grad_b_skip_scratch[bi_5 + k]  += d_skip;
        grad_b_main_scratch[bi_5 + k]  += d_main;
        grad_b_gate_scratch[bi_5 + k]  += d_gate_lin;
    }
    __syncthreads();

    // Pass 2: per-(k, m_out=m) compute d_eta_1[k, m] + grad_w_{main,gate}.
    // 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;

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

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

    // 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]
    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).
    // 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) {
            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. 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]
    // 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) {
            grad_w_skip_scratch[bi_5_h + (long long)k * HIDDEN_H + i]
                += s_d_skip[k] * h_bi_i;

            float contrib = s_d_skip[k] * w_skip[k * HIDDEN_H + i];
            #pragma unroll
            for (int mm = 0; mm < HEAD_MID_H; ++mm) {
                contrib += s_d_z1[k * HEAD_MID_H + mm]
                         * w1[((long long)(k * HEAD_MID_H) + mm) * HIDDEN_H + i];
            }
            acc_grad_h += s_lambda[k] * contrib;
        }
        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_grad_h + carry;
    }
}
