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

/* =====================================================================
 * Mamba2 selective scan — ml-alpha (Phase 1d.1)
 *
 * Forward + analytical backward for the SSM scan used by the supervised
 * snapshot-stream pipeline. Purpose-built for the supervised path — no
 * ISV adaptive signals, no per-position temporal_weight, no NULL-pointer
 * branches.
 *
 * Layout conventions (row-major):
 *   a_proj          : [N, K, state_d]    per-step state-update gates (pre-sigmoid)
 *   b_proj          : [N, K, state_d]    per-step state-input contributions
 *   w_c             : [sh2, state_d]     output mix weights (per-hidden-channel)
 *   h_s2            : [N, sh2]           residual input added to scan output
 *   h_enriched      : [N, sh2]           OUTPUT of forward: h_s2 + (W_c · final_state)
 *   d_h_enriched    : [N, sh2]           upstream gradient flowing into h_enriched
 *
 * Backward scratch (per-channel, atomicAdd-free):
 *   d_a_per_channel : [N, sh2, K, state_d]  each (i, j) thread writes its
 *                                            unique slot; reduced across j by
 *                                            `mamba2_alpha_reduce_d_proj`
 *   d_b_per_channel : [N, sh2, K, state_d]  symmetric
 *   d_w_c_per_sample: [N, sh2, state_d]     reduced across N by
 *                                            `mamba2_alpha_reduce_d_w_c`
 *
 * Backward outputs (after reductions):
 *   d_a_proj        : [N, K, state_d]    sum over j of d_a_per_channel
 *   d_b_proj        : [N, K, state_d]    sum over j of d_b_per_channel
 *   d_w_c           : [sh2, state_d]     sum over N of d_w_c_per_sample
 *   d_h_s2          : [N, sh2]           identity passthrough of d_h_enriched
 *
 * Kernel constraints (must hold at every launch):
 *   state_d ≤ 32     (register array `float x[32]`)
 *   K ≤ 96           (local-memory array `float x_hist[96 * 32]` = 12 KiB / thread —
 *                     spills to per-thread DRAM-backed local memory but L2-cached;
 *                     chosen so K ≥ longest practical horizon for tick-level
 *                     supervised heads while keeping the backward replay buffer
 *                     comfortably below the L2 working-set ceiling)
 *
 * Rust constructor enforces both via Mamba2BlockConfig::validate.
 * ===================================================================== */

#include <cuda_runtime.h>

// state_d capacity: raised from 16 → 32 (per-thread register array
// `float x[32]` = 128 bytes; backward replay cache `x_hist[K*32]` up
// to 12 KiB/thread of local memory at K=96). L40S/H100 register file
// (256 KiB/SM) handles this without occupancy collapse for our block
// dim (32-128 threads/block). 32-wide state empirically validated as
// the next step up the capacity ladder from 16; further increases
// will need a different storage scheme (smem-tiled state).
#define MAMBA2_ALPHA_MAX_STATE_D 32
#define MAMBA2_ALPHA_MAX_K       96

/* ---------------------------------------------------------------------
 * Forward scan: one thread per (sample i, hidden-channel j).
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_scan_fwd(
    const float* __restrict__ a_proj,
    const float* __restrict__ b_proj,
    const float* __restrict__ w_c,
    const float* __restrict__ h_s2,
    float* __restrict__ h_enriched,
    int N,
    int K,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    float x[MAMBA2_ALPHA_MAX_STATE_D];
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) x[s] = 0.0f;

    for (int t = 0; t < K; t++) {
        long long base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float gate = 1.0f / (1.0f + expf(-a_proj[base + s]));
            x[s] = gate * x[s] + b_proj[base + s];
        }
    }

    float ctx = 0.0f;
    for (int s = 0; s < state_d; s++) {
        ctx += w_c[(long long)j * state_d + s] * x[s];
    }
    h_enriched[(long long)i * sh2 + j] = h_s2[(long long)i * sh2 + j] + ctx;
}


/* ---------------------------------------------------------------------
 * Backward scan: one thread per (sample i, hidden-channel j).
 *
 * Each thread writes to its UNIQUE slot in the per-channel scratch buffers
 * `d_a_per_channel[i, j, t, s]` and `d_b_per_channel[i, j, t, s]`. No
 * atomicAdd; cross-channel reduction is done by `mamba2_alpha_reduce_d_proj`
 * below. The `d_w_c_per_sample[i, j, s]` slot is also (i, j)-unique;
 * reduced across N by `mamba2_alpha_reduce_d_w_c`.
 *
 * `d_h_s2[i, j]` is the identity passthrough of `d_h_enriched[i, j]`.
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_scan_bwd(
    const float* __restrict__ a_proj,
    const float* __restrict__ b_proj,
    const float* __restrict__ d_h_enriched,
    const float* __restrict__ w_c,
    float* __restrict__ d_a_per_channel,   // [N, sh2, K, state_d]
    float* __restrict__ d_b_per_channel,   // [N, sh2, K, state_d]
    float* __restrict__ d_w_c_per_sample,  // [N, sh2, state_d]
    float* __restrict__ d_h_s2,            // [N, sh2]
    int N,
    int K,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    /* Replay the forward state path; cache x[t][s] for the reverse scan. */
    float x_hist[MAMBA2_ALPHA_MAX_K * MAMBA2_ALPHA_MAX_STATE_D];
    float x[MAMBA2_ALPHA_MAX_STATE_D];
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) x[s] = 0.0f;

    for (int t = 0; t < K; t++) {
        long long base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float gate = 1.0f / (1.0f + expf(-a_proj[base + s]));
            x[s] = gate * x[s] + b_proj[base + s];
            x_hist[t * MAMBA2_ALPHA_MAX_STATE_D + s] = x[s];
        }
    }

    /* Initial d_state at t = K-1: gradient flowing back through the dot
     * product `ctx[i,j] = sum_s w_c[j,s] * x[s]` is d_h_ij * w_c[j,s]. */
    float d_state[MAMBA2_ALPHA_MAX_STATE_D];
    float d_h_ij = d_h_enriched[(long long)i * sh2 + j];
    for (int s = 0; s < state_d; s++) {
        d_state[s] = d_h_ij * w_c[(long long)j * state_d + s];
    }

    /* W_c gradient contribution from this (i, j):
     *   d_w_c[j, s] += d_h_ij * x_final[s]
     * Each (i, j) thread writes its UNIQUE slot in d_w_c_per_sample[i, j, s];
     * reduction across i happens in mamba2_alpha_reduce_d_w_c. */
    long long w_c_slot = ((long long)i * sh2 + j) * state_d;
    for (int s = 0; s < state_d; s++) {
        d_w_c_per_sample[w_c_slot + s] =
            d_h_ij * x_hist[(K - 1) * MAMBA2_ALPHA_MAX_STATE_D + s];
    }

    /* Reverse scan. At each step t (going K-1 → 0):
     *   d_b[t, s]  = d_state[s]                            (b is added linearly)
     *   d_a[t, s]  = d_state[s] * x_prev[s] * σ'(a[t, s])  (a feeds the gate)
     *   d_state[s] = d_state[s] * gate[t, s]               (propagate to t-1)
     *
     * Per-channel scratch slot for (i, j, t, s):
     *   d_a_per_channel[((i * sh2 + j) * K + t) * state_d + s]
     */
    long long per_chan_base = ((long long)i * sh2 + j) * (long long)K * state_d;
    for (int t = K - 1; t >= 0; t--) {
        long long fwd_base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float a_raw = a_proj[fwd_base + s];
            float gate = 1.0f / (1.0f + expf(-a_raw));
            float sig_deriv = gate * (1.0f - gate);

            float x_prev = (t == 0) ? 0.0f : x_hist[(t - 1) * MAMBA2_ALPHA_MAX_STATE_D + s];

            long long slot = per_chan_base + (long long)t * state_d + s;
            d_b_per_channel[slot] = d_state[s];
            d_a_per_channel[slot] = d_state[s] * x_prev * sig_deriv;

            /* Propagate to t-1. */
            d_state[s] = d_state[s] * gate;
        }
    }

    /* d_h_s2 is identity passthrough (h_s2 added linearly to h_enriched). */
    d_h_s2[(long long)i * sh2 + j] = d_h_ij;
}


/* ---------------------------------------------------------------------
 * Reduction kernel: sum d_a_per_channel[N, sh2, K, state_d] across j →
 *                   d_a_proj[N, K, state_d].
 *
 * Same kernel handles d_b reduction (call twice with different
 * input/output pointers). Avoids code duplication.
 *
 * Grid: (N, K, ceil(state_d / 32)), Block: 32
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_reduce_d_proj(
    const float* __restrict__ d_per_channel,  // [N, sh2, K, state_d]
    float* __restrict__ d_proj,               // [N, K, state_d]
    int N,
    int K,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int t = blockIdx.y;
    int s = blockIdx.z * blockDim.x + threadIdx.x;
    if (i >= N || t >= K || s >= state_d) return;

    /* Sum over j: d_proj[i, t, s] = sum_j d_per_channel[i, j, t, s] */
    float sum = 0.0f;
    for (int j = 0; j < sh2; j++) {
        long long slot = (((long long)i * sh2 + j) * K + t) * state_d + s;
        sum += d_per_channel[slot];
    }
    d_proj[((long long)i * K + t) * state_d + s] = sum;
}


/* ---------------------------------------------------------------------
 * Per-step variants — supervise the SSM at EVERY timestep, not just the
 * final position. Used by ml-alpha PerceptionTrainer.
 *
 * Forward writes per-step h_enriched_seq[N, K, sh2]; backward accepts
 * d_h_enriched_seq[N, K, sh2] and accumulates gradient injections at
 * every step before propagating the recurrent d_state through the
 * gate chain. Per-step semantics:
 *
 *   h_enriched_seq[i, t, j] = h_s2[i, j] + sum_s w_c[j, s] * x[i, t, s]
 *
 *   d_h_s2[i, j]     = sum_t d_h_enriched_seq[i, t, j]      (residual added every step)
 *   d_w_c[j, s]      = sum_i sum_t d_h_enriched_seq[i, t, j] * x[i, t, s]
 *   d_state[s] at t  = (carried d_state from t+1) + sum_j d_h_enriched_seq[i, t, j] * w_c[j, s]
 *
 * Backward kernel: thread per (i, j) injects its own j-channel
 * contribution at each step; the per-channel reductions across j
 * (mamba2_alpha_reduce_d_proj / _d_w_c) work UNCHANGED — same scratch
 * shapes, same launch configs.
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_scan_fwd_seq(
    const float* __restrict__ a_proj,        // [N, K, state_d]
    const float* __restrict__ b_proj,        // [N, K, state_d]
    const float* __restrict__ w_c,           // [sh2, state_d]
    const float* __restrict__ h_s2,          // [N, sh2]
    float* __restrict__ h_enriched_seq,      // [N, K, sh2] — written at every step
    int N,
    int K,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    float x[MAMBA2_ALPHA_MAX_STATE_D];
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) x[s] = 0.0f;

    const float h_s2_ij = h_s2[(long long)i * sh2 + j];

    for (int t = 0; t < K; t++) {
        long long fwd_base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float gate = 1.0f / (1.0f + expf(-a_proj[fwd_base + s]));
            x[s] = gate * x[s] + b_proj[fwd_base + s];
        }
        float ctx = 0.0f;
        for (int s = 0; s < state_d; s++) {
            ctx += w_c[(long long)j * state_d + s] * x[s];
        }
        h_enriched_seq[(((long long)i * K) + t) * sh2 + j] = h_s2_ij + ctx;
    }
}


/* ---------------------------------------------------------------------
 * SINGLE-STEP forward scan — CRT Phase A0.5 incremental SSM.
 *
 * Same arithmetic as mamba2_alpha_scan_fwd_seq with K=1, but READS the
 * recurrent SSM register state from `x_state[N, sh2, state_d]` and WRITES
 * the post-step state back. Lets the caller advance the SSM by one
 * snapshot at a time across many launches without re-running over a
 * K-window each time (which is what scan_fwd_seq with K=1 would do —
 * it'd reset x[] to zero on every call and never advance).
 *
 * The dedicated kernel is needed because the existing scan_fwd_seq
 * unconditionally zero-initialises its register-array state at entry
 * (line 253-255 in this file). Calling it with K=1 would discard prior
 * state every launch.
 *
 *   x_state          : [N, sh2, state_d]    PERSISTENT SSM register state
 *                                           (one thread per (i, j) owns
 *                                           state_d floats); read at start,
 *                                           written back at end.
 *   a_proj           : [N, state_d]         per-step gate (single snapshot)
 *   b_proj           : [N, state_d]         per-step input (single snapshot)
 *   w_c              : [sh2, state_d]       output mix
 *   h_s2             : [N, sh2]             residual (typically zeros)
 *   h_out            : [N, sh2]             OUTPUT enriched hidden state
 *
 * Launch:  grid=(N, ceil_div(sh2, blockDim.x), 1) block=(32-128, 1, 1)
 *          — same launch shape as scan_fwd_seq with sh2 channels.
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_scan_fwd_step(
    float* __restrict__ x_state,             // [N, sh2, state_d] in/out
    const float* __restrict__ a_proj,        // [N, state_d]
    const float* __restrict__ b_proj,        // [N, state_d]
    const float* __restrict__ w_c,           // [sh2, state_d]
    const float* __restrict__ h_s2,          // [N, sh2]
    float* __restrict__ h_out,               // [N, sh2]
    int N,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    /* Load this (i, j) thread's SSM register state from DRAM. */
    float x[MAMBA2_ALPHA_MAX_STATE_D];
    long long state_base = ((long long)i * sh2 + j) * state_d;
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) x[s] = 0.0f;
    for (int s = 0; s < state_d; s++) {
        x[s] = x_state[state_base + s];
    }

    /* Advance by one step: x[s] = sigmoid(a) * x[s] + b. */
    long long ab_base = (long long)i * state_d;
    for (int s = 0; s < state_d; s++) {
        float gate = 1.0f / (1.0f + expf(-a_proj[ab_base + s]));
        x[s] = gate * x[s] + b_proj[ab_base + s];
    }

    /* Compute output contraction: h_out[i, j] = h_s2[i, j] + sum_s w_c[j, s] * x[s]. */
    float ctx = 0.0f;
    for (int s = 0; s < state_d; s++) {
        ctx += w_c[(long long)j * state_d + s] * x[s];
    }
    h_out[(long long)i * sh2 + j] = h_s2[(long long)i * sh2 + j] + ctx;

    /* Write post-step state back to DRAM. */
    for (int s = 0; s < state_d; s++) {
        x_state[state_base + s] = x[s];
    }
}


extern "C" __global__ void mamba2_alpha_scan_bwd_seq(
    const float* __restrict__ a_proj,            // [N, K, state_d]
    const float* __restrict__ b_proj,            // [N, K, state_d]
    const float* __restrict__ d_h_enriched_seq,  // [N, K, sh2]
    const float* __restrict__ w_c,               // [sh2, state_d]
    float* __restrict__ d_a_per_channel,         // [N, sh2, K, state_d]
    float* __restrict__ d_b_per_channel,         // [N, sh2, K, state_d]
    float* __restrict__ d_w_c_per_sample,        // [N, sh2, state_d]
    float* __restrict__ d_h_s2,                  // [N, sh2]
    int N,
    int K,
    int sh2,
    int state_d
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    /* Replay forward state, caching x[t][s]. */
    float x_hist[MAMBA2_ALPHA_MAX_K * MAMBA2_ALPHA_MAX_STATE_D];
    float x[MAMBA2_ALPHA_MAX_STATE_D];
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) x[s] = 0.0f;

    for (int t = 0; t < K; t++) {
        long long fwd_base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float gate = 1.0f / (1.0f + expf(-a_proj[fwd_base + s]));
            x[s] = gate * x[s] + b_proj[fwd_base + s];
            x_hist[t * MAMBA2_ALPHA_MAX_STATE_D + s] = x[s];
        }
    }

    /* Zero d_state and d_w_c_per_sample (the latter accumulates across t). */
    float d_state[MAMBA2_ALPHA_MAX_STATE_D];
    #pragma unroll
    for (int s = 0; s < MAMBA2_ALPHA_MAX_STATE_D; s++) d_state[s] = 0.0f;
    long long w_c_slot = ((long long)i * sh2 + j) * state_d;
    for (int s = 0; s < state_d; s++) d_w_c_per_sample[w_c_slot + s] = 0.0f;

    float d_h_s2_sum = 0.0f;

    /* Reverse scan: at each step t (going K-1 → 0), inject this step's
     * gradient into d_state via W_c BEFORE recording d_a/d_b and
     * propagating through the gate. */
    long long per_chan_base = ((long long)i * sh2 + j) * (long long)K * state_d;
    for (int t = K - 1; t >= 0; t--) {
        long long fwd_base = ((long long)i * K + t) * state_d;
        float d_h_ij_t = d_h_enriched_seq[(((long long)i * K) + t) * sh2 + j];
        d_h_s2_sum += d_h_ij_t;

        /* W_c grad contribution at step t: d_w_c[j,s] += d_h_ij_t * x[t,s]. */
        for (int s = 0; s < state_d; s++) {
            d_w_c_per_sample[w_c_slot + s] +=
                d_h_ij_t * x_hist[t * MAMBA2_ALPHA_MAX_STATE_D + s];
        }
        /* Inject step-t gradient into d_state via W_c. */
        for (int s = 0; s < state_d; s++) {
            d_state[s] += d_h_ij_t * w_c[(long long)j * state_d + s];
        }
        /* Record d_b[t,s], d_a[t,s] from the CURRENT (post-injection) d_state. */
        for (int s = 0; s < state_d; s++) {
            float a_raw = a_proj[fwd_base + s];
            float gate = 1.0f / (1.0f + expf(-a_raw));
            float sig_deriv = gate * (1.0f - gate);
            float x_prev = (t == 0) ? 0.0f : x_hist[(t - 1) * MAMBA2_ALPHA_MAX_STATE_D + s];

            long long slot = per_chan_base + (long long)t * state_d + s;
            d_b_per_channel[slot] = d_state[s];
            d_a_per_channel[slot] = d_state[s] * x_prev * sig_deriv;

            d_state[s] = d_state[s] * gate;
        }
    }
    d_h_s2[(long long)i * sh2 + j] = d_h_s2_sum;
}


/* ---------------------------------------------------------------------
 * Phase 1d.4 backtest — per-trade PnL kernel.
 *
 * For each (threshold τ, sequence n), apply the trading rule:
 *   pred_dev = probs[n] - 0.5
 *   if |pred_dev| > τ:
 *     direction = sign(pred_dev) in {-1, +1}
 *     pnl       = direction * (price_kt[n] - price_t[n]) - cost
 *     trade     = 1
 *   else:
 *     pnl       = 0
 *     trade     = 0
 *
 * Output is row-major `[T, N]` for both pnl and trade. Allows a single
 * kernel launch to evaluate the full threshold sweep in one pass.
 *
 * Grid:  (T, ceil(N / 256)), Block: 256
 * --------------------------------------------------------------------- */
extern "C" __global__ void backtest_per_trade_pnl(
    const float* __restrict__ probs,        /* [N] stacker sigmoid output ∈ [0, 1] */
    const float* __restrict__ prices_t,     /* [N] mid-price at sequence end-bar */
    const float* __restrict__ prices_kt,    /* [N] mid-price at end-bar + horizon */
    const float* __restrict__ thresholds,   /* [T] confidence thresholds */
    float cost,                              /* round-trip cost in price units */
    float* __restrict__ pnl_out,            /* [T, N] per-trade PnL (0 if no trade) */
    int*   __restrict__ trade_out,          /* [T, N] 1 if trade, 0 otherwise */
    int N,
    int T
) {
    int t_idx = blockIdx.x;
    int n_idx = blockIdx.y * blockDim.x + threadIdx.x;
    if (t_idx >= T || n_idx >= N) return;

    float prob   = probs[n_idx];
    float thresh = thresholds[t_idx];
    float dev    = prob - 0.5f;
    float abs_d  = fabsf(dev);

    long long slot = (long long)t_idx * N + n_idx;
    if (abs_d > thresh) {
        float dir = (dev > 0.0f) ? 1.0f : -1.0f;
        pnl_out[slot]   = dir * (prices_kt[n_idx] - prices_t[n_idx]) - cost;
        trade_out[slot] = 1;
    } else {
        pnl_out[slot]   = 0.0f;
        trade_out[slot] = 0;
    }
}


/* ---------------------------------------------------------------------
 * Phase 1d.4 backtest — block tree-reduce one row of [T, N] data.
 *
 * Per-threshold reduction: sum across N to produce a [T]-shaped output.
 * Avoids atomicAdd via the standard block-tree shared-memory reduction.
 *
 * Grid:  (T, 1), Block: 256
 * Shared: 256 floats
 * --------------------------------------------------------------------- */
extern "C" __global__ void backtest_sum_reduce_f32(
    const float* __restrict__ input,   /* [T, N] */
    float*       __restrict__ output,  /* [T] */
    int N,
    int T
) {
    __shared__ float sdata[256];
    int t_idx = blockIdx.x;
    if (t_idx >= T) return;

    int tid = threadIdx.x;
    long long base = (long long)t_idx * N;

    /* Each thread accumulates its slice of the N axis. */
    float local = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        local += input[base + i];
    }
    sdata[tid] = local;
    __syncthreads();

    /* Block tree-reduce in shared memory. */
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    if (tid == 0) output[t_idx] = sdata[0];
}


/* ---------------------------------------------------------------------
 * Phase 1d.4 backtest — block tree-reduce SQUARED values along N.
 *
 * Computes `sum_i(input[t, i]^2)` per threshold t. Used for variance
 * computation: var = sum_sq / count - mean^2 (Welford-equivalent for
 * a single pass; n is large so f32 numerical error is acceptable).
 *
 * Grid:  (T, 1), Block: 256
 * --------------------------------------------------------------------- */
extern "C" __global__ void backtest_sum_squared_reduce_f32(
    const float* __restrict__ input,   /* [T, N] */
    float*       __restrict__ output,  /* [T] */
    int N,
    int T
) {
    __shared__ float sdata[256];
    int t_idx = blockIdx.x;
    if (t_idx >= T) return;

    int tid = threadIdx.x;
    long long base = (long long)t_idx * N;

    float local = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float v = input[base + i];
        local += v * v;
    }
    sdata[tid] = local;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    if (tid == 0) output[t_idx] = sdata[0];
}


/* ---------------------------------------------------------------------
 * Phase 1d.4 backtest — sum reduce for INT data (trade counts).
 *
 * Same algorithm as float reduce but operates on int32 input. Counts
 * the number of trades per threshold.
 *
 * Grid:  (T, 1), Block: 256
 * --------------------------------------------------------------------- */
extern "C" __global__ void backtest_sum_reduce_i32(
    const int* __restrict__ input,   /* [T, N] */
    int*       __restrict__ output,  /* [T] */
    int N,
    int T
) {
    __shared__ int sdata[256];
    int t_idx = blockIdx.x;
    if (t_idx >= T) return;

    int tid = threadIdx.x;
    long long base = (long long)t_idx * N;

    int local = 0;
    for (int i = tid; i < N; i += blockDim.x) {
        local += input[base + i];
    }
    sdata[tid] = local;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    if (tid == 0) output[t_idx] = sdata[0];
}


/* ---------------------------------------------------------------------
 * AdamW step — bias-corrected decoupled-weight-decay update for one
 * parameter tensor. Identical to ml-core's adamw_update kernel, included
 * here so ml-alpha's cubin is self-contained (no cross-crate cubin loads).
 *
 *   grad_scale: 1.0 disables gradient clipping; <1.0 applied as a scale.
 *   t:          1-indexed training step (for bias correction).
 *
 * Grid:  ceil(n / 256), Block: 256.
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_adamw_step(
    float* __restrict__ param,
    const float* __restrict__ grad,
    float* __restrict__ m,
    float* __restrict__ v,
    float lr,
    float beta1,
    float beta2,
    float epsilon,
    float weight_decay,
    float grad_scale,
    int t,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float g = grad[i] * grad_scale;
    float p = param[i];

    /* Decoupled weight decay. */
    p = p * (1.0f - lr * weight_decay);

    /* Moment updates. */
    float mi = beta1 * m[i] + (1.0f - beta1) * g;
    float vi = beta2 * v[i] + (1.0f - beta2) * g * g;
    m[i] = mi;
    v[i] = vi;

    /* Bias correction. */
    float bc1 = 1.0f - powf(beta1, (float)t);
    float bc2 = 1.0f - powf(beta2, (float)t);
    float m_hat = mi / bc1;
    float v_hat = vi / bc2;

    /* Parameter update. */
    p = p - lr * m_hat / (sqrtf(v_hat) + epsilon);
    param[i] = p;
}


/// Variant of `mamba2_alpha_adamw_step` that reads BOTH `grad_scale` and
/// `t` (step counter) from DEVICE pointers — companion to the
/// GPU-resident grad-clip pipeline in `grad_norm.cu` AND the CUDA Graph
/// capture path. Eliminates the host scalars that would otherwise be
/// baked into the captured graph (replays would freeze grad_scale to
/// its first-step value and step counter would stop advancing → wrong
/// bias correction). Use `mamba2_alpha_increment_step_counter` once per
/// training step to advance the counter inside the captured region.
extern "C" __global__ void mamba2_alpha_adamw_step_devscale(
    float* __restrict__ param,
    const float* __restrict__ grad,
    float* __restrict__ m,
    float* __restrict__ v,
    float lr,
    float beta1,
    float beta2,
    float epsilon,
    float weight_decay,
    const float* __restrict__ grad_scale_ptr,
    const int* __restrict__ step_ptr,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float g = grad[i] * grad_scale_ptr[0];
    float p = param[i];

    /* Decoupled weight decay. */
    p = p * (1.0f - lr * weight_decay);

    /* Moment updates. */
    float mi = beta1 * m[i] + (1.0f - beta1) * g;
    float vi = beta2 * v[i] + (1.0f - beta2) * g * g;
    m[i] = mi;
    v[i] = vi;

    /* Bias correction. */
    int t = step_ptr[0];
    float bc1 = 1.0f - powf(beta1, (float)t);
    float bc2 = 1.0f - powf(beta2, (float)t);
    float m_hat = mi / bc1;
    float v_hat = vi / bc2;

    /* Parameter update. */
    p = p - lr * m_hat / (sqrtf(v_hat) + epsilon);
    param[i] = p;
}

/// Increment the device-resident step counter for Mamba2 AdamW.
/// Single-thread kernel; launch once per training step inside the
/// captured graph, AFTER all `mamba2_alpha_adamw_step_devscale` launches
/// that read the current step value.
extern "C" __global__ void mamba2_alpha_increment_step_counter(int* __restrict__ step_ptr) {
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        step_ptr[0] += 1;
    }
}


/* ---------------------------------------------------------------------
 * Reduction kernel: sum d_w_c_per_sample[N, sh2, state_d] across N →
 *                   d_w_c[sh2, state_d].
 *
 * Grid: (sh2, ceil(state_d / 32)), Block: 32
 * --------------------------------------------------------------------- */
extern "C" __global__ void mamba2_alpha_reduce_d_w_c(
    const float* __restrict__ d_w_c_per_sample,  // [N, sh2, state_d]
    float* __restrict__ d_w_c,                   // [sh2, state_d]
    int N,
    int sh2,
    int state_d
) {
    int j = blockIdx.x;
    int s = blockIdx.y * blockDim.x + threadIdx.x;
    if (j >= sh2 || s >= state_d) return;

    /* Sum over i: d_w_c[j, s] = sum_i d_w_c_per_sample[i, j, s] */
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += d_w_c_per_sample[((long long)i * sh2 + j) * state_d + s];
    }
    d_w_c[(long long)j * state_d + s] = sum;
}