diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d803382a1..a0aec3bde 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -48,8 +48,6 @@ use cudarc::driver::{ use cudarc::nvrtc::Ptx; use tracing::info; -use super::gpu_weights::{BranchingWeightSetBf16, DuelingWeightSetBf16}; - use crate::MLError; use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; @@ -238,13 +236,17 @@ pub struct GpuDqnTrainer { f32_to_bf16_kernel: CudaFunction, bf16_to_f32_kernel: CudaFunction, - // ── BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ── - // Allocated lazily on first train_step when weight sets are available. - // Synced after each Adam step (online) and after each EMA update (target). - online_dueling_bf16: Option, - online_branching_bf16: Option, - target_dueling_bf16: Option, - target_branching_bf16: Option, + // ── Flat BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ── + // Single contiguous BF16 buffer per network (online + target), same GOFF_* layout + // as the F32 flat buffers. One f32_to_bf16_kernel launch converts the entire buffer + // instead of 20 per-tensor launches. Forward kernels read at precomputed byte offsets. + bf16_params_buf: CudaSlice, // [TOTAL_PARAMS] flat online BF16 + bf16_target_params_buf: CudaSlice, // [TOTAL_PARAMS] flat target BF16 + bf16_mirrors_initialized: bool, + /// Precomputed byte offsets into flat BF16 buffers for each of the 20 weight tensors. + /// Layout matches GOFF_* order: w_s1, b_s1, w_s2, b_s2, ..., w_bu2, b_bu2. + /// Each offset is in bytes (element offset * sizeof(u16)). + bf16_goff_byte_offsets: [u64; 20], // ── Batch input buffers (uploaded per step) ───────────────────── states_buf: CudaSlice, // [B, STATE_DIM] @@ -275,6 +277,7 @@ pub struct GpuDqnTrainer { // ── Backward / Adam buffers ───────────────────────────────────── grad_buf: CudaSlice, // [TOTAL_PARAMS] gradient accumulator params_buf: CudaSlice, // [TOTAL_PARAMS] flat online parameters + target_params_buf: CudaSlice, // [TOTAL_PARAMS] flat target parameters (EMA) m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm @@ -286,6 +289,7 @@ pub struct GpuDqnTrainer { adam_step: i32, total_params: usize, params_initialized: bool, + target_params_initialized: bool, // ── Shared memory size ────────────────────────────────────────── shmem_bytes: usize, @@ -468,11 +472,30 @@ impl GpuDqnTrainer { // ── Allocate backward / Adam buffers ──────────────────────── let grad_buf = alloc_f32(&stream, total_params, "grad_buf")?; let params_buf = alloc_f32(&stream, total_params, "params_buf")?; + let target_params_buf = alloc_f32(&stream, total_params, "target_params_buf")?; let m_buf = alloc_f32(&stream, total_params, "adam_m")?; let v_buf = alloc_f32(&stream, total_params, "adam_v")?; let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?; let t_buf = alloc_i32(&stream, 1, "adam_t")?; + // ── Allocate flat BF16 weight mirror buffers ────────────────── + // One contiguous BF16 buffer per network; single f32_to_bf16_kernel + // launch converts the entire flat F32 buffer instead of 20 per-tensor + // launches. Forward kernels read at precomputed GOFF byte offsets. + let bf16_params_buf = alloc_u16(&stream, total_params, "bf16_params")?; + let bf16_target_params_buf = alloc_u16(&stream, total_params, "bf16_target_params")?; + + // Precompute byte offsets into flat BF16 buffers (GOFF_* layout). + let param_sizes = compute_param_sizes(&config); + let mut bf16_goff_byte_offsets = [0_u64; 20]; + { + let mut offset = 0_u64; + for i in 0..20 { + bf16_goff_byte_offsets[i] = offset; + offset += (param_sizes[i] as u64) * (std::mem::size_of::() as u64); + } + } + // ── Allocate consolidated transfer buffers ───────────────── // Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights let upload_staging_len = b * config.state_dim * 2 + b * 4; // 2*B*SD + 4*B @@ -512,7 +535,7 @@ impl GpuDqnTrainer { + b * num_branches * config.num_atoms * 2 + b * 2 + 1) * std::mem::size_of::(); - let optim_bytes = total_params * 4 * std::mem::size_of::(); // grad + params + m + v + let optim_bytes = total_params * 5 * std::mem::size_of::(); // grad + params + target_params + m + v info!( batch_size = b, @@ -536,10 +559,10 @@ impl GpuDqnTrainer { per_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, - online_dueling_bf16: None, - online_branching_bf16: None, - target_dueling_bf16: None, - target_branching_bf16: None, + bf16_params_buf, + bf16_target_params_buf, + bf16_mirrors_initialized: false, + bf16_goff_byte_offsets, states_buf, next_states_buf, actions_buf, @@ -560,6 +583,7 @@ impl GpuDqnTrainer { q_out_buf, grad_buf, params_buf, + target_params_buf, m_buf, v_buf, grad_norm_buf, @@ -567,6 +591,7 @@ impl GpuDqnTrainer { adam_step: 0, total_params, params_initialized: false, + target_params_initialized: false, shmem_bytes, training_graph: None, upload_staging_buf, @@ -600,78 +625,112 @@ impl GpuDqnTrainer { // BF16 weight mirror management // ═══════════════════════════════════════════════════════════════════ - /// Ensure BF16 weight mirrors are allocated and synced from F32 originals. + /// Ensure flat BF16 weight mirrors are synced from F32 originals. /// - /// Called lazily on first `train_step()` or `forward_loss()`. Allocates - /// 4 BF16 mirror sets (online + target, dueling + branching) and syncs - /// initial F32 weights into them via the `f32_to_bf16_kernel`. + /// Called lazily on first `train_step()` or `forward_loss()`. Performs + /// initial F32 → BF16 conversion via 2 kernel launches (1 online + 1 target) + /// over the flat parameter buffers (same GOFF_* layout). fn ensure_bf16_mirrors( &mut self, - online_d: &DuelingWeightSet, - online_b: &BranchingWeightSet, - target_d: &DuelingWeightSet, - target_b: &BranchingWeightSet, + _online_d: &DuelingWeightSet, + _online_b: &BranchingWeightSet, + _target_d: &DuelingWeightSet, + _target_b: &BranchingWeightSet, ) -> Result<(), MLError> { - if self.online_dueling_bf16.is_some() { + if self.bf16_mirrors_initialized { return Ok(()); } - // Allocate BF16 mirrors from F32 weight sets - let mut od_bf16 = DuelingWeightSetBf16::alloc_from(online_d, &self.stream)?; - let mut ob_bf16 = BranchingWeightSetBf16::alloc_from(online_b, &self.stream)?; - let mut td_bf16 = DuelingWeightSetBf16::alloc_from(target_d, &self.stream)?; - let mut tb_bf16 = BranchingWeightSetBf16::alloc_from(target_b, &self.stream)?; + // Single kernel launch: flat F32 params_buf → flat BF16 bf16_params_buf + self.launch_flat_bf16_convert_online()?; + // Single kernel launch: flat F32 target_params_buf → flat BF16 bf16_target_params_buf + self.launch_flat_bf16_convert_target()?; - // Initial F32 → BF16 sync - od_bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?; - ob_bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?; - td_bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?; - tb_bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?; + self.bf16_mirrors_initialized = true; - self.online_dueling_bf16 = Some(od_bf16); - self.online_branching_bf16 = Some(ob_bf16); - self.target_dueling_bf16 = Some(td_bf16); - self.target_branching_bf16 = Some(tb_bf16); - - info!("GpuDqnTrainer: BF16 weight mirrors allocated and synced (4 sets, 40 tensors)"); + info!( + total_params = self.total_params, + "GpuDqnTrainer: flat BF16 weight mirrors synced (2 kernel launches, 0 per-tensor)" + ); Ok(()) } - /// Sync online BF16 mirrors from F32 weight tensors. + /// Single fused F32 → BF16 conversion: params_buf → bf16_params_buf. + fn launch_flat_bf16_convert_online(&self) -> Result<(), MLError> { + let n = self.total_params; + if n == 0 { return Ok(()); } + let blocks = ((n + 255) / 256) as u32; + let cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + self.stream + .launch_builder(&self.f32_to_bf16_kernel) + .arg(&self.params_buf) + .arg(&self.bf16_params_buf) + .arg(&(n as i32)) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("flat f32_to_bf16 (online): {e}")))?; + } + Ok(()) + } + + /// Single fused F32 → BF16 conversion: target_params_buf → bf16_target_params_buf. + fn launch_flat_bf16_convert_target(&self) -> Result<(), MLError> { + let n = self.total_params; + if n == 0 { return Ok(()); } + let blocks = ((n + 255) / 256) as u32; + let cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + self.stream + .launch_builder(&self.f32_to_bf16_kernel) + .arg(&self.target_params_buf) + .arg(&self.bf16_target_params_buf) + .arg(&(n as i32)) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("flat f32_to_bf16 (target): {e}")))?; + } + Ok(()) + } + + /// Compute 20 raw device pointers into a flat BF16 buffer at GOFF_* offsets. + /// + /// Returns `[ptr_w_s1, ptr_b_s1, ptr_w_s2, ..., ptr_w_bu2, ptr_b_bu2]`. + /// Each pointer is the base of the flat buffer + the precomputed byte offset. + fn bf16_weight_ptrs(&self, bf16_buf: &CudaSlice) -> [u64; 20] { + let base = raw_device_ptr_u16(bf16_buf, &self.stream); + let mut ptrs = [0_u64; 20]; + for i in 0..20 { + ptrs[i] = base + self.bf16_goff_byte_offsets[i]; + } + ptrs + } + + /// Sync online BF16 mirrors from the flat F32 `params_buf`. /// /// Called inside the CUDA Graph after `unflatten_online_weights()` so that /// the next graph replay's forward kernel reads updated BF16 weights. - /// 12 conversion kernel launches (one per dueling tensor) + 8 (branching). + /// Single kernel launch over the entire flat buffer (was 20 per-tensor). fn sync_online_bf16( - &mut self, - online_d: &DuelingWeightSet, - online_b: &BranchingWeightSet, + &self, + _online_d: &DuelingWeightSet, + _online_b: &BranchingWeightSet, ) -> Result<(), MLError> { - if let Some(ref mut bf16) = self.online_dueling_bf16 { - bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?; - } - if let Some(ref mut bf16) = self.online_branching_bf16 { - bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?; - } - Ok(()) + self.launch_flat_bf16_convert_online() } - /// Sync target BF16 mirrors from F32 weight tensors. + /// Sync target BF16 mirrors from the flat F32 `target_params_buf`. /// /// Called after `target_ema_update()` so that the next forward kernel - /// reads updated BF16 target weights. - fn sync_target_bf16( - &mut self, - target_d: &DuelingWeightSet, - target_b: &BranchingWeightSet, - ) -> Result<(), MLError> { - if let Some(ref mut bf16) = self.target_dueling_bf16 { - bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?; - } - if let Some(ref mut bf16) = self.target_branching_bf16 { - bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?; - } - Ok(()) + /// reads updated BF16 target weights. Single kernel launch (was 20 per-tensor). + fn sync_target_bf16(&self) -> Result<(), MLError> { + self.launch_flat_bf16_convert_target() } // ═══════════════════════════════════════════════════════════════════ @@ -1065,8 +1124,8 @@ impl GpuDqnTrainer { &mut self, states: &CudaSlice, batch_size: usize, - online_dueling: &DuelingWeightSet, - online_branching: &BranchingWeightSet, + _online_dueling: &DuelingWeightSet, + _online_branching: &BranchingWeightSet, ) -> Result<&CudaSlice, MLError> { if batch_size > self.config.batch_size { return Err(MLError::ModelError(format!( @@ -1075,17 +1134,12 @@ impl GpuDqnTrainer { ))); } - // Ensure BF16 mirrors are initialized (lazy on first call). - // forward_only_q only needs online weights but ensure_bf16_mirrors - // requires all 4 sets. Create dummy target refs if needed. - if self.online_dueling_bf16.is_none() { - // Allocate BF16 mirrors for online weights only - let mut od_bf16 = DuelingWeightSetBf16::alloc_from(online_dueling, &self.stream)?; - let mut ob_bf16 = BranchingWeightSetBf16::alloc_from(online_branching, &self.stream)?; - od_bf16.sync_from_f32(online_dueling, &self.f32_to_bf16_kernel, &self.stream)?; - ob_bf16.sync_from_f32(online_branching, &self.f32_to_bf16_kernel, &self.stream)?; - self.online_dueling_bf16 = Some(od_bf16); - self.online_branching_bf16 = Some(ob_bf16); + // Ensure flat BF16 online buffer is synced from params_buf (lazy on first call). + // forward_only_q only needs online weights — single kernel launch over + // the flat buffer if not yet initialized. + if !self.bf16_mirrors_initialized { + self.launch_flat_bf16_convert_online()?; + self.bf16_mirrors_initialized = true; } // Upload states to trainer's states_buf via DtoD copy @@ -1104,10 +1158,8 @@ impl GpuDqnTrainer { /// Argument order matches `dqn_forward_only_kernel` in `dqn_training_kernel.cu`: /// 1 input + 20 online BF16 weights + 1 output + 1 config = 23 args. fn launch_forward_only(&self, batch_size: usize) -> Result<(), MLError> { - let od = self.online_dueling_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 online dueling mirrors not initialized".into()))?; - let ob = self.online_branching_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 online branching mirrors not initialized".into()))?; + // Compute 20 raw device pointers into the flat BF16 online buffer + let op = self.bf16_weight_ptrs(&self.bf16_params_buf); let batch_size_i32 = batch_size as i32; @@ -1118,33 +1170,34 @@ impl GpuDqnTrainer { }; // Safety: argument order matches the extern "C" kernel signature exactly. - // All CudaSlice lifetimes are valid (owned by self). Grid/block = 1 warp per sample. + // All pointers are into the pre-allocated bf16_params_buf (owned by self). + // Grid/block = 1 warp per sample. unsafe { self.stream .launch_builder(&self.forward_only_kernel) // ── Input states (1) ────────────────────────────────── .arg(&self.states_buf) - // ── Online network BF16 weights (20) ───────────────── - .arg(&od.w_s1) - .arg(&od.b_s1) - .arg(&od.w_s2) - .arg(&od.b_s2) - .arg(&od.w_v1) - .arg(&od.b_v1) - .arg(&od.w_v2) - .arg(&od.b_v2) - .arg(&od.w_a1) - .arg(&od.b_a1) - .arg(&od.w_a2) - .arg(&od.b_a2) - .arg(&ob.w_bo1) - .arg(&ob.b_bo1) - .arg(&ob.w_bo2) - .arg(&ob.b_bo2) - .arg(&ob.w_bu1) - .arg(&ob.b_bu1) - .arg(&ob.w_bu2) - .arg(&ob.b_bu2) + // ── Online network BF16 weights (20) — views into flat bf16_params_buf ─ + .arg(&op[0]) // w_s1 + .arg(&op[1]) // b_s1 + .arg(&op[2]) // w_s2 + .arg(&op[3]) // b_s2 + .arg(&op[4]) // w_v1 + .arg(&op[5]) // b_v1 + .arg(&op[6]) // w_v2 + .arg(&op[7]) // b_v2 + .arg(&op[8]) // w_a1 (branch 0 FC) + .arg(&op[9]) // b_a1 + .arg(&op[10]) // w_a2 (branch 0 out) + .arg(&op[11]) // b_a2 + .arg(&op[12]) // w_bo1 (branch 1 FC) + .arg(&op[13]) // b_bo1 + .arg(&op[14]) // w_bo2 (branch 1 out) + .arg(&op[15]) // b_bo2 + .arg(&op[16]) // w_bu1 (branch 2 FC) + .arg(&op[17]) // b_bu1 + .arg(&op[18]) // w_bu2 (branch 2 out) + .arg(&op[19]) // b_bu2 // ── Q-value output (1) ───────────────────────────────── .arg(&self.q_out_buf) // ── Config (1) ────────────────────────────────────────── @@ -1279,8 +1332,8 @@ impl GpuDqnTrainer { // are stable — the graph replays writes to the same addresses. self.unflatten_online_weights(online_d, online_b)?; - // ── 5. Sync online BF16 mirrors from updated F32 tensors ───── - // 20 conversion kernel launches (capturable in CUDA Graph). + // ── 5. Sync online BF16 mirror from updated flat params_buf ───── + // Single f32_to_bf16_kernel launch (capturable in CUDA Graph). // Next graph replay's forward kernel reads updated BF16 weights. self.sync_online_bf16(online_d, online_b)?; @@ -1298,12 +1351,10 @@ impl GpuDqnTrainer { pub fn invalidate_training_graph(&mut self) { self.training_graph = None; self.params_initialized = false; - // Drop BF16 mirrors — they'll be reallocated + synced on next train_step. - // This ensures the graph captures fresh BF16 pointers if weight sets change. - self.online_dueling_bf16 = None; - self.online_branching_bf16 = None; - self.target_dueling_bf16 = None; - self.target_branching_bf16 = None; + // Reset BF16 mirrors flag — they'll be re-synced on next train_step. + // Flat BF16 buffers are pre-allocated (no reallocation needed); only the + // content needs refreshing via a single f32_to_bf16_kernel launch each. + self.bf16_mirrors_initialized = false; } /// Invalidate cached state after external weight modifications. @@ -1463,22 +1514,19 @@ impl GpuDqnTrainer { // Kernel launch methods // ═══════════════════════════════════════════════════════════════════ - /// Launch the forward+loss kernel with BF16 weight mirrors. + /// Launch the forward+loss kernel with flat BF16 weight buffers. /// /// Argument order matches `dqn_forward_loss_kernel` in `dqn_training_kernel.cu`: /// 6 batch data + 20 online BF16 weights + 20 target BF16 weights + 8 activation saves /// + 3 outputs + 2 config = 59 args. /// - /// BF16 mirrors must be populated before this call (via `ensure_bf16_mirrors`). + /// BF16 flat buffers must be synced before this call (via `ensure_bf16_mirrors`). + /// Weight pointers are raw u64 device addresses into the pre-allocated flat + /// `bf16_params_buf` / `bf16_target_params_buf` at precomputed GOFF_* offsets. fn launch_forward_loss(&self) -> Result<(), MLError> { - let od = self.online_dueling_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 online dueling mirrors not initialized".into()))?; - let ob = self.online_branching_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 online branching mirrors not initialized".into()))?; - let td = self.target_dueling_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 target dueling mirrors not initialized".into()))?; - let tb = self.target_branching_bf16.as_ref() - .ok_or_else(|| MLError::ModelError("BF16 target branching mirrors not initialized".into()))?; + // Compute 20 raw device pointers for online and target flat BF16 buffers + let op = self.bf16_weight_ptrs(&self.bf16_params_buf); + let tp = self.bf16_weight_ptrs(&self.bf16_target_params_buf); let b = self.config.batch_size; let batch_size_i32 = b as i32; @@ -1491,8 +1539,7 @@ impl GpuDqnTrainer { }; // Safety: argument order matches the extern "C" kernel signature exactly. - // All CudaSlice lifetimes are valid (owned by self). - // BF16 weight pointers (CudaSlice) match kernel's __nv_bfloat16* params. + // All pointers are into pre-allocated flat BF16 buffers (owned by self). // Grid/block dimensions match kernel expectations (1 warp per sample). unsafe { self.stream @@ -1504,51 +1551,48 @@ impl GpuDqnTrainer { .arg(&self.rewards_buf) .arg(&self.dones_buf) .arg(&self.is_weights_buf) - // ── Online network BF16 weights (20) ───────────────── - .arg(&od.w_s1) - .arg(&od.b_s1) - .arg(&od.w_s2) - .arg(&od.b_s2) - .arg(&od.w_v1) - .arg(&od.b_v1) - .arg(&od.w_v2) - .arg(&od.b_v2) - // Branch 0 (exposure) — DuelingWeightSet advantage slot - .arg(&od.w_a1) - .arg(&od.b_a1) - .arg(&od.w_a2) - .arg(&od.b_a2) - // Branch 1 (order) — BranchingWeightSet - .arg(&ob.w_bo1) - .arg(&ob.b_bo1) - .arg(&ob.w_bo2) - .arg(&ob.b_bo2) - // Branch 2 (urgency) — BranchingWeightSet - .arg(&ob.w_bu1) - .arg(&ob.b_bu1) - .arg(&ob.w_bu2) - .arg(&ob.b_bu2) - // ── Target network BF16 weights (20) ───────────────── - .arg(&td.w_s1) - .arg(&td.b_s1) - .arg(&td.w_s2) - .arg(&td.b_s2) - .arg(&td.w_v1) - .arg(&td.b_v1) - .arg(&td.w_v2) - .arg(&td.b_v2) - .arg(&td.w_a1) - .arg(&td.b_a1) - .arg(&td.w_a2) - .arg(&td.b_a2) - .arg(&tb.w_bo1) - .arg(&tb.b_bo1) - .arg(&tb.w_bo2) - .arg(&tb.b_bo2) - .arg(&tb.w_bu1) - .arg(&tb.b_bu1) - .arg(&tb.w_bu2) - .arg(&tb.b_bu2) + // ── Online network BF16 weights (20) — views into flat bf16_params_buf ─ + .arg(&op[0]) // w_s1 + .arg(&op[1]) // b_s1 + .arg(&op[2]) // w_s2 + .arg(&op[3]) // b_s2 + .arg(&op[4]) // w_v1 + .arg(&op[5]) // b_v1 + .arg(&op[6]) // w_v2 + .arg(&op[7]) // b_v2 + .arg(&op[8]) // w_a1 (branch 0 FC) + .arg(&op[9]) // b_a1 + .arg(&op[10]) // w_a2 (branch 0 out) + .arg(&op[11]) // b_a2 + .arg(&op[12]) // w_bo1 (branch 1 FC) + .arg(&op[13]) // b_bo1 + .arg(&op[14]) // w_bo2 (branch 1 out) + .arg(&op[15]) // b_bo2 + .arg(&op[16]) // w_bu1 (branch 2 FC) + .arg(&op[17]) // b_bu1 + .arg(&op[18]) // w_bu2 (branch 2 out) + .arg(&op[19]) // b_bu2 + // ── Target network BF16 weights (20) — views into flat bf16_target_params_buf ─ + .arg(&tp[0]) // w_s1 + .arg(&tp[1]) // b_s1 + .arg(&tp[2]) // w_s2 + .arg(&tp[3]) // b_s2 + .arg(&tp[4]) // w_v1 + .arg(&tp[5]) // b_v1 + .arg(&tp[6]) // w_v2 + .arg(&tp[7]) // b_v2 + .arg(&tp[8]) // w_a1 + .arg(&tp[9]) // b_a1 + .arg(&tp[10]) // w_a2 + .arg(&tp[11]) // b_a2 + .arg(&tp[12]) // w_bo1 + .arg(&tp[13]) // b_bo1 + .arg(&tp[14]) // w_bo2 + .arg(&tp[15]) // b_bo2 + .arg(&tp[16]) // w_bu1 + .arg(&tp[17]) // b_bu1 + .arg(&tp[18]) // w_bu2 + .arg(&tp[19]) // b_bu2 // ── Saved activations (8) ─────────────────────────── .arg(&self.save_h_s1) .arg(&self.save_h_s2) @@ -1833,86 +1877,146 @@ impl GpuDqnTrainer { // GPU-native Polyak EMA target update // ═══════════════════════════════════════════════════════════════════ + /// Copy 20 individual target weight tensors into the flat `target_params_buf`. + /// + /// Same GOFF_* layout as `flatten_online_weights()`. + /// Pure device-to-device copies -- zero host roundtrip. + fn flatten_target_weights( + &self, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + let sizes = compute_param_sizes(&self.config); + let dst_base = raw_device_ptr(&self.target_params_buf, &self.stream); + + let slices: [&CudaSlice; 20] = [ + &target_d.w_s1, &target_d.b_s1, + &target_d.w_s2, &target_d.b_s2, + &target_d.w_v1, &target_d.b_v1, + &target_d.w_v2, &target_d.b_v2, + &target_d.w_a1, &target_d.b_a1, + &target_d.w_a2, &target_d.b_a2, + &target_b.w_bo1, &target_b.b_bo1, + &target_b.w_bo2, &target_b.b_bo2, + &target_b.w_bu1, &target_b.b_bu1, + &target_b.w_bu2, &target_b.b_bu2, + ]; + + let mut byte_offset: u64 = 0; + for (i, slice) in slices.iter().enumerate() { + let num_bytes = sizes[i] * std::mem::size_of::(); + let src = raw_device_ptr(slice, &self.stream); + dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten_target")?; + byte_offset += num_bytes as u64; + } + + Ok(()) + } + + /// Copy flat `target_params_buf` back to 20 individual target weight tensors. + /// + /// Called after the fused EMA kernel to scatter updated flat target weights + /// back into the individual `DuelingWeightSet` + `BranchingWeightSet` tensors. + /// Pure device-to-device copies -- zero host roundtrip. + fn unflatten_target_weights( + &self, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + let sizes = compute_param_sizes(&self.config); + let src_base = raw_device_ptr(&self.target_params_buf, &self.stream); + + let slices: [&CudaSlice; 20] = [ + &target_d.w_s1, &target_d.b_s1, + &target_d.w_s2, &target_d.b_s2, + &target_d.w_v1, &target_d.b_v1, + &target_d.w_v2, &target_d.b_v2, + &target_d.w_a1, &target_d.b_a1, + &target_d.w_a2, &target_d.b_a2, + &target_b.w_bo1, &target_b.b_bo1, + &target_b.w_bo2, &target_b.b_bo2, + &target_b.w_bu1, &target_b.b_bu1, + &target_b.w_bu2, &target_b.b_bu2, + ]; + + let mut byte_offset: u64 = 0; + for (i, slice) in slices.iter().enumerate() { + let num_bytes = sizes[i] * std::mem::size_of::(); + let dst = raw_device_ptr(slice, &self.stream); + dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?; + byte_offset += num_bytes as u64; + } + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // GPU-native Polyak EMA target update (fused single kernel) + // ═══════════════════════════════════════════════════════════════════ + /// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` /// - /// Updates all 20 target weight tensors in-place from the corresponding - /// online weight tensors using the EMA kernel. Runs OUTSIDE the captured - /// CUDA Graph — device pointers are stable so the graph stays valid. + /// Fused single-kernel update over flat parameter buffers. On first call, + /// flattens target weights into `target_params_buf` (same GOFF_* layout as + /// `params_buf`). Then launches ONE EMA kernel over the entire flat buffer + /// instead of 20 per-tensor launches. After the kernel, scatters the updated + /// flat target weights back to individual tensors via `unflatten_target_weights()`. /// - /// Eliminates the reverse-sync → CPU Polyak → forward-sync round-trip - /// that previously required 120 D2D copies + 120 Candle ops per step. + /// Runs OUTSIDE the captured CUDA Graph -- device pointers are stable so + /// the graph stays valid. pub fn target_ema_update( &mut self, - online_d: &DuelingWeightSet, - online_b: &BranchingWeightSet, + _online_d: &DuelingWeightSet, + _online_b: &BranchingWeightSet, target_d: &mut DuelingWeightSet, target_b: &mut BranchingWeightSet, tau: f32, ) -> Result<(), MLError> { - let sizes = compute_param_sizes(&self.config); - // Sync stream and clear any stale errors from graph capture phase. // cudarc stores errors from cuStreamWaitEvent on disabled events during // graph capture. check_err() consumes them so bind_to_thread() succeeds. unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } let _ = self.stream.context().check_err(); - // Paired (target, online) slices in GOFF_* order (20 pairs) - let pairs: [(&CudaSlice, &CudaSlice); 20] = [ - (&target_d.w_s1, &online_d.w_s1), - (&target_d.b_s1, &online_d.b_s1), - (&target_d.w_s2, &online_d.w_s2), - (&target_d.b_s2, &online_d.b_s2), - (&target_d.w_v1, &online_d.w_v1), - (&target_d.b_v1, &online_d.b_v1), - (&target_d.w_v2, &online_d.w_v2), - (&target_d.b_v2, &online_d.b_v2), - (&target_d.w_a1, &online_d.w_a1), - (&target_d.b_a1, &online_d.b_a1), - (&target_d.w_a2, &online_d.w_a2), - (&target_d.b_a2, &online_d.b_a2), - (&target_b.w_bo1, &online_b.w_bo1), - (&target_b.b_bo1, &online_b.b_bo1), - (&target_b.w_bo2, &online_b.w_bo2), - (&target_b.b_bo2, &online_b.b_bo2), - (&target_b.w_bu1, &online_b.w_bu1), - (&target_b.b_bu1, &online_b.b_bu1), - (&target_b.w_bu2, &online_b.w_bu2), - (&target_b.b_bu2, &online_b.b_bu2), - ]; - - for (i, (target_slice, online_slice)) in pairs.iter().enumerate() { - let n = sizes[i] as i32; - if sizes[i] == 0 { - continue; - } - let blocks = ((sizes[i] + 255) / 256) as u32; - let launch_cfg = LaunchConfig { - grid_dim: (blocks, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }; - - // Use raw device pointers — graph capture left stale events on - // weight CudaSlices, making cudarc's launch_builder.arg() fail. - let t_ptr = raw_device_ptr(target_slice, &self.stream); - let o_ptr = raw_device_ptr(online_slice, &self.stream); - unsafe { - self.stream - .launch_builder(&self.ema_kernel) - .arg(&t_ptr) - .arg(&o_ptr) - .arg(&tau) - .arg(&n) - .launch(launch_cfg) - .map_err(|e| { - MLError::ModelError(format!("dqn_ema_kernel launch[{i}]: {e}")) - })?; - } + // First call: flatten target weights into flat buffer (20 DtoD copies, once only). + // Subsequent calls reuse the flat buffer which is kept in sync by unflatten at end. + if !self.target_params_initialized { + self.flatten_target_weights(target_d, target_b)?; + self.target_params_initialized = true; } + // Single fused EMA kernel over all parameters at once. + // params_buf contains the latest online weights (kept in sync by + // unflatten_online_weights after each Adam step). + let n = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + let t_ptr = raw_device_ptr(&self.target_params_buf, &self.stream); + let o_ptr = raw_device_ptr(&self.params_buf, &self.stream); + unsafe { + self.stream + .launch_builder(&self.ema_kernel) + .arg(&t_ptr) + .arg(&o_ptr) + .arg(&tau) + .arg(&n) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_ema_kernel flat launch: {e}")) + })?; + } + + // Scatter flat target buffer back to individual weight tensors + self.unflatten_target_weights(target_d, target_b)?; + // Sync target BF16 mirrors from updated F32 target weights - self.sync_target_bf16(target_d, target_b)?; + // Single kernel launch over flat target_params_buf → bf16_target_params_buf + self.sync_target_bf16()?; Ok(()) } @@ -2262,6 +2366,13 @@ fn raw_device_ptr_u32(slice: &CudaSlice, stream: &CudaStream) -> u64 { ptr } +/// Extract raw CUDA device pointer from a `CudaSlice` (BF16 weight buffers). +fn raw_device_ptr_u16(slice: &CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr(stream); + let _no_drop = std::mem::ManuallyDrop::new(guard); + ptr +} + /// Async device-to-device memcpy with error context. pub(crate) fn dtod_copy( dst: u64,