From c95abbf3daa89c764848f32c7ba9c37f5a675329 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 00:25:31 +0200 Subject: [PATCH] =?UTF-8?q?feat(cuda):=20GPU-resident=20PER=20kernels=20?= =?UTF-8?q?=E2=80=94=20push,=20sample,=20update,=20tree=5Frebuild?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 4 + crates/ml-alpha/cuda/rl_per_push.cu | 181 ++++++++++++++++++ crates/ml-alpha/cuda/rl_per_sample.cu | 118 ++++++++++++ crates/ml-alpha/cuda/rl_per_tree_rebuild.cu | 44 +++++ .../ml-alpha/cuda/rl_per_update_priority.cu | 61 ++++++ 5 files changed, 408 insertions(+) create mode 100644 crates/ml-alpha/cuda/rl_per_push.cu create mode 100644 crates/ml-alpha/cuda/rl_per_sample.cu create mode 100644 crates/ml-alpha/cuda/rl_per_tree_rebuild.cu create mode 100644 crates/ml-alpha/cuda/rl_per_update_priority.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 1e03c74c0..d6d04bd45 100644 --- a/crates/ml-alpha/build.rs +++ b/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) ]; diff --git a/crates/ml-alpha/cuda/rl_per_push.cu b/crates/ml-alpha/cuda/rl_per_push.cu new file mode 100644 index 000000000..71a62a226 --- /dev/null +++ b/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; +} diff --git a/crates/ml-alpha/cuda/rl_per_sample.cu b/crates/ml-alpha/cuda/rl_per_sample.cu new file mode 100644 index 000000000..a229bdd69 --- /dev/null +++ b/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]; +} diff --git a/crates/ml-alpha/cuda/rl_per_tree_rebuild.cu b/crates/ml-alpha/cuda/rl_per_tree_rebuild.cu new file mode 100644 index 000000000..ce74b8156 --- /dev/null +++ b/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; + } +} diff --git a/crates/ml-alpha/cuda/rl_per_update_priority.cu b/crates/ml-alpha/cuda/rl_per_update_priority.cu new file mode 100644 index 000000000..0e2888c16 --- /dev/null +++ b/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; + } +}