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) <noreply@anthropic.com>

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)

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;
+    }
+}

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];
+}