foxhunt/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu

// bce_loss_multi_horizon.cu — base-weight-only multi-horizon BCE.
//
// Single source of truth for the multi-horizon BCE. Loss is the
// base-weight-scaled per-horizon mean of unweighted BCE — no Kendall
// σ damping (2026-05-18). Per-horizon prioritization is done outside
// this kernel by the ISV-driven λ controller scaling the trunk
// grad_h contribution in heads_grn_bwd (per `pearl_adam_normalizes_loss_weights`
// — affect the gradient, not the loss aggregate).
//
// Math:
//
//   raw_bce_h  = Σ_{i in horizon h, m_i = 1} L_i        (unweighted sum)
//   count_h    = #{i in horizon h : m_i = 1}
//   mean_bce_h = raw_bce_h / count_h
//   w_h        = base_weight_h                          (typically 1)
//   total_loss = Σ_h [ w_h * mean_bce_h ]
//
// Gradients:
//   d L / d p_i           = m_i * (w_h / count_h) * (p − y) / (p (1 − p))
//   d L / d log_sigma_h   = 0  (σ is a sidecar — emitted by
//                               horizon_ema_and_lambda for telemetry,
//                               not used in the loss path)
//
// Per `feedback_nvidia_grade_perf_for_kernels`:
//   - Warp-shuffle reduction (`__shfl_xor_sync`), NOT block tree-reduce.
//   - One `__syncthreads` for the cross-warp aggregate, total = O(1) barriers.
//   - Inactive lanes contribute 0 via ternary; never divergent shuffles.
//   - Hard-coded N_HORIZONS_BCE = 5 → per-warp 5-element accumulator stays in registers.
//
// Block layout: grid = (1), block = (256) = 8 warps. Single block over
// the full [n_pos, n_horizons] grid is fine — we never have > 256 ×
// (per-thread fan-in) positions in the trainer hot loop (max K × B =
// 64 × 8 = 512, well within stride-loop reach).

#define BCS_N_HORIZONS 5
#define BCS_BLOCK      256
#define BCS_N_WARPS    (BCS_BLOCK / 32)   // == 8

__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}

__device__ __forceinline__ int warp_reduce_sum_int(int v) {
    #pragma unroll
    for (int s = 16; s > 0; s >>= 1) {
        v += __shfl_xor_sync(0xffffffff, v, s);
    }
    return v;
}

extern "C" __global__ void bce_multi_horizon_forward_backward(
    const float* __restrict__ probs,         // [n_pos * n_horizons]
    const float* __restrict__ labels,        // [n_pos * n_horizons]
    const float* __restrict__ base_weights,  // [n_horizons]  (per-horizon λ from ISV; constants OK)
    const float* __restrict__ log_sigma_h,   // [n_horizons]  (learnable Kendall σ logarithm)
    int n_pos,
    int n_horizons,
    float* __restrict__ loss_out,            // [1]                  — total Kendall loss
    float* __restrict__ mean_bce_per_h_out,  // [n_horizons]         — unweighted per-horizon mean (ISV signal)
    float* __restrict__ grad_probs,          // [n_pos * n_horizons] — ∂L/∂p
    int*   __restrict__ valid_count_out,     // [1]                  — Σ_i m_i (unweighted)
    float* __restrict__ d_log_sigma_h        // [n_horizons]         — ∂L/∂log_sigma_h
) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int total = n_pos * n_horizons;

    // Per-thread per-horizon partials in registers.
    float local_raw[BCS_N_HORIZONS];
    float local_cnt[BCS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < BCS_N_HORIZONS; ++h) {
        local_raw[h] = 0.0f;
        local_cnt[h] = 0.0f;
    }
    int local_valid = 0;

    // Pass 1 — per-thread strided accumulation. Same loop also zeros
    // grad_probs at masked positions so the second pass can skip those.
    for (int i = tid; i < total; i += BCS_BLOCK) {
        const float y = labels[i];
        if (isnan(y)) {
            grad_probs[i] = 0.0f;
            continue;
        }
        const int h = i % n_horizons;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        const float L = -(y * logf(p) + (1.0f - y) * logf(1.0f - p));
        local_raw[h] += L;
        local_cnt[h] += 1.0f;
        local_valid  += 1;
    }

    // Warp-shuffle reduce each of the 5 horizon sums + the per-horizon
    // count + the global valid count. Each lane in warp gets the full
    // warp-sum at the end (no need to broadcast manually).
    float warp_raw[BCS_N_HORIZONS];
    float warp_cnt[BCS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < BCS_N_HORIZONS; ++h) {
        warp_raw[h] = warp_reduce_sum(local_raw[h]);
        warp_cnt[h] = warp_reduce_sum(local_cnt[h]);
    }
    int warp_valid = warp_reduce_sum_int(local_valid);

    // Cross-warp reduce — only lane 0 of each warp writes its partial.
    __shared__ float s_raw[BCS_N_WARPS * BCS_N_HORIZONS];
    __shared__ float s_cnt[BCS_N_WARPS * BCS_N_HORIZONS];
    __shared__ int   s_valid[BCS_N_WARPS];
    if (lane == 0) {
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            s_raw[warp * BCS_N_HORIZONS + h] = warp_raw[h];
            s_cnt[warp * BCS_N_HORIZONS + h] = warp_cnt[h];
        }
        s_valid[warp] = warp_valid;
    }
    __syncthreads();

    // Warp 0 finishes: reduce N_WARPS partials. All 32 lanes participate;
    // inactive lanes (tid ≥ N_WARPS) contribute 0 via ternary (NOT a
    // divergent shuffle).
    __shared__ float mean_bce_h_shared[BCS_N_HORIZONS];
    __shared__ float w_eff_shared[BCS_N_HORIZONS];
    __shared__ int valid_count_shared;

    if (warp == 0) {
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            float v_raw = (lane < BCS_N_WARPS) ? s_raw[lane * BCS_N_HORIZONS + h] : 0.0f;
            float v_cnt = (lane < BCS_N_WARPS) ? s_cnt[lane * BCS_N_HORIZONS + h] : 0.0f;
            v_raw = warp_reduce_sum(v_raw);
            v_cnt = warp_reduce_sum(v_cnt);
            if (lane == 0) {
                const float c_h = v_cnt;
                const float mean_bce = (c_h > 0.0f) ? v_raw / c_h : 0.0f;
                mean_bce_h_shared[h] = mean_bce;
                mean_bce_per_h_out[h] = mean_bce;

                // σ sidecar (2026-05-18): Kendall σ damped every horizon
                // by ~30% (σ ≈ √mean_bce ≈ 0.84 → w_h ≈ 0.71) and biased
                // the encoder away from h6000 — empirically -2pp mean_auc
                // vs Phase 1+2+3 across multiple architectures. The
                // Kendall framework's assumption ("damp the noisier
                // task") inverts what we need ("learn the harder
                // horizon harder"). σ is now emitted by
                // horizon_ema_and_lambda for telemetry only and does
                // NOT participate in the BCE forward weighting or
                // backward gradient. `log_sigma_h` arg is preserved
                // for kernel-signature stability.
                const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
                const float w_h = bw;
                w_eff_shared[h] = (c_h > 0.0f) ? w_h / c_h : 0.0f;
                d_log_sigma_h[h] = 0.0f;
            }
        }
        int v_valid = (lane < BCS_N_WARPS) ? s_valid[lane] : 0;
        v_valid = warp_reduce_sum_int(v_valid);
        if (lane == 0) {
            valid_count_shared = v_valid;
            valid_count_out[0] = v_valid;
        }
    }
    __syncthreads();

    // Total loss — small serial reduce by thread 0 over 5 horizons.
    // σ-Kendall regularizer (log σ_h term) intentionally absent — see
    // sidecar note above. Loss is unweighted per-horizon BCE sum,
    // scaled only by `base_weights[h]` (typically uniform [1,1,1,1,1]).
    if (tid == 0) {
        float total_loss = 0.0f;
        #pragma unroll
        for (int h = 0; h < BCS_N_HORIZONS; ++h) {
            const float bw = (base_weights != nullptr) ? base_weights[h] : 1.0f;
            total_loss += bw * mean_bce_h_shared[h];
        }
        loss_out[0] = total_loss;
    }
    __syncthreads();

    // Pass 2 — per-element grad. Coalesced strided access. w_eff_shared
    // is read-only at this point, broadcast freely across threads.
    for (int i = tid; i < total; i += BCS_BLOCK) {
        const float y = labels[i];
        if (isnan(y)) {
            continue;  // grad was already zeroed in pass 1
        }
        const int h = i % n_horizons;
        const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f);
        grad_probs[i] = w_eff_shared[h] * (p - y) / (p * (1.0f - p));
    }
}