perf(dqn): GPU searchsorted kernel for PER sampling via CustomOp2
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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::<f32>()?;
|
||||
let targets = s2.as_slice::<f32>()?;
|
||||
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<i64> = 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<Result<cudarc::nvrtc::Ptx, String>> = 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<f32> from both storages (zero-copy)
|
||||
let cumsum_slice = s1.as_cuda_slice::<f32>()?;
|
||||
let cumsum_view = cumsum_slice.slice(l1.start_offset()..);
|
||||
let targets_slice = s2.as_cuda_slice::<f32>()?;
|
||||
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::<i64>(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<Self, MLError> {
|
||||
@@ -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<GpuBatch, MLError> {
|
||||
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::<f32>()?;
|
||||
let targets_vec = targets.to_vec1::<f32>()?;
|
||||
|
||||
let mut indices_vec: Vec<i64> = 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::<f32>()?;
|
||||
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<GpuBatch, MLError> {
|
||||
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::<f32>()?;
|
||||
|
||||
// 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::<f32>()?;
|
||||
let targets_vec = targets.to_vec1::<f32>()?;
|
||||
|
||||
let mut rank_indices_vec: Vec<i64> = 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::<f32>()?;
|
||||
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,
|
||||
|
||||
39
crates/ml-dqn/src/searchsorted_kernel.cu
Normal file
39
crates/ml-dqn/src/searchsorted_kernel.cu
Normal file
@@ -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;
|
||||
}
|
||||
Reference in New Issue
Block a user