From 9ad7eb257c508ef09425b97b7f350d1169aef0ab Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 01:41:15 +0200 Subject: [PATCH] =?UTF-8?q?feat(cuda):=20coalesced=20PER=20push=20?= =?UTF-8?q?=E2=80=94=20ring=20+=20flush=20kernels=20(Grid=3DB,=20Block=3D1?= =?UTF-8?q?28)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rl_per_push_ring: 128 threads copy h_t in one coalesced 512-byte transaction. Thread 0 writes scalars + flush_flag. rl_per_push_flush: block 0 prefix-sums flush_flags, signals via volatile ready flag. All flushing blocks do coalesced 128-thread h_t + h_tp1 + scalars write to replay. No atomicAdd. Expected: 124μs → ~15μs per step (8× faster). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/rl_per_push_flush.cu | 200 ++++++++++++++++++++++ crates/ml-alpha/cuda/rl_per_push_ring.cu | 109 ++++++++++++ 2 files changed, 309 insertions(+) create mode 100644 crates/ml-alpha/cuda/rl_per_push_flush.cu create mode 100644 crates/ml-alpha/cuda/rl_per_push_ring.cu diff --git a/crates/ml-alpha/cuda/rl_per_push_flush.cu b/crates/ml-alpha/cuda/rl_per_push_flush.cu new file mode 100644 index 000000000..980b84834 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_per_push_flush.cu @@ -0,0 +1,200 @@ +/** + * rl_per_push_flush.cu — Coordinated coalesced replay write. + * + * Second kernel in the split rl_per_push pipeline. After rl_per_push_ring has + * written data to the n-step ring and set flush_flags[b] for completed n-step + * transitions, this kernel coordinates slot allocation and performs the actual + * coalesced write to the replay buffer. + * + * Architecture: + * - Grid=(b_size, 1, 1), Block=(128, 1, 1). One block per batch element. + * - Block 0 runs the prefix-sum to assign write slots (single-threaded, O(B)). + * - All other blocks spin on a volatile global-memory ready flag (no atomics). + * - Once ready, each flushing block does a coalesced 128-thread copy of h_t + * and h_tp1, then thread 0 computes the n-step return and writes scalars + + * priority + resets the ring count. + * + * Coordination: block-0 prefix-sum + volatile ready flag. + * No atomicAdd — slot assignment is deterministic serial scan (feedback_no_atomicadd). + * Pre-compiled cubin only (feedback_no_nvrtc). + * + * Launch: CUDA_HOME=/usr/local/cuda nvcc -cubin -arch sm_86 -O3 \ + * --generate-line-info rl_per_push_flush.cu -o rl_per_push_flush.cubin + */ + +#define HIDDEN_DIM 128 +#define N_STEP_MAX 16 +#define SCALARS_PER_TRANSITION 7 + +/* ISV slot indices for RL config */ +#define RL_GAMMA_INDEX 400 +#define RL_N_STEP_INDEX 403 + +extern "C" __global__ void rl_per_push_flush( + /* ── Flush coordination ── */ + const int* __restrict__ flush_flags, /* [B] from ring kernel: 1=flush, 0=skip */ + int* __restrict__ write_offsets, /* [B+2] output: per-batch offsets, base, ready flag */ + + /* ── N-step ring (read source) ── */ + const float* __restrict__ nstep_scalars, /* [B * N_STEP_MAX * 3]: (r_scaled, r_raw, done) per ring slot */ + const float* __restrict__ nstep_h_t, /* [B * N_STEP_MAX * HIDDEN_DIM] */ + const unsigned int* __restrict__ nstep_write_idx, /* [B] ring write cursor */ + const unsigned int* __restrict__ nstep_count, /* [B] items in ring */ + + /* ── Current step data (for h_tp1) ── */ + const float* __restrict__ h_tp1_current, /* [B * HIDDEN_DIM] */ + const float* __restrict__ log_pi_old, /* [B] */ + const int* __restrict__ actions, /* [B] */ + const float* __restrict__ isv, /* ISV array */ + + /* ── Replay storage (write destination) ── */ + float* __restrict__ replay_h_t, /* [capacity * HIDDEN_DIM] */ + float* __restrict__ replay_h_tp1, /* [capacity * HIDDEN_DIM] */ + float* __restrict__ replay_scalars, /* [capacity * SCALARS_PER_TRANSITION] */ + float* __restrict__ priority_tree, /* [2 * capacity] (leaf layer at offset capacity) */ + unsigned int* __restrict__ write_head, /* [1] circular write cursor */ + unsigned int* __restrict__ replay_len, /* [1] current buffer occupancy */ + float* __restrict__ max_priority, /* [1] running max priority for new inserts */ + + /* ── N-step count reset (written on flush) ── */ + unsigned int* __restrict__ nstep_count_out, /* [B] — set to 0 for flushed batches */ + + int b_size, + int capacity +) +{ + /* ==================================================================== + * PHASE 1: Block 0 coordination (thread 0 only). + * + * Computes prefix-sum of flush_flags to assign contiguous write slots, + * advances write_head and replay_len, stores base for other blocks, + * and signals ready via volatile write. + * ==================================================================== */ + if (blockIdx.x == 0 && threadIdx.x == 0) { + unsigned int base = write_head[0]; + + /* Exclusive prefix-sum: write_offsets[i] = number of flushes before batch i */ + int total = 0; + for (int i = 0; i < b_size; ++i) { + write_offsets[i] = total; + total += flush_flags[i]; + } + + /* Advance write head and replay length */ + if (total > 0) { + write_head[0] = (base + (unsigned int)total) % (unsigned int)capacity; + unsigned int len = replay_len[0] + (unsigned int)total; + if (len > (unsigned int)capacity) len = (unsigned int)capacity; + replay_len[0] = len; + } + + /* Store base write head for other blocks at slot [b_size] */ + write_offsets[b_size] = (int)base; + + /* Fence: ensure all writes above are globally visible before signaling */ + __threadfence(); + + /* Signal ready at slot [b_size + 1] */ + write_offsets[b_size + 1] = 1; + __threadfence(); + } + + /* ==================================================================== + * PHASE 2: All blocks wait for ready signal. + * + * Non-block-0 blocks spin on a volatile read of the ready flag. + * Block 0 already wrote it, so it proceeds immediately. + * ==================================================================== */ + if (threadIdx.x == 0 && blockIdx.x != 0) { + volatile int* ready = (volatile int*)&write_offsets[b_size + 1]; + while (*ready == 0) { + /* spin — volatile load, no atomic */ + } + } + __syncthreads(); + + /* ==================================================================== + * PHASE 3: Coalesced replay write (128 threads per block). + * + * Each flushing block writes: + * - h_t[HIDDEN_DIM]: from oldest ring entry (coalesced, all 128 threads) + * - h_tp1[HIDDEN_DIM]: from current step (coalesced, all 128 threads) + * - scalars[7]: n-step return + metadata (thread 0 only) + * - priority: max_priority into leaf layer (thread 0 only) + * - reset: nstep_count_out[b] = 0 (thread 0 only) + * ==================================================================== */ + const int b = blockIdx.x; + + /* Early exit: nothing to flush for this batch */ + if (flush_flags[b] == 0) return; + + /* Compute target replay slot */ + const unsigned int base = (unsigned int)write_offsets[b_size]; + const unsigned int my_offset = (unsigned int)write_offsets[b]; + const unsigned int my_slot = (base + my_offset) % (unsigned int)capacity; + const int tid = threadIdx.x; + + /* Read n-step config from ISV */ + int n_step_raw = (int)isv[RL_N_STEP_INDEX]; + const int n_step = (n_step_raw < 1) ? 1 : ((n_step_raw > N_STEP_MAX) ? N_STEP_MAX : n_step_raw); + const float gamma = isv[RL_GAMMA_INDEX]; + + /* Determine oldest ring entry to read h_t from */ + const unsigned int count = nstep_count[b]; + const int actual_n = ((int)count < n_step) ? (int)count : n_step; + const unsigned int oldest_ring = (nstep_write_idx[b] + (unsigned int)N_STEP_MAX - (unsigned int)actual_n) + % (unsigned int)N_STEP_MAX; + + /* Coalesced h_t copy: oldest ring entry -> replay (128 threads, 1 float each) */ + replay_h_t[my_slot * HIDDEN_DIM + tid] = + nstep_h_t[(b * N_STEP_MAX + oldest_ring) * HIDDEN_DIM + tid]; + + /* Coalesced h_tp1 copy: current hidden state -> replay (128 threads, 1 float each) */ + replay_h_tp1[my_slot * HIDDEN_DIM + tid] = + h_tp1_current[b * HIDDEN_DIM + tid]; + + /* Thread 0: compute n-step return, write scalars, set priority, reset count */ + if (tid == 0) { + /* N-step discounted return computation */ + float r_scaled = 0.0f; + float r_raw = 0.0f; + float gamma_pow = 1.0f; + float n_step_gamma = 1.0f; + int any_done = 0; + + for (int k = 0; k < actual_n; ++k) { + unsigned int rk = (oldest_ring + (unsigned int)k) % (unsigned int)N_STEP_MAX; + int sc_idx = (b * N_STEP_MAX + (int)rk) * 3; + r_scaled += gamma_pow * nstep_scalars[sc_idx + 0]; + r_raw += gamma_pow * nstep_scalars[sc_idx + 1]; + if (nstep_scalars[sc_idx + 2] > 0.5f) { + any_done = 1; + n_step_gamma = 0.0f; + break; + } + gamma_pow *= gamma; + } + if (!any_done) { + n_step_gamma = gamma_pow; /* gamma^actual_n for bootstrap */ + } + + /* Write transition scalars to replay buffer */ + int sc_out = (int)my_slot * SCALARS_PER_TRANSITION; + replay_scalars[sc_out + 0] = (float)actions[b]; /* action index */ + replay_scalars[sc_out + 1] = r_scaled; /* n-step scaled return */ + replay_scalars[sc_out + 2] = r_raw; /* n-step raw return */ + replay_scalars[sc_out + 3] = (float)any_done; /* terminal flag */ + replay_scalars[sc_out + 4] = log_pi_old[b]; /* log pi(a|s) at collection */ + replay_scalars[sc_out + 5] = n_step_gamma; /* discount for bootstrap */ + replay_scalars[sc_out + 6] = r_raw; /* raw return (duplicate for diagnostics) */ + + /* Set priority to current max (standard PER new-insert convention). + * priority_tree leaf layer starts at offset capacity. */ + 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 (ring consumed) */ + nstep_count_out[b] = 0; + } +} diff --git a/crates/ml-alpha/cuda/rl_per_push_ring.cu b/crates/ml-alpha/cuda/rl_per_push_ring.cu new file mode 100644 index 000000000..ab4248aa6 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_per_push_ring.cu @@ -0,0 +1,109 @@ +/* ===================================================================== + * rl_per_push_ring.cu — Coalesced n-step ring write kernel + * + * Writes the current step's hidden state, scaled reward, raw reward, + * and done flag into a per-batch n-step ring buffer. The ring write + * for h_t is fully coalesced: 128 threads each write one float, + * producing a single 512-byte transaction per batch element. + * + * This is the first of two kernels splitting the old rl_per_push + * (Grid=1, Block=B, 124us/call). The second kernel (rl_per_flush) + * handles the actual PER priority computation and replay-buffer push + * for batches whose flush_flag is set. + * + * Launch config: Grid=(b_size, 1, 1), Block=(128, 1, 1). + * One block per batch element. 128 threads = HIDDEN_DIM. + * + * Memory access pattern: + * h_t read: coalesced 512B (128 threads x 4B, contiguous in b) + * h_t write: coalesced 512B (128 threads x 4B, ring slot contiguous) + * scalars: thread-0 only, 3 floats (negligible) + * ring meta: thread-0 only, 2 uint32 reads + 2 writes + * + * No atomicAdd. No shared memory. No nvrtc (pre-compiled cubin). + * ===================================================================== */ + +#define HIDDEN_DIM 128 +#define N_STEP_MAX 16 +#define RL_GAMMA_INDEX 400 +#define RL_N_STEP_INDEX 403 + +extern "C" __global__ void rl_per_push_ring( + const float* __restrict__ h_t_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, /* ISV vector (shared, read-only)*/ + /* N-step ring buffers (per batch element) */ + float* __restrict__ nstep_scalars, /* [B, N_STEP_MAX, 3] */ + float* __restrict__ nstep_h_t, /* [B, N_STEP_MAX, 128] */ + unsigned int* __restrict__ nstep_write_idx, /* [B] */ + unsigned int* __restrict__ nstep_count, /* [B] */ + /* Output: flush decision per batch element */ + int* __restrict__ flush_flags, /* [B] */ + int b_size +) +{ + const int b = blockIdx.x; + const int tid = threadIdx.x; + + if (b >= b_size) return; + + /* ------------------------------------------------------------------ + * Read ring metadata (same value for every thread in the block). + * No shared memory needed — all 128 threads read the same two scalars + * from global memory; the L1 cache serves subsequent reads from the + * same cache line at zero cost. + * ------------------------------------------------------------------ */ + const unsigned int ring_idx = nstep_write_idx[b]; + const unsigned int count = nstep_count[b]; + + /* ------------------------------------------------------------------ + * ISV-driven n_step: clamp to [1, N_STEP_MAX]. + * Thread 0 reads the ISV slot; other threads don't need the value + * until the flush decision, which only thread 0 computes. + * ------------------------------------------------------------------ */ + int n_step = N_STEP_MAX; /* default if ISV read is out of range */ + if (tid == 0) { + float n_step_f = isv[RL_N_STEP_INDEX]; + int n_step_i = (int)n_step_f; + n_step = (n_step_i < 1) ? 1 : ((n_step_i > N_STEP_MAX) ? N_STEP_MAX : n_step_i); + } + + /* ------------------------------------------------------------------ + * Coalesced h_t copy: each thread writes one float. + * Source: h_t_current[b * HIDDEN_DIM + tid] (contiguous per batch) + * Dest: nstep_h_t[(b * N_STEP_MAX + ring_idx) * HIDDEN_DIM + tid] + * + * 128 threads x 4 bytes = 512 bytes = one coalesced transaction. + * ------------------------------------------------------------------ */ + const long long src_offset = (long long)b * HIDDEN_DIM + tid; + const long long dst_offset = ((long long)b * N_STEP_MAX + ring_idx) * HIDDEN_DIM + tid; + nstep_h_t[dst_offset] = h_t_current[src_offset]; + + /* ------------------------------------------------------------------ + * Thread 0: write scalar ring entries and advance ring pointer. + * ------------------------------------------------------------------ */ + if (tid == 0) { + /* Scalar layout: [rewards_scaled, rewards_raw, dones] per ring slot */ + const long long scalar_base = ((long long)b * N_STEP_MAX + ring_idx) * 3; + nstep_scalars[scalar_base + 0] = rewards_scaled[b]; + nstep_scalars[scalar_base + 1] = rewards_raw[b]; + nstep_scalars[scalar_base + 2] = dones[b]; + + /* Advance ring write pointer (modular) */ + nstep_write_idx[b] = (ring_idx + 1) % N_STEP_MAX; + + /* Increment count (saturates at N_STEP_MAX — ring is full) */ + const unsigned int new_count = (count < (unsigned int)N_STEP_MAX) + ? count + 1 + : (unsigned int)N_STEP_MAX; + nstep_count[b] = new_count; + + /* Flush decision: ring is full OR episode terminated */ + const int should_flush = ((int)new_count >= n_step) || (dones[b] > 0.5f); + flush_flags[b] = should_flush ? 1 : 0; + } +}