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 ≤ 16     (register array `float x[16]`)
 *   K ≤ 32           (local-memory array `float x_hist[32 * 16]` ≤ 2 KiB / thread)
 *
 * Rust constructor enforces both via Mamba2BlockConfig::validate.
 * ===================================================================== */

#include <cuda_runtime.h>

#define MAMBA2_ALPHA_MAX_STATE_D 16
#define MAMBA2_ALPHA_MAX_K       32

/* ---------------------------------------------------------------------
 * 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;
}


/* ---------------------------------------------------------------------
 * 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;
}


/* ---------------------------------------------------------------------
 * 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;
}