diff --git a/crates/ml-alpha/src/data/loader.rs b/crates/ml-alpha/src/data/loader.rs index f2d515e9e..5559e15e8 100644 --- a/crates/ml-alpha/src/data/loader.rs +++ b/crates/ml-alpha/src/data/loader.rs @@ -714,43 +714,48 @@ impl MultiHorizonLoader { "upload_to_gpu: pre-converting snapshots", ); - let mut all_raws: Vec = Vec::with_capacity(total_snaps); let mut file_offsets: Vec = Vec::with_capacity(n_files); let mut file_sizes: Vec = Vec::with_capacity(n_files); let mut offset: usize = 0; for lf in &self.files_loaded { file_offsets.push(offset as i32); file_sizes.push(lf.snapshots.len() as i32); - for (idx, snap) in lf.snapshots.iter().enumerate() { - let prev = if idx == 0 { snap } else { &lf.snapshots[idx - 1] }; - all_raws.push(convert(snap, prev, lf.regime_full[idx])); - } offset += lf.snapshots.len(); } - debug_assert_eq!(all_raws.len(), total_snaps); - // ── 2. Upload snapshots as raw bytes ───────────────────────────── - // Reinterpret Vec as &[u8] for the upload. The GPU - // kernel reads the same struct layout via its `Mbp10Raw` definition. + // ── 2. Upload snapshots per-file (streaming) ──────────────────── + // Allocate the full GPU buffer upfront, then convert + upload one + // file at a time. Peak host memory = one file's worth of converted + // snapshots (~1-2GB) instead of all files (~10GB). let total_bytes = total_snaps * MBP10_RAW_INPUT_BYTES; - let snap_bytes: &[u8] = unsafe { - std::slice::from_raw_parts( - all_raws.as_ptr() as *const u8, - total_bytes, - ) - }; let snapshots_d = stream .alloc_zeros::(total_bytes) .context("upload_to_gpu: snapshots alloc")?; - unsafe { - cudarc::driver::sys::cuMemcpyHtoD_v2( - snapshots_d.raw_ptr(), - snap_bytes.as_ptr() as *const std::ffi::c_void, - total_bytes, + let mut gpu_offset_bytes: usize = 0; + for (f_idx, lf) in self.files_loaded.iter().enumerate() { + let n = lf.snapshots.len(); + let mut file_raws: Vec = Vec::with_capacity(n); + for (idx, snap) in lf.snapshots.iter().enumerate() { + let prev = if idx == 0 { snap } else { &lf.snapshots[idx - 1] }; + file_raws.push(convert(snap, prev, lf.regime_full[idx])); + } + let file_bytes = n * MBP10_RAW_INPUT_BYTES; + unsafe { + cudarc::driver::sys::cuMemcpyHtoD_v2( + snapshots_d.raw_ptr() + gpu_offset_bytes as u64, + file_raws.as_ptr() as *const std::ffi::c_void, + file_bytes, + ); + } + gpu_offset_bytes += file_bytes; + tracing::info!( + file = f_idx, + n_snapshots = n, + gpu_offset_mb = gpu_offset_bytes / (1024 * 1024), + "upload_to_gpu: file uploaded", ); + drop(file_raws); } - // Drop the host-side Vec now — it can be multi-GB. - drop(all_raws); tracing::info!( gpu_bytes = total_bytes, "upload_to_gpu: snapshots uploaded",