75f83888ca
Mixed-precision loss/grad kernels: - MSE + C51 loss: float softmax/projection/TD-error (prevents bf16 exp overflow) - MSE + C51 grad: float arithmetic + bf16 range clamp before atomicAdd - Shared memory: float (4 bytes/elem) for numerically stable reductions - Bias kernels: float add+clamp ±500 (prevents bf16 Inf cascade between layers) - Noisy bias kernel: same float clamping fast_isnan/fast_isinf (ROOT CAUSE FIX): - nvcc --use_fast_math implies --no-nans → isnan()/isinf() compiled to false - ALL NaN guards across ALL kernels were dead code - Added bit-pattern IEEE 754 checks to common_device_functions.cuh - Replaced isnan/isinf in 7 kernel files (21 occurrences) - ml-dqn build.rs: all kernels now get common header (no more standalone) f32 PER IS-weights: - GpuBatchSlices.weights: CudaSlice<u16> → CudaSlice<f32> - GpuBatch.weights: GpuTensor → CudaSlice<f32> - Loss/grad kernel signatures: const __nv_bfloat16* → const float* - Upload path: separate f32 memcpy instead of bf16 staging - Eliminates bf16 overflow in IS-weight storage CUTLASS padding: - pad32() helper: round up to next multiple of 32 - 6 value-logit buffers: pad32(num_atoms) (51 → 64) - 6 branch-logit buffers: +32*3 padding per branch 895/895 unit tests, 8/9 smoke tests pass. 50-epoch convergence: NaN at step ~100-200 — backward pass produces NaN gradients within the CUDA graph replay (same atomic execution as Adam). Root cause: bf16 backward GemmEx inputs can overflow. Needs mixed-precision backward pass (same pattern as loss kernels) or f32 gradient output buffers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1164 lines
55 KiB
Rust
1164 lines
55 KiB
Rust
#![allow(unsafe_code)] // Required for CUDA kernel launches
|
|
|
|
//! GPU-Resident Replay Buffer -- pure cudarc `CudaSlice` internals.
|
|
//!
|
|
//! PER sampling uses an O(log n) GPU segment tree instead of O(n) prefix sum.
|
|
//! The segment tree is a flat `CudaSlice<f32>` of size `2 * capacity_pow2`:
|
|
//! root at index 1, leaves at `capacity_pow2..2*capacity_pow2`. Internal
|
|
//! nodes store sums of children. Sampling is parallel root-to-leaf traversal
|
|
//! with Philox RNG (1 kernel launch), replacing the old 5-kernel pipeline
|
|
//! (pow_alpha + prefix_sum + threshold_gen + searchsorted + i64_cast).
|
|
//!
|
|
//! Output `GpuBatchSlices` wraps gathered data as raw `CudaSlice` buffers
|
|
//! for downstream neural network consumption.
|
|
|
|
use std::sync::{Arc, OnceLock};
|
|
|
|
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)]
|
|
pub struct GpuBatchSlices {
|
|
pub states: CudaSlice<u16>, // [batch_size * state_dim] bf16 on GPU
|
|
pub next_states: CudaSlice<u16>, // [batch_size * state_dim] bf16 on GPU
|
|
pub actions: CudaSlice<u32>, // [batch_size] u32 on GPU
|
|
pub rewards: CudaSlice<u16>, // [batch_size] bf16 on GPU
|
|
pub dones: CudaSlice<u16>, // [batch_size] bf16 on GPU (0.0/1.0)
|
|
pub weights: CudaSlice<f32>, // [batch_size] f32 on GPU (IS weights — f32 to avoid bf16 overflow → Inf → NaN)
|
|
pub indices: CudaSlice<u32>, // [batch_size] u32 on GPU (buffer indices)
|
|
/// Episode IDs for sampled transitions `[batch_size]` i32 on GPU.
|
|
/// Used by HER Future/Final strategies for episode-aware donor selection.
|
|
pub episode_ids: Option<CudaSlice<i32>>,
|
|
pub batch_size: usize,
|
|
pub state_dim: usize,
|
|
}
|
|
|
|
impl GpuBatchSlices {
|
|
/// Convert to `GpuBatch` (GpuTensor-based) via GPU-only bf16->f32 cast kernels.
|
|
///
|
|
/// Zero CPU downloads -- all type conversions happen on GPU via custom CUDA kernels.
|
|
pub fn into_gpu_batch(self, stream: &Arc<CudaStream>) -> Result<crate::replay_buffer_type::GpuBatch, MLError> {
|
|
let bs = self.batch_size;
|
|
let sd = self.state_dim;
|
|
|
|
let kernels = get_cast_kernels(stream)?;
|
|
|
|
// bf16 states -> f32 via GPU cast kernel (zero host roundtrip)
|
|
let states = bf16_slice_to_gpu_tensor_gpu(&self.states, vec![bs, sd], stream, kernels)?;
|
|
let next_states = bf16_slice_to_gpu_tensor_gpu(&self.next_states, vec![bs, sd], stream, kernels)?;
|
|
|
|
// u32 actions -> GpuTensor via GPU cast kernel (zero host roundtrip)
|
|
let actions = u32_slice_to_gpu_tensor_gpu(&self.actions, vec![bs], stream, kernels)?;
|
|
// u32 indices: DtoD clone (stays u32, no bf16 conversion)
|
|
let indices = {
|
|
let mut dst = stream.alloc_zeros::<u32>(bs)
|
|
.map_err(|e| MLError::ModelError(format!("indices clone alloc: {e}")))?;
|
|
let nbytes = bs * std::mem::size_of::<u32>();
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst.raw_ptr(), self.indices.raw_ptr(), nbytes, stream.cu_stream(),
|
|
).map_err(|e| MLError::ModelError(format!("indices DtoD: {e}")))?;
|
|
}
|
|
dst
|
|
};
|
|
|
|
// bf16 slices -> bf16 GpuTensor via DtoD reinterpret (zero cast needed, already bf16)
|
|
let rewards = bf16_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?;
|
|
let dones = bf16_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?;
|
|
|
|
// IS-weights stay as f32 — no bf16 conversion (bf16 overflows to Inf for large weights)
|
|
let weights = dtod_clone_f32(stream, &self.weights, bs, "w_f32")?;
|
|
|
|
Ok(crate::replay_buffer_type::GpuBatch {
|
|
states,
|
|
actions,
|
|
rewards,
|
|
next_states,
|
|
dones,
|
|
weights,
|
|
indices,
|
|
episode_ids: self.episode_ids,
|
|
})
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Cast kernel cache for bf16->f32 and u32->f32 GPU-only conversions
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct CastKernels {
|
|
bf16_to_f32: CudaFunction,
|
|
u32_to_f32: CudaFunction,
|
|
f32_to_u32: CudaFunction,
|
|
f32_to_bf16: CudaFunction,
|
|
}
|
|
|
|
static CAST_KERNELS: OnceLock<Result<CastKernels, String>> = OnceLock::new();
|
|
|
|
fn get_cast_kernels(stream: &Arc<CudaStream>) -> Result<&'static CastKernels, MLError> {
|
|
let result = CAST_KERNELS.get_or_init(|| {
|
|
static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cast_kernels.cubin"));
|
|
let ctx = stream.context();
|
|
let module = ctx.load_cubin(CUBIN.to_vec())
|
|
.map_err(|e| format!("cast cubin load: {e}"))?;
|
|
let ld = |n: &str| -> Result<CudaFunction, String> {
|
|
module.load_function(n).map_err(|e| format!("{n}: {e}"))
|
|
};
|
|
Ok(CastKernels {
|
|
bf16_to_f32: ld("bf16_to_f32_cast")?,
|
|
u32_to_f32: ld("u32_to_f32_cast")?,
|
|
f32_to_u32: ld("f32_idx_to_u32")?,
|
|
f32_to_bf16: ld("f32_to_bf16_cast")?,
|
|
})
|
|
});
|
|
match result {
|
|
Ok(k) => Ok(k),
|
|
Err(e) => Err(MLError::ModelError(e.clone())),
|
|
}
|
|
}
|
|
|
|
/// Public API for f32->u32 GPU cast (used by `replay_buffer_type` for index conversion).
|
|
pub struct F32ToU32Caster {
|
|
kernel: &'static CudaFunction,
|
|
bf16_to_f32_kernel: &'static CudaFunction,
|
|
}
|
|
|
|
impl F32ToU32Caster {
|
|
/// Cast bf16-encoded indices to u32 on GPU (zero CPU download).
|
|
///
|
|
/// Internally does bf16->f32->u32 via two GPU kernel launches.
|
|
pub fn cast_f32_to_u32(
|
|
&self,
|
|
src: &CudaSlice<half::bf16>,
|
|
n: usize,
|
|
stream: &Arc<CudaStream>,
|
|
) -> Result<CudaSlice<u32>, MLError> {
|
|
// Step 1: bf16 -> f32 via bf16_to_f32_cast kernel
|
|
let f32_buf = stream.alloc_zeros::<f32>(n)
|
|
.map_err(|e| MLError::ModelError(format!("bf16->f32 alloc: {e}")))?;
|
|
let ni = n as i32;
|
|
// SAFETY: f32_buf, src are valid device allocations of at least n elements.
|
|
unsafe {
|
|
stream.launch_builder(self.bf16_to_f32_kernel)
|
|
.arg(&f32_buf).arg(src).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("bf16_to_f32 kernel: {e}")))?;
|
|
}
|
|
// Step 2: f32 -> u32 via f32_idx_to_u32 kernel
|
|
let out = stream.alloc_zeros::<u32>(n)
|
|
.map_err(|e| MLError::ModelError(format!("f32->u32 alloc: {e}")))?;
|
|
// SAFETY: out, f32_buf are valid device allocations of at least n elements.
|
|
unsafe {
|
|
stream.launch_builder(self.kernel)
|
|
.arg(&out).arg(&f32_buf).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("f32_to_u32 kernel: {e}")))?;
|
|
}
|
|
Ok(out)
|
|
}
|
|
}
|
|
|
|
/// Get the f32->u32 GPU caster (compiles kernel on first call, cached thereafter).
|
|
pub fn get_cast_kernels_f32_to_u32(stream: &Arc<CudaStream>) -> Result<F32ToU32Caster, MLError> {
|
|
let kernels = get_cast_kernels(stream)?;
|
|
Ok(F32ToU32Caster { kernel: &kernels.f32_to_u32, bf16_to_f32_kernel: &kernels.bf16_to_f32 })
|
|
}
|
|
|
|
/// Convert a bf16 `CudaSlice<u16>` to `GpuTensor` via DtoD reinterpret (zero CPU download).
|
|
///
|
|
/// `half::bf16` has the same binary repr as `u16`, so this is a bitwise copy.
|
|
fn bf16_slice_to_gpu_tensor_gpu(
|
|
src: &CudaSlice<u16>,
|
|
shape: Vec<usize>,
|
|
stream: &Arc<CudaStream>,
|
|
_kernels: &CastKernels,
|
|
) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
|
|
let n = src.len();
|
|
// Allocate bf16 buffer and DtoD copy (u16 and bf16 are the same bits)
|
|
let mut out = stream.alloc_zeros::<half::bf16>(n)
|
|
.map_err(|e| MLError::ModelError(format!("bf16 reinterpret alloc: {e}")))?;
|
|
let num_bytes = n * std::mem::size_of::<u16>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = src.device_ptr(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
let dst_ptr = {
|
|
let (ptr, guard) = out.device_ptr_mut(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
// SAFETY: u16 and half::bf16 have identical binary representations.
|
|
// Both buffers are valid device allocations of at least num_bytes.
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
|
|
).map_err(|e| MLError::ModelError(format!("bf16 reinterpret dtod: {e}")))?;
|
|
}
|
|
GpuTensor::new(out, shape)
|
|
}
|
|
|
|
/// Convert a u32 `CudaSlice<u32>` to bf16 `GpuTensor` via GPU cast kernels (zero CPU download).
|
|
///
|
|
/// Two-step: u32 -> f32 -> bf16 via GPU kernels.
|
|
fn u32_slice_to_gpu_tensor_gpu(
|
|
src: &CudaSlice<u32>,
|
|
shape: Vec<usize>,
|
|
stream: &Arc<CudaStream>,
|
|
kernels: &CastKernels,
|
|
) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
|
|
let n = src.len();
|
|
let ni = n as i32;
|
|
// Step 1: u32 -> f32
|
|
let f32_buf = stream.alloc_zeros::<f32>(n)
|
|
.map_err(|e| MLError::ModelError(format!("u32->f32 alloc: {e}")))?;
|
|
// SAFETY: f32_buf, src are valid device allocations of at least n elements.
|
|
unsafe {
|
|
stream.launch_builder(&kernels.u32_to_f32)
|
|
.arg(&f32_buf).arg(src).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("u32_to_f32 kernel: {e}")))?;
|
|
}
|
|
// Step 2: f32 -> bf16
|
|
let out = stream.alloc_zeros::<u16>(n)
|
|
.map_err(|e| MLError::ModelError(format!("f32->bf16 alloc: {e}")))?;
|
|
// SAFETY: out, f32_buf are valid device allocations of at least n elements.
|
|
unsafe {
|
|
stream.launch_builder(&kernels.f32_to_bf16)
|
|
.arg(&out).arg(&f32_buf).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("f32_to_bf16 kernel: {e}")))?;
|
|
}
|
|
// Reinterpret u16 -> bf16 via DtoD copy
|
|
let mut bf16_buf = stream.alloc_zeros::<half::bf16>(n)
|
|
.map_err(|e| MLError::ModelError(format!("bf16 reinterpret alloc: {e}")))?;
|
|
let num_bytes = n * std::mem::size_of::<u16>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = out.device_ptr(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
let dst_ptr = {
|
|
let (ptr, guard) = bf16_buf.device_ptr_mut(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
// SAFETY: u16 and half::bf16 have identical binary representations.
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
|
|
).map_err(|e| MLError::ModelError(format!("u32->bf16 dtod: {e}")))?;
|
|
}
|
|
GpuTensor::new(bf16_buf, shape)
|
|
}
|
|
|
|
/// Convert a f32 `CudaSlice<f32>` to bf16 `GpuTensor` via GPU cast kernel (zero CPU download).
|
|
#[allow(dead_code)]
|
|
fn f32_slice_to_gpu_tensor_gpu(
|
|
src: &CudaSlice<f32>,
|
|
shape: Vec<usize>,
|
|
stream: &Arc<CudaStream>,
|
|
kernels: &CastKernels,
|
|
) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
|
|
use ml_core::cuda_autograd::GpuTensor;
|
|
|
|
let n = src.len();
|
|
let ni = n as i32;
|
|
// f32 -> bf16 (stored as u16)
|
|
let out_u16 = stream.alloc_zeros::<u16>(n)
|
|
.map_err(|e| MLError::ModelError(format!("f32->bf16 alloc: {e}")))?;
|
|
// SAFETY: out_u16, src are valid device allocations of at least n elements.
|
|
unsafe {
|
|
stream.launch_builder(&kernels.f32_to_bf16)
|
|
.arg(&out_u16).arg(src).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("f32_to_bf16 kernel: {e}")))?;
|
|
}
|
|
// Reinterpret u16 -> bf16 via DtoD copy
|
|
let mut bf16_buf = stream.alloc_zeros::<half::bf16>(n)
|
|
.map_err(|e| MLError::ModelError(format!("bf16 reinterpret alloc: {e}")))?;
|
|
let num_bytes = n * std::mem::size_of::<u16>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = out_u16.device_ptr(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
let dst_ptr = {
|
|
let (ptr, guard) = bf16_buf.device_ptr_mut(stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
// SAFETY: u16 and half::bf16 have identical binary representations.
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
|
|
).map_err(|e| MLError::ModelError(format!("f32->bf16 dtod: {e}")))?;
|
|
}
|
|
GpuTensor::new(bf16_buf, shape)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Compiled kernel cache
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct ReplayKernels {
|
|
scatter_insert_f32: CudaFunction,
|
|
scatter_insert_u32: CudaFunction,
|
|
scatter_insert_bf16: CudaFunction,
|
|
#[allow(dead_code)]
|
|
gather_f32: CudaFunction,
|
|
gather_u32: CudaFunction,
|
|
gather_bf16_rows: CudaFunction,
|
|
gather_bf16: CudaFunction,
|
|
is_weights_f32: CudaFunction,
|
|
normalize_weights_f32: CudaFunction,
|
|
fill_from_gpu_f32: CudaFunction,
|
|
reduce_max_f32: CudaFunction,
|
|
i64_to_u32: CudaFunction,
|
|
f32_to_bf16_cast: CudaFunction,
|
|
max_of_two_f32: CudaFunction,
|
|
// Segment tree kernels (replace prefix_sum + searchsorted + pow_alpha pipeline)
|
|
seg_tree_update: CudaFunction,
|
|
seg_tree_insert: CudaFunction,
|
|
seg_tree_sample: CudaFunction,
|
|
seg_tree_gather_prios: 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 segment tree kernels cubin
|
|
static ST_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/seg_tree_kernel.cubin"));
|
|
let st_mod = ctx.load_cubin(ST_CUBIN.to_vec())
|
|
.map_err(|e| MLError::ModelError(format!("seg_tree cubin load: {e}")))?;
|
|
|
|
Ok(Self {
|
|
scatter_insert_f32: ld("scatter_insert_f32")?,
|
|
scatter_insert_u32: ld("scatter_insert_u32")?,
|
|
scatter_insert_bf16: ld("scatter_insert_bf16")?,
|
|
gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?,
|
|
gather_bf16_rows: ld("gather_bf16_rows")?,
|
|
gather_bf16: ld("gather_bf16")?,
|
|
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")?,
|
|
f32_to_bf16_cast: ld("f32_to_bf16_cast")?,
|
|
max_of_two_f32: ld("max_of_two_f32")?,
|
|
seg_tree_update: st_mod.load_function("seg_tree_update")
|
|
.map_err(|e| MLError::ModelError(format!("st_update fn: {e}")))?,
|
|
seg_tree_insert: st_mod.load_function("seg_tree_insert")
|
|
.map_err(|e| MLError::ModelError(format!("st_insert fn: {e}")))?,
|
|
seg_tree_sample: st_mod.load_function("seg_tree_sample")
|
|
.map_err(|e| MLError::ModelError(format!("st_sample fn: {e}")))?,
|
|
seg_tree_gather_prios: st_mod.load_function("seg_tree_gather_prios")
|
|
.map_err(|e| MLError::ModelError(format!("st_gather fn: {e}")))?,
|
|
})
|
|
}
|
|
}
|
|
|
|
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<u16>, next_states: CudaSlice<u16>,
|
|
actions: CudaSlice<u32>, rewards: CudaSlice<u16>,
|
|
dones: CudaSlice<u16>, 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,
|
|
// Segment tree: flat array [2 * capacity_pow2], root at [1], leaves at [cap_pow2..2*cap_pow2]
|
|
seg_tree: CudaSlice<f32>,
|
|
capacity_pow2: usize,
|
|
// Pre-allocated PER sampling buffers (zero cuMemAlloc after warmup)
|
|
sample_indices_i64: CudaSlice<i64>,
|
|
sample_indices_u32: CudaSlice<u32>,
|
|
sample_states: CudaSlice<u16>,
|
|
sample_next_states: CudaSlice<u16>,
|
|
sample_actions: CudaSlice<u32>,
|
|
sample_rewards: CudaSlice<u16>,
|
|
sample_dones: CudaSlice<u16>,
|
|
sample_priorities: CudaSlice<f32>,
|
|
sample_weights: CudaSlice<f32>,
|
|
sample_max_weight: CudaSlice<f32>,
|
|
sample_episode_ids: CudaSlice<i32>,
|
|
total_sum_buf: CudaSlice<f32>,
|
|
rng_step: u32,
|
|
}
|
|
|
|
impl Drop for GpuReplayBuffer {
|
|
fn drop(&mut self) {
|
|
// Synchronize stream before CudaSlice fields drop.
|
|
// The segment tree, 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 = a16(stream, cap * sd, "s")?;
|
|
let ns = a16(stream, cap * sd, "ns")?;
|
|
let a = a32u(stream, cap, "a")?;
|
|
let r = a16(stream, cap, "r")?;
|
|
let d = a16(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}")))?;
|
|
|
|
// Segment tree: flat array [2 * capacity_pow2], all zeros initially.
|
|
// Leaves at [cap_pow2..2*cap_pow2], root at [1]. Power-of-2 for balanced tree.
|
|
let cap_pow2 = cap.next_power_of_two();
|
|
let seg = a32f(stream, 2 * cap_pow2, "seg_tree")?;
|
|
|
|
// 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 = a16(stream, mbs * sd, "s_states")?;
|
|
let sns = a16(stream, mbs * sd, "s_nstates")?;
|
|
let sa = a32u(stream, mbs, "s_act")?;
|
|
let sr = a16(stream, mbs, "s_rew")?;
|
|
let sdn = a16(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")?;
|
|
|
|
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,
|
|
seg_tree: seg, capacity_pow2: cap_pow2,
|
|
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,
|
|
rng_step: 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.seg_tree).map_err(|e| MLError::ModelError(format!("seg_tree 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 }
|
|
|
|
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) };
|
|
let el = eff * sd;
|
|
let mut b_s = a16(&self.stream, el, "ib_s")?;
|
|
let mut b_n = a16(&self.stream, el, "ib_n")?;
|
|
let ni = el as i32;
|
|
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..) };
|
|
// SAFETY: b_s, ss, b_n, sn are valid device allocations. Kernel indices are bounds-checked.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
|
|
.arg(&mut b_s).arg(&ss).arg(&ni).launch(lcfg(el))
|
|
.map_err(|e| MLError::ModelError(format!("cast s: {e}")))?;
|
|
}
|
|
// SAFETY: same as above for next-state cast.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
|
|
.arg(&mut b_n).arg(&sn).arg(&ni).launch(lcfg(el))
|
|
.map_err(|e| MLError::ModelError(format!("cast n: {e}")))?;
|
|
}
|
|
let (ci, cpi, sdi, bsi) = (self.write_cursor as i32, cap as i32, sd as i32, eff as i32);
|
|
// SAFETY: states, b_s, next_states, b_n are valid device allocations. Indices within capacity.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.states).arg(&b_s).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
|
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s: {e}")))?;
|
|
}
|
|
// SAFETY: same as above for next-state scatter.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.next_states).arg(&b_n).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
|
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc n: {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..) };
|
|
// SAFETY: actions, rewards, dones, priorities are valid device allocations within capacity.
|
|
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}")))?;
|
|
}
|
|
// Cast f32 rewards/dones to bf16 (u16) then scatter insert as bf16 (state_dim=1)
|
|
let one_i = 1_i32;
|
|
let eff_i = eff as i32;
|
|
let mut b_r = a16(&self.stream, eff, "ib_r")?;
|
|
// SAFETY: b_r, sr are valid device allocations of at least eff elements.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
|
|
.arg(&mut b_r).arg(&sr).arg(&eff_i).launch(lcfg(eff))
|
|
.map_err(|e| MLError::ModelError(format!("cast r: {e}")))?;
|
|
}
|
|
// SAFETY: rewards buffer (bf16) valid, b_r cast above.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.rewards).arg(&b_r).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
|
|
}
|
|
let mut b_d = a16(&self.stream, eff, "ib_d")?;
|
|
// SAFETY: b_d, sd2 are valid device allocations of at least eff elements.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
|
|
.arg(&mut b_d).arg(&sd2).arg(&eff_i).launch(lcfg(eff))
|
|
.map_err(|e| MLError::ModelError(format!("cast d: {e}")))?;
|
|
}
|
|
// SAFETY: dones buffer (bf16) valid, b_d cast above.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.dones).arg(&b_d).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
|
|
}
|
|
let mut pt = a32f(&self.stream, eff, "pt")?;
|
|
// SAFETY: GPU-only broadcast of max_priority scalar (zero CPU readback).
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.fill_from_gpu_f32)
|
|
.arg(&mut pt).arg(&self.max_priority).arg(&bsi)
|
|
.launch(lcfg(eff))
|
|
.map_err(|e| MLError::ModelError(format!("fill mp: {e}")))?;
|
|
}
|
|
// SAFETY: priorities buffer valid, pt filled above.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
|
.arg(&self.priorities).arg(&pt).arg(&ci).arg(&cpi).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc p: {e}")))?;
|
|
}
|
|
// Update segment tree leaves for inserted indices.
|
|
// Build index array on CPU (small: batch_size <= 1024 typically) and upload.
|
|
let insert_indices: Vec<u32> = (0..eff)
|
|
.map(|j| ((self.write_cursor + j) % cap) as u32)
|
|
.collect();
|
|
let mut idx_buf = a32u(&self.stream, eff, "ib_idx")?;
|
|
self.stream.memcpy_htod(&insert_indices, &mut idx_buf)
|
|
.map_err(|e| MLError::ModelError(format!("ib idx htod: {e}")))?;
|
|
let al = self.config.alpha;
|
|
let cap_pow2_i = self.capacity_pow2 as i32;
|
|
// SAFETY: seg_tree, idx_buf, pt are valid. seg_tree has 2*cap_pow2 elements.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.seg_tree_insert)
|
|
.arg(&self.seg_tree).arg(&idx_buf).arg(&pt).arg(&al)
|
|
.arg(&cap_pow2_i).arg(&bsi)
|
|
.launch(lcfg(eff))
|
|
.map_err(|e| MLError::ModelError(format!("st insert: {e}")))?;
|
|
}
|
|
// Write episode IDs for inserted positions: simple sequential IDs (slot index).
|
|
// For flat replay buffers, each buffer slot is its own "episode" (L=1).
|
|
// HER Future/Final strategies use these to identify episode boundaries.
|
|
let ep_ids_host: Vec<i32> = (0..eff)
|
|
.map(|j| ((self.write_cursor + j) % cap) as i32)
|
|
.collect();
|
|
let mut ep_buf = a32i(&self.stream, eff, "ib_ep")?;
|
|
self.stream.memcpy_htod(&ep_ids_host, &mut ep_buf)
|
|
.map_err(|e| MLError::ModelError(format!("ep htod: {e}")))?;
|
|
// SAFETY: episode_ids and ep_buf are valid device allocations. Reinterpret i32 as u32 (same size).
|
|
unsafe {
|
|
let ep_dst = &*(&self.episode_ids as *const CudaSlice<i32> as *const CudaSlice<u32>);
|
|
let ep_src = &*(&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/rewards/dones are already `CudaSlice<half::bf16>`.
|
|
///
|
|
/// Skips the f32→bf16 cast for states (already bf16), and uses DtoD copy
|
|
/// for rewards/dones via scatter_insert_bf16 (reinterpreted as bf16 scatter).
|
|
pub fn insert_batch_bf16(
|
|
&mut self,
|
|
sf: &CudaSlice<half::bf16>,
|
|
nf: &CudaSlice<half::bf16>,
|
|
ac: &CudaSlice<u32>,
|
|
rw: &CudaSlice<half::bf16>,
|
|
dn: &CudaSlice<half::bf16>,
|
|
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) };
|
|
let el = eff * sd;
|
|
|
|
// States are already bf16 — skip cast, scatter directly
|
|
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);
|
|
// SAFETY: states, ss, next_states, sn are valid bf16 device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.states).arg(&ss).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
|
|
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s bf16: {e}")))?;
|
|
}
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.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 bf16: {e}")))?;
|
|
}
|
|
// Actions (u32) — same path
|
|
let sa = if off > 0 { ac.slice(off..) } else { ac.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}")))?;
|
|
}
|
|
// Rewards/dones: both input and storage are bf16 — scatter directly (state_dim=1)
|
|
{
|
|
let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
|
|
let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
|
|
let one_i = 1_i32;
|
|
// SAFETY: rewards/dones buffers are u16 (bf16), rw/dn CudaSlice<half::bf16> same repr.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r bf16: {e}")))?;
|
|
}
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
|
|
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d bf16: {e}")))?;
|
|
}
|
|
}
|
|
// Priorities — same as insert_batch
|
|
let mut pt = a32f(&self.stream, eff, "pt")?;
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.fill_from_gpu_f32)
|
|
.arg(&mut pt).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(&pt).arg(&ci).arg(&cpi).arg(&bsi)
|
|
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc p: {e}")))?;
|
|
}
|
|
let insert_indices: Vec<u32> = (0..eff)
|
|
.map(|j| ((self.write_cursor + j) % cap) as u32)
|
|
.collect();
|
|
let mut idx_buf = a32u(&self.stream, eff, "ib_idx")?;
|
|
self.stream.memcpy_htod(&insert_indices, &mut idx_buf)
|
|
.map_err(|e| MLError::ModelError(format!("ib idx htod: {e}")))?;
|
|
let al = self.config.alpha;
|
|
let cap_pow2_i = self.capacity_pow2 as i32;
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.seg_tree_insert)
|
|
.arg(&self.seg_tree).arg(&idx_buf).arg(&pt).arg(&al)
|
|
.arg(&cap_pow2_i).arg(&bsi)
|
|
.launch(lcfg(eff))
|
|
.map_err(|e| MLError::ModelError(format!("st insert: {e}")))?;
|
|
}
|
|
// Episode IDs
|
|
let ep_ids_host: Vec<i32> = (0..eff)
|
|
.map(|j| ((self.write_cursor + j) % cap) as i32)
|
|
.collect();
|
|
let mut ep_buf = a32i(&self.stream, eff, "ib_ep")?;
|
|
self.stream.memcpy_htod(&ep_ids_host, &mut 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 = &*(&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<GpuBatchSlices, 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;
|
|
let cap_pow2_i = self.capacity_pow2 as i32;
|
|
let buf_size_i = n as i32;
|
|
|
|
// Step 1: Segment tree sample — parallel root-to-leaf traversal with Philox RNG.
|
|
// O(log n) per sample instead of O(n) prefix sum + O(log n) binary search.
|
|
// Replaces: pow_alpha + prefix_sum + threshold_gen + searchsorted.
|
|
self.rng_step = self.rng_step.wrapping_add(1);
|
|
let seed = self.rng_step;
|
|
// SAFETY: seg_tree, sample_indices_i64 are valid device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.seg_tree_sample)
|
|
.arg(&self.seg_tree).arg(&mut self.sample_indices_i64)
|
|
.arg(&seed).arg(&cap_pow2_i).arg(&bsi).arg(&buf_size_i)
|
|
.launch(lcfg(batch_size))
|
|
.map_err(|e| MLError::ModelError(format!("st sample: {e}")))?;
|
|
}
|
|
|
|
// Step 2: 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.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.gather_bf16_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: {e}")))?;
|
|
}
|
|
// SAFETY: same as above for next_states gather.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.gather_bf16_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: {e}")))?;
|
|
}
|
|
// SAFETY: same context, actions buffer valid.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.gather_u32)
|
|
.arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi)
|
|
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?;
|
|
}
|
|
// SAFETY: same context, rewards buffer valid.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.gather_bf16)
|
|
.arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi)
|
|
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?;
|
|
}
|
|
// SAFETY: same context, dones buffer valid.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.gather_bf16)
|
|
.arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi)
|
|
.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>);
|
|
self.stream.launch_builder(&self.kernels.gather_u32)
|
|
.arg(ep_dst).arg(ep_src).arg(&self.sample_indices_i64).arg(&bsi)
|
|
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g ep: {e}")))?;
|
|
}
|
|
|
|
// Step 4: gather sampled priority^alpha from segment tree leaves for IS weights.
|
|
// SAFETY: sample_priorities, seg_tree, sample_indices_i64 are valid device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.seg_tree_gather_prios)
|
|
.arg(&mut self.sample_priorities).arg(&self.seg_tree)
|
|
.arg(&self.sample_indices_i64).arg(&cap_pow2_i).arg(&bsi)
|
|
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g sp: {e}")))?;
|
|
}
|
|
|
|
// Step 5: DtoD copy tree[1] (root = total sum) -> total_sum_buf for IS weights
|
|
{
|
|
let tree_root = self.seg_tree.slice(1..2);
|
|
let num_bytes = std::mem::size_of::<f32>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = tree_root.device_ptr(&self.stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
let dst_ptr = {
|
|
let (ptr, guard) = self.total_sum_buf.device_ptr_mut(&self.stream);
|
|
let _no_drop = std::mem::ManuallyDrop::new(guard);
|
|
ptr
|
|
};
|
|
// SAFETY: src_ptr and dst_ptr are valid device pointers. num_bytes = sizeof(f32).
|
|
unsafe {
|
|
cudarc::driver::result::memcpy_dtod_async(
|
|
dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
|
|
).map_err(|e| MLError::ModelError(format!("ts dtod: {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}")))?;
|
|
}
|
|
|
|
// Return DtoD clones of pre-allocated slices sized to actual batch_size.
|
|
// The caller owns the returned GpuBatchSlices (consumed by into_gpu_batch),
|
|
// so we DtoD-clone the relevant portions. All copies are async on the stream.
|
|
let ep_ids = dtod_clone_i32(&self.stream, &self.sample_episode_ids, batch_size, "o_ep")?;
|
|
|
|
// IS weights stay f32 — bf16 overflows to Inf for large weights, causing NaN loss.
|
|
let weights_f32 = dtod_clone_f32(&self.stream, &self.sample_weights, batch_size, "o_w_f32")?;
|
|
|
|
Ok(GpuBatchSlices {
|
|
states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?,
|
|
next_states: dtod_clone_u16(&self.stream, &self.sample_next_states, batch_size * sd, "o_n")?,
|
|
actions: dtod_clone_u32(&self.stream, &self.sample_actions, batch_size, "o_act")?,
|
|
rewards: dtod_clone_u16(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
|
|
dones: dtod_clone_u16(&self.stream, &self.sample_dones, batch_size, "o_d")?,
|
|
weights: weights_f32,
|
|
indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?,
|
|
episode_ids: Some(ep_ids),
|
|
batch_size,
|
|
state_dim: sd,
|
|
})
|
|
}
|
|
|
|
/// Update priorities from GPU-resident index and `td_error` `CudaSlices`.
|
|
///
|
|
/// Uses the segment tree update kernel: computes `(|td|^alpha + eps)`,
|
|
/// writes to priorities buffer AND tree leaves, propagates sums to root.
|
|
/// O(log n) per update instead of O(n) prefix sum rebuild.
|
|
pub fn update_priorities_gpu(
|
|
&mut self,
|
|
indices: &CudaSlice<u32>,
|
|
td_errors: &CudaSlice<half::bf16>,
|
|
bs: usize,
|
|
) -> Result<(), MLError> {
|
|
let _nvtx = NvtxRange::new("per_update_priorities");
|
|
if bs == 0 { return Ok(()); }
|
|
// Convert bf16 td_errors to f32 for the CUDA kernel (which reads float*)
|
|
let kernels = get_cast_kernels(&self.stream)?;
|
|
let td_errors_f32 = {
|
|
let n = bs;
|
|
let ni = n as i32;
|
|
let out = self.stream.alloc_zeros::<f32>(n)
|
|
.map_err(|e| MLError::ModelError(format!("td bf16->f32 alloc: {e}")))?;
|
|
// SAFETY: out, td_errors are valid device allocations of at least n elements.
|
|
unsafe {
|
|
self.stream.launch_builder(&kernels.bf16_to_f32)
|
|
.arg(&out).arg(td_errors).arg(&ni)
|
|
.launch(lcfg(n))
|
|
.map_err(|e| MLError::ModelError(format!("td bf16_to_f32: {e}")))?;
|
|
}
|
|
out
|
|
};
|
|
let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
|
|
let cap_pow2_i = self.capacity_pow2 as i32;
|
|
let mut bm = a32f(&self.stream, 1, "bm")?;
|
|
self.stream.memset_zeros(&mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?;
|
|
// seg_tree_update: writes priorities[], tree leaves, propagates, atomicMax.
|
|
// SAFETY: seg_tree, priorities, bm, indices, td_errors_f32 are valid device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.seg_tree_update)
|
|
.arg(&self.seg_tree).arg(&self.priorities).arg(&mut bm)
|
|
.arg(indices).arg(&td_errors_f32)
|
|
.arg(&al).arg(&ep).arg(&cap_pow2_i).arg(&bsi)
|
|
.launch(lcfg(bs)).map_err(|e| MLError::ModelError(format!("st update: {e}")))?;
|
|
}
|
|
self.pending_max_priority = Some(match self.pending_max_priority.take() {
|
|
Some(prev) => {
|
|
let mut r = a32f(&self.stream, 1, "mm")?;
|
|
// SAFETY: r, prev, bm are valid 1-element device allocations on same context.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.max_of_two_f32)
|
|
.arg(&mut r).arg(&prev).arg(&bm)
|
|
.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}")))?;
|
|
}
|
|
r
|
|
}
|
|
None => bm,
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
pub fn flush_max_priority(&mut self) -> Result<(), MLError> {
|
|
if let Some(pend) = self.pending_max_priority.take() {
|
|
// GPU-only max of pending vs current (zero CPU readback)
|
|
let mut result = a32f(&self.stream, 1, "fm")?;
|
|
// SAFETY: result, pend, max_priority are valid 1-element device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.max_of_two_f32)
|
|
.arg(&mut result).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}")))?;
|
|
}
|
|
// DtoD copy result into self.max_priority
|
|
let num_bytes = std::mem::size_of::<f32>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = result.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
|
|
};
|
|
// SAFETY: dst_ptr and src_ptr are valid device pointers. num_bytes = sizeof(f32).
|
|
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}")))?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn apply_max_priority_scalar(&mut self, mp: f32) -> Result<(), MLError> {
|
|
if mp > 0.0 {
|
|
// Upload CPU scalar to GPU, then GPU-only max comparison (zero download)
|
|
let mut mp_gpu = a32f(&self.stream, 1, "amp")?;
|
|
self.stream.memcpy_htod(&[mp], &mut mp_gpu)
|
|
.map_err(|e| MLError::ModelError(format!("amp htod: {e}")))?;
|
|
let mut result = a32f(&self.stream, 1, "amr")?;
|
|
// SAFETY: result, mp_gpu, max_priority are valid 1-element device allocations.
|
|
unsafe {
|
|
self.stream.launch_builder(&self.kernels.max_of_two_f32)
|
|
.arg(&mut result).arg(&mp_gpu).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}")))?;
|
|
}
|
|
// DtoD copy result into self.max_priority
|
|
let num_bytes = std::mem::size_of::<f32>();
|
|
let src_ptr = {
|
|
let (ptr, guard) = result.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
|
|
};
|
|
// SAFETY: dst_ptr and src_ptr are valid device pointers. num_bytes = sizeof(f32).
|
|
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<u16> { &self.states }
|
|
pub const fn next_states_slice(&self) -> &CudaSlice<u16> { &self.next_states }
|
|
pub const fn actions_slice(&self) -> &CudaSlice<u32> { &self.actions }
|
|
pub const fn rewards_slice(&self) -> &CudaSlice<u16> { &self.rewards }
|
|
pub const fn dones_slice(&self) -> &CudaSlice<u16> { &self.dones }
|
|
pub const fn priorities_slice(&self) -> &CudaSlice<f32> { &self.priorities }
|
|
|
|
/// Sample proportional indices and IS weights as GPU-resident `CudaSlices`.
|
|
///
|
|
/// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<f32>)` on GPU (weights are f32).
|
|
/// Callers process indices on GPU -- zero CPU download.
|
|
pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<f32>), MLError> {
|
|
let b = self.sample_proportional(bs)?;
|
|
Ok((b.indices, b.weights))
|
|
}
|
|
|
|
}
|
|
|
|
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 a16(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<u16>, MLError> {
|
|
s.alloc_zeros::<u16>(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}")))
|
|
}
|
|
|
|
/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
|
|
fn dtod_clone_f32(s: &Arc<CudaStream>, src: &CudaSlice<f32>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
|
|
let mut dst = a32f(s, n, nm)?;
|
|
let sv = src.slice(..n);
|
|
s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
|
|
Ok(dst)
|
|
}
|
|
|
|
/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
|
|
fn dtod_clone_u32(s: &Arc<CudaStream>, src: &CudaSlice<u32>, n: usize, nm: &str) -> Result<CudaSlice<u32>, MLError> {
|
|
let mut dst = a32u(s, n, nm)?;
|
|
let sv = src.slice(..n);
|
|
s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
|
|
Ok(dst)
|
|
}
|
|
|
|
/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
|
|
fn dtod_clone_i32(s: &Arc<CudaStream>, src: &CudaSlice<i32>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLError> {
|
|
let mut dst = a32i(s, n, nm)?;
|
|
let sv = src.slice(..n);
|
|
s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
|
|
Ok(dst)
|
|
}
|
|
|
|
/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
|
|
fn dtod_clone_u16(s: &Arc<CudaStream>, src: &CudaSlice<u16>, n: usize, nm: &str) -> Result<CudaSlice<u16>, MLError> {
|
|
let mut dst = a16(s, n, nm)?;
|
|
let sv = src.slice(..n);
|
|
s.memcpy_dtod(&sv, &mut dst).map_err(|e| MLError::ModelError(format!("dtod {nm}: {e}")))?;
|
|
Ok(dst)
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
}
|