From efdf355e8a357464e70d71bd8730875d46a840a5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 19:33:33 +0200 Subject: [PATCH] feat: add regime_branch_gate, adaptive_atom_positions, mamba2_temporal_scan CUDA kernels Three new GPU kernels for OOS performance enhancement: - regime_branch_gate: learned per-branch importance weights from regime features - adaptive_atom_positions: softmax-normalized learnable C51 atom spacing - mamba2_temporal_scan + mamba2_update_history: selective SSM for temporal context Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/build.rs | 1 + .../src/cuda_pipeline/experience_kernels.cu | 152 ++++++++++++++++++ .../cuda_pipeline/mamba2_temporal_kernel.cu | 88 ++++++++++ 3 files changed, 241 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index ef49a407a..9b5a11afb 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -80,6 +80,7 @@ fn main() { "trade_stats_kernel.cu", "backward_kernels.cu", "iqn_cvar_kernel.cu", + "mamba2_temporal_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 2847fed33..d0c8c737c 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -3984,3 +3984,155 @@ extern "C" __global__ void qlstm_train_step( } } } + +/* ================================================================== */ +/* Kernel: regime_branch_gate */ +/* ================================================================== */ + +/** + * Regime-conditioned branch gating: learned per-branch importance weights. + * + * regime_input = [ADX, CUSUM, Q_gap, atom_utilization] (4 features) + * branch_importance[4] = softmax(W_regime @ regime_input + b_regime) + * Q_gated[b, a] = branch_importance[branch_of(a)] * Q_in[b, a] + * + * Grid: ceil(N / 256), Block: 256. One thread per sample. + */ +extern "C" __global__ void regime_branch_gate( + float* __restrict__ q_values, + const float* __restrict__ states, + const float* __restrict__ q_gap, + const float* __restrict__ atom_util, + const float* __restrict__ w_regime, + const float* __restrict__ b_regime, + int N, + int state_dim, + int b0_size, + int b1_size, + int b2_size, + int b3_size +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + int total_actions = b0_size + b1_size + b2_size + b3_size; + + /* Extract regime features from state vector */ + float adx = states[(long long)i * state_dim + 40]; + float cusum = states[(long long)i * state_dim + 41]; + float qgap_i = q_gap[i]; + float util_i = atom_util[0]; + + float regime_input[4] = { adx, cusum, qgap_i, util_i }; + + /* Compute W_regime @ regime_input + b_regime */ + float logits[4]; + for (int d = 0; d < 4; d++) { + float val = b_regime[d]; + for (int k = 0; k < 4; k++) { + val += w_regime[d * 4 + k] * regime_input[k]; + } + logits[d] = val; + } + + /* Softmax over 4 branches */ + float max_logit = logits[0]; + for (int d = 1; d < 4; d++) + max_logit = fmaxf(max_logit, logits[d]); + + float sum_exp = 0.0f; + float importance[4]; + for (int d = 0; d < 4; d++) { + importance[d] = expf(logits[d] - max_logit); + sum_exp += importance[d]; + } + for (int d = 0; d < 4; d++) + importance[d] /= sum_exp; + + /* Scale Q-values by branch importance */ + float* q_row = q_values + (long long)i * total_actions; + int offset = 0; + int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size }; + for (int d = 0; d < 4; d++) { + for (int a = 0; a < branch_sizes[d]; a++) { + q_row[offset + a] *= importance[d]; + } + offset += branch_sizes[d]; + } +} + +/* ================================================================== */ +/* Kernel: adaptive_atom_positions */ +/* ================================================================== */ + +/** + * Compute adaptive C51 atom positions from learnable spacing parameters. + * + * spacing = softmax(spacing_raw) + * positions[j] = v_min + (v_max - v_min) * cumsum(spacing)[j] + * + * Single-block kernel: num_atoms=51 fits in one block. + * Grid: (1,1,1), Block: (256,1,1). Shared memory: 2 * num_atoms * sizeof(float). + */ +extern "C" __global__ void adaptive_atom_positions( + const float* __restrict__ spacing_raw, + float* __restrict__ positions_out, + int num_atoms, + float v_min, + float v_max +) { + int tid = threadIdx.x; + extern __shared__ float shmem[]; + float* s_softmax = shmem; + float* s_cumsum = shmem + num_atoms; + + /* Load spacing_raw into shared memory */ + if (tid < num_atoms) { + s_softmax[tid] = spacing_raw[tid]; + } + __syncthreads(); + + /* Find max for numerical stability (single thread, num_atoms=51 is tiny) */ + float max_val; + if (tid == 0) { + max_val = s_softmax[0]; + for (int j = 1; j < num_atoms; j++) + max_val = fmaxf(max_val, s_softmax[j]); + s_cumsum[0] = max_val; /* store max temporarily */ + } + __syncthreads(); + max_val = s_cumsum[0]; + __syncthreads(); + + /* Compute softmax */ + if (tid < num_atoms) { + s_softmax[tid] = expf(s_softmax[tid] - max_val); + } + __syncthreads(); + + /* Sum for normalization */ + if (tid == 0) { + float sum = 0.0f; + for (int j = 0; j < num_atoms; j++) + sum += s_softmax[j]; + float inv_sum = 1.0f / fmaxf(sum, 1e-8f); + for (int j = 0; j < num_atoms; j++) + s_softmax[j] *= inv_sum; + } + __syncthreads(); + + /* Inclusive prefix sum (sequential, num_atoms=51) */ + if (tid == 0) { + s_cumsum[0] = s_softmax[0]; + for (int j = 1; j < num_atoms; j++) + s_cumsum[j] = s_cumsum[j-1] + s_softmax[j]; + } + __syncthreads(); + + /* Write positions: positions[0] = v_min, positions[j>0] = v_min + range * cumsum[j-1] */ + if (tid < num_atoms) { + float range = v_max - v_min; + float frac = (tid == 0) ? 0.0f : s_cumsum[tid - 1]; + positions_out[tid] = v_min + range * frac; + } +} diff --git a/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu new file mode 100644 index 000000000..f95bebb44 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu @@ -0,0 +1,88 @@ +/** + * Mamba2 selective scan for temporal context from a rolling history buffer. + * + * Implements the selective state space mechanism: + * A_t = sigmoid(W_A @ h_t) -- forget gate (input-dependent) + * B_t = W_B @ h_t -- input projection + * x_t = A_t * x_{t-1} + B_t * h_t -- recurrent scan + * temporal_context = (W_C @ h_t) * x_t + * h_enriched = h_s2 + temporal_context -- residual addition + * + * Grid: ceil(N / 256), Block: 256. One thread per sample. + */ +extern "C" __global__ void mamba2_temporal_scan( + const float* __restrict__ h_history, + const float* __restrict__ h_s2, + float* __restrict__ h_enriched, + const float* __restrict__ w_a, + const float* __restrict__ w_b, + const float* __restrict__ w_c, + int N, + int K, + int sh2, + int state_d +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + const float* h_current = h_s2 + (long long)i * sh2; + + /* Initialize SSM state x[state_d] to zero */ + float x[16]; /* state_d <= 16 */ + for (int s = 0; s < state_d; s++) + x[s] = 0.0f; + + /* Sequential scan over K history steps */ + for (int t = 0; t < K; t++) { + const float* h_t = h_history + ((long long)i * K + t) * sh2; + + for (int s = 0; s < state_d; s++) { + float a_val = 0.0f; + float b_val = 0.0f; + for (int j = 0; j < sh2; j++) { + a_val += w_a[(long long)j * state_d + s] * h_t[j]; + b_val += w_b[(long long)j * state_d + s] * h_t[j]; + } + float gate = 1.0f / (1.0f + expf(-a_val)); + + x[s] = gate * x[s] + b_val; + } + } + + /* Compute temporal context and residual addition */ + float* out = h_enriched + (long long)i * sh2; + for (int j = 0; j < sh2; j++) { + float ctx = 0.0f; + for (int s = 0; s < state_d; s++) { + ctx += w_c[(long long)j * state_d + s] * x[s]; + } + out[j] = h_current[j] + ctx; + } +} + +/** + * Update rolling history buffer: shift left by 1, insert current h_s2 at position K-1. + * + * Grid: ceil(N * SH2 / 256), Block: 256. + */ +extern "C" __global__ void mamba2_update_history( + float* __restrict__ h_history, + const float* __restrict__ h_s2, + int N, + int K, + int sh2 +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = N * sh2; + if (idx >= total) return; + + int i = idx / sh2; + int j = idx % sh2; + + float* base = h_history + (long long)i * K * sh2; + for (int t = 0; t < K - 1; t++) { + base[(long long)t * sh2 + j] = base[(long long)(t + 1) * sh2 + j]; + } + + base[(long long)(K - 1) * sh2 + j] = h_s2[(long long)i * sh2 + j]; +}