From 0f45537e6a6cece768d6d2b3a96138120b07ffa9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 9 Mar 2026 08:33:28 +0100 Subject: [PATCH] perf(dqn): GPU searchsorted kernel for PER sampling via CustomOp2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eliminates the CPU roundtrip in PER binary search: previously the cumsum tensor (~400KB for 100K buffer) was downloaded to CPU, searched in a loop, then indices uploaded back. Now a CUDA kernel runs one thread per target with O(log n) binary search — zero DMA. - Add searchsorted_kernel.cu (40-line CUDA binary search kernel) - Implement SearchSorted as Candle CustomOp2 with CPU fallback - Wire into sample_proportional() and sample_rank_based() - Fix all clippy warnings in gpu_replay_buffer.rs (const fn, shadow, doc backticks, to_owned, div_ceil, module_name_repetitions) - Fix wildcard_enum_match_arm in mixed_precision.rs Co-Authored-By: Claude Opus 4.6 --- crates/ml-core/src/mixed_precision.rs | 2 +- crates/ml-dqn/src/gpu_replay_buffer.rs | 286 ++++++++++++++++------- crates/ml-dqn/src/searchsorted_kernel.cu | 39 ++++ 3 files changed, 244 insertions(+), 83 deletions(-) create mode 100644 crates/ml-dqn/src/searchsorted_kernel.cu diff --git a/crates/ml-core/src/mixed_precision.rs b/crates/ml-core/src/mixed_precision.rs index 88dbce9ef..62ac1da6b 100644 --- a/crates/ml-core/src/mixed_precision.rs +++ b/crates/ml-core/src/mixed_precision.rs @@ -25,7 +25,7 @@ use crate::MLError; pub fn align_dim_for_tensor_cores(dim: usize, device: &Device) -> usize { match device { Device::Cuda(_) => (dim + 7) & !7, - _ => dim, + Device::Cpu | Device::Metal(_) => dim, } } diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index ebd196280..ce7e008aa 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -1,15 +1,162 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launch in SearchSorted + //! GPU-Resident Replay Buffer with parallel prefix-sum PER sampling //! //! Stores all experience data as contiguous GPU tensors in a ring buffer. -//! Sampling uses prefix-sum + binary search (proportional) or radix sort -//! (rank-based) — all on GPU with zero CPU involvement. +//! Sampling uses prefix-sum + GPU searchsorted (proportional) or radix sort +//! (rank-based) — all on GPU with zero CPU involvement for the binary search. -use candle_core::{Device, DType, Tensor}; +use candle_core::{CpuStorage, Device, DType, Layout, Shape, Tensor}; use ml_core::MLError; use crate::mixed_precision::training_dtype; use crate::replay_buffer_type::GpuBatch; +// --------------------------------------------------------------------------- +// SearchSorted CustomOp2 — GPU binary search via CUDA kernel +// --------------------------------------------------------------------------- + +/// Candle `CustomOp2` that performs GPU-native binary search (searchsorted). +/// +/// Given a sorted cumulative-sum tensor `[n]` and a targets tensor `[batch_size]`, +/// finds the insertion index for each target via binary search. +/// +/// On CUDA: compiles and launches a custom CUDA kernel (one thread per target, +/// O(log n) work per thread). Zero CPU-GPU data transfer. +/// +/// On CPU: falls back to Rust's `binary_search_by` (identical to the previous +/// implementation, used for testing). +struct SearchSorted { + /// Number of active elements in the sorted data (may be < tensor capacity). + n: usize, +} + +impl candle_core::CustomOp2 for SearchSorted { + fn name(&self) -> &'static str { + "searchsorted" + } + + /// CPU fallback: binary search each target in the cumsum slice. + fn cpu_fwd( + &self, + s1: &CpuStorage, + l1: &Layout, + s2: &CpuStorage, + l2: &Layout, + ) -> candle_core::Result<(CpuStorage, Shape)> { + let cumsum = s1.as_slice::()?; + let targets = s2.as_slice::()?; + let n = self.n; + let offset1 = l1.start_offset(); + let offset2 = l2.start_offset(); + let batch_size = l2.shape().elem_count(); + + let mut indices: Vec = Vec::with_capacity(batch_size); + let end = (offset1 + n).min(cumsum.len()); + let search_slice = cumsum.get(offset1..end).unwrap_or(&[]); + for t in 0..batch_size { + let target = targets.get(offset2 + t).copied().unwrap_or(0.0); + let idx = match search_slice.binary_search_by(|v| { + v.partial_cmp(&target).unwrap_or(std::cmp::Ordering::Less) + }) { + Ok(pos) => pos, + Err(pos) => pos.min(n.saturating_sub(1)), + }; + indices.push(idx as i64); + } + + Ok((CpuStorage::I64(indices), Shape::from_dims(&[batch_size]))) + } + + /// GPU path: launch CUDA searchsorted kernel (zero CPU-GPU data transfer). + #[cfg(feature = "cuda")] + fn cuda_fwd( + &self, + s1: &candle_core::CudaStorage, + l1: &Layout, + s2: &candle_core::CudaStorage, + l2: &Layout, + ) -> candle_core::Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::cudarc; + use cudarc::driver::{LaunchConfig, PushKernelArg}; + use std::sync::OnceLock; + + // Compile PTX once per process (NVRTC is expensive, ~100ms) + static SEARCHSORTED_PTX: OnceLock> = OnceLock::new(); + let ptx_result = SEARCHSORTED_PTX.get_or_init(|| { + let src = include_str!("searchsorted_kernel.cu"); + cudarc::nvrtc::compile_ptx(src) + .map_err(|e| format!("searchsorted CUDA kernel compilation failed: {e}")) + }); + let ptx = ptx_result.as_ref().map_err(|e| { + candle_core::Error::Cuda(e.clone().into()) + })?; + + let n = self.n; + let batch_size = l2.shape().elem_count(); + let dev = &s1.device; + let stream = dev.cuda_stream(); + + // Extract CudaSlice from both storages (zero-copy) + let cumsum_slice = s1.as_cuda_slice::()?; + let cumsum_view = cumsum_slice.slice(l1.start_offset()..); + let targets_slice = s2.as_cuda_slice::()?; + let targets_view = targets_slice.slice(l2.start_offset()..); + + // Load module + function (cheap after PTX compile) + let module = stream.context().load_module(ptx.clone()).map_err(|e| { + candle_core::Error::Cuda(format!("searchsorted module load: {e}").into()) + })?; + let func = module.load_function("searchsorted_kernel").map_err(|e| { + candle_core::Error::Cuda(format!("searchsorted function load: {e}").into()) + })?; + + // Allocate output buffer (i64 for Candle's index_select) + let mut out_indices = stream.alloc_zeros::(batch_size).map_err(|e| { + candle_core::Error::Cuda(format!("alloc searchsorted output: {e}").into()) + })?; + + // Launch: 1 thread per target, 256 threads per block + let threads_per_block = 256_u32; + let blocks = (batch_size as u32).div_ceil(threads_per_block); + let config = LaunchConfig { + grid_dim: (blocks.max(1), 1, 1), + block_dim: (threads_per_block, 1, 1), + shared_mem_bytes: 0, + }; + let n_i32 = n as i32; + let bs_i32 = batch_size as i32; + + // Safety: kernel parameter order matches searchsorted_kernel.cu exactly. + // cumsum_view has >= n elements, targets_view has >= batch_size elements, + // out_indices has batch_size elements. All slices are valid GPU memory. + unsafe { + stream + .launch_builder(&func) + .arg(&cumsum_view) + .arg(&targets_view) + .arg(&mut out_indices) + .arg(&n_i32) + .arg(&bs_i32) + .launch(config) + .map_err(|e| { + candle_core::Error::Cuda( + format!("searchsorted kernel launch: {e}").into(), + ) + })?; + } + + // Wrap output CudaSlice back into CudaStorage + let out_storage = candle_core::CudaStorage::wrap_cuda_slice(out_indices, dev.clone()); + Ok((out_storage, Shape::from_dims(&[batch_size]))) + } +} + +// --------------------------------------------------------------------------- +// GpuReplayBuffer +// --------------------------------------------------------------------------- + /// Configuration for GPU replay buffer +#[allow(clippy::module_name_repetitions)] #[derive(Debug, Clone)] pub struct GpuReplayBufferConfig { pub capacity: usize, @@ -51,9 +198,9 @@ impl GpuReplayBuffer { /// Create a new GPU replay buffer with pre-allocated storage. /// /// All tensors are allocated on the given device (CPU or CUDA) at creation. - /// For 100K capacity × 51 state_dim, this uses ~47 MB VRAM. + /// For 100K capacity × 51 `state_dim`, this uses ~47 MB VRAM. /// Maximum VRAM budget for the replay buffer (4 GB). - /// Prevents accidental OOM from misconfigured capacity/state_dim. + /// Prevents accidental OOM from misconfigured `capacity`/`state_dim`. const MAX_BYTES: usize = 4 * 1024 * 1024 * 1024; pub fn new(config: GpuReplayBufferConfig, device: &Device) -> Result { @@ -65,10 +212,12 @@ impl GpuReplayBuffer { // actions: cap * 4 (u32), rewards + dones + priorities: 3 * cap * 4 let bytes_needed = 2 * cap * sdim * 4 + 4 * cap * 4; if bytes_needed > Self::MAX_BYTES { + #[allow(clippy::integer_division)] + let needed_mb = bytes_needed / (1024 * 1024); + #[allow(clippy::integer_division)] + let limit_mb = Self::MAX_BYTES / (1024 * 1024); return Err(MLError::ModelError(format!( - "GPU replay buffer would need {} MB (limit {} MB). Reduce capacity or state_dim.", - bytes_needed / (1024 * 1024), - Self::MAX_BYTES / (1024 * 1024), + "GPU replay buffer would need {needed_mb} MB (limit {limit_mb} MB). Reduce capacity or state_dim.", ))); } @@ -97,22 +246,22 @@ impl GpuReplayBuffer { } /// Number of experiences currently in the buffer. - pub fn len(&self) -> usize { + pub const fn len(&self) -> usize { self.size } /// Maximum number of experiences the buffer can hold. - pub fn capacity(&self) -> usize { + pub const fn capacity(&self) -> usize { self.config.capacity } /// Whether the buffer is empty. - pub fn is_empty(&self) -> bool { + pub const fn is_empty(&self) -> bool { self.size == 0 } /// Whether the buffer has enough experiences to sample a batch. - pub fn can_sample(&self, batch_size: usize) -> bool { + pub const fn can_sample(&self, batch_size: usize) -> bool { self.size >= batch_size } @@ -121,18 +270,17 @@ impl GpuReplayBuffer { if self.config.beta_annealing_steps == 0 { return self.config.beta_max; } - let progress = (self.current_step as f32) / (self.config.beta_annealing_steps as f32); - let progress = progress.min(1.0); + let progress = ((self.current_step as f32) / (self.config.beta_annealing_steps as f32)).min(1.0); self.config.beta_start + (self.config.beta_max - self.config.beta_start) * progress } /// Step beta annealing counter. - pub fn step(&mut self) { + pub const fn step(&mut self) { self.current_step = self.current_step.saturating_add(1); } /// Clear all experiences. Resets cursor and size but keeps allocations. - pub fn clear(&mut self) { + pub const fn clear(&mut self) { self.write_cursor = 0; self.size = 0; self.max_priority = 1.0; @@ -140,22 +288,22 @@ impl GpuReplayBuffer { } /// Reference to the underlying device. - pub fn device(&self) -> &Device { + pub const fn device(&self) -> &Device { &self.device } /// PER alpha exponent. - pub fn alpha(&self) -> f32 { + pub const fn alpha(&self) -> f32 { self.config.alpha } /// Priority epsilon floor. - pub fn epsilon(&self) -> f32 { + pub const fn epsilon(&self) -> f32 { self.config.epsilon } /// State dimensionality. - pub fn state_dim(&self) -> usize { + pub const fn state_dim(&self) -> usize { self.config.state_dim } @@ -180,8 +328,8 @@ impl GpuReplayBuffer { } // Cast incoming states to match buffer dtype (e.g. F32 → BF16 on CUDA) - let states = &states.to_dtype(self.states.dtype())?; - let next_states = &next_states.to_dtype(self.next_states.dtype())?; + let cast_states = &states.to_dtype(self.states.dtype())?; + let cast_next_states = &next_states.to_dtype(self.next_states.dtype())?; let cap = self.config.capacity; let cursor = self.write_cursor; @@ -195,8 +343,8 @@ impl GpuReplayBuffer { if cursor + batch_size <= cap { // Contiguous write — no wrapping - self.states = self.states.slice_scatter(states, 0, cursor)?; - self.next_states = self.next_states.slice_scatter(next_states, 0, cursor)?; + self.states = self.states.slice_scatter(cast_states, 0, cursor)?; + self.next_states = self.next_states.slice_scatter(cast_next_states, 0, cursor)?; self.actions = self.actions.slice_scatter(actions, 0, cursor)?; self.rewards = self.rewards.slice_scatter(rewards, 0, cursor)?; self.dones = self.dones.slice_scatter(dones, 0, cursor)?; @@ -207,8 +355,8 @@ impl GpuReplayBuffer { let head_len = batch_size - tail_len; // Tail: cursor..capacity - let s_tail = states.narrow(0, 0, tail_len)?; - let ns_tail = next_states.narrow(0, 0, tail_len)?; + let s_tail = cast_states.narrow(0, 0, tail_len)?; + let ns_tail = cast_next_states.narrow(0, 0, tail_len)?; let a_tail = actions.narrow(0, 0, tail_len)?; let r_tail = rewards.narrow(0, 0, tail_len)?; let d_tail = dones.narrow(0, 0, tail_len)?; @@ -222,8 +370,8 @@ impl GpuReplayBuffer { self.priorities = self.priorities.slice_scatter(&p_tail, 0, cursor)?; // Head: 0..head_len - let s_head = states.narrow(0, tail_len, head_len)?; - let ns_head = next_states.narrow(0, tail_len, head_len)?; + let s_head = cast_states.narrow(0, tail_len, head_len)?; + let ns_head = cast_next_states.narrow(0, tail_len, head_len)?; let a_head = actions.narrow(0, tail_len, head_len)?; let r_head = rewards.narrow(0, tail_len, head_len)?; let d_head = dones.narrow(0, tail_len, head_len)?; @@ -242,13 +390,14 @@ impl GpuReplayBuffer { Ok(()) } - /// Proportional PER sampling via cumulative sum + binary search. + /// Proportional PER sampling via cumulative sum + GPU searchsorted. /// /// Returns a `GpuBatch` with all tensors on the same device as the buffer. /// IS weights are computed as `(N * p_i / sum)^(-beta) / max_weight`. /// - /// Binary search requires CPU access to the cumsum, so cumsum and targets - /// are pulled to CPU (read-only, no write-back). All other ops are GPU-native. + /// On CUDA: binary search runs entirely on GPU via a custom CUDA kernel + /// (`SearchSorted` `CustomOp2`). Zero CPU-GPU data transfer for the search. + /// On CPU: falls back to Rust's `binary_search_by` (for testing). pub fn sample_proportional(&self, batch_size: usize) -> Result { if !self.can_sample(batch_size) { return Err(MLError::ModelError(format!( @@ -270,31 +419,17 @@ impl GpuReplayBuffer { if total_sum <= 0.0 || !total_sum.is_finite() { return Err(MLError::ModelError( - "Priority sum is zero or non-finite".to_string(), + "Priority sum is zero or non-finite".to_owned(), )); } // Generate uniform random targets in [0, total_sum) let targets = Tensor::rand(0.0_f32, total_sum, &[batch_size], &self.device)?; - // Binary search: cumsum + targets pulled to CPU (read-only, no write-back). - // Candle lacks a GPU searchsorted kernel, so CPU binary search is required. - let cumsum_vec = cumsum.to_vec1::()?; - let targets_vec = targets.to_vec1::()?; - - let mut indices_vec: Vec = Vec::with_capacity(batch_size); - for &target in &targets_vec { - let idx = match cumsum_vec.binary_search_by(|v| { - v.partial_cmp(&target).unwrap_or(std::cmp::Ordering::Less) - }) { - Ok(i) => i, - Err(i) => i.min(n - 1), - }; - indices_vec.push(idx as i64); - } - - // Single CPU→GPU transfer for indices; u32 cast on GPU (avoids second from_vec) - let indices_for_select = Tensor::from_vec(indices_vec, &[batch_size], &self.device)?; + // GPU searchsorted: dispatches to CUDA kernel or CPU fallback automatically. + // On CUDA: zero CPU-GPU data transfer (eliminates ~400KB cumsum download). + // On CPU: equivalent to previous binary_search_by implementation. + let indices_for_select = cumsum.apply_op2_no_bwd(&targets, &SearchSorted { n })?; let indices = indices_for_select.to_dtype(DType::U32)?; // Gather experiences at sampled indices @@ -311,13 +446,13 @@ impl GpuReplayBuffer { let probs = (sampled_prios.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_t)?; let neg_beta = Tensor::full(-self.current_beta(), &[batch_size], &self.device)?; - let weights = probs.pow(&neg_beta)?; + let raw_weights = probs.pow(&neg_beta)?; // Normalize weights by max - let max_weight = weights.max(0)?; + let max_weight = raw_weights.max(0)?; let max_val = max_weight.to_vec0::()?; - let weights = if max_val > 0.0 { - weights.broadcast_div(&max_weight)? + let norm_weights = if max_val > 0.0 { + raw_weights.broadcast_div(&max_weight)? } else { Tensor::ones(&[batch_size], DType::F32, &self.device)? }; @@ -328,7 +463,7 @@ impl GpuReplayBuffer { rewards, next_states, dones, - weights, + weights: norm_weights, indices, }) } @@ -336,7 +471,7 @@ impl GpuReplayBuffer { /// Rank-based PER sampling via sort + rank probabilities. /// /// Sorts priorities descending, computes `P(i) = 1/rank(i)^alpha`, - /// normalizes, then samples via cumulative sum + binary search. + /// normalizes, then samples via cumulative sum + GPU searchsorted. /// Rank tensor generated on-device via arange; i64→u32 cast on GPU. pub fn sample_rank_based(&self, batch_size: usize) -> Result { if !self.can_sample(batch_size) { @@ -352,12 +487,12 @@ impl GpuReplayBuffer { let active_prios = self.priorities.narrow(0, 0, n)?; // Reshape to [1, n] for sort_last_dim, then squeeze back let prios_2d = active_prios.reshape(&[1, n])?; - let (_sorted, sort_indices) = prios_2d.sort_last_dim(false)?; // descending - let sort_indices = sort_indices.squeeze(0)?.to_dtype(DType::I64)?; // [n] as i64 + let (_sorted, sort_indices_2d) = prios_2d.sort_last_dim(false)?; // descending + let sort_indices = sort_indices_2d.squeeze(0)?.to_dtype(DType::I64)?; // [n] as i64 // Rank probabilities: P(rank) = 1/rank^alpha, rank 1-based // GPU-native: arange on device instead of CPU Vec allocation - let ranks_tensor = Tensor::arange(1u32, (n + 1) as u32, &self.device)? + let ranks_tensor = Tensor::arange(1_u32, (n + 1) as u32, &self.device)? .to_dtype(DType::F32)?; let alpha_tensor = Tensor::full(self.config.alpha, &[n], &self.device)?; let rank_probs = ranks_tensor.pow(&alpha_tensor)?.recip()?; // 1/rank^alpha @@ -366,24 +501,11 @@ impl GpuReplayBuffer { let cumsum = rank_probs.cumsum(0)?; let total_sum = cumsum.narrow(0, n - 1, 1)?.squeeze(0)?.to_vec0::()?; - // Binary search: cumsum + targets pulled to CPU (read-only, no write-back) + // GPU searchsorted for rank index lookup let targets = Tensor::rand(0.0_f32, total_sum, &[batch_size], &self.device)?; - let cumsum_vec = cumsum.to_vec1::()?; - let targets_vec = targets.to_vec1::()?; - - let mut rank_indices_vec: Vec = Vec::with_capacity(batch_size); - for &target in &targets_vec { - let idx = match cumsum_vec.binary_search_by(|v| { - v.partial_cmp(&target).unwrap_or(std::cmp::Ordering::Less) - }) { - Ok(i) => i, - Err(i) => i.min(n - 1), - }; - rank_indices_vec.push(idx as i64); - } + let rank_idx_tensor = cumsum.apply_op2_no_bwd(&targets, &SearchSorted { n })?; // Map rank indices back to original buffer indices via sort_indices - let rank_idx_tensor = Tensor::from_vec(rank_indices_vec, &[batch_size], &self.device)?; let original_indices = sort_indices.index_select(&rank_idx_tensor, 0)?; // GPU-native i64→u32 cast (eliminates GPU→CPU→GPU roundtrip for type conversion) @@ -403,12 +525,12 @@ impl GpuReplayBuffer { let probs = (sampled_rank_probs.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_t)?; let neg_beta = Tensor::full(-self.current_beta(), &[batch_size], &self.device)?; - let weights = probs.pow(&neg_beta)?; + let raw_weights = probs.pow(&neg_beta)?; - let max_weight = weights.max(0)?; + let max_weight = raw_weights.max(0)?; let max_val = max_weight.to_vec0::()?; - let weights = if max_val > 0.0 { - weights.broadcast_div(&max_weight)? + let norm_weights = if max_val > 0.0 { + raw_weights.broadcast_div(&max_weight)? } else { Tensor::ones(&[batch_size], DType::F32, &self.device)? }; @@ -419,7 +541,7 @@ impl GpuReplayBuffer { rewards, next_states, dones, - weights, + weights: norm_weights, indices: indices_u32, }) } @@ -427,9 +549,9 @@ impl GpuReplayBuffer { /// Update priorities from GPU-resident TD errors. /// /// Computes `new_priority = |td_error|^alpha + epsilon` entirely via tensor ops. - /// Uses scatter_add with a delta trick to write new priorities without any - /// capacity-sized CPU roundtrip: gather old values, compute delta, scatter_add. - /// Only a single scalar (max_priority) is read back to CPU. + /// Uses `scatter_add` with a delta trick to write new priorities without any + /// capacity-sized CPU roundtrip: gather old values, compute delta, `scatter_add`. + /// Only a single scalar (`max_priority`) is read back to CPU. pub fn update_priorities_gpu( &mut self, indices: &Tensor, diff --git a/crates/ml-dqn/src/searchsorted_kernel.cu b/crates/ml-dqn/src/searchsorted_kernel.cu new file mode 100644 index 000000000..c967517e3 --- /dev/null +++ b/crates/ml-dqn/src/searchsorted_kernel.cu @@ -0,0 +1,39 @@ +// GPU searchsorted kernel for PER (Prioritized Experience Replay) sampling. +// +// Each thread performs binary search for one target value in a sorted +// cumulative-sum array. Replaces the CPU roundtrip in proportional and +// rank-based PER sampling: eliminates ~400KB DMA per sample call +// (100K cumsum download) and CPU binary search loop. +// +// Grid: ceil(batch_size / 256) blocks × 256 threads. + +extern "C" __global__ void searchsorted_kernel( + const float* __restrict__ cumsum, // [n] sorted cumulative sums (on GPU) + const float* __restrict__ targets, // [batch_size] random search targets (on GPU) + long long* __restrict__ out_indices, // [batch_size] output indices (i64 for Candle) + int n, // number of elements in cumsum + int batch_size // number of targets to search +) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= batch_size) return; + + float target = targets[tid]; + + // Binary search: find first i where cumsum[i] >= target (upper bound) + int lo = 0; + int hi = n; + while (lo < hi) { + int mid = lo + ((hi - lo) >> 1); + if (cumsum[mid] < target) { + lo = mid + 1; + } else { + hi = mid; + } + } + + // Clamp to valid range [0, n-1] + if (lo >= n) lo = n - 1; + if (lo < 0) lo = 0; + + out_indices[tid] = (long long)lo; +}