feat(cuda): GPU-resident PER kernels — push, sample, update, tree_rebuild

4 kernels for all-device prioritized experience replay:
- rl_per_push: n-step accumulation + single-block prefix-sum for
  write_head coordination (no atomicAdd)
- rl_per_sample: stratified proportional sampling via top-down
  sum-tree walk with xorshift32 PRNG + inline gather
- rl_per_update_priority: write |TD|^α to leaves + shared-mem
  block-wide max reduction
- rl_per_tree_rebuild: bottom-up parallel scan with __threadfence
  between levels (15 passes for capacity=32768, no atomics)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/build.rs

@@ -101,6 +101,10 @@ const KERNELS: &[&str] = &[
     "rl_sample_noise",               // CUDA graph prereq: device-side factored noise f(rand) for NoisyLinear; replaces host ChaCha8 + mapped-pinned upload
     "rl_write_u64",                  // CUDA graph prereq: single-thread u64 scalar write for device-resident ts_ns (graph-captured kernels read via pointer)
     "rl_fused_reward_pipeline",      // P3: 7→1 fused per-batch reward extraction + shaping + EMA + outcome update
+    "rl_per_push",                   // GPU PER: n-step accumulation + circular write with prefix-sum coordination
+    "rl_per_sample",                 // GPU PER: stratified proportional sampling via sum-tree walk + gather
+    "rl_per_update_priority",        // GPU PER: write |TD|^α to tree leaves + block-wide max reduction
+    "rl_per_tree_rebuild",           // GPU PER: bottom-up parallel sum-tree rebuild (no atomics)
     "rl_increment_step",             // device-resident step counter bump (ISV[548] += 1.0); graph-safe prereq — removes scalar current_step from all downstream kernel args
     "rl_fused_controllers",          // fused kernel: 10 RL ISV controllers in one launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda) — saves 9 kernel launch overheads (~40-80μs/step)
 ];

crates/ml-alpha/cuda/rl_per_push.cu

@@ -0,0 +1,181 @@
+/* =====================================================================
+ * rl_per_push.cu — GPU-resident PER: push transitions into replay buffer
+ *
+ * Grid=(1,1,1), Block=(b_size,1,1).
+ *
+ * Each thread handles one batch element's n-step accumulation. When a
+ * thread needs to flush (n_step reached OR done), it signals via shared
+ * memory. Thread 0 prefix-sums the flush-flags to assign contiguous
+ * write_head slots. No atomicAdd — coordination is via shared memory
+ * and __syncthreads.
+ *
+ * ISV reads:
+ *   isv[RL_GAMMA_INDEX=400]   — discount factor γ
+ *   isv[RL_N_STEP_INDEX=403]  — n-step horizon (clamped 1..N_STEP_MAX)
+ *   isv[RL_PER_ALPHA_INDEX=405] — (unused here, used by update_priority)
+ * ===================================================================== */
+
+#define HIDDEN_DIM             128
+#define N_STEP_MAX             16
+#define SCALARS_PER_TRANSITION 7
+#define RL_GAMMA_INDEX         400
+#define RL_N_STEP_INDEX        403
+
+extern "C" __global__ void rl_per_push(
+    /* Current step data */
+    const float* __restrict__ h_t_current,      /* [B, 128] */
+    const float* __restrict__ h_tp1_current,    /* [B, 128] */
+    const float* __restrict__ rewards_scaled,   /* [B] */
+    const float* __restrict__ rewards_raw,      /* [B] */
+    const float* __restrict__ dones,            /* [B] */
+    const float* __restrict__ log_pi_old,       /* [B] */
+    const int*   __restrict__ actions,          /* [B] */
+    const float* __restrict__ isv,
+    /* Replay storage */
+    float* __restrict__ replay_h_t,             /* [capacity, 128] */
+    float* __restrict__ replay_h_tp1,           /* [capacity, 128] */
+    float* __restrict__ replay_scalars,         /* [capacity, 7] */
+    float* __restrict__ priority_tree,          /* [2 * capacity] leaves at [cap..2*cap) */
+    unsigned int* __restrict__ write_head,      /* [1] */
+    unsigned int* __restrict__ replay_len,      /* [1] */
+    float* __restrict__ max_priority,           /* [1] */
+    /* N-step ring per batch element */
+    float* __restrict__ nstep_scalars,          /* [B, N_STEP_MAX, 3] (scaled_r, raw_r, done) */
+    float* __restrict__ nstep_h_t,             /* [B, N_STEP_MAX, 128] */
+    unsigned int* __restrict__ nstep_write_idx, /* [B] */
+    unsigned int* __restrict__ nstep_count,     /* [B] */
+    int b_size,
+    int capacity
+)
+{
+    /* Shared memory for flush coordination */
+    extern __shared__ int smem[];
+    /* Layout: [0..b_size) = flush_flag, [b_size..2*b_size) = write_offset */
+    int* flush_flag  = smem;
+    int* write_offset = smem + b_size;
+
+    const int b = threadIdx.x;
+    if (b >= b_size) return;
+
+    /* ── Read ISV ─────────────────────────────────────────────────────── */
+    int n_step = (int)isv[RL_N_STEP_INDEX];
+    if (n_step < 1)          n_step = 1;
+    if (n_step > N_STEP_MAX) n_step = N_STEP_MAX;
+    const float gamma = isv[RL_GAMMA_INDEX];
+
+    /* ── Write current transition into n-step ring ────────────────────── */
+    const unsigned int ring_idx = nstep_write_idx[b];
+    const int ring_base_scalars = (b * N_STEP_MAX + ring_idx) * 3;
+    nstep_scalars[ring_base_scalars + 0] = rewards_scaled[b];
+    nstep_scalars[ring_base_scalars + 1] = rewards_raw[b];
+    nstep_scalars[ring_base_scalars + 2] = dones[b];
+
+    const int ring_base_h = (b * N_STEP_MAX + ring_idx) * HIDDEN_DIM;
+    for (int i = 0; i < HIDDEN_DIM; ++i) {
+        nstep_h_t[ring_base_h + i] = h_t_current[b * HIDDEN_DIM + i];
+    }
+
+    nstep_write_idx[b] = (ring_idx + 1) % N_STEP_MAX;
+    unsigned int count = nstep_count[b] + 1;
+    nstep_count[b] = count;
+
+    /* ── Decide whether to flush ──────────────────────────────────────── */
+    const int should_flush = ((int)count >= n_step) || (dones[b] > 0.5f);
+    flush_flag[b] = should_flush ? 1 : 0;
+    __syncthreads();
+
+    /* ── Thread 0: prefix-sum flush flags → assign write slots ────────── */
+    if (b == 0) {
+        int total = 0;
+        for (int i = 0; i < b_size; ++i) {
+            write_offset[i] = total;
+            total += flush_flag[i];
+        }
+        /* Advance write_head by total flushes */
+        if (total > 0) {
+            unsigned int base = write_head[0];
+            /* Store the base for all threads to read */
+            write_offset[b_size - 1 + 1] = 0; /* reuse trick: store base in flush_flag[b_size] area */
+            /* Actually we store base differently — just use write_offset encoding */
+            /* We'll encode base_write_head into smem[2*b_size] */
+            smem[2 * b_size] = (int)base;
+
+            /* Update write_head */
+            write_head[0] = (base + (unsigned int)total) % (unsigned int)capacity;
+
+            /* Update replay_len */
+            unsigned int cur_len = replay_len[0];
+            unsigned int new_len = cur_len + (unsigned int)total;
+            if (new_len > (unsigned int)capacity) new_len = (unsigned int)capacity;
+            replay_len[0] = new_len;
+        }
+    }
+    __syncthreads();
+
+    /* ── Each flushing thread writes to its assigned slot ─────────────── */
+    if (!should_flush) return;
+
+    const unsigned int base_write = (unsigned int)smem[2 * b_size];
+    const unsigned int my_slot = (base_write + (unsigned int)write_offset[b]) % (unsigned int)capacity;
+
+    /* Compute n-step return (iterate ring from oldest) */
+    const int actual_n = ((int)count < n_step) ? (int)count : n_step;
+    const unsigned int oldest_ring = (nstep_write_idx[b] + N_STEP_MAX - (unsigned int)actual_n) % N_STEP_MAX;
+
+    float n_step_return_scaled = 0.0f;
+    float n_step_return_raw = 0.0f;
+    float gamma_power = 1.0f;
+    float n_step_gamma = 1.0f;
+    int any_done = 0;
+
+    for (int k = 0; k < actual_n; ++k) {
+        const unsigned int ring_k = (oldest_ring + (unsigned int)k) % N_STEP_MAX;
+        const int sc_base = (b * N_STEP_MAX + ring_k) * 3;
+        const float r_s = nstep_scalars[sc_base + 0];
+        const float r_r = nstep_scalars[sc_base + 1];
+        const float d   = nstep_scalars[sc_base + 2];
+
+        n_step_return_scaled += gamma_power * r_s;
+        n_step_return_raw    += gamma_power * r_r;
+        gamma_power *= gamma;
+
+        if (d > 0.5f) {
+            any_done = 1;
+            /* Truncate: γⁿ = 0 for bootstrap when done encountered */
+            n_step_gamma = 0.0f;
+            break;
+        }
+    }
+    if (!any_done) {
+        n_step_gamma = gamma_power; /* γ^actual_n */
+    }
+
+    /* Copy oldest h_t from ring → replay_h_t[my_slot] */
+    const int oldest_h_base = (b * N_STEP_MAX + oldest_ring) * HIDDEN_DIM;
+    for (int i = 0; i < HIDDEN_DIM; ++i) {
+        replay_h_t[my_slot * HIDDEN_DIM + i] = nstep_h_t[oldest_h_base + i];
+    }
+
+    /* Copy h_tp1_current[b] → replay_h_tp1[my_slot] */
+    for (int i = 0; i < HIDDEN_DIM; ++i) {
+        replay_h_tp1[my_slot * HIDDEN_DIM + i] = h_tp1_current[b * HIDDEN_DIM + i];
+    }
+
+    /* Write scalars: [action_f32, scaled_reward, raw_reward, done, log_pi_old, n_step_gamma, n_step_return_raw] */
+    const int sc_out = my_slot * SCALARS_PER_TRANSITION;
+    replay_scalars[sc_out + 0] = (float)actions[b];
+    replay_scalars[sc_out + 1] = n_step_return_scaled;
+    replay_scalars[sc_out + 2] = n_step_return_raw;
+    replay_scalars[sc_out + 3] = (float)any_done;
+    replay_scalars[sc_out + 4] = log_pi_old[b];
+    replay_scalars[sc_out + 5] = n_step_gamma;
+    replay_scalars[sc_out + 6] = n_step_return_raw;
+
+    /* Set priority leaf = max_priority (or 1.0 if max==0) */
+    float max_p = max_priority[0];
+    if (max_p < 1e-9f) max_p = 1.0f;
+    priority_tree[capacity + (int)my_slot] = max_p;
+
+    /* Reset n-step count for this batch element */
+    nstep_count[b] = 0;
+}

crates/ml-alpha/cuda/rl_per_sample.cu

@@ -0,0 +1,118 @@
+/* =====================================================================
+ * rl_per_sample.cu — GPU-resident PER: proportional sampling via sum-tree
+ *
+ * Grid=(b_size), Block=(1). One thread per sample.
+ *
+ * Each thread samples a leaf from the priority sum-tree using stratified
+ * sampling (segment per thread) with xorshift32 PRNG, then gathers the
+ * transition data from replay storage.
+ *
+ * ISV reads: none (alpha used only at priority-update time)
+ * ===================================================================== */
+
+#define HIDDEN_DIM             128
+#define SCALARS_PER_TRANSITION 7
+
+/* xorshift32 PRNG — same pattern used throughout the codebase */
+__device__ __forceinline__ unsigned int xorshift32(unsigned int* state) {
+    unsigned int x = *state;
+    x ^= x << 13;
+    x ^= x >> 17;
+    x ^= x << 5;
+    *state = x;
+    return x;
+}
+
+extern "C" __global__ void rl_per_sample(
+    const float* __restrict__ priority_tree,     /* [2 * capacity] */
+    const float* __restrict__ replay_h_t,        /* [capacity, 128] */
+    const float* __restrict__ replay_h_tp1,      /* [capacity, 128] */
+    const float* __restrict__ replay_scalars,    /* [capacity, 7] */
+    const unsigned int* __restrict__ replay_len, /* [1] */
+    const float* __restrict__ isv,
+    unsigned int* __restrict__ prng_state,       /* [B] */
+    /* Outputs */
+    float* __restrict__ sampled_h_t,             /* [B, 128] */
+    float* __restrict__ sampled_h_tp1,           /* [B, 128] */
+    float* __restrict__ sampled_rewards,         /* [B] */
+    float* __restrict__ sampled_dones,           /* [B] */
+    float* __restrict__ sampled_log_pi_old,      /* [B] */
+    float* __restrict__ sampled_n_step_gammas,   /* [B] */
+    int*   __restrict__ sampled_actions,         /* [B] */
+    unsigned int* __restrict__ sample_indices,   /* [B] */
+    int b_size,
+    int capacity
+)
+{
+    const int b = blockIdx.x * blockDim.x + threadIdx.x;
+    if (b >= b_size) return;
+
+    const unsigned int len = replay_len[0];
+    const float total_priority = priority_tree[1]; /* root */
+
+    /* Guard: empty or zero-priority replay */
+    if (total_priority < 1e-9f || len == 0) {
+        for (int i = 0; i < HIDDEN_DIM; ++i) {
+            sampled_h_t[b * HIDDEN_DIM + i] = 0.0f;
+            sampled_h_tp1[b * HIDDEN_DIM + i] = 0.0f;
+        }
+        sampled_rewards[b] = 0.0f;
+        sampled_dones[b] = 0.0f;
+        sampled_log_pi_old[b] = 0.0f;
+        sampled_n_step_gammas[b] = 0.0f;
+        sampled_actions[b] = 0;
+        sample_indices[b] = 0;
+        return;
+    }
+
+    /* ── Seed PRNG if cold ────────────────────────────────────────────── */
+    if (prng_state[b] == 0) {
+        prng_state[b] = (unsigned int)(b + 1) * 2654435761u;
+    }
+
+    /* ── Stratified sampling: draw u in [b*segment, (b+1)*segment) ────── */
+    const float segment = total_priority / (float)b_size;
+    unsigned int rng = xorshift32(&prng_state[b]);
+    const float u_frac = (float)(rng & 0x00FFFFFFu) / (float)0x01000000u; /* [0, 1) */
+    float u = ((float)b + u_frac) * segment;
+
+    /* Clamp to avoid floating-point overshoot */
+    if (u >= total_priority) u = total_priority - 1e-6f;
+    if (u < 0.0f) u = 0.0f;
+
+    /* ── Walk tree top-down ───────────────────────────────────────────── */
+    int idx = 1;
+    while (idx < capacity) {
+        const int left = 2 * idx;
+        const float left_val = priority_tree[left];
+        if (u <= left_val) {
+            idx = left;
+        } else {
+            u -= left_val;
+            idx = left + 1;
+        }
+    }
+    const int leaf = idx - capacity;
+
+    /* Clamp leaf to valid range */
+    const int safe_leaf = (leaf >= 0 && leaf < (int)len) ? leaf : 0;
+    sample_indices[b] = (unsigned int)safe_leaf;
+
+    /* ── Gather h_t ──────────────────────────────────────────────────── */
+    for (int i = 0; i < HIDDEN_DIM; ++i) {
+        sampled_h_t[b * HIDDEN_DIM + i] = replay_h_t[safe_leaf * HIDDEN_DIM + i];
+    }
+
+    /* ── Gather h_tp1 ────────────────────────────────────────────────── */
+    for (int i = 0; i < HIDDEN_DIM; ++i) {
+        sampled_h_tp1[b * HIDDEN_DIM + i] = replay_h_tp1[safe_leaf * HIDDEN_DIM + i];
+    }
+
+    /* ── Unpack scalars ──────────────────────────────────────────────── */
+    const int sc_base = safe_leaf * SCALARS_PER_TRANSITION;
+    sampled_actions[b]        = (int)replay_scalars[sc_base + 0];
+    sampled_rewards[b]        = replay_scalars[sc_base + 1]; /* n-step scaled return */
+    sampled_dones[b]          = replay_scalars[sc_base + 3];
+    sampled_log_pi_old[b]     = replay_scalars[sc_base + 4];
+    sampled_n_step_gammas[b]  = replay_scalars[sc_base + 5];
+}

crates/ml-alpha/cuda/rl_per_tree_rebuild.cu

@@ -0,0 +1,44 @@
+/* =====================================================================
+ * rl_per_tree_rebuild.cu — GPU-resident PER: bottom-up parallel sum-tree rebuild
+ *
+ * Grid=(128), Block=(256). Total 32768 threads = capacity.
+ *
+ * Rebuilds the entire sum-tree from leaf priorities (at indices
+ * [capacity..2*capacity)) up to the root (index 1). Each level is
+ * processed in parallel with __threadfence() device-wide barriers
+ * between levels.
+ *
+ * No atomicAdd — each internal node is written by exactly one thread.
+ * ===================================================================== */
+
+extern "C" __global__ void rl_per_tree_rebuild(
+    float* __restrict__ priority_tree, /* [2 * capacity] */
+    int capacity
+)
+{
+    const int tid_global = blockIdx.x * blockDim.x + threadIdx.x;
+    const int total_threads = gridDim.x * blockDim.x;
+
+    /* Bottom-up: level 0 = parents of leaves, level (log2(cap)-1) = root */
+    /* At level L, there are capacity >> (L+1) internal nodes.            */
+    /* Node indices at level L: [capacity >> (L+1) .. capacity >> L)      */
+
+    int nodes_at_level = capacity >> 1; /* level 0: cap/2 nodes */
+    int start = nodes_at_level;         /* first node index at this level */
+
+    while (nodes_at_level >= 1) {
+        /* Each thread processes multiple nodes via grid-stride loop */
+        for (int i = tid_global; i < nodes_at_level; i += total_threads) {
+            const int node = start + i;
+            priority_tree[node] = priority_tree[2 * node] + priority_tree[2 * node + 1];
+        }
+
+        /* Device-wide fence: all writes at this level visible before
+         * any thread reads them at the next level */
+        __threadfence();
+
+        /* Move up one level */
+        nodes_at_level >>= 1;
+        start >>= 1;
+    }
+}

crates/ml-alpha/cuda/rl_per_update_priority.cu

@@ -0,0 +1,61 @@
+/* =====================================================================
+ * rl_per_update_priority.cu — GPU-resident PER: update leaf priorities
+ *
+ * Grid=(1), Block=(b_size). After a training step, writes new priorities
+ * (from absolute TD-errors) into the sum-tree leaf positions.
+ *
+ * No atomicAdd — max_priority reduction uses block-wide tree-reduce in
+ * shared memory.
+ *
+ * ISV reads:
+ *   isv[RL_PER_ALPHA_INDEX=405] — prioritization exponent α
+ * ===================================================================== */
+
+#define RL_PER_ALPHA_INDEX 405
+
+extern "C" __global__ void rl_per_update_priority(
+    const unsigned int* __restrict__ sample_indices, /* [B] */
+    const float* __restrict__ td_errors,             /* [B] absolute values */
+    float* __restrict__ priority_tree,               /* [2 * capacity] */
+    float* __restrict__ max_priority,                /* [1] */
+    const float* __restrict__ isv,
+    int b_size,
+    int capacity
+)
+{
+    extern __shared__ float smax[];
+
+    const int b = threadIdx.x;
+    if (b >= b_size) return;
+
+    /* ── Read alpha from ISV ──────────────────────────────────────────── */
+    const float alpha = isv[RL_PER_ALPHA_INDEX];
+
+    /* ── Compute priority = (|td_error| + ε)^α ───────────────────────── */
+    const float td_abs = fabsf(td_errors[b]);
+    const float priority = powf(td_abs + 1e-6f, alpha);
+
+    /* ── Write priority to tree leaf ──────────────────────────────────── */
+    const int leaf = (int)sample_indices[b];
+    priority_tree[capacity + leaf] = priority;
+
+    /* ── Block-wide tree-reduce to find max priority ──────────────────── */
+    smax[b] = priority;
+    __syncthreads();
+
+    /* Power-of-two tree reduce (handles non-power-of-two b_size) */
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (b < stride && (b + stride) < b_size) {
+            const float other = smax[b + stride];
+            if (other > smax[b]) smax[b] = other;
+        }
+        __syncthreads();
+    }
+
+    /* Thread 0 writes the block max to global max_priority */
+    if (b == 0) {
+        const float old_max = max_priority[0];
+        const float new_max = smax[0];
+        max_priority[0] = (new_max > old_max) ? new_max : old_max;
+    }
+}