diff --git a/crates/ml-dqn/build.rs b/crates/ml-dqn/build.rs index 96dce7d07..74486d572 100644 --- a/crates/ml-dqn/build.rs +++ b/crates/ml-dqn/build.rs @@ -46,7 +46,7 @@ fn main() { "residual_kernels.cu", "cast_kernels.cu", "replay_buffer_kernels.cu", - "seg_tree_kernel.cu", + "per_kernels.cu", ]; let standalone_kernels: [&str; 0] = []; diff --git a/crates/ml-dqn/src/per_kernels.cu b/crates/ml-dqn/src/per_kernels.cu new file mode 100644 index 000000000..1486a0337 --- /dev/null +++ b/crates/ml-dqn/src/per_kernels.cu @@ -0,0 +1,257 @@ +// PER (Prioritized Experience Replay) GPU kernels — decoupled lookback scan. +// +// Architecture: flat priority array + single-pass O(N) prefix sum +// (Merrill & Garland, 2016) + binary search sampling. +// +// Four kernels: +// 1. per_update_pa — priority update from BF16 TD errors +// 2. per_insert_pa — insert raw priorities at max_priority +// 3. per_prefix_scan — decoupled lookback inclusive prefix sum +// 4. per_sample — proportional sampling via binary search + +// --------------------------------------------------------------------------- +// 1. per_update_pa: priority leaf update from BF16 TD errors +// Reads bf16 TD errors directly (no separate cast kernel). +// Computes p = |td|^alpha + eps, writes raw priority, computes pa = p^alpha, +// writes to priorities_pa. atomicMax for batch max. +// Launch: grid=(ceil(batch_size/256)), block=(256). +// --------------------------------------------------------------------------- +extern "C" __global__ void per_update_pa( + float* __restrict__ priorities_pa, + float* __restrict__ priorities, + float* __restrict__ batch_max, + const unsigned int* __restrict__ indices, + const __nv_bfloat16* __restrict__ td_errors_bf16, + float alpha, float epsilon, + int capacity, int batch_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + unsigned int idx = indices[i]; + if (idx >= (unsigned int)capacity) return; + + float td = fabsf(__bfloat162float(td_errors_bf16[i])); + float new_prio = powf(td, alpha) + epsilon; + if (new_prio < epsilon) new_prio = epsilon; + if (new_prio > 1e6f) new_prio = 1e6f; + + priorities[idx] = new_prio; + + int ival = __float_as_int(new_prio); + atomicMax((int*)batch_max, ival); + + float pa = powf(new_prio, alpha); + priorities_pa[idx] = pa; +} + +// --------------------------------------------------------------------------- +// 2. per_insert_pa: insert priorities at max_priority +// Takes raw priority values, computes pa = priority^alpha, writes to +// priorities_pa. +// Launch: grid=(ceil(batch_size/256)), block=(256). +// --------------------------------------------------------------------------- +extern "C" __global__ void per_insert_pa( + float* __restrict__ priorities_pa, + const unsigned int* __restrict__ indices, + const float* __restrict__ raw_priorities, + float alpha, + int capacity, int batch_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + unsigned int idx = indices[i]; + if (idx >= (unsigned int)capacity) return; + + float pa = powf(raw_priorities[i], alpha); + priorities_pa[idx] = pa; +} + +// --------------------------------------------------------------------------- +// 3. per_prefix_scan: decoupled lookback inclusive prefix sum +// (Merrill & Garland, 2016). Single-pass O(N) scan. +// +// Status flags per tile: INVALID(0), AGGREGATE(1), PREFIX(2). +// Stored in separate arrays for f32 precision. +// +// Launch: grid=(ceil(size/256)), block=(256), shared_mem=256*4. +// --------------------------------------------------------------------------- +#define PER_TILE_SIZE 256 + +#define PER_STATE_INVALID 0u +#define PER_STATE_AGGREGATE 1u +#define PER_STATE_PREFIX_AVAILABLE 2u + +extern "C" __global__ void per_prefix_scan( + const float* __restrict__ input, + float* __restrict__ output, + volatile unsigned int* __restrict__ tile_state, + volatile float* __restrict__ tile_aggregate, + volatile float* __restrict__ tile_prefix, + int size) +{ + extern __shared__ float sdata[]; + + int tile_id = blockIdx.x; + int tid = threadIdx.x; + int gid = tile_id * PER_TILE_SIZE + tid; + + // Load into shared memory + float val = (gid < size) ? input[gid] : 0.0f; + sdata[tid] = val; + __syncthreads(); + + // Intra-tile inclusive scan (Hillis-Steele) + for (int offset = 1; offset < PER_TILE_SIZE; offset <<= 1) { + float add = (tid >= offset) ? sdata[tid - offset] : 0.0f; + __syncthreads(); + sdata[tid] += add; + __syncthreads(); + } + + float tile_total = sdata[PER_TILE_SIZE - 1]; + + // Decoupled lookback: thread 0 publishes state and computes prefix + if (tid == 0) { + if (tile_id == 0) { + // First tile: prefix = 0, total is already the prefix sum + tile_aggregate[0] = tile_total; + tile_prefix[0] = tile_total; + __threadfence(); + tile_state[0] = PER_STATE_PREFIX_AVAILABLE; + __threadfence(); + } else { + // Publish aggregate first + tile_aggregate[tile_id] = tile_total; + __threadfence(); + tile_state[tile_id] = PER_STATE_AGGREGATE; + __threadfence(); + + // Lookback: scan preceding tiles to compute exclusive prefix + float exclusive_prefix = 0.0f; + int look = tile_id - 1; + + while (look >= 0) { + unsigned int state; + do { + state = tile_state[look]; + } while (state == PER_STATE_INVALID); + + if (state == PER_STATE_PREFIX_AVAILABLE) { + exclusive_prefix += tile_prefix[look]; + break; + } else { + exclusive_prefix += tile_aggregate[look]; + look--; + } + } + + // Publish our prefix (inclusive: exclusive_prefix + tile_total) + tile_prefix[tile_id] = exclusive_prefix + tile_total; + __threadfence(); + tile_state[tile_id] = PER_STATE_PREFIX_AVAILABLE; + __threadfence(); + + // Store exclusive prefix in shared memory for broadcast + sdata[PER_TILE_SIZE - 1] = exclusive_prefix; + } + } + + // All threads wait for thread 0 to finish lookback + __syncthreads(); + + // Add exclusive prefix to all elements (tile 0 has prefix = 0) + float prefix = (tile_id == 0) ? 0.0f : sdata[PER_TILE_SIZE - 1]; + + // Reload local scan value (was overwritten for tile_id > 0) + float scan_val; + if (tile_id == 0) { + scan_val = sdata[tid]; + } else { + // sdata[PER_TILE_SIZE-1] was overwritten with exclusive_prefix; + // we need to reconstruct. Re-read from sdata before it was overwritten. + // Actually, only sdata[PER_TILE_SIZE-1] was overwritten by thread 0. + // For tid < PER_TILE_SIZE-1, sdata[tid] is still valid. + // For tid == PER_TILE_SIZE-1, the scan result = tile_total (saved above). + if (tid == PER_TILE_SIZE - 1) { + scan_val = tile_total; + } else { + scan_val = sdata[tid]; + } + } + + if (gid < size) { + output[gid] = prefix + scan_val; + } +} + +// --------------------------------------------------------------------------- +// Philox single-round helper (counter-based PRNG) +// --------------------------------------------------------------------------- +__device__ __forceinline__ +unsigned int philox_single(unsigned int counter, unsigned int key) { + unsigned int hi, lo; + lo = counter * 0xD2511F53u; + hi = __umulhi(counter, 0xD2511F53u); + for (int r = 0; r < 10; r++) { + unsigned int t = hi ^ key; + hi = lo * 0xCD9E8D57u; + lo = __umulhi(lo, 0xCD9E8D57u); + lo ^= t; + key += 0x9E3779B9u; + } + return lo ^ hi; +} + +// --------------------------------------------------------------------------- +// 4. per_sample: proportional sampling via binary search on prefix sum +// Each thread generates a Philox random threshold in [0, total_sum) and +// binary searches prefix_sum[0..size]. Also gathers priorities_pa[index] +// for IS weight computation. Thread 0 writes total_sum to output. +// Launch: grid=(ceil(batch_size/256)), block=(256). +// --------------------------------------------------------------------------- +extern "C" __global__ void per_sample( + const float* __restrict__ prefix_sum, + const float* __restrict__ priorities_pa, + long long* __restrict__ out_indices, + float* __restrict__ out_priorities, + float* __restrict__ total_sum_out, + unsigned int seed, + int size, int batch_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + + // Thread 0 writes total sum + if (i == 0) { + total_sum_out[0] = (size > 0) ? prefix_sum[size - 1] : 0.0f; + } + + if (i >= batch_size) return; + + float total_sum = (size > 0) ? prefix_sum[size - 1] : 1.0f; + + // Philox RNG: generate uniform in [0, total_sum) + unsigned int bits = philox_single((unsigned int)i, seed); + float u = (float)(bits >> 8) * (1.0f / 16777216.0f); + float threshold = u * total_sum; + + // Binary search: find smallest index where prefix_sum[index] > threshold + int lo = 0; + int hi = size - 1; + while (lo < hi) { + int mid = lo + (hi - lo) / 2; + if (prefix_sum[mid] <= threshold) { + lo = mid + 1; + } else { + hi = mid; + } + } + + // Clamp to valid range + if (lo >= size) lo = size - 1; + if (lo < 0) lo = 0; + + out_indices[i] = (long long)lo; + out_priorities[i] = priorities_pa[lo]; +}