gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(cuda): coalesced PER push — ring + flush kernels (Grid=B, Block=128)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 01:41:15 +02:00
parent be3d5ceeae
commit 9ad7eb257c
2 changed files with 309 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

200
crates/ml-alpha/cuda/rl_per_push_flush.cu Normal file
View File

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

109
crates/ml-alpha/cuda/rl_per_push_ring.cu Normal file
View File

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