diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index ecd592612..12ba633d0 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -70,11 +70,12 @@ //! ## Phase status //! //! This module provides the building blocks for Phase 2 Task 2. The full -//! cuBLAS backward path requires a standalone C51 loss kernel that emits -//! `dL/d_logits` for each branch. Until that kernel exists, the existing -//! `dqn_backward_kernel` (atomicAdd path) remains active. The functions -//! here are wired via `GpuDqnTrainer::launch_cublas_backward()` with a -//! TODO gate that enables them once the C51 loss kernel is ready. +//! The cuBLAS backward path consumes `dL/d_logits` produced by the +//! standalone C51 loss kernel (`c51_loss_batched` in +//! `c51_loss_kernel.cu`) and is wired via +//! `GpuDqnTrainer::launch_cublas_backward()`, which is called from the +//! fused training step. The earlier 1-warp-per-sample +//! `dqn_backward_kernel` (atomicAdd path) has been removed. use std::collections::HashMap; use std::mem::ManuallyDrop; @@ -1453,9 +1454,10 @@ impl CublasBackwardSet { /// Full backward pass through the branching dueling network via cuBLAS SGEMM. /// - /// Requires pre-computed `dL/d_logits` for each head from the C51 loss kernel. - /// Until the standalone C51 loss kernel exists (Phase 2 Task 2), this function - /// is gated behind a TODO comment in `GpuDqnTrainer::launch_cublas_backward`. + /// Consumes pre-computed `dL/d_logits` for each head emitted by + /// `c51_loss_batched` (`c51_loss_kernel.cu`), and is invoked by + /// `GpuDqnTrainer::launch_cublas_backward` during the fused training + /// step. /// /// ## Backward order (reverse of forward) /// diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 752f30e5e..de66a57c6 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1732,10 +1732,20 @@ pub struct GpuDqnTrainer { risk_adam_step: i32, // ── ISV scratch scalars (pinned device-mapped, written by GPU kernels) ── - // TODO(isv): Wire C51 loss kernel to write mean |TD-error| here + // Zero-initialised scratch. No writer currently emits batch mean + // |TD-error| into this slot, so `isv_signals[2]` (the TD-error EMA + // in `isv_signal_update`) accumulates a constant-zero signal and + // effectively decays to zero. Downstream consumers treat ISV[2] as + // an unavailable signal. Keeping the buffer allocated preserves the + // `isv_signal_update` kernel ABI — a future C51-loss writeback can + // populate it without a signature change. td_error_scratch_pinned: *mut f32, td_error_scratch_dev_ptr: u64, - // TODO(isv): Wire ensemble aggregate kernel to write mean variance here + // Zero-initialised scratch for mean ensemble variance. Same + // situation as `td_error_scratch_pinned`: no kernel currently + // writes into it, so `isv_signals[3]` (ensemble-variance EMA) and + // `isv_signals[4]` (variance velocity) stay at zero. Buffer kept + // for `isv_signal_update` ABI stability. ensemble_var_scratch_pinned: *mut f32, ensemble_var_scratch_dev_ptr: u64, reward_scratch_pinned: *mut f32, @@ -12642,9 +12652,16 @@ impl GpuDqnTrainer { /// Deterministic loss reduction: sequential sum of per_sample_loss → total_loss / batch_size. /// Grid=(1,1,1), Block=(1,1,1). Must stay single-thread inside graph_mega on Hopper. + /// + /// Also emits the batch mean |TD-error| into `td_error_scratch_dev_ptr` + /// by running the same reduction kernel over `td_errors_buf`. This feeds + /// ISV[2] (the TD-error EMA in `isv_signal_update`), which downstream + /// consumers previously saw as a constant-zero signal because the scratch + /// was never written. fn launch_loss_reduce(&self, total_loss_dev_ptr: u64) -> Result<(), MLError> { let b = self.config.batch_size as i32; let per_sample_loss_buf_ptr = self.per_sample_loss_buf.raw_ptr(); + let td_errors_buf_ptr = self.td_errors_buf.raw_ptr(); unsafe { self.stream .launch_builder(&self.c51_loss_reduce_kernel) @@ -12657,6 +12674,19 @@ impl GpuDqnTrainer { shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!("c51_loss_reduce launch: {e}")))?; + /* Second launch: batch mean |TD-error| into the ISV scratch. + * Generic mean-reduction kernel, reused verbatim. */ + self.stream + .launch_builder(&self.c51_loss_reduce_kernel) + .arg(&td_errors_buf_ptr) + .arg(&self.td_error_scratch_dev_ptr) + .arg(&b) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("c51_td_error_reduce launch: {e}")))?; } Ok(()) }