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

// output_smoothness.cu — per-horizon temporal smoothness regularizer.
//
// Penalises (p[k] - p[k-1])² for adjacent positions within each sequence,
// weighted per-horizon by λ[h]. λ[h] scales with horizon length so slow
// horizons (h6000) are forced to change much more slowly than fast ones
// (h30). Forward loss is reported per-horizon (raw, pre-λ) for telemetry
// + lumped (λ-weighted) total.
//
// Per `feedback_nvidia_grade_perf_for_kernels.md`:
//   - Warp-shuffle reduction; no atomicAdd.
//   - Single block over the [K, B, N_HORIZONS] grid; stride loop per
//     thread covers all elements.
//   - No host branches, no divergent shuffles (inactive lanes contribute
//     0 via ternary).
//   - N_HORIZONS horizon accumulators kept in registers (BCS-style).
//
// Per `feedback_no_atomicadd.md`: single-block tree-reduce (no atomics).
// Per `pearl_one_unbounded_signal_per_reward.md`: smoothness contributes
// a bounded (≤ 1 per pair after sigmoid) additive term — no
// unboundedness concerns.
//
// Post-2026-05-22 horizon-rebase: OS_N_HORIZONS = 3 (matches Rust-side
// `crates/ml-alpha/src/heads.rs::N_HORIZONS`).

#define OS_N_HORIZONS 3
#define OS_BLOCK      256
#define OS_N_WARPS    (OS_BLOCK / 32)   // 8

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

// Smoothness forward + backward.
//
// Inputs:
//   probs           [K * B * N_HORIZONS]   row-major, index = k*B*N_HORIZONS + b*N_HORIZONS + h
//   lambda_per_h    [N_HORIZONS]           λ[h] per horizon (already includes base × ratio)
//   K, B            (scalars)
//
// Outputs:
//   loss_out                  [1]            Σ_h λ[h] · raw_h           — total loss
//   loss_raw_per_h            [N_HORIZONS]   raw mean-sq-diff per horizon (no λ) — telemetry
//   grad_probs                [K*B*N_HORIZONS]  ACCUMULATED (+=) into existing BCE grad
//
// The gradient at position j for the same (b, h):
//   grad += (2 λ[h] / N_pairs) · Δ(j, b, h)
//
// where Δ has the 3-branch boundary structure documented in the plan.
//
// Single block, OS_BLOCK threads. Each thread strides over the K*B*N_HORIZONS
// flat range computing its own Δ(j, b, h) for the gradient pass.
// Forward loss is accumulated per-horizon in registers, then warp/block
// reduced; only thread 0 writes the N_HORIZONS raw outputs + the λ-weighted total.
extern "C" __global__ void output_smoothness_loss_and_grad(
    const float* __restrict__ probs,            // [K * B * N_HORIZONS]
    const float* __restrict__ lambda_per_h,     // [N_HORIZONS]
    int K,
    int B,
    float* __restrict__ loss_out,               // [1]
    float* __restrict__ loss_raw_per_h,         // [N_HORIZONS]
    float* __restrict__ grad_probs              // [K * B * N_HORIZONS]   (+=)
) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int total = K * B * OS_N_HORIZONS;
    const int stride_k = B * OS_N_HORIZONS;     // stride between k and k+1 at fixed (b,h); B*N_HORIZONS
    const int n_pairs = (K - 1) * B;
    // N_pairs == 0 when K==1; in that case we still need to zero outputs.
    const float inv_n_pairs = (n_pairs > 0) ? (1.0f / (float)n_pairs) : 0.0f;

    // ── Forward: per-thread per-horizon raw accumulator ──
    // Anchor each (p[k]-p[k-1])² pair on its RIGHT element (k ≥ 1) so
    // each pair is counted exactly once across all threads.
    float local_raw[OS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < OS_N_HORIZONS; ++h) local_raw[h] = 0.0f;

    for (int i = tid; i < total; i += OS_BLOCK) {
        const int h = i % OS_N_HORIZONS;
        // Decompose i into (k, b, h). k = i / (B*5); b = (i / 5) % B.
        const int k = i / stride_k;
        // Only positions k ≥ 1 anchor a (p[k]-p[k-1])² pair.
        if (k >= 1) {
            const float p_cur  = probs[i];
            const float p_prev = probs[i - stride_k];
            const float d = p_cur - p_prev;
            local_raw[h] += d * d;
        }
    }

    // Warp reduce per-horizon partial sums.
    float warp_raw[OS_N_HORIZONS];
    #pragma unroll
    for (int h = 0; h < OS_N_HORIZONS; ++h) {
        warp_raw[h] = os_warp_reduce_sum(local_raw[h]);
    }

    __shared__ float s_raw[OS_N_WARPS * OS_N_HORIZONS];
    if (lane == 0) {
        #pragma unroll
        for (int h = 0; h < OS_N_HORIZONS; ++h) {
            s_raw[warp * OS_N_HORIZONS + h] = warp_raw[h];
        }
    }
    __syncthreads();

    // Warp 0 finishes the cross-warp reduce.
    __shared__ float s_lambda[OS_N_HORIZONS];
    if (warp == 0) {
        #pragma unroll
        for (int h = 0; h < OS_N_HORIZONS; ++h) {
            float v = (lane < OS_N_WARPS) ? s_raw[lane * OS_N_HORIZONS + h] : 0.0f;
            v = os_warp_reduce_sum(v);
            if (lane == 0) {
                const float raw_h = v * inv_n_pairs;
                loss_raw_per_h[h] = raw_h;
                s_lambda[h] = lambda_per_h[h];
                // s_raw repurposed below as λ-weighted scratch.
                s_raw[h] = s_lambda[h] * raw_h;
            }
        }
    }
    __syncthreads();

    // Thread 0 sums the N_HORIZONS λ-weighted terms into loss_out.
    if (tid == 0) {
        float total_loss = 0.0f;
        #pragma unroll
        for (int h = 0; h < OS_N_HORIZONS; ++h) total_loss += s_raw[h];
        loss_out[0] = total_loss;
    }
    __syncthreads();

    // ── Backward: per-thread compute Δ(j, b, h) and accumulate into grad_probs ──
    // Same stride loop. Three boundary cases:
    //   j = 0       : Δ = p[0]   − p[1]
    //   j = K−1     : Δ = p[K−1] − p[K−2]
    //   interior    : Δ = 2·p[j] − p[j−1] − p[j+1]
    // Scale: factor = 2 · λ[h] / N_pairs.
    //
    // s_lambda is now populated for all threads.
    for (int i = tid; i < total; i += OS_BLOCK) {
        const int h = i % OS_N_HORIZONS;
        const int k = i / stride_k;
        const float lam = s_lambda[h];
        const float factor = 2.0f * lam * inv_n_pairs;
        const float p_j = probs[i];
        const bool has_left  = (k >= 1);
        const bool has_right = (k <= K - 2);
        // Branchless boundary: at k=0 we read self for the "left" slot
        // and at k=K-1 we read self for the "right" slot; both indices
        // stay in-bounds and the lm/rm multipliers below zero the
        // spurious contribution. Every lane executes both loads — no
        // thread divergence on boundaries.
        const int idx_left  = has_left  ? (i - stride_k) : i;
        const int idx_right = has_right ? (i + stride_k) : i;
        const float p_jm1 = probs[idx_left];
        const float p_jp1 = probs[idx_right];
        // Δ(j, b, h):
        //   left contribution  =  (p[j] - p[j-1])   if has_left  else 0
        //   right contribution = -(p[j+1] - p[j])   if has_right else 0
        // Sum = 2·p[j] − p[j-1] − p[j+1]  for interior;
        //       p[j] − p[j+1]             for j==0      (has_left=F, has_right=T)
        //       p[j] − p[j-1]             for j==K-1    (has_left=T,  has_right=F)
        // Use float multiplies to encode has_left/has_right (1.0 / 0.0).
        const float lm = has_left  ? 1.0f : 0.0f;
        const float rm = has_right ? 1.0f : 0.0f;
        const float delta = (lm + rm) * p_j - lm * p_jm1 - rm * p_jp1;
        grad_probs[i] += factor * delta;
    }
}