feat: per_kernels.cu — decoupled lookback scan + binary search sampling

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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] = [];

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];
+}