diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 895fad73a..39ed07fd8 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -524,20 +524,17 @@ impl GpuIqnHead { (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos(); } } - cos_features = stream.clone_htod(&cos_feat_host).map_err(|e| { - MLError::ModelError(format!("IQN htod cos_features ({} f32): {e}", d * n)) - })?; + cos_features = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &cos_feat_host) + .map_err(|e| MLError::ModelError(format!("IQN cos_features upload via pinned ({} f32): {e}", d * n)))?; let mut tiled = Vec::with_capacity(b * n); for _ in 0..b { tiled.extend_from_slice(fixed); } - online_taus = stream.clone_htod(&tiled).map_err(|e| { - MLError::ModelError(format!("IQN htod online_taus ({} f32): {e}", b * n)) - })?; - target_taus = stream.clone_htod(&tiled).map_err(|e| { - MLError::ModelError(format!("IQN htod target_taus ({} f32): {e}", b * n)) - })?; + online_taus = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &tiled) + .map_err(|e| MLError::ModelError(format!("IQN online_taus upload via pinned ({} f32): {e}", b * n)))?; + target_taus = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &tiled) + .map_err(|e| MLError::ModelError(format!("IQN target_taus upload via pinned ({} f32): {e}", b * n)))?; } let branch_actions = stream.alloc_zeros::(b * 4).map_err(|e| { MLError::ModelError(format!("IQN alloc branch_actions: {e}")) @@ -2149,10 +2146,9 @@ fn init_iqn_xavier_weights( debug_assert_eq!(offset, total, "IQN weight layout mismatch"); - // Upload to GPU with cuBLAS tile padding (zeros past total_params) - let params_buf = stream.clone_htod(&weights).map_err(|e| { - MLError::ModelError(format!("IQN htod online_params ({} f32): {e}", total + cublas_pad)) - })?; + // Upload to GPU via mapped-pinned staging (zeros past total_params remain in `weights` Vec) + let params_buf = super::mapped_pinned::clone_to_device_f32_via_pinned(stream, &weights) + .map_err(|e| MLError::ModelError(format!("IQN online_params upload via pinned ({} f32): {e}", total + cublas_pad)))?; Ok(params_buf) } @@ -2193,17 +2189,17 @@ impl Drop for GpuIqnHead { } } -/// Clone a CudaSlice by copying device → host → device. +/// Clone a `CudaSlice` via a single DtoD copy (no host crossing). +/// +/// Replaces the legacy device→host→device round-trip that violated both +/// `feedback_gpu_cpu_roundtrip.md` (no DtoH) and +/// `feedback_no_htod_htoh_only_mapped_pinned.md` (no HtoD). Used at +/// constructor time to seed the IQN target network from the online params. fn clone_cuda_slice( stream: &Arc, src: &CudaSlice, n: usize, ) -> Result, MLError> { - let host = stream.clone_dtoh(src).map_err(|e| { - MLError::ModelError(format!("IQN clone dtoh ({n} f32): {e}")) - })?; - let dst = stream.clone_htod(&host).map_err(|e| { - MLError::ModelError(format!("IQN clone htod ({n} f32): {e}")) - })?; - Ok(dst) + super::clone_cuda_slice_f32(src, stream) + .map_err(|e| MLError::ModelError(format!("IQN clone DtoD ({n} f32): {e}"))) } diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 9db3fffd7..1d7f7d770 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -109,3 +109,6 @@ ### Fix 8 — Migrate weight-upload sites in attention/tlob/iql/weights (4 sites) `gpu_weights.rs` (1 site: `RmsNormWeightSet::ones` constructor), `gpu_attention.rs` (1 site: attention params constructor), `gpu_tlob.rs` (1 site: TLOB params constructor — `mod tests` blocks intentionally untouched per scope), `gpu_iql_trainer.rs` (2 sites: per_sample_support seed, IQL params init). All `htod_f32` calls rewritten to `mapped_pinned::upload_f32_via_pinned`; the gpu_weights.rs ones-init uses `clone_to_device_f32_via_pinned` since the destination is freshly allocated. + +### Fix 9 — Migrate `gpu_iqn_head.rs` constructors + DtoD-only target clone (5 sites) +`gpu_iqn_head.rs`: 4 `stream.clone_htod` calls (cos_features precompute, online_taus + target_taus tile broadcast, init_iqn_xavier_weights upload) rewritten to `mapped_pinned::clone_to_device_f32_via_pinned`. Additionally the helper `clone_cuda_slice` was replaced — it previously did device→host (`clone_dtoh`) → host→device (`clone_htod`) round-trip, double-violating both the no-DtoH and no-HtoD rules. Now uses `clone_cuda_slice_f32` (single async DtoD), the only callsite being the target-network seed at constructor time (`target_params = clone_cuda_slice(&online_params)`).