diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 0b64038ec..a86292c7c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -4882,6 +4882,40 @@ impl GpuDqnTrainer { self.cql_logit_grad_kernel.is_some() && self.config.cql_alpha > 0.0 } + /// F8/G5: Scale the C51-contributed portion of `grad_buf` by the adaptive c51_budget. + /// + /// Must be called AFTER `submit_forward_ops_main` (which runs C51 backward and writes + /// C51 gradients into `grad_buf`), and BEFORE any CQL/IQN SAXPY accumulations. + /// + /// The composite effect is: + /// grad = c51_budget × C51_grad + cql_budget × CQL_grad + iqn_budget × IQN_grad + /// + /// Uses `scale_f32_ungraphed` (loaded from a separate CUmodule) so it is safe to + /// call outside CUDA graph capture without corrupting the forward_child graph state. + pub fn apply_c51_budget_scale(&mut self, c51_budget: f32) -> Result<(), MLError> { + // No-op when budget ≈ 1.0 (common at health=0 — avoids unnecessary kernel launch). + if (c51_budget - 1.0).abs() < 1e-6 { + return Ok(()); + } + let grad_ptr = self.ptrs.grad_buf; + let n = self.total_params as i32; + let blocks = ((self.total_params as u32 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.scale_f32_ungraphed) + .arg(&grad_ptr) + .arg(&c51_budget) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("c51_budget_scale: {e}")))?; + } + Ok(()) + } + /// Add CQL gradient scratch to grad_buf via plain SAXPY. /// /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index e9356bf73..c6b958e64 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1159,7 +1159,14 @@ impl FusedTrainingCtx { // component balance. Safety is still handled by the single global clip in Adam. // B4/G5: Compute adaptive budgets once per step from ISV health/regime signals. - let (_c51_budget, iqn_budget, cql_budget, _ens_budget) = self.compute_adaptive_budgets(); + let (c51_budget, iqn_budget, cql_budget, _ens_budget) = self.compute_adaptive_budgets(); + + // F8/G5: Scale C51 contribution in grad_buf by c51_budget. + // C51 backward already wrote into grad_buf during submit_forward_ops_main (Phase 2). + // This scale MUST run before CQL/IQN SAXPYs below so only C51's portion is scaled. + // Order: c51_budget × C51_grad (here), then cql_budget × CQL_grad, iqn_budget × IQN_grad. + self.trainer.apply_c51_budget_scale(c51_budget) + .map_err(|e| anyhow::anyhow!("c51_budget_scale: {e}"))?; // DIAGNOSTIC: sync between each aux op to find hanging kernel