ab94ce2a49
Stage 1+2 of #162 (CUDA Graph capture of training step). The AdamW kernels previously took the step counter as a host scalar arg, which gets baked into kernel args at CUDA Graph capture time — replays would freeze the counter and produce wrong bias-correction values. Both AdamW variants now read the step from a device pointer, advanced by a tiny 1-thread `increment_counter` kernel that goes inside the captured region. Each replay correctly increments and observes the new step value. Kernel changes: adamw_step.cu: - adamw_step: int step → const int* step_ptr - adamw_increment_counter: new, +=1 on step_ptr[0] mamba2_alpha_kernel.cu: - mamba2_alpha_adamw_step_devscale: int t → const int* step_ptr - mamba2_alpha_increment_step_counter: new Rust changes: trainer/optim.rs (AdamW): - host `step: i32` → device `step_count_d: CudaSlice<i32>` - step(): launch increment kernel BEFORE adamw kernel; both read device counter via pointer arg. - step_count(): test-only accessor, mapped-pinned readback (sync). mamba2_block.rs (Mamba2AdamW): - kept host `step_count: i32` for legacy paths (`step`, `step_from_buffers`) which aren't capture-compatible anyway (host grad-norm dtoh, host scalar grad_scale). - added device `step_count_d: CudaSlice<i32>` for the production gpu_clip path; advances via `kernel_increment_step` kernel inside the captured region. - adamw_apply_devscale: `t: i32` → `step_d: &CudaSlice<i32>`. Validation: - 4 adamw_invariants tests pass (step_count_increments specifically exercises the device counter). - 10 mamba2_block lib tests pass (training_loop_decreases_loss exercises legacy host-counter path). - Synthetic overfit smoke: initial=0.25 → final=0.0006 (matches pre-refactor trajectory bit-for-bit-equivalent). Stage 3+4 (capture brackets + first-call-capture / subsequent-replay state machine in step_batched) follows in the next commit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
646 lines
24 KiB
Plaintext
646 lines
24 KiB
Plaintext
/* =====================================================================
|
||
* 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 ≤ 96 (local-memory array `float x_hist[96 * 16]` = 6 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>
|
||
|
||
#define MAMBA2_ALPHA_MAX_STATE_D 16
|
||
#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;
|
||
}
|
||
}
|
||
|
||
|
||
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;
|
||
}
|