diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 1edb6912b..6990d3d6d 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -2936,3 +2936,40 @@ extern "C" __global__ void selectivity_backward( } atomicAdd(&d_sel_params[SH2], d_z); /* db */ } + +/** + * Phase 1: Compute mean of Q-values into a scratch scalar. + * Grid: (1, 1, 1), Block: (256, 1, 1). + */ +extern "C" __global__ void q_mean_reduce( + const float* __restrict__ q_values, + float* __restrict__ mean_out, + int total) +{ + __shared__ float sdata[256]; + float sum = 0.0f; + for (int i = threadIdx.x; i < total; i += 256) { + sum += q_values[i]; + } + sdata[threadIdx.x] = sum; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) mean_out[0] = sdata[0] / (float)total; +} + +/** + * Phase 2: Subtract mean from all Q-values in-place. + * Grid: ceil(total/256), Block: 256. + */ +extern "C" __global__ void q_mean_subtract( + float* __restrict__ q_values, + const float* __restrict__ mean_val, + int total) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + q_values[i] -= mean_val[0]; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a38722cdf..ee2e6ea32 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1253,6 +1253,48 @@ impl GpuDqnTrainer { self.q_coord_buf.raw_ptr() } + /// Center Q-values in q_out_buf by subtracting their global mean. + /// + /// Two-phase GPU operation (no CPU roundtrip): + /// Phase 1 (q_mean_reduce): [B * total_actions] → mean scalar in q_mean_scratch + /// Phase 2 (q_mean_subtract): q_out_buf[i] -= q_mean_scratch[0] in-place + /// + /// Prevents bootstrapping drift (Q-mean 0→+0.47 over 18 epochs observed in + /// train-vpb4w). Must be called AFTER compute_expected_q, BEFORE launch_q_attention. + pub(crate) fn launch_q_mean_center(&self, batch_size: usize) -> Result<(), MLError> { + let total = (batch_size * self.total_actions()) as i32; + let q_ptr = self.q_out_buf.raw_ptr(); + let scratch_ptr = self.q_mean_scratch.raw_ptr(); + // Phase 1: reduce all Q-values to their mean + unsafe { + self.stream.launch_builder(&self.q_mean_reduce_kernel) + .arg(&q_ptr) + .arg(&scratch_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("q_mean_reduce: {e}")))?; + } + // Phase 2: subtract mean from every Q-value + let blocks = ((total as u32 + 255) / 256).max(1); + unsafe { + self.stream.launch_builder(&self.q_mean_subtract_kernel) + .arg(&q_ptr) + .arg(&scratch_ptr) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("q_mean_subtract: {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, @@ -4777,6 +4819,9 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("compute_expected_q (q_stats): {e}")))?; } + // Q-mean centering: removes bootstrapping drift before attention and stats. + self.launch_q_mean_center(batch_size)?; + // 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)?; diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2eb062733..c4a1def9f 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1983,6 +1983,14 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("launch_q_attention: {e}")) } + /// Center Q-values in q_out_buf by subtracting their global mean. + /// Prevents bootstrapping drift. Must be called after compute_expected_q, + /// before launch_q_attention. + pub(crate) fn launch_q_mean_center(&self, batch_size: usize) -> Result<()> { + self.trainer.launch_q_mean_center(batch_size) + .map_err(|e| anyhow::anyhow!("launch_q_mean_center: {e}")) + } + /// Selectivity gate forward: sigmoid(dot(W_sel, h_s2) + b_sel) per sample. pub(crate) fn launch_selectivity_forward(&self, batch_size: usize) -> Result<()> { self.trainer.launch_selectivity_forward(batch_size)