gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(htod): migrate gpu_iqn_head.rs to mapped-pinned + DtoD-only clone (5 sites)

Browse Source 
- 4 stream.clone_htod sites in IQN constructor: cos_features precompute,
  online_taus + target_taus tile broadcast, init_iqn_xavier_weights
  upload — all rewritten to mapped_pinned::clone_to_device_f32_via_pinned.
- clone_cuda_slice helper rewritten to a single DtoD copy via the
  existing super::clone_cuda_slice_f32 helper. Previously did
  device→host→device, double-violating both no-DtoH and no-HtoD rules.
  Sole callsite is constructor seeding of target params from online.

Audit row appended (Fix 9) in docs/dqn-gpu-hot-path-audit.md.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 21:01:51 +02:00
parent f997f1fa78
commit f77ebeb214
2 changed files with 20 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -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::<i32>(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<f32> by copying device → host → device.
/// Clone a `CudaSlice<f32>` 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<CudaStream>,
src: &CudaSlice<f32>,
n: usize,
) -> Result<CudaSlice<f32>, 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}")))
}

3
docs/dqn-gpu-hot-path-audit.md
View File

@@ -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)`).