diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index a97c2d87e..3f38271ca 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -298,6 +298,73 @@ impl DqnGpuData { }) } + /// Upload DQN training data to GPU from separate contiguous arrays. + /// + /// Accepts features, targets, and OFI as separate slices matching the fxcache + /// flat binary layout — no tuple allocation or per-bar heap Vec required. + /// OFI is uploaded only when at least one value is non-zero. + pub fn upload_slices( + features: &[[f64; 42]], + targets: &[[f64; 4]], + ofi: &[[f64; 8]], + stream: &Arc, + ) -> Result { + let num_bars = features.len(); + if num_bars == 0 { + return Err(MLError::ModelError("Empty training data".to_owned())); + } + + let feature_dim = 42; + let target_dim = 4; + + let estimated_bytes = estimate_vram_bytes(num_bars * (feature_dim + target_dim)); + if estimated_bytes > MAX_UPLOAD_BYTES { + return Err(MLError::ModelError(format!( + "Training data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} bars)", + estimated_bytes as f64 / 1_073_741_824.0, + num_bars, + ))); + } + + let flat_features: Vec = features + .iter() + .flat_map(|f| f.iter().map(|&v| v as f32)) + .collect(); + + let flat_targets: Vec = targets + .iter() + .flat_map(|t| t.iter().map(|&v| v as f32)) + .collect(); + + let features_gpu = clone_htod_f32_to_bf16(stream, &flat_features)?; + let targets_gpu = clone_htod_f32_to_bf16(stream, &flat_targets)?; + + let ofi_features = if !ofi.is_empty() && ofi.iter().any(|row| row.iter().any(|&v| v != 0.0)) { + let mut flat_ofi = Vec::with_capacity(num_bars * 8); + for i in 0..num_bars { + if i < ofi.len() { + for &v in ofi[i].iter() { + flat_ofi.push(v as f32); + } + } else { + flat_ofi.extend_from_slice(&[0.0_f32; 8]); + } + } + Some(clone_htod_f32_to_bf16(stream, &flat_ofi)?) + } else { + None + }; + + Ok(Self { + features: features_gpu, + targets: targets_gpu, + ofi_features, + num_bars, + feature_dim, + aligned_state_dim: None, + }) + } + /// Upload OFI features (8 dims per bar) from MBP-10 order book data. /// /// Must be called after `upload()` with a slice matching `num_bars` length.