gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cleanup: wire td_error ISV scratch + rewrite stale TODOs declaratively

Browse Source 
Part A of pre-L40S cleanup.

1. Wire td_error batch mean into ISV scratch (gpu_dqn_trainer.rs):
   `launch_loss_reduce` now runs the generic `c51_loss_reduce` kernel a
   second time over `td_errors_buf` into `td_error_scratch_dev_ptr`.
   ISV[2] (TD-error EMA in `isv_signal_update`) was previously reading
   a zero-initialised scratch and accumulated a constant-zero signal.
   This was a genuinely missing kernel writeback — the C51 loss kernel
   was already emitting per-sample |TD-error| into `td_errors_buf`
   (c51_loss_kernel.cu:1096), it just wasn't being batch-reduced.

   Reuses the existing `c51_loss_reduce` (generic mean-reduction, single
   block, deterministic) rather than adding a new kernel — no new CUDA
   surface, no ABI change.

2. Remove 2 stale TODOs from batched_backward.rs docstrings that
   described a migration that's actually complete:
   - Module docstring said "dqn_backward_kernel (atomicAdd path)
     remains active" — the atomicAdd kernel has been removed; cuBLAS
     backward is wired via launch_cublas_backward.
   - `backward_full` docstring said "gated behind TODO" — the function
     is actively called from the fused training step.

3. Rewrite 2 ISV scratch field comments as declarative: td_error_scratch
   is now wired (as per change 1); ensemble_var_scratch remains
   zero-initialised and its comment honestly describes that consumers
   (ISV[3] and [4]) treat it as unavailable. Per feedback_no_todo_fixme.md,
   replaces the TODO(isv) markers with declarative descriptions of
   current behaviour. Future wiring is tracked in the plan, not in code
   aspirational markers.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 08:21:03 +02:00
parent d9fee6ef8d
commit 7f92fa242c
2 changed files with 42 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -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)
///

34
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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(())
}