diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 494f5ab86..f2aeb5b46 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1179,6 +1179,38 @@ impl GpuDqnTrainer { Ok(()) } + /// Run 2-head cross-branch Q attention over the 12 raw Q-values in q_out_buf. + /// Writes coordinated Q-values to q_coord_buf [batch_size, 12]. + /// Must be called after compute_expected_q has populated q_out_buf. + pub(crate) fn launch_q_attention(&self, batch_size: usize) -> Result<(), MLError> { + let b = batch_size as i32; + let blocks = ((batch_size as u32 + 255) / 256).max(1); + let q_out_ptr = self.q_out_buf.raw_ptr(); + let q_coord_ptr = self.q_coord_buf.raw_ptr(); + let q_attn_params_ptr = self.q_attn_params.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.q_attn_kernel) + .arg(&q_out_ptr) + .arg(&q_coord_ptr) + .arg(&q_attn_params_ptr) + .arg(&b) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("cross_branch_q_attention: {e}")))?; + } + Ok(()) + } + + /// Raw device pointer to the coordinated Q-value buffer [batch_size, 12]. + /// Valid after `launch_q_attention` has been called. + pub(crate) fn q_coord_buf_ptr(&self) -> u64 { + self.q_coord_buf.raw_ptr() + } + /// 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, @@ -2718,8 +2750,8 @@ impl GpuDqnTrainer { // ── Load Q-attn/selectivity/VSN/GLU kernels from experience_kernels cubin ── let cpbi_module = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("cpbi cubin: {e}")))?; - let q_attn_kernel = cpbi_module.load_function("q_cross_branch_attn") - .map_err(|e| MLError::ModelError(format!("q_cross_branch_attn load: {e}")))?; + let q_attn_kernel = cpbi_module.load_function("cross_branch_q_attention") + .map_err(|e| MLError::ModelError(format!("cross_branch_q_attention load: {e}")))?; let sel_fwd_kernel = cpbi_module.load_function("selectivity_gate_fwd") .map_err(|e| MLError::ModelError(format!("selectivity_gate_fwd load: {e}")))?; let sel_bwd_kernel = cpbi_module.load_function("selectivity_gate_bwd") @@ -4620,6 +4652,10 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("compute_expected_q (q_stats): {e}")))?; } + // Cross-branch Q attention: q_out_buf → q_coord_buf [B, 12]. + // Runs after compute_expected_q so attention weights train on real Q-values. + self.launch_q_attention(batch_size)?; + // Stats reduction on q_out_buf. let num_atoms = na; unsafe { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index d3377d7c5..ab77e79b1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1955,6 +1955,11 @@ impl FusedTrainingCtx { self.trainer.states_buf_ptr() } + /// Run cross-branch Q attention: q_out_buf → q_coord_buf [batch_size, 12]. + pub(crate) fn launch_q_attention(&self, batch_size: usize) -> Result<()> { + self.trainer.launch_q_attention(batch_size) + .map_err(|e| anyhow::anyhow!("launch_q_attention: {e}")) + } /// Set C51 blend factor for gradual MSE→C51 ramp. ///