diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 932c6684d..ffc979061 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1111,26 +1111,12 @@ impl GpuDqnTrainer { } } - // ── 2. Cast IQN d_h_s2 (bf16) → bw_d_h_s2 (f32) ───────────────── - // IQN head produces bf16 gradient. Cast to f32 for the backward scratch. + // ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ──────────── + // IQN d_h_s2_buf is f32 (atomicAdd in backward kernel). No cast needed. { - let src = iqn_d_h_s2_ptr; - let dst = self.ptrs.bw_d_h_s2; - let n_elems = (b * sh2) as i32; - let blocks = ((b * sh2 + 255) / 256) as u32; - unsafe { - self.stream - .launch_builder(&self.bf16_to_f32_kernel) - .arg(&src) - .arg(&dst) - .arg(&n_elems) - .launch(LaunchConfig { - grid_dim: (blocks, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }) - .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 bf16→f32: {e}")))?; - } + let n_bytes = b * sh2 * std::mem::size_of::(); + dtod_copy(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, &self.stream, 0, "iqn_d_h_s2_copy") + .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 f32 DtoD: {e}")))?; } // ── 3. ReLU mask: bw_d_h_s2 *= (save_h_s2 > 0) (f32 dx, bf16 act) ──