diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 6f90d9358..186be0786 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -5,7 +5,7 @@ //! pre-uploaded buffers instead of creating per-step copies. //! //! All modules operate on raw `CudaSlice` via cudarc. -//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice` buffers. +//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice` buffers. use std::sync::Arc; use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; @@ -52,72 +52,74 @@ pub fn align_to_tensor_cores(dim: usize) -> usize { /// Estimate VRAM usage in bytes for pre-uploaded f32 data. pub fn estimate_vram_bytes(num_elements: usize) -> usize { - num_elements * std::mem::size_of::() + num_elements * std::mem::size_of::() } // --------------------------------------------------------------------------- -// BF16 ↔ F32 host-side conversion helpers +// F32 host ↔ device transfer helpers // --------------------------------------------------------------------------- -// All GPU buffers are CudaSlice. Host code works in f32. -// These helpers bridge the boundary at every HtoD / DtoH transfer. +// All GPU buffers are CudaSlice. Host code works in f32. +// Direct transfers — no bf16 conversion. -/// Convert an f32 slice to a Vec of bf16 for GPU upload. -#[inline] -pub fn f32_slice_to_bf16(src: &[f32]) -> Vec { - src.iter().map(|&x| half::bf16::from_f32(x)).collect() -} - -/// Upload f32 host data into an existing `CudaSlice`. -/// -/// Converts f32 → bf16 on the host, then calls `memcpy_htod`. -pub fn htod_f32_to_bf16( +/// Upload f32 host data into an existing `CudaSlice`. +pub fn htod_f32( stream: &Arc, src: &[f32], - dst: &mut CudaSlice, + dst: &mut CudaSlice, ) -> Result<(), MLError> { - let bf16_data = f32_slice_to_bf16(src); - stream.memcpy_htod(&bf16_data, dst) - .map_err(|e| MLError::ModelError(format!("htod_f32_to_bf16: {e}"))) + stream.memcpy_htod(src, dst) + .map_err(|e| MLError::ModelError(format!("htod_f32: {e}"))) } -/// Upload f32 host data to a new `CudaSlice` (clone_htod equivalent). -pub fn clone_htod_f32_to_bf16( +/// Upload f32 host data to a new `CudaSlice` (clone_htod). +pub fn clone_htod_f32( stream: &Arc, src: &[f32], -) -> Result, MLError> { - let bf16_data = f32_slice_to_bf16(src); - stream.clone_htod(&bf16_data) - .map_err(|e| MLError::ModelError(format!("clone_htod_f32_to_bf16: {e}"))) +) -> Result, MLError> { + stream.clone_htod(src) + .map_err(|e| MLError::ModelError(format!("clone_htod_f32: {e}"))) } -/// Download `CudaSlice` (or `CudaView`) to an f32 host buffer. +/// Download `CudaSlice` (or `CudaView`) to an f32 host buffer. /// -/// Allocates a temporary bf16 host buffer, downloads, then converts to f32. -pub fn dtoh_bf16_to_f32>( +/// Uses raw cuStreamSynchronize + cuMemcpyDtoH_v2 for CUDA-graph safety +/// (avoids cudarc event tracking which can deadlock inside captured graphs). +pub fn dtoh_f32>( stream: &Arc, src: &Src, dst: &mut [f32], ) -> Result<(), MLError> { let n = dst.len(); - let mut bf16_buf = vec![half::bf16::ZERO; n]; - let num_bytes = n * std::mem::size_of::(); + let num_bytes = n * std::mem::size_of::(); let (src_ptr, _guard) = src.device_ptr(stream); // Sync + raw DtoH to avoid cudarc event tracking (graph-safe) #[allow(unsafe_code)] unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); cudarc::driver::sys::cuMemcpyDtoH_v2( - bf16_buf.as_mut_ptr().cast(), + dst.as_mut_ptr().cast(), src_ptr, num_bytes, ); } - for (d, &s) in dst.iter_mut().zip(bf16_buf.iter()) { - *d = s.to_f32(); - } Ok(()) } +// Legacy aliases — delegate to the f32 functions above. +// Keeps the 188 existing call sites compiling during incremental migration. +#[inline] +pub fn htod_f32_to_bf16(stream: &Arc, src: &[f32], dst: &mut CudaSlice) -> Result<(), MLError> { + htod_f32(stream, src, dst) +} +#[inline] +pub fn clone_htod_f32_to_bf16(stream: &Arc, src: &[f32]) -> Result, MLError> { + clone_htod_f32(stream, src) +} +#[inline] +pub fn dtoh_bf16_to_f32>(stream: &Arc, src: &Src, dst: &mut [f32]) -> Result<(), MLError> { + dtoh_f32(stream, src, dst) +} + /// Compute optimal (grid_dim, block_dim) for a 1-D kernel launch. /// /// Selects the largest block size that divides 32 (warp size) and does not @@ -167,21 +169,21 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3 /// The fused experience collector kernel (4490 lines, branching+C51+NoisyNets) // compile_ptx_for_device re-export removed: all kernels now use precompiled cubins. -/// Clone a `CudaSlice` via DtoD memcpy. +/// Clone a `CudaSlice` via DtoD memcpy. /// /// The caller owns the returned slice independently of the source. pub fn clone_cuda_slice_f32( - src: &CudaSlice, + src: &CudaSlice, stream: &Arc, -) -> Result, crate::MLError> { +) -> Result, crate::MLError> { let n = src.len(); - let mut dst = stream.alloc_zeros::(n).map_err(|e| { + let mut dst = stream.alloc_zeros::(n).map_err(|e| { crate::MLError::ModelError(format!("clone_cuda_slice_f32 alloc: {e}")) })?; { let (src_ptr, _src_guard) = src.device_ptr(stream); let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream); - let num_bytes = n * std::mem::size_of::(); + let num_bytes = n * std::mem::size_of::(); #[allow(unsafe_code)] unsafe { cudarc::driver::result::memcpy_dtod_async( @@ -218,16 +220,16 @@ pub fn clone_cuda_slice_u32( /// Pre-uploaded GPU training data for DQN trainer. /// /// Holds market features [N * 42] and target prices [N * 4] as GPU-resident -/// `CudaSlice` buffers. Portfolio features (3 dims) are computed per-bar +/// `CudaSlice` buffers. Portfolio features (3 dims) are computed per-bar /// by the trainer and concatenated on-device. OFI features (8 dims from MBP-10 /// order book) are optionally uploaded and appended to the state vector. pub struct DqnGpuData { /// Market features [num_bars * 42] on GPU (f32, row-major) - pub features: CudaSlice, + pub features: CudaSlice, /// Target prices [num_bars * 4] on GPU (f32, row-major) - pub targets: CudaSlice, + pub targets: CudaSlice, /// OFI features [num_bars * 8] on GPU (f32), from MBP-10 data - pub ofi_features: Option>, + pub ofi_features: Option>, /// Number of training bars pub num_bars: usize, /// Feature dimension (42) @@ -415,19 +417,19 @@ impl DqnGpuData { /// D2D subrange copy: extract `count` f32 elements at `offset` into a new `CudaSlice`. pub(crate) fn d2d_subrange( - src: &CudaSlice, + src: &CudaSlice, offset: usize, count: usize, stream: &Arc, - ) -> Result, MLError> { - let mut dst = stream.alloc_zeros::(count).map_err(|e| { + ) -> Result, MLError> { + let mut dst = stream.alloc_zeros::(count).map_err(|e| { MLError::ModelError(format!("d2d_subrange alloc: {e}")) })?; { let src_view = src.slice(offset..offset + count); let (src_ptr, _sg) = src_view.device_ptr(stream); let (dst_ptr, _dg) = dst.device_ptr_mut(stream); - let nbytes = count * std::mem::size_of::(); + let nbytes = count * std::mem::size_of::(); #[allow(unsafe_code)] unsafe { cudarc::driver::result::memcpy_dtod_async( @@ -439,7 +441,7 @@ impl DqnGpuData { } /// Get market features for a single bar as a [42] `CudaSlice` (D2D copy from GPU buffer). - pub fn bar_features(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_features(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { if bar_idx >= self.num_bars { return Err(MLError::ModelError(format!("Bar {bar_idx} out of range (num_bars={})", self.num_bars))); } @@ -447,13 +449,13 @@ impl DqnGpuData { } /// Get market features for a batch of bars as a [batch_size * 42] `CudaSlice`. - pub fn batch_features(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_features(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); Self::d2d_subrange(&self.features, start * self.feature_dim, count * self.feature_dim, stream) } /// Get target prices for a single bar as a [4] `CudaSlice`. - pub fn bar_targets(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_targets(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { if bar_idx >= self.num_bars { return Err(MLError::ModelError(format!("Bar {bar_idx} target out of range (num_bars={})", self.num_bars))); } @@ -461,16 +463,16 @@ impl DqnGpuData { } /// Get target prices for a batch as [batch_size * 4] `CudaSlice`. - pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); Self::d2d_subrange(&self.targets, start * 4, count * 4, stream) } - /// Get target values for a bar as a GPU-resident [4] `CudaSlice`. + /// Get target values for a bar as a GPU-resident [4] `CudaSlice`. /// /// Tensor order: [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw]. /// Stays on GPU. Callers that need CPU scalars should extract at their own boundary. - pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { self.bar_targets(bar_idx, stream) } @@ -483,13 +485,13 @@ impl DqnGpuData { bar_idx: usize, portfolio_features: &[f32; 3], stream: &Arc, - ) -> Result, MLError> { + ) -> Result, MLError> { let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 }; let raw_dim = self.feature_dim + 3 + ofi_dim; let final_dim = self.aligned_state_dim.unwrap_or(raw_dim); // Allocate output buffer (zeros for padding region) - let mut dst = stream.alloc_zeros::(final_dim) + let mut dst = stream.alloc_zeros::(final_dim) .map_err(|e| MLError::ModelError(format!("build_state alloc: {e}")))?; // DtoD copy: market features -> dst[0..feature_dim] @@ -518,7 +520,7 @@ impl DqnGpuData { count: usize, portfolio_features: &[f32; 3], stream: &Arc, - ) -> Result, MLError> { + ) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); if count == 0 { return Err(MLError::ModelError("Empty batch for state construction".to_owned())); @@ -528,7 +530,7 @@ impl DqnGpuData { let final_dim = self.aligned_state_dim.unwrap_or(raw_dim); // Allocate output buffer (zeros for padding region) - let mut dst = stream.alloc_zeros::(count * final_dim) + let mut dst = stream.alloc_zeros::(count * final_dim) .map_err(|e| MLError::ModelError(format!("build_batch alloc: {e}")))?; // Upload portfolio features once (3 scalars) @@ -574,8 +576,8 @@ impl DqnGpuData { /// DtoD copy `count` f32 elements from `src` (offset 0) into `dst` at element offset. #[allow(unsafe_code)] fn dtod_copy_into( - src: &CudaSlice, - dst: &mut CudaSlice, + src: &CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, count: usize, stream: &Arc, @@ -589,9 +591,9 @@ impl DqnGpuData { /// buffer. The source is sliced via `CudaSlice::slice()` for safety. #[allow(unsafe_code)] fn dtod_copy_into_at_offset( - src: &CudaSlice, + src: &CudaSlice, src_offset_elems: usize, - dst: &mut CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, count: usize, stream: &Arc, @@ -599,12 +601,12 @@ impl DqnGpuData { if count == 0 { return Ok(()); } - let nbytes = count * std::mem::size_of::(); + let nbytes = count * std::mem::size_of::(); let src_view = src.slice(src_offset_elems..src_offset_elems + count); let (src_ptr, _sg) = src_view.device_ptr(stream); // Get base pointer of dst, then offset by dst_offset_elems * sizeof(f32) let (dst_base, _dg) = dst.device_ptr_mut(stream); - let dst_ptr = dst_base + (dst_offset_elems * std::mem::size_of::()) as u64; + let dst_ptr = dst_base + (dst_offset_elems * std::mem::size_of::()) as u64; unsafe { cudarc::driver::result::memcpy_dtod_async( dst_ptr, src_ptr, nbytes, stream.cu_stream(), @@ -620,7 +622,7 @@ impl DqnGpuData { #[allow(unsafe_code)] fn htod_copy_into( src: &[f32], - dst: &mut CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, stream: &Arc, ) -> Result<(), MLError> { @@ -709,7 +711,7 @@ impl GpuBufferPool { } } - // Upload the used slice to GPU via f32→bf16 conversion + clone_htod. + // Upload the used slice to GPU via clone_htod. let features = clone_htod_f32_to_bf16(stream, &self.feature_buf[..feat_len])?; let targets = clone_htod_f32_to_bf16(stream, &self.target_buf[..targ_len])?; @@ -731,10 +733,10 @@ impl GpuBufferPool { /// Pre-uploaded GPU training data for PPO trainer. /// -/// Holds all rollout states as a single [N * state_dim] `CudaSlice`. +/// Holds all rollout states as a single [N * state_dim] `CudaSlice`. pub struct PpoGpuData { /// All states [num_steps * state_dim] on GPU (f32, row-major) - pub states: CudaSlice, + pub states: CudaSlice, /// Number of training steps pub num_steps: usize, /// State dimension @@ -751,7 +753,7 @@ impl std::fmt::Debug for PpoGpuData { } impl PpoGpuData { - /// Upload PPO market data to GPU as a single contiguous `CudaSlice`. + /// Upload PPO market data to GPU as a single contiguous `CudaSlice`. pub fn upload(market_data: &[Vec], stream: &Arc) -> Result { let num_steps = market_data.len(); if num_steps == 0 { @@ -796,7 +798,7 @@ impl PpoGpuData { } /// Get a single step's state as a [state_dim] `CudaSlice` (D2D copy from GPU buffer). - pub fn step_state(&self, step_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn step_state(&self, step_idx: usize, stream: &Arc) -> Result, MLError> { if step_idx >= self.num_steps { return Err(MLError::ModelError(format!("Step {step_idx} out of range (num_steps={})", self.num_steps))); } @@ -804,7 +806,7 @@ impl PpoGpuData { } /// Get a batch of states as [batch_size * state_dim] `CudaSlice`. - pub fn batch_states(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_states(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_steps.saturating_sub(start)); DqnGpuData::d2d_subrange(&self.states, start * self.state_dim, count * self.state_dim, stream) } @@ -850,11 +852,10 @@ mod tests { let batch = gpu_data.batch_features(0, 10, &stream).expect("batch"); assert_eq!(batch.len(), 10 * 42); - // Download target values and verify on CPU (bf16 -> f32 conversion) + // Download target values and verify on CPU let targets_gpu = gpu_data.bar_target_values(0, &stream).expect("targets"); - let mut targets_bf16 = vec![half::bf16::ZERO; 4]; - stream.memcpy_dtoh(&targets_gpu, &mut targets_bf16).expect("DtoH"); // test readback - let targets_host: Vec = targets_bf16.iter().map(|v| v.to_f32()).collect(); + let mut targets_host = vec![0.0_f32; 4]; + dtoh_f32(&stream, &targets_gpu, &mut targets_host).expect("DtoH"); assert!((targets_host.first().copied().unwrap_or(0.0) - 100.0).abs() < 0.5); } @@ -962,10 +963,10 @@ mod tests { #[test] fn test_vram_estimation() { - // 300K bars × 55 features × 2 bytes (BF16) = ~33 MB + // 300K bars × 55 features × 4 bytes (f32) = ~66 MB let bytes = estimate_vram_bytes(300_000 * 55); - assert!(bytes < 50_000_000); // Under 50 MB - assert!(bytes > 30_000_000); // Over 30 MB + assert!(bytes < 100_000_000); // Under 100 MB + assert!(bytes > 60_000_000); // Over 60 MB } #[test] @@ -1001,13 +1002,11 @@ mod tests { let batch = gpu_data.build_batch_states(10, 32, &portfolio, &stream).expect("batch"); assert_eq!(batch.len(), 32 * 45); // 42 market + 3 portfolio - // Download and verify portfolio features in row 0 (bf16 -> f32 conversion) - let mut host_bf16 = vec![half::bf16::ZERO; 32 * 45]; - stream.memcpy_dtoh(&batch, &mut host_bf16).expect("DtoH"); // test readback - let host: Vec = host_bf16.iter().map(|v| v.to_f32()).collect(); + // Download and verify portfolio features in row 0 + let mut host = vec![0.0_f32; 32 * 45]; + dtoh_f32(&stream, &batch, &mut host).expect("DtoH"); // Row 0: features[42..45] should be portfolio [0.95, 0.5, 0.0001] - // BF16 precision: relax tolerance let pf = &host[42..45]; assert!((pf[0] - 0.95).abs() < 0.02, "portfolio[0]={}", pf[0]); assert!((pf[1] - 0.5).abs() < 0.02, "portfolio[1]={}", pf[1]); @@ -1072,11 +1071,10 @@ mod tests { let g2 = pool.upload_dqn(&data2, &stream).expect("upload2"); assert_eq!(g2.num_bars, 200); - // Download target values and verify (bf16 -> f32 conversion) + // Download target values and verify let t = g2.bar_target_values(0, &stream).expect("targets"); - let mut host_bf16 = vec![half::bf16::ZERO; 4]; - stream.memcpy_dtoh(&t, &mut host_bf16).expect("DtoH"); // test readback - let host: Vec = host_bf16.iter().map(|v| v.to_f32()).collect(); + let mut host = vec![0.0_f32; 4]; + dtoh_f32(&stream, &t, &mut host).expect("DtoH"); assert!((host[0] - 200.0).abs() < 1.0); } @@ -1231,17 +1229,14 @@ mod tests { let advantages_host: Vec = (0..total).map(|i| i as f32 * 0.1 - 2.0).collect(); let returns_host: Vec = (0..total).map(|i| i as f32 * 0.5).collect(); - let states_bf16: Vec = states_host.iter().map(|&v| half::bf16::from_f32(v)).collect(); - let log_probs_bf16: Vec = vec![half::bf16::from_f32(-1.5); total]; - let advantages_bf16: Vec = advantages_host.iter().map(|&v| half::bf16::from_f32(v)).collect(); - let returns_bf16: Vec = returns_host.iter().map(|&v| half::bf16::from_f32(v)).collect(); + let log_probs_host: Vec = vec![-1.5_f32; total]; let batch = PpoExperienceBatch { - states: stream.clone_htod(&states_bf16).unwrap(), + states: stream.clone_htod(&states_host).unwrap(), actions: stream.clone_htod(&(0..total).map(|i| (i % 45) as i32).collect::>()).unwrap(), - log_probs: stream.clone_htod(&log_probs_bf16).unwrap(), - advantages: stream.clone_htod(&advantages_bf16).unwrap(), - returns: stream.clone_htod(&returns_bf16).unwrap(), + log_probs: stream.clone_htod(&log_probs_host).unwrap(), + advantages: stream.clone_htod(&advantages_host).unwrap(), + returns: stream.clone_htod(&returns_host).unwrap(), done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::>()).unwrap(), n_episodes, timesteps, @@ -1322,17 +1317,15 @@ mod tests { } let returns_host: Vec = advantages_host.iter().map(|&a| a + 5.0).collect(); - let states_bf16: Vec = vec![half::bf16::ZERO; total * state_dim]; - let log_probs_bf16: Vec = vec![half::bf16::from_f32(-1.0); total]; - let advantages_bf16: Vec = advantages_host.iter().map(|&v| half::bf16::from_f32(v)).collect(); - let returns_bf16: Vec = returns_host.iter().map(|&v| half::bf16::from_f32(v)).collect(); + let states_f32: Vec = vec![0.0_f32; total * state_dim]; + let log_probs_f32: Vec = vec![-1.0_f32; total]; let batch = PpoExperienceBatch { - states: stream.clone_htod(&states_bf16).unwrap(), + states: stream.clone_htod(&states_f32).unwrap(), actions: stream.clone_htod(&vec![1_i32; total]).unwrap(), - log_probs: stream.clone_htod(&log_probs_bf16).unwrap(), - advantages: stream.clone_htod(&advantages_bf16).unwrap(), - returns: stream.clone_htod(&returns_bf16).unwrap(), + log_probs: stream.clone_htod(&log_probs_f32).unwrap(), + advantages: stream.clone_htod(&advantages_host).unwrap(), + returns: stream.clone_htod(&returns_host).unwrap(), done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::>()).unwrap(), n_episodes, timesteps, @@ -1347,7 +1340,7 @@ mod tests { let r = returns.get(i).copied().unwrap_or(0.0); let a = advantages.get(i).copied().unwrap_or(0.0); assert!( - (r - a - 5.0).abs() < 0.1, // BF16 precision: values round-trip with ~1% error + (r - a - 5.0).abs() < 0.1, // f32 precision: values round-trip cleanly "Value at {} should be ~5.0, got {}", i, r - a