diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4ee57248f..e8d768038 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -389,7 +389,7 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT cfg.adv_h, // [9] b_b0fc cfg.branch_0_size * cfg.num_atoms * cfg.adv_h, // [10] w_b0out cfg.branch_0_size * cfg.num_atoms, // [11] b_b0out - cfg.adv_h * cfg.shared_h2, // [12] w_b1fc + cfg.adv_h * (cfg.shared_h2 + 3), // [12] w_b1fc (direction-conditioned: SH2+3 input) cfg.adv_h, // [13] b_b1fc cfg.branch_1_size * cfg.num_atoms * cfg.adv_h, // [14] w_b1out cfg.branch_1_size * cfg.num_atoms, // [15] b_b1out @@ -459,6 +459,8 @@ struct CachedPtrs { save_h_b1: u64, save_h_b2: u64, save_h_b3: u64, + mag_concat_buf: u64, + d_mag_concat_buf: u64, bw_d_h_s1: u64, bw_d_h_s2: u64, bw_d_h_v: u64, @@ -588,6 +590,16 @@ pub struct GpuDqnTrainer { save_h_b1: CudaSlice, // [B, ADV_H] save_h_b2: CudaSlice, // [B, ADV_H] save_h_b3: CudaSlice, // [B, ADV_H] + + /// Magnitude branch concat input [B, SH2+3]: [h_s2; Q_dir] for direction conditioning. + mag_concat_buf: CudaSlice, + /// Backward scratch for magnitude concat dX [B, SH2+3]. + d_mag_concat_buf: CudaSlice, + /// Kernel: mag_concat_qdir — builds [h_s2; Q_dir] concat. + mag_concat_kernel: CudaFunction, + /// Kernel: strided_accumulate — extracts d_h_s2 from d_mag_concat in backward. + strided_accumulate_kernel: CudaFunction, + save_current_lp: CudaSlice, // [B, NUM_BRANCHES(4), NUM_ATOMS] save_projected: CudaSlice, // [B, NUM_BRANCHES(4), NUM_ATOMS] @@ -1083,6 +1095,71 @@ impl GpuDqnTrainer { self.branch_scales_ptr = ptr; } + /// Build [h_s2; Q_dir] concat buffer for magnitude branch conditioning. + /// Must be called AFTER direction branch forward (branch 0 logits are ready). + pub(crate) fn launch_mag_concat(&self, v_logits_ptr: u64, b_logits_ptr: u64) -> Result<(), MLError> { + let b = self.config.batch_size; + let sh2 = self.config.shared_h2 as i32; + let na = self.config.num_atoms as i32; + let b0 = self.config.branch_0_size as i32; + let n = b as i32; + let blocks = ((b as u32 + 255) / 256).max(1); + let h_s2_ptr = self.ptrs.save_h_s2; + let concat_ptr = self.ptrs.mag_concat_buf; + let support_ptr = self.per_sample_support_ptr; + unsafe { + self.stream + .launch_builder(&self.mag_concat_kernel) + .arg(&h_s2_ptr) + .arg(&v_logits_ptr) + .arg(&b_logits_ptr) + .arg(&concat_ptr) + .arg(&n) + .arg(&sh2) + .arg(&na) + .arg(&b0) + .arg(&support_ptr) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("mag_concat_qdir: {e}")))?; + } + Ok(()) + } + + /// Accumulate first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2]. + pub(crate) fn accumulate_d_h_s2_from_concat( + &self, + d_concat_ptr: u64, + d_h_s2_ptr: u64, + batch: usize, + beta: f32, + ) -> Result<(), MLError> { + let total = (batch * self.config.shared_h2) as i32; + let blocks = ((total as u32 + 255) / 256).max(1); + let sh2 = self.config.shared_h2 as i32; + let sh2_plus_3 = (self.config.shared_h2 + 3) as i32; + unsafe { + self.stream + .launch_builder(&self.strided_accumulate_kernel) + .arg(&d_concat_ptr) + .arg(&d_h_s2_ptr) + .arg(&sh2) + .arg(&sh2_plus_3) + .arg(&total) + .arg(&beta) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("strided_accumulate: {e}")))?; + } + Ok(()) + } + /// Update IQN readiness from the current IQN loss. /// Readiness ramps 0→1 as loss decreases from initial value. /// Uses adaptive-rate EMA (same pattern as eval_v_range). @@ -2370,6 +2447,9 @@ impl GpuDqnTrainer { let save_h_b1 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b1")?; let save_h_b2 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b2")?; let save_h_b3 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b3")?; + let mag_concat_dim = config.shared_h2 + 3; + let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?; + let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?; let save_current_lp = alloc_f32( &stream, b * num_branches * config.num_atoms, @@ -2591,6 +2671,16 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?; info!("GpuDqnTrainer: expected_q + q_stats kernels compiled"); + // ── Load mag_concat and strided_accumulate from experience_kernels cubin ── + let exp_module_for_mag = stream.context() + .load_cubin(EXPECTED_Q_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("exp cubin for mag_concat: {e}")))?; + let mag_concat_kernel = exp_module_for_mag.load_function("mag_concat_qdir") + .map_err(|e| MLError::ModelError(format!("mag_concat_qdir load: {e}")))?; + let strided_accumulate_kernel = exp_module_for_mag.load_function("strided_accumulate") + .map_err(|e| MLError::ModelError(format!("strided_accumulate load: {e}")))?; + info!("GpuDqnTrainer: mag_concat_qdir + strided_accumulate kernels loaded"); + // ── Compile CQL penalty kernel (if enabled) ────────────────────── let cql_logit_grad_kernel = if config.cql_alpha > 0.0 { match compile_cql_logit_grad_kernel(&stream) { @@ -2908,6 +2998,8 @@ impl GpuDqnTrainer { save_h_b1: save_h_b1.raw_ptr(), save_h_b2: save_h_b2.raw_ptr(), save_h_b3: save_h_b3.raw_ptr(), + mag_concat_buf: mag_concat_buf.raw_ptr(), + d_mag_concat_buf: d_mag_concat_buf.raw_ptr(), bw_d_h_s1: bw_d_h_s1.raw_ptr(), bw_d_h_s2: bw_d_h_s2.raw_ptr(), bw_d_h_v: bw_d_h_v.raw_ptr(), @@ -3060,6 +3152,10 @@ impl GpuDqnTrainer { save_h_b1, save_h_b2, save_h_b3, + mag_concat_buf, + d_mag_concat_buf, + mag_concat_kernel, + strided_accumulate_kernel, save_current_lp, save_projected, per_sample_loss_buf,