diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 97f776f25..bd2f33ab1 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -199,7 +199,7 @@ extern "C" __global__ void c51_loss_batched( const float* __restrict__ curiosity_errors, /* [B] per-sample curiosity prediction error (0 = no penalty) */ float curiosity_q_penalty_lambda, /* scaling factor: gamma *= 1/(1 + lambda * error) */ - float gamma, + const float* __restrict__ gamma_buf, /* [B] per-sample effective gamma */ int batch_size, int num_atoms, const float* __restrict__ per_sample_support, /* [B, 3] per-sample: [v_min, v_max, delta_z] */ @@ -331,10 +331,10 @@ extern "C" __global__ void c51_loss_batched( * gamma down so the Bellman target approaches the immediate reward, * preventing overconfident extrapolation into novel states. * gamma_eff = gamma / (1 + lambda * prediction_error) */ - float gamma_eff = gamma; + float gamma_eff = gamma_buf[sample_id]; if (curiosity_q_penalty_lambda > 0.0f) { float cur_err = curiosity_errors[sample_id]; - gamma_eff = gamma / (1.0f + curiosity_q_penalty_lambda * cur_err); + gamma_eff = gamma_buf[sample_id] / (1.0f + curiosity_q_penalty_lambda * cur_err); } float total_ce = 0.0f; diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 6bf1eba86..0486af8af 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -4930,3 +4930,17 @@ extern "C" __global__ void isv_forward( gamma_raw += w_gamma[k] * embedding[k]; gamma_mod_out[0] = 0.5f + 0.5f / (1.0f + expf(-gamma_raw)); } + +/* ================================================================== */ +/* Kernel: fill_gamma_buf — broadcast scalar gamma to per-sample buf */ +/* ================================================================== */ +extern "C" __global__ void fill_gamma_buf( + float* __restrict__ gamma_buf, /* [B] output */ + float base_gamma, + const float* __restrict__ gamma_mod, /* [1] from isv_forward */ + int B +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= B) return; + gamma_buf[i] = base_gamma * gamma_mod[0]; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4c0eb878e..6b8af79a5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1338,6 +1338,7 @@ pub struct GpuDqnTrainer { // ── ISV forward outputs ── isv_forward_kernel: CudaFunction, + fill_gamma_buf_kernel: CudaFunction, isv_embedding_buf: CudaSlice, // [ISV_DIM] = [8] branch_gate_buf: CudaSlice, // [4] gamma_mod_buf: CudaSlice, // [1] @@ -2067,6 +2068,28 @@ impl GpuDqnTrainer { Ok(()) } + /// Broadcast base_gamma * gamma_mod[0] into gamma_buf [B] for per-sample C51 loss. + pub(crate) fn fill_gamma_buf(&self) -> Result<(), MLError> { + let base_gamma = self.adaptive_gamma.powi(self.config.n_steps as i32); + let blocks = ((self.config.batch_size as u32 + 255) / 256).max(1); + let b = self.config.batch_size as i32; + let gamma_mod_ptr = self.gamma_mod_buf.raw_ptr(); + unsafe { + self.stream.launch_builder(&self.fill_gamma_buf_kernel) + .arg(&self.gamma_buf) + .arg(&base_gamma) + .arg(&gamma_mod_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!("fill_gamma_buf: {e}")))?; + } + Ok(()) + } + /// LR scale factor for new components: min(1.0, step / 500). fn new_component_lr_scale(&self) -> f32 { (self.new_component_warmup_step as f32 / NEW_COMPONENT_WARMUP_STEPS as f32).min(1.0) @@ -4160,7 +4183,9 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("isv_signal_update load: {e}")))?; let isv_forward_kernel = exp_module_for_mag.load_function("isv_forward") .map_err(|e| MLError::ModelError(format!("isv_forward load: {e}")))?; - info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward kernels loaded"); + let fill_gamma_buf_kernel = exp_module_for_mag.load_function("fill_gamma_buf") + .map_err(|e| MLError::ModelError(format!("fill_gamma_buf load: {e}")))?; + info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward + fill_gamma_buf kernels loaded"); // ── G5: Epistemic-gated magnitude — pinned var_ema threshold ─ let (var_ema_pinned, var_ema_dev_ptr) = { @@ -5654,6 +5679,7 @@ impl GpuDqnTrainer { lagged_td_error_dev_ptr, isv_signal_update_kernel, isv_forward_kernel, + fill_gamma_buf_kernel, isv_embedding_buf, branch_gate_buf, gamma_mod_buf, @@ -6292,6 +6318,7 @@ impl GpuDqnTrainer { self.launch_mse_loss()?; self.launch_loss_reduce(self.mse_loss_dev_ptr)?; self.launch_mse_grad_to_scratch()?; + self.fill_gamma_buf()?; self.launch_c51_loss()?; self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; @@ -7545,6 +7572,7 @@ impl GpuDqnTrainer { self.q_divergence_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), ); } + self.fill_gamma_buf()?; self.launch_c51_loss()?; self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; @@ -7697,6 +7725,7 @@ impl GpuDqnTrainer { self.q_divergence_dev_ptr, 0, std::mem::size_of::(), self.stream.cu_stream(), ); } + self.fill_gamma_buf()?; self.launch_c51_loss()?; self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; @@ -8169,9 +8198,9 @@ impl GpuDqnTrainer { let on_next_b2_ptr = on_next_b1_ptr + (b * b1 * na * f32_sz) as u64; let on_next_b3_ptr = on_next_b2_ptr + (b * b2 * na * f32_sz) as u64; - // N-step returns: use gamma^n for the Bellman projection. - // The experience collector pre-computes R_n = sum(gamma^i * r_i). - let gamma = self.adaptive_gamma.powi(self.config.n_steps as i32); + // N-step returns: gamma_buf [B] is pre-filled by fill_gamma_buf kernel + // (base_gamma^n * gamma_mod[0] per sample). + let gamma_buf_ptr = self.gamma_buf.raw_ptr(); let batch_i32 = b as i32; let na_i32 = na as i32; let b0_i32 = b0 as i32; @@ -8222,7 +8251,7 @@ impl GpuDqnTrainer { .arg(&self.curiosity_error_buf) .arg(&self.config.curiosity_q_penalty_lambda) // ── Config (8 — per_sample_support replaces v_min+v_max) ── - .arg(&gamma) + .arg(&gamma_buf_ptr) .arg(&batch_i32) .arg(&na_i32) .arg(&self.per_sample_support_ptr)