29d58efca2
Eliminated 5 per-step GPU allocations: - insert_prio_buf, insert_idx_buf, insert_ep_buf: 3x cuMemAlloc per insert_batch/insert_batch_bf16 call (up to 4M elements each) - scratch_f32: reusable 1-element temp for flush_max_priority and apply_max_priority_scalar (was 2x a32f per call) - update_batch_spare: pre-allocated spare for None→Some priority swap, replenished by flush_max_priority (was 1x a32f per epoch start) cuMemAlloc is a synchronizing driver call that stalls the GPU pipeline. With 8192 batch_size and multiple inserts per epoch, this was measurable overhead on H100. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
930 lines
45 KiB
Rust
930 lines
45 KiB
Rust
#![allow(unsafe_code)] // Required for CUDA kernel launches
|
||
|
||
//! GPU-Resident Replay Buffer -- pure cudarc `CudaSlice` internals.
|
||
//!
|
||
//! PER sampling uses a prefix sum over `priorities_pa` (priority^alpha).
|
||
//! `per_prefix_scan` builds the prefix sum, `per_sample` does uniform
|
||
//! threshold search with Philox RNG.
|
||
//!
|
||
//! Output `GpuBatchPtrs` wraps gathered data as raw `CudaSlice` buffers
|
||
//! for downstream neural network consumption.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||
use ml_core::nvtx::NvtxRange;
|
||
use ml_core::MLError;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// GPU batch output (CudaSlice-based, no Candle Tensor dependency)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Pre-built GPU batch for training. All fields are raw `CudaSlice` on GPU.
|
||
#[allow(missing_debug_implementations)]
|
||
/// Zero-alloc batch: raw GPU pointers to pre-allocated sample buffers.
|
||
/// Valid until the next `sample_proportional` call (buffers are reused).
|
||
/// No cuMemAlloc, no DtoD clone — pure pointer pass-through.
|
||
pub struct GpuBatchPtrs {
|
||
pub states_ptr: u64,
|
||
pub next_states_ptr: u64,
|
||
pub actions_ptr: u64,
|
||
pub rewards_ptr: u64,
|
||
pub dones_ptr: u64,
|
||
pub weights_ptr: u64,
|
||
pub indices_ptr: u64,
|
||
pub episode_ids_ptr: u64,
|
||
pub batch_size: usize,
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Compiled kernel cache
|
||
// ---------------------------------------------------------------------------
|
||
|
||
struct ReplayKernels {
|
||
scatter_insert_f32: CudaFunction,
|
||
scatter_insert_u32: CudaFunction,
|
||
gather_f32: CudaFunction,
|
||
gather_u32: CudaFunction,
|
||
/// #30 F32 row gather for f32 state storage.
|
||
gather_f32_rows: CudaFunction,
|
||
is_weights_f32: CudaFunction,
|
||
normalize_weights_f32: CudaFunction,
|
||
fill_from_gpu_f32: CudaFunction,
|
||
reduce_max_f32: CudaFunction,
|
||
i64_to_u32: CudaFunction,
|
||
max_of_two_f32: CudaFunction,
|
||
// PER prefix sum kernels
|
||
per_update_pa: CudaFunction,
|
||
per_insert_pa: CudaFunction,
|
||
per_prefix_scan: CudaFunction,
|
||
per_sample: CudaFunction,
|
||
}
|
||
|
||
impl ReplayKernels {
|
||
fn compile(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
|
||
let ctx = stream.context();
|
||
|
||
// Load precompiled replay buffer kernels cubin
|
||
static RB_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/replay_buffer_kernels.cubin"));
|
||
let rb_mod = ctx.load_cubin(RB_CUBIN.to_vec())
|
||
.map_err(|e| MLError::ModelError(format!("rb cubin load: {e}")))?;
|
||
let ld = |n: &str| -> Result<CudaFunction, MLError> {
|
||
rb_mod.load_function(n).map_err(|e| MLError::ModelError(format!("load {n}: {e}")))
|
||
};
|
||
|
||
// Load precompiled PER prefix sum kernels cubin
|
||
static PER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_kernels.cubin"));
|
||
let per_mod = ctx.load_cubin(PER_CUBIN.to_vec())
|
||
.map_err(|e| MLError::ModelError(format!("per cubin load: {e}")))?;
|
||
let per_ld = |n: &str| -> Result<CudaFunction, MLError> {
|
||
per_mod.load_function(n).map_err(|e| MLError::ModelError(format!("per {n}: {e}")))
|
||
};
|
||
|
||
Ok(Self {
|
||
scatter_insert_f32: ld("scatter_insert_f32")?,
|
||
scatter_insert_u32: ld("scatter_insert_u32")?,
|
||
gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?,
|
||
gather_f32_rows: ld("gather_f32_rows")?,
|
||
is_weights_f32: ld("is_weights_f32")?,
|
||
normalize_weights_f32: ld("normalize_weights_f32")?,
|
||
fill_from_gpu_f32: ld("fill_from_gpu_f32")?,
|
||
reduce_max_f32: ld("reduce_max_f32")?,
|
||
i64_to_u32: ld("i64_to_u32")?,
|
||
max_of_two_f32: ld("max_of_two_f32")?,
|
||
per_update_pa: per_ld("per_update_pa")?,
|
||
per_insert_pa: per_ld("per_insert_pa")?,
|
||
per_prefix_scan: per_ld("per_prefix_scan")?,
|
||
per_sample: per_ld("per_sample")?,
|
||
})
|
||
}
|
||
}
|
||
|
||
fn lcfg(n: usize) -> LaunchConfig {
|
||
let t = 256_u32; let b = (n as u32).div_ceil(t);
|
||
LaunchConfig { grid_dim: (b.max(1), 1, 1), block_dim: (t, 1, 1), shared_mem_bytes: 0 }
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Config + Buffer
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[allow(clippy::module_name_repetitions)]
|
||
#[derive(Debug, Clone)]
|
||
pub struct GpuReplayBufferConfig {
|
||
pub capacity: usize, pub state_dim: usize, pub alpha: f32,
|
||
pub beta_start: f32, pub beta_max: f32, pub beta_annealing_steps: usize,
|
||
pub epsilon: f32, pub max_memory_bytes: usize,
|
||
/// Maximum batch size for pre-allocated sampling buffers. Defaults to 1024.
|
||
pub max_batch_size: usize,
|
||
}
|
||
|
||
pub struct GpuReplayBuffer {
|
||
config: GpuReplayBufferConfig,
|
||
stream: Arc<CudaStream>,
|
||
kernels: ReplayKernels,
|
||
states: CudaSlice<f32>, next_states: CudaSlice<f32>,
|
||
actions: CudaSlice<u32>, rewards: CudaSlice<f32>,
|
||
dones: CudaSlice<f32>, priorities: CudaSlice<f32>,
|
||
/// Episode IDs per buffer slot `[capacity]` i32 on GPU.
|
||
/// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`.
|
||
episode_ids: CudaSlice<i32>,
|
||
write_cursor: usize, size: usize,
|
||
max_priority: CudaSlice<f32>,
|
||
pending_max_priority: Option<CudaSlice<f32>>,
|
||
current_step: usize,
|
||
// PER prefix sum buffers
|
||
priorities_pa: CudaSlice<f32>,
|
||
prefix_sum: CudaSlice<f32>,
|
||
scan_tile_state: CudaSlice<u32>,
|
||
scan_tile_aggregate: CudaSlice<f32>,
|
||
scan_tile_prefix: CudaSlice<f32>,
|
||
scan_tile_counter: CudaSlice<u32>,
|
||
// Pre-allocated PER sampling buffers (zero cuMemAlloc after warmup)
|
||
sample_indices_i64: CudaSlice<i64>,
|
||
sample_indices_u32: CudaSlice<u32>,
|
||
sample_states: CudaSlice<f32>,
|
||
sample_next_states: CudaSlice<f32>,
|
||
sample_actions: CudaSlice<u32>,
|
||
sample_rewards: CudaSlice<f32>,
|
||
sample_dones: CudaSlice<f32>,
|
||
sample_priorities: CudaSlice<f32>,
|
||
sample_weights: CudaSlice<f32>,
|
||
sample_max_weight: CudaSlice<f32>,
|
||
sample_episode_ids: CudaSlice<i32>,
|
||
total_sum_buf: CudaSlice<f32>,
|
||
// Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step)
|
||
update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator
|
||
update_max_merge: CudaSlice<f32>, // [1] merge pending + batch max
|
||
update_batch_spare: Option<CudaSlice<f32>>, // [1] spare for None→Some swap (zero alloc)
|
||
// Pre-allocated insert scratch buffers (zero cuMemAlloc per insert)
|
||
insert_prio_buf: CudaSlice<f32>, // [insert_scratch_cap] priority broadcast
|
||
insert_idx_buf: CudaSlice<u32>, // [insert_scratch_cap] per_insert_pa indices
|
||
insert_ep_buf: CudaSlice<i32>, // [insert_scratch_cap] episode IDs
|
||
insert_scratch_cap: usize, // current allocation size
|
||
scratch_f32: CudaSlice<f32>, // [1] reusable temp for flush/apply_max
|
||
rng_step: u32,
|
||
/// Size of the most recent `sample_proportional` call. Used by `sample_weights_ref` /
|
||
/// `sample_indices_ref` to return correctly-sized views (pre-allocated buffers are
|
||
/// `max_batch_size` large; only the first `last_batch_size` elements are valid).
|
||
last_batch_size: usize,
|
||
}
|
||
|
||
impl Drop for GpuReplayBuffer {
|
||
fn drop(&mut self) {
|
||
// Synchronize stream before CudaSlice fields drop.
|
||
// PER sampling and priority update kernels may have pending work.
|
||
// Without sync, cuMemFree races with async kernel writes.
|
||
#[allow(unsafe_code)]
|
||
unsafe {
|
||
cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
|
||
}
|
||
}
|
||
}
|
||
|
||
impl GpuReplayBuffer {
|
||
pub fn new(config: GpuReplayBufferConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
|
||
let (cap, sd) = (config.capacity, config.state_dim);
|
||
let mbs = if config.max_batch_size == 0 { 1024 } else { config.max_batch_size };
|
||
let need = 2 * cap * sd * 2 + 5 * cap * 4;
|
||
if need > config.max_memory_bytes {
|
||
#[allow(clippy::integer_division)]
|
||
return Err(MLError::ModelError(format!(
|
||
"GPU replay buffer needs {} MB (limit {} MB)",
|
||
need / (1024 * 1024), config.max_memory_bytes / (1024 * 1024),
|
||
)));
|
||
}
|
||
let k = ReplayKernels::compile(stream)?;
|
||
let s = a32f(stream, cap * sd, "s")?;
|
||
let ns = a32f(stream, cap * sd, "ns")?;
|
||
let a = a32u(stream, cap, "a")?;
|
||
let r = a32f(stream, cap, "r")?;
|
||
let d = a32f(stream, cap, "d")?;
|
||
let p = a32f(stream, cap, "p")?;
|
||
let mut mp = a32f(stream, 1, "mp")?;
|
||
stream.memcpy_htod(&[1.0_f32], &mut mp).map_err(|e| MLError::ModelError(format!("mp: {e}")))?;
|
||
|
||
// PER prefix sum buffers
|
||
let priorities_pa = a32f(stream, cap, "priorities_pa")?;
|
||
let prefix_sum = a32f(stream, cap, "prefix_sum")?;
|
||
let tile_size = 256_usize;
|
||
let num_scan_tiles = cap.div_ceil(tile_size);
|
||
let scan_tile_state = stream.alloc_zeros::<u32>(num_scan_tiles)
|
||
.map_err(|e| MLError::ModelError(format!("alloc scan_tile_state: {e}")))?;
|
||
let scan_tile_aggregate = a32f(stream, num_scan_tiles, "scan_tile_agg")?;
|
||
let scan_tile_prefix = a32f(stream, num_scan_tiles, "scan_tile_pfx")?;
|
||
let scan_tile_counter = stream.alloc_zeros::<u32>(1)
|
||
.map_err(|e| MLError::ModelError(format!("alloc scan_tile_counter: {e}")))?;
|
||
|
||
// Pre-allocate PER sampling buffers (zero cuMemAlloc after warmup)
|
||
let si64 = stream.alloc_zeros::<i64>(mbs).map_err(|e| MLError::ModelError(format!("alloc s_i64: {e}")))?;
|
||
let su32 = a32u(stream, mbs, "s_idx")?;
|
||
let ss = a32f(stream, mbs * sd, "s_states")?;
|
||
let sns = a32f(stream, mbs * sd, "s_nstates")?;
|
||
let sa = a32u(stream, mbs, "s_act")?;
|
||
let sr = a32f(stream, mbs, "s_rew")?;
|
||
let sdn = a32f(stream, mbs, "s_done")?;
|
||
let sp = a32f(stream, mbs, "s_pri")?;
|
||
let sw = a32f(stream, mbs, "s_wt")?;
|
||
let smw = a32f(stream, 1, "s_mw")?;
|
||
let tsb = a32f(stream, 1, "ts_buf")?;
|
||
let ep_ids = a32i(stream, cap, "episode_ids")?;
|
||
let s_ep = a32i(stream, mbs, "s_episode_ids")?;
|
||
|
||
// Pre-allocate PER update scratch buffers (zero cuMemAlloc per step)
|
||
let update_batch_max = a32f(stream, 1, "update_batch_max")?;
|
||
let update_max_merge = a32f(stream, 1, "update_max_merge")?;
|
||
let update_batch_spare = a32f(stream, 1, "update_batch_spare")?;
|
||
|
||
// Pre-allocate insert scratch buffers (zero cuMemAlloc per insert)
|
||
// Initial size = max_batch_size; grows if insert_batch receives larger batches.
|
||
let isc = mbs;
|
||
let insert_prio_buf = a32f(stream, isc, "ins_prio")?;
|
||
let insert_idx_buf = a32u(stream, isc, "ins_idx")?;
|
||
let insert_ep_buf = a32i(stream, isc, "ins_ep")?;
|
||
let scratch_f32 = a32f(stream, 1, "scratch")?;
|
||
|
||
Ok(Self {
|
||
config, stream: Arc::clone(stream), kernels: k,
|
||
states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
|
||
episode_ids: ep_ids,
|
||
write_cursor: 0, size: 0, max_priority: mp,
|
||
pending_max_priority: None, current_step: 0,
|
||
priorities_pa, prefix_sum,
|
||
scan_tile_state, scan_tile_aggregate, scan_tile_prefix, scan_tile_counter,
|
||
sample_indices_i64: si64,
|
||
sample_indices_u32: su32, sample_states: ss,
|
||
sample_next_states: sns, sample_actions: sa,
|
||
sample_rewards: sr, sample_dones: sdn,
|
||
sample_priorities: sp, sample_weights: sw,
|
||
sample_episode_ids: s_ep,
|
||
sample_max_weight: smw, total_sum_buf: tsb,
|
||
update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare),
|
||
insert_prio_buf, insert_idx_buf, insert_ep_buf, insert_scratch_cap: isc,
|
||
scratch_f32,
|
||
rng_step: 0,
|
||
last_batch_size: 0,
|
||
})
|
||
}
|
||
|
||
pub const fn len(&self) -> usize { self.size }
|
||
pub const fn capacity(&self) -> usize { self.config.capacity }
|
||
pub const fn is_empty(&self) -> bool { self.size == 0 }
|
||
pub const fn can_sample(&self, bs: usize) -> bool { self.size >= bs }
|
||
pub fn current_beta(&self) -> f32 {
|
||
if self.config.beta_annealing_steps == 0 { return self.config.beta_max; }
|
||
let p = ((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) * p
|
||
}
|
||
pub const fn step(&mut self) { self.current_step = self.current_step.saturating_add(1); }
|
||
pub fn clear(&mut self) -> Result<(), MLError> {
|
||
self.write_cursor = 0; self.size = 0;
|
||
self.stream.memcpy_htod(&[1.0_f32], &mut self.max_priority).map_err(|e| MLError::ModelError(format!("{e}")))?;
|
||
self.stream.memset_zeros(&mut self.priorities_pa).map_err(|e| MLError::ModelError(format!("priorities_pa clear: {e}")))?;
|
||
self.stream.memset_zeros(&mut self.prefix_sum).map_err(|e| MLError::ModelError(format!("prefix_sum clear: {e}")))?;
|
||
self.pending_max_priority = None; self.current_step = 0; Ok(())
|
||
}
|
||
pub const fn stream(&self) -> &Arc<CudaStream> { &self.stream }
|
||
pub const fn alpha(&self) -> f32 { self.config.alpha }
|
||
pub const fn epsilon(&self) -> f32 { self.config.epsilon }
|
||
pub const fn state_dim(&self) -> usize { self.config.state_dim }
|
||
|
||
/// Grow insert scratch buffers if `eff` exceeds current allocation.
|
||
/// Only triggers cuMemAlloc on the first large insert — subsequent calls reuse.
|
||
fn ensure_insert_scratch(&mut self, eff: usize) -> Result<(), MLError> {
|
||
if eff <= self.insert_scratch_cap { return Ok(()); }
|
||
self.insert_prio_buf = a32f(&self.stream, eff, "ins_prio")?;
|
||
self.insert_idx_buf = a32u(&self.stream, eff, "ins_idx")?;
|
||
self.insert_ep_buf = a32i(&self.stream, eff, "ins_ep")?;
|
||
self.insert_scratch_cap = eff;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn insert_batch(&mut self, sf: &CudaSlice<f32>, nf: &CudaSlice<f32>,
|
||
ac: &CudaSlice<u32>, rw: &CudaSlice<f32>, dn: &CudaSlice<f32>, bs: usize,
|
||
) -> Result<(), MLError> {
|
||
if bs == 0 { return Ok(()); }
|
||
let (cap, sd) = (self.config.capacity, self.config.state_dim);
|
||
let (eff, off) = if bs > cap { (cap, bs - cap) } else { (bs, 0) };
|
||
self.ensure_insert_scratch(eff)?;
|
||
let el = eff * sd;
|
||
// #30: f32 states — scatter insert directly (no bf16 cast)
|
||
let ss = if off > 0 { sf.slice(off * sd..) } else { sf.slice(0..) };
|
||
let sn = if off > 0 { nf.slice(off * sd..) } else { nf.slice(0..) };
|
||
let (ci, cpi, sdi, bsi) = (self.write_cursor as i32, cap as i32, sd as i32, eff as i32);
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.states).arg(&ss).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s f32: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.next_states).arg(&sn).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc n f32: {e}")))?;
|
||
}
|
||
let sa = if off > 0 { ac.slice(off..) } else { ac.slice(0..) };
|
||
let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
|
||
let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_u32)
|
||
.arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
|
||
}
|
||
// Broadcast max_priority into pre-allocated prio buffer
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.fill_from_gpu_f32)
|
||
.arg(&mut self.insert_prio_buf).arg(&self.max_priority).arg(&bsi)
|
||
.launch(lcfg(eff))
|
||
.map_err(|e| MLError::ModelError(format!("fill mp: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.priorities).arg(&self.insert_prio_buf).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc p: {e}")))?;
|
||
}
|
||
// Write priority^alpha to priorities_pa using pre-allocated index buffer
|
||
let insert_indices: Vec<u32> = (0..eff)
|
||
.map(|j| ((self.write_cursor + j) % cap) as u32)
|
||
.collect();
|
||
self.stream.memcpy_htod(&insert_indices, &mut self.insert_idx_buf)
|
||
.map_err(|e| MLError::ModelError(format!("ib idx htod: {e}")))?;
|
||
let al = self.config.alpha;
|
||
let cap_i = self.config.capacity as i32;
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.per_insert_pa)
|
||
.arg(&self.priorities_pa).arg(&self.insert_idx_buf).arg(&self.insert_prio_buf).arg(&al)
|
||
.arg(&cap_i).arg(&bsi)
|
||
.launch(lcfg(eff))
|
||
.map_err(|e| MLError::ModelError(format!("per_insert_pa: {e}")))?;
|
||
}
|
||
// Write episode IDs using pre-allocated buffer
|
||
let ep_ids_host: Vec<i32> = (0..eff)
|
||
.map(|j| ((self.write_cursor + j) % cap) as i32)
|
||
.collect();
|
||
self.stream.memcpy_htod(&ep_ids_host, &mut self.insert_ep_buf)
|
||
.map_err(|e| MLError::ModelError(format!("ep htod: {e}")))?;
|
||
unsafe {
|
||
let ep_dst = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
||
let ep_src = &*(&self.insert_ep_buf as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_u32)
|
||
.arg(ep_dst).arg(ep_src).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc ep: {e}")))?;
|
||
}
|
||
|
||
self.write_cursor = (self.write_cursor + eff) % cap;
|
||
self.size = (self.size + eff).min(cap);
|
||
Ok(())
|
||
}
|
||
|
||
/// Insert a batch where states are already `CudaSlice<half::bf16>` and
|
||
/// rewards/dones are `CudaSlice<f32>` (no bf16 NaN risk).
|
||
///
|
||
/// Skips the f32→bf16 cast for states (already bf16). Rewards/dones
|
||
/// scatter-insert as f32 directly.
|
||
/// #30: Accept f32 states from the experience collector.
|
||
/// States are stored as f32 in the replay buffer (no bf16 truncation).
|
||
pub fn insert_batch_bf16(
|
||
&mut self,
|
||
sf: &CudaSlice<f32>, // f32 states from experience collector
|
||
nf: &CudaSlice<f32>, // f32 next_states
|
||
ac: &CudaSlice<u32>,
|
||
rw: &CudaSlice<f32>,
|
||
dn: &CudaSlice<f32>,
|
||
bs: usize,
|
||
) -> Result<(), MLError> {
|
||
if bs == 0 { return Ok(()); }
|
||
let (cap, sd) = (self.config.capacity, self.config.state_dim);
|
||
let (eff, off) = if bs > cap { (cap, bs - cap) } else { (bs, 0) };
|
||
self.ensure_insert_scratch(eff)?;
|
||
let el = eff * sd;
|
||
|
||
// #30: f32 states — scatter insert directly (full precision storage)
|
||
let ss = if off > 0 { sf.slice(off * sd..) } else { sf.slice(0..) };
|
||
let sn = if off > 0 { nf.slice(off * sd..) } else { nf.slice(0..) };
|
||
let (ci, cpi, sdi, bsi) = (self.write_cursor as i32, cap as i32, sd as i32, eff as i32);
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.states).arg(&ss).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s f32: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.next_states).arg(&sn).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc n f32: {e}")))?;
|
||
}
|
||
let sa = if off > 0 { ac.slice(off..) } else { ac.slice(0..) };
|
||
{
|
||
let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
|
||
let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_u32)
|
||
.arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r f32: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d f32: {e}")))?;
|
||
}
|
||
}
|
||
// Broadcast max_priority into pre-allocated prio buffer
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.fill_from_gpu_f32)
|
||
.arg(&mut self.insert_prio_buf).arg(&self.max_priority).arg(&bsi)
|
||
.launch(lcfg(eff))
|
||
.map_err(|e| MLError::ModelError(format!("fill mp: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||
.arg(&self.priorities).arg(&self.insert_prio_buf).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc p: {e}")))?;
|
||
}
|
||
// Write priority^alpha to priorities_pa using pre-allocated buffers
|
||
let insert_indices: Vec<u32> = (0..eff)
|
||
.map(|j| ((self.write_cursor + j) % cap) as u32)
|
||
.collect();
|
||
self.stream.memcpy_htod(&insert_indices, &mut self.insert_idx_buf)
|
||
.map_err(|e| MLError::ModelError(format!("ib idx htod: {e}")))?;
|
||
let al = self.config.alpha;
|
||
let cap_i = self.config.capacity as i32;
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.per_insert_pa)
|
||
.arg(&self.priorities_pa).arg(&self.insert_idx_buf).arg(&self.insert_prio_buf).arg(&al)
|
||
.arg(&cap_i).arg(&bsi)
|
||
.launch(lcfg(eff))
|
||
.map_err(|e| MLError::ModelError(format!("per_insert_pa: {e}")))?;
|
||
}
|
||
// Episode IDs using pre-allocated buffer
|
||
let ep_ids_host: Vec<i32> = (0..eff)
|
||
.map(|j| ((self.write_cursor + j) % cap) as i32)
|
||
.collect();
|
||
self.stream.memcpy_htod(&ep_ids_host, &mut self.insert_ep_buf)
|
||
.map_err(|e| MLError::ModelError(format!("ep htod: {e}")))?;
|
||
unsafe {
|
||
let ep_dst = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
||
let ep_src = &*(&self.insert_ep_buf as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
||
self.stream.launch_builder(&self.kernels.scatter_insert_u32)
|
||
.arg(ep_dst).arg(ep_src).arg(&ci).arg(&cpi).arg(&bsi)
|
||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc ep: {e}")))?;
|
||
}
|
||
|
||
self.write_cursor = (self.write_cursor + eff) % cap;
|
||
self.size = (self.size + eff).min(cap);
|
||
Ok(())
|
||
}
|
||
|
||
pub fn sample_proportional(&mut self, batch_size: usize) -> Result<GpuBatchPtrs, MLError> {
|
||
let _nvtx = NvtxRange::new("per_sample_proportional");
|
||
if !self.can_sample(batch_size) {
|
||
return Err(MLError::ModelError(format!("Cannot sample {batch_size} from {}", self.size)));
|
||
}
|
||
let mbs = if self.config.max_batch_size == 0 { 1024 } else { self.config.max_batch_size };
|
||
if batch_size > mbs {
|
||
return Err(MLError::ModelError(format!(
|
||
"batch_size {batch_size} exceeds max_batch_size {mbs}"
|
||
)));
|
||
}
|
||
let (n, sd) = (self.size, self.config.state_dim);
|
||
let nb = -self.current_beta();
|
||
let bsi = batch_size as i32;
|
||
|
||
// Step 1: Prefix scan of priorities_pa[0..size]
|
||
// Zero tile state + counter before each scan.
|
||
self.stream.memset_zeros(&mut self.scan_tile_state)
|
||
.map_err(|e| MLError::ModelError(format!("zero scan state: {e}")))?;
|
||
self.stream.memset_zeros(&mut self.scan_tile_counter)
|
||
.map_err(|e| MLError::ModelError(format!("zero scan counter: {e}")))?;
|
||
|
||
let size_i = n as i32;
|
||
let actual_tiles = n.div_ceil(256);
|
||
let num_tiles_i = actual_tiles as i32;
|
||
// Limit grid to max resident blocks to prevent spin-wait deadlock.
|
||
// Each block grabs tiles dynamically via atomicAdd on tile_counter.
|
||
let max_resident_blocks = 512_u32; // H100: 132 SMs × ~4 blocks/SM
|
||
let scan_blocks = (actual_tiles as u32).min(max_resident_blocks);
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.per_prefix_scan)
|
||
.arg(&self.priorities_pa)
|
||
.arg(&mut self.prefix_sum)
|
||
.arg(&self.scan_tile_state)
|
||
.arg(&self.scan_tile_aggregate)
|
||
.arg(&self.scan_tile_prefix)
|
||
.arg(&self.scan_tile_counter)
|
||
.arg(&size_i)
|
||
.arg(&num_tiles_i)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (scan_blocks.max(1), 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 256 * 4,
|
||
})
|
||
.map_err(|e| MLError::ModelError(format!("per_prefix_scan: {e}")))?;
|
||
}
|
||
|
||
// Step 2: Sample via binary search + gather priorities_pa
|
||
self.rng_step = self.rng_step.wrapping_add(1);
|
||
let seed = self.rng_step;
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.per_sample)
|
||
.arg(&self.prefix_sum)
|
||
.arg(&self.priorities_pa)
|
||
.arg(&mut self.sample_indices_i64)
|
||
.arg(&mut self.sample_priorities)
|
||
.arg(&mut self.total_sum_buf)
|
||
.arg(&seed)
|
||
.arg(&size_i)
|
||
.arg(&bsi)
|
||
.launch(lcfg(batch_size))
|
||
.map_err(|e| MLError::ModelError(format!("per_sample: {e}")))?;
|
||
}
|
||
|
||
// Step 3: i64 -> u32 indices for output
|
||
// SAFETY: sample_indices_u32 and sample_indices_i64 are valid device allocations.
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.i64_to_u32)
|
||
.arg(&mut self.sample_indices_u32).arg(&self.sample_indices_i64).arg(&bsi)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("i2u: {e}")))?;
|
||
}
|
||
|
||
// Step 3: gather into pre-allocated buffers
|
||
let sdi = sd as i32;
|
||
// SAFETY: sample_states, states, sample_indices_i64 are valid device allocations.
|
||
// #30: gather f32 rows (full-precision state storage)
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.gather_f32_rows)
|
||
.arg(&mut self.sample_states).arg(&self.states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g s f32: {e}")))?;
|
||
}
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.gather_f32_rows)
|
||
.arg(&mut self.sample_next_states).arg(&self.next_states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi)
|
||
.launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g n f32: {e}")))?;
|
||
}
|
||
// SAFETY: same context, actions buffer valid.
|
||
unsafe {
|
||
let cap_i32 = self.capacity() as i32;
|
||
self.stream.launch_builder(&self.kernels.gather_u32)
|
||
.arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?;
|
||
}
|
||
// SAFETY: same context, rewards buffer valid (f32).
|
||
unsafe {
|
||
let cap_i32 = self.capacity() as i32;
|
||
self.stream.launch_builder(&self.kernels.gather_f32)
|
||
.arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?;
|
||
}
|
||
// SAFETY: same context, dones buffer valid (f32).
|
||
unsafe {
|
||
let cap_i32 = self.capacity() as i32;
|
||
self.stream.launch_builder(&self.kernels.gather_f32)
|
||
.arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?;
|
||
}
|
||
|
||
// Step 3b: gather episode_ids for HER strategies.
|
||
// Reinterpret i32 buffers as u32 (same bit width) for the gather_u32 kernel.
|
||
// SAFETY: episode_ids and sample_episode_ids are CudaSlice<i32>, same layout as u32.
|
||
// The gather kernel copies raw bytes, so reinterpret is safe for same-size types.
|
||
unsafe {
|
||
let ep_src = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
||
let ep_dst = &mut *(&mut self.sample_episode_ids as *mut CudaSlice<i32> as *mut CudaSlice<u32>);
|
||
let cap_ep = self.capacity() as i32;
|
||
self.stream.launch_builder(&self.kernels.gather_u32)
|
||
.arg(ep_dst).arg(ep_src).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_ep)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g ep: {e}")))?;
|
||
}
|
||
|
||
// Step 6: IS weights via GPU-resident total_sum (zero CPU readback)
|
||
// is_weights_f32 reads total_sum from GPU pointer.
|
||
// SAFETY: sample_weights, sample_priorities, total_sum_buf are valid device allocations.
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.is_weights_f32)
|
||
.arg(&mut self.sample_weights).arg(&self.sample_priorities).arg(&self.total_sum_buf)
|
||
.arg(&nb).arg(&(n as i32)).arg(&bsi)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("isw: {e}")))?;
|
||
}
|
||
|
||
// Step 7: reduce max weight
|
||
self.stream.memset_zeros(&mut self.sample_max_weight)
|
||
.map_err(|e| MLError::ModelError(format!("mw zero: {e}")))?;
|
||
let rt = 256_u32.min(batch_size as u32).max(1);
|
||
let rb = (batch_size as u32).div_ceil(rt);
|
||
// SAFETY: sample_weights, sample_max_weight are valid device allocations.
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.reduce_max_f32)
|
||
.arg(&self.sample_weights).arg(&mut self.sample_max_weight).arg(&bsi)
|
||
.launch(LaunchConfig { grid_dim: (rb.max(1),1,1), block_dim: (rt,1,1), shared_mem_bytes: rt*4 })
|
||
.map_err(|e| MLError::ModelError(format!("rm: {e}")))?;
|
||
}
|
||
|
||
// Step 8: normalize weights by max
|
||
// SAFETY: sample_weights, sample_max_weight are valid. Normalization divides element-wise.
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.normalize_weights_f32)
|
||
.arg(&mut self.sample_weights).arg(&self.sample_max_weight).arg(&bsi)
|
||
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("nw: {e}")))?;
|
||
}
|
||
|
||
// Record the batch size so that sample_weights_ref / sample_indices_ref return
|
||
// correctly-sized views (the pre-allocated buffers are max_batch_size elements,
|
||
// but only the first `batch_size` are valid after this sample).
|
||
self.last_batch_size = batch_size;
|
||
|
||
// Zero-alloc: return raw pointers to pre-allocated sample buffers.
|
||
// Gather kernels already wrote into self.sample_* buffers above.
|
||
// No DtoD clone, no cuMemAlloc — pure GPU-resident references.
|
||
Ok(GpuBatchPtrs {
|
||
states_ptr: self.sample_states.raw_ptr(),
|
||
next_states_ptr: self.sample_next_states.raw_ptr(),
|
||
actions_ptr: self.sample_actions.raw_ptr(),
|
||
rewards_ptr: self.sample_rewards.raw_ptr(),
|
||
dones_ptr: self.sample_dones.raw_ptr(),
|
||
weights_ptr: self.sample_weights.raw_ptr(),
|
||
indices_ptr: self.sample_indices_u32.raw_ptr(),
|
||
episode_ids_ptr: self.sample_episode_ids.raw_ptr(),
|
||
batch_size,
|
||
state_dim: sd,
|
||
})
|
||
}
|
||
|
||
/// Reference to pre-allocated sample indices buffer (valid after sample_proportional).
|
||
///
|
||
/// Returns a view sized to the most recent `sample_proportional` batch size.
|
||
/// The underlying buffer is `max_batch_size` elements; only `last_batch_size` are valid.
|
||
pub fn sample_indices_ref(&self) -> cudarc::driver::CudaView<'_, u32> {
|
||
self.sample_indices_u32.slice(..self.last_batch_size)
|
||
}
|
||
|
||
/// Reference to pre-allocated sample episode IDs buffer (valid after sample_proportional).
|
||
pub fn sample_episode_ids_ref(&self) -> &CudaSlice<i32> { &self.sample_episode_ids }
|
||
|
||
/// Reference to pre-allocated sample weights buffer (valid after sample_proportional).
|
||
///
|
||
/// Returns a view sized to the most recent `sample_proportional` batch size.
|
||
/// The underlying buffer is `max_batch_size` elements; only `last_batch_size` are valid.
|
||
pub fn sample_weights_ref(&self) -> cudarc::driver::CudaView<'_, f32> {
|
||
self.sample_weights.slice(..self.last_batch_size)
|
||
}
|
||
|
||
pub fn update_priorities_gpu_raw(&mut self, _indices_ptr: u64, td_errors: &CudaSlice<half::bf16>, bs: usize, ext_stream: Option<&Arc<CudaStream>>) -> Result<(), MLError> {
|
||
let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice<u32>) };
|
||
self.update_priorities_gpu(idx_slice, td_errors, bs, ext_stream)
|
||
}
|
||
|
||
/// Update PER priorities from TD errors using the prefix-sum kernel.
|
||
///
|
||
/// `per_update_pa` reads bf16 td_errors directly (no cast needed),
|
||
/// computes `(|td|^alpha + eps)`, and writes to both `priorities` and
|
||
/// `priorities_pa`. Prefix sum is rebuilt lazily before sampling.
|
||
///
|
||
/// If `ext_stream` is provided, all GPU work runs on that stream instead of
|
||
/// the replay buffer's own stream. This avoids cross-stream synchronization
|
||
/// when the TD errors were computed on a different stream (e.g., the trainer's
|
||
/// main stream). All pre-allocated scratch buffers are used — zero per-step
|
||
/// GPU allocations.
|
||
pub fn update_priorities_gpu(
|
||
&mut self,
|
||
indices: &CudaSlice<u32>,
|
||
td_errors: &CudaSlice<half::bf16>,
|
||
bs: usize,
|
||
ext_stream: Option<&Arc<CudaStream>>,
|
||
) -> Result<(), MLError> {
|
||
if bs == 0 { return Ok(()); }
|
||
let stream_owned = ext_stream.cloned().unwrap_or_else(|| Arc::clone(&self.stream));
|
||
let stream = &stream_owned;
|
||
|
||
let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
|
||
let cap_i = self.config.capacity as i32;
|
||
|
||
stream.memset_zeros(&mut self.update_batch_max)
|
||
.map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
|
||
|
||
// per_update_pa: reads bf16 td_errors directly, writes priorities + priorities_pa
|
||
unsafe {
|
||
stream.launch_builder(&self.kernels.per_update_pa)
|
||
.arg(&self.priorities_pa)
|
||
.arg(&self.priorities)
|
||
.arg(&self.update_batch_max)
|
||
.arg(indices)
|
||
.arg(td_errors)
|
||
.arg(&al).arg(&ep).arg(&cap_i).arg(&bsi)
|
||
.launch(lcfg(bs))
|
||
.map_err(|e| MLError::ModelError(format!("per_update_pa: {e}")))?;
|
||
}
|
||
|
||
// Merge batch max into pending (unchanged logic)
|
||
match self.pending_max_priority.take() {
|
||
Some(prev) => {
|
||
unsafe {
|
||
stream.launch_builder(&self.kernels.max_of_two_f32)
|
||
.arg(&self.update_max_merge).arg(&prev).arg(&self.update_batch_max)
|
||
.launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
|
||
.map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?;
|
||
}
|
||
self.pending_max_priority = Some(std::mem::replace(&mut self.update_max_merge, prev));
|
||
}
|
||
None => {
|
||
// Rotate: batch_max → pending, spare → batch_max (zero alloc)
|
||
// spare was pre-allocated at construction; replenished by flush_max_priority.
|
||
if let Some(spare) = self.update_batch_spare.take() {
|
||
self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, spare));
|
||
} else {
|
||
// Fallback: first epoch before flush replenishes spare
|
||
let fresh = a32f(stream, 1, "bm_epoch")?;
|
||
self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, fresh));
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub fn flush_max_priority(&mut self) -> Result<(), MLError> {
|
||
if let Some(pend) = self.pending_max_priority.take() {
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.max_of_two_f32)
|
||
.arg(&mut self.scratch_f32).arg(&pend).arg(&self.max_priority)
|
||
.launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
|
||
.map_err(|e| MLError::ModelError(format!("flush max_of_two: {e}")))?;
|
||
}
|
||
let num_bytes = std::mem::size_of::<f32>();
|
||
let src_ptr = {
|
||
let (ptr, guard) = self.scratch_f32.device_ptr(&self.stream);
|
||
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
||
ptr
|
||
};
|
||
let dst_ptr = {
|
||
let (ptr, guard) = self.max_priority.device_ptr_mut(&self.stream);
|
||
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
||
ptr
|
||
};
|
||
unsafe {
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
|
||
).map_err(|e| MLError::ModelError(format!("flush dtod: {e}")))?;
|
||
}
|
||
// Replenish spare with the consumed pend buffer (zero alloc next epoch)
|
||
self.update_batch_spare = Some(pend);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub fn apply_max_priority_scalar(&mut self, mp: f32) -> Result<(), MLError> {
|
||
if mp > 0.0 {
|
||
self.stream.memcpy_htod(&[mp], &mut self.scratch_f32)
|
||
.map_err(|e| MLError::ModelError(format!("amp htod: {e}")))?;
|
||
// max(scratch_f32=mp, max_priority) → update_max_merge
|
||
unsafe {
|
||
self.stream.launch_builder(&self.kernels.max_of_two_f32)
|
||
.arg(&mut self.update_max_merge).arg(&self.scratch_f32).arg(&self.max_priority)
|
||
.launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
|
||
.map_err(|e| MLError::ModelError(format!("amp max: {e}")))?;
|
||
}
|
||
let num_bytes = std::mem::size_of::<f32>();
|
||
let src_ptr = {
|
||
let (ptr, guard) = self.update_max_merge.device_ptr(&self.stream);
|
||
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
||
ptr
|
||
};
|
||
let dst_ptr = {
|
||
let (ptr, guard) = self.max_priority.device_ptr_mut(&self.stream);
|
||
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
||
ptr
|
||
};
|
||
unsafe {
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
|
||
).map_err(|e| MLError::ModelError(format!("amp dtod: {e}")))?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Raw `CudaSlice` accessors for direct GPU consumption.
|
||
pub const fn states_slice(&self) -> &CudaSlice<f32> { &self.states }
|
||
pub const fn next_states_slice(&self) -> &CudaSlice<f32> { &self.next_states }
|
||
pub const fn actions_slice(&self) -> &CudaSlice<u32> { &self.actions }
|
||
pub const fn rewards_slice(&self) -> &CudaSlice<f32> { &self.rewards }
|
||
pub const fn dones_slice(&self) -> &CudaSlice<f32> { &self.dones }
|
||
pub const fn priorities_slice(&self) -> &CudaSlice<f32> { &self.priorities }
|
||
|
||
}
|
||
|
||
impl std::fmt::Debug for GpuReplayBuffer {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.debug_struct("GpuReplayBuffer")
|
||
.field("capacity", &self.config.capacity).field("size", &self.size)
|
||
.field("state_dim", &self.config.state_dim)
|
||
.field("write_cursor", &self.write_cursor).finish()
|
||
}
|
||
}
|
||
|
||
fn a32f(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
|
||
s.alloc_zeros::<f32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
|
||
}
|
||
fn a32u(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<u32>, MLError> {
|
||
s.alloc_zeros::<u32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
|
||
}
|
||
fn a32i(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLError> {
|
||
s.alloc_zeros::<i32>(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}")))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn make_stream() -> Arc<CudaStream> {
|
||
cudarc::driver::CudaContext::new(0)
|
||
.expect("CUDA required")
|
||
.new_stream()
|
||
.expect("CUDA stream")
|
||
}
|
||
|
||
#[test]
|
||
fn test_creation() {
|
||
let c = GpuReplayBufferConfig { capacity: 1000, state_dim: 48, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 100_000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 256 };
|
||
let b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
|
||
assert_eq!(b.len(), 0); assert_eq!(b.capacity(), 1000); assert!(b.is_empty());
|
||
}
|
||
#[test]
|
||
fn test_beta() {
|
||
let c = GpuReplayBufferConfig { capacity: 100, state_dim: 4, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 64 };
|
||
let mut b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
|
||
assert!((b.current_beta() - 0.4).abs() < 1e-6);
|
||
for _ in 0..500 { b.step(); }
|
||
assert!(b.current_beta() > 0.4 && b.current_beta() < 1.0);
|
||
for _ in 0..600 { b.step(); }
|
||
assert!((b.current_beta() - 1.0).abs() < 1e-6);
|
||
}
|
||
#[test]
|
||
fn test_clear() {
|
||
let c = GpuReplayBufferConfig { capacity: 100, state_dim: 4, alpha: 0.6, beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, max_memory_bytes: 4<<30, max_batch_size: 64 };
|
||
let mut b = GpuReplayBuffer::new(c, &make_stream()).expect("buf");
|
||
b.step(); b.clear().expect("clear");
|
||
assert_eq!(b.len(), 0); assert_eq!(b.current_step, 0);
|
||
}
|
||
|
||
/// Regression test: prefix scan must use `size` tiles, not `capacity` tiles.
|
||
/// With capacity=1M and size=256, the old code scanned 3907 tiles instead of 1.
|
||
/// On H100 with capacity=24.7M this caused a hang (96K tiles × decoupled lookback).
|
||
#[test]
|
||
#[ignore] // Requires CUDA GPU
|
||
fn test_per_prefix_scan_large_capacity() {
|
||
let stream = make_stream();
|
||
let cap = 1_000_000;
|
||
let sd = 4;
|
||
let bs = 64;
|
||
let c = GpuReplayBufferConfig {
|
||
capacity: cap, state_dim: sd, alpha: 0.6,
|
||
beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000,
|
||
epsilon: 1e-6, max_memory_bytes: 8 << 30, max_batch_size: bs,
|
||
};
|
||
let mut b = GpuReplayBuffer::new(c, &stream).expect("buf");
|
||
|
||
// Insert only 256 experiences into a 1M-slot buffer (ratio 3906:1)
|
||
let n = 256;
|
||
let s = stream.alloc_zeros::<f32>(n * sd).unwrap();
|
||
let ns = stream.alloc_zeros::<f32>(n * sd).unwrap();
|
||
let a = stream.alloc_zeros::<u32>(n).unwrap();
|
||
let mut r = stream.alloc_zeros::<f32>(n).unwrap();
|
||
let d = stream.alloc_zeros::<f32>(n).unwrap();
|
||
// Give non-zero rewards so priorities are non-zero after update
|
||
let rewards_host: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0) * 0.01).collect();
|
||
stream.memcpy_htod(&rewards_host, &mut r).unwrap();
|
||
b.insert_batch(&s, &ns, &a, &r, &d, n).unwrap();
|
||
assert_eq!(b.len(), n);
|
||
|
||
// sample_proportional triggers per_prefix_scan — must complete, not hang
|
||
let batch = b.sample_proportional(bs).expect("sample must not hang");
|
||
assert_eq!(batch.batch_size, bs);
|
||
|
||
// Verify all sampled indices are within [0, size)
|
||
let mut host_idx = vec![0u32; bs];
|
||
unsafe {
|
||
cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream());
|
||
}
|
||
stream.memcpy_dtoh(&b.sample_indices_u32.slice(..bs), &mut host_idx).unwrap();
|
||
for (i, &idx) in host_idx.iter().enumerate() {
|
||
assert!(
|
||
(idx as usize) < n,
|
||
"index {i} = {idx} exceeds buffer size {n}"
|
||
);
|
||
}
|
||
}
|
||
}
|