feat(dqn): Plan 1 Task 9 — atoms_update GPU kernel + AtomsMonitor + delete CPU paths
atoms_update_kernel.cu: 4-block kernel (grid=(4,1,1), block=(256,1,1)) replaces the CPU loop that launched adaptive_atom_positions 4× per branch. Each block reads ISV v-range slots [23..31) and spacing_raw params[52..56); computes softmax + cumsum; writes atom_positions_buf[branch * num_atoms]. Same formula as the existing adaptive_atom_positions kernel in experience_kernels.cu. Deleted recompute_atom_positions and warm_start_atom_positions from GpuDqnTrainer and their fused_training.rs delegates. step_atom_positions now calls launch_atoms_update. training_loop.rs: recompute_atom_positions → launch_atoms_update; warm_start_atom_positions + compute_reward_quantiles block removed. AtomsMonitor reads ISV[V_CENTER_DIR_INDEX=23] as representative scalar; diagnose() exposes all 8 v-range slots [23..31). state_reset_registry.rs "follow-up task" descriptions updated for all 4 adaptive slots. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -88,6 +88,7 @@ fn main() {
|
||||
"epsilon_update_kernel.cu",
|
||||
"gamma_update_kernel.cu",
|
||||
"kelly_cap_update_kernel.cu",
|
||||
"atoms_update_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
104
crates/ml/src/cuda_pipeline/atoms_update_kernel.cu
Normal file
104
crates/ml/src/cuda_pipeline/atoms_update_kernel.cu
Normal file
@@ -0,0 +1,104 @@
|
||||
/* atoms_update — C51 adaptive atom position recompute, GPU-driven (Plan 1 Task 9).
|
||||
*
|
||||
* Replaces the CPU-side `recompute_atom_positions` loop which launched the
|
||||
* existing `adaptive_atom_positions` kernel once per branch. Here, all 4
|
||||
* branches are dispatched in a single launch: grid = (4, 1, 1),
|
||||
* block = (256, 1, 1), shared memory = 2 * num_atoms * sizeof(float).
|
||||
* blockIdx.x selects the branch (0=dir, 1=mag, 2=ord, 3=urg).
|
||||
*
|
||||
* Reads:
|
||||
* spacing_raw[b] — learnable softmax spacing params for branch b,
|
||||
* accessed via pre-resolved device pointers passed as args.
|
||||
* isv_signals — ISV bus (per-branch v-range at slots 23+2*b / 24+2*b).
|
||||
*
|
||||
* Writes:
|
||||
* positions_out — atom grid [4 * num_atoms], branch b at offset b*num_atoms.
|
||||
*
|
||||
* Formula (same as adaptive_atom_positions in experience_kernels.cu):
|
||||
* spacing = softmax(spacing_raw)
|
||||
* positions[j] = v_min + (v_max - v_min) * cumsum(spacing)[j]
|
||||
* with v_min = isv[23+2*b] - isv[24+2*b], v_max = isv[23+2*b] + isv[24+2*b].
|
||||
*/
|
||||
|
||||
extern "C" __global__ void atoms_update(
|
||||
unsigned long long spacing_ptr_b0, /* device ptr to spacing_raw[dir] */
|
||||
unsigned long long spacing_ptr_b1, /* device ptr to spacing_raw[mag] */
|
||||
unsigned long long spacing_ptr_b2, /* device ptr to spacing_raw[ord] */
|
||||
unsigned long long spacing_ptr_b3, /* device ptr to spacing_raw[urg] */
|
||||
float* __restrict__ positions_out, /* [4 * num_atoms] */
|
||||
int num_atoms,
|
||||
const float* __restrict__ isv_signals
|
||||
) {
|
||||
int branch = (int)blockIdx.x;
|
||||
if (branch >= 4) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
extern __shared__ float shmem[];
|
||||
float* s_softmax = shmem;
|
||||
float* s_cumsum = shmem + num_atoms;
|
||||
|
||||
/* Per-branch v-range from ISV: centre at 23+2*b, half at 24+2*b */
|
||||
float v_center = isv_signals[23 + 2 * branch];
|
||||
float v_half = isv_signals[24 + 2 * branch];
|
||||
float v_min = v_center - v_half;
|
||||
float v_max = v_center + v_half;
|
||||
|
||||
/* Resolve spacing_raw pointer for this branch */
|
||||
const float* spacing_raw;
|
||||
switch (branch) {
|
||||
case 0: spacing_raw = (const float*)spacing_ptr_b0; break;
|
||||
case 1: spacing_raw = (const float*)spacing_ptr_b1; break;
|
||||
case 2: spacing_raw = (const float*)spacing_ptr_b2; break;
|
||||
default: spacing_raw = (const float*)spacing_ptr_b3; break;
|
||||
}
|
||||
|
||||
/* Load spacing_raw into shared memory */
|
||||
if (tid < num_atoms) {
|
||||
s_softmax[tid] = spacing_raw[tid];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* Find max for numerical stability (single thread, num_atoms <= 256 is tiny) */
|
||||
if (tid == 0) {
|
||||
float max_val = s_softmax[0];
|
||||
for (int j = 1; j < num_atoms; j++)
|
||||
max_val = fmaxf(max_val, s_softmax[j]);
|
||||
s_cumsum[0] = max_val; /* store max temporarily */
|
||||
}
|
||||
__syncthreads();
|
||||
float max_val = s_cumsum[0];
|
||||
__syncthreads();
|
||||
|
||||
/* Compute softmax */
|
||||
if (tid < num_atoms) {
|
||||
s_softmax[tid] = expf(s_softmax[tid] - max_val);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* Sum for normalization */
|
||||
if (tid == 0) {
|
||||
float sum = 0.0f;
|
||||
for (int j = 0; j < num_atoms; j++)
|
||||
sum += s_softmax[j];
|
||||
float inv_sum = 1.0f / fmaxf(sum, 1e-8f);
|
||||
for (int j = 0; j < num_atoms; j++)
|
||||
s_softmax[j] *= inv_sum;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* Inclusive prefix sum (sequential, tid=0 only) */
|
||||
if (tid == 0) {
|
||||
s_cumsum[0] = s_softmax[0];
|
||||
for (int j = 1; j < num_atoms; j++)
|
||||
s_cumsum[j] = s_cumsum[j-1] + s_softmax[j];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* Write positions: positions[0] = v_min, positions[j>0] = v_min + range * cumsum[j-1] */
|
||||
float* out = positions_out + (long long)branch * num_atoms;
|
||||
if (tid < num_atoms) {
|
||||
float range = v_max - v_min;
|
||||
float frac = (tid == 0) ? 0.0f : s_cumsum[tid - 1];
|
||||
out[tid] = v_min + range * frac;
|
||||
}
|
||||
}
|
||||
@@ -85,6 +85,8 @@ static EPSILON_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/e
|
||||
static GAMMA_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gamma_update_kernel.cubin"));
|
||||
/// Plan 1 Task 11: GPU-driven effective Kelly cap.
|
||||
static KELLY_CAP_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/kelly_cap_update_kernel.cubin"));
|
||||
/// Plan 1 Task 9: GPU-driven C51 atom position recompute (4-branch, single launch).
|
||||
static ATOMS_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/atoms_update_kernel.cubin"));
|
||||
|
||||
/// Mamba2 temporal scan configuration.
|
||||
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
|
||||
@@ -228,8 +230,8 @@ pub const SHARPE_EMA_INDEX: usize = 22;
|
||||
/// ISV v-range unification bundle (spec 2026-04-23): per-branch Q-support
|
||||
/// centre + half-width, 2 slots per branch (direction=0, magnitude=1,
|
||||
/// order=2, urgency=3). Written by `update_eval_v_range` from the per-branch
|
||||
/// Q-stats EMA state; read by `adaptive_atom_positions` (atom grid),
|
||||
/// `warm_start_atom_positions` (quantile clamp) and `update_per_sample_support`
|
||||
/// Q-stats EMA state; read by the `atoms_update` GPU kernel (`launch_atoms_update`,
|
||||
/// Plan 1 Task 9) and `update_per_sample_support`
|
||||
/// (loss kernel projection range). Bootstrap values at construction + fold
|
||||
/// reset: `v_center = 0`, `v_half = (config.v_max - config.v_min) / 2` — at
|
||||
/// epoch 1 this reproduces the pre-spec `[config.v_min, config.v_max]`
|
||||
@@ -1416,6 +1418,9 @@ pub struct GpuDqnTrainer {
|
||||
/// Plan 1 Task 11: GPU-driven effective Kelly cap. Cold-path (per-epoch).
|
||||
/// Reads ISV[12] + portfolio_state[n_envs, PS_STRIDE]; writes ISV[KELLY_CAP_EFF_INDEX=44].
|
||||
kelly_cap_update_kernel: CudaFunction,
|
||||
/// Plan 1 Task 9: GPU-driven C51 atom position recompute. Cold-path (per-epoch + SGD step).
|
||||
/// Reads ISV v-range slots [23..31) + spacing_raw params; writes atom_positions_buf.
|
||||
atoms_update_kernel: CudaFunction,
|
||||
/// Host-side cached per-component norms populated at epoch boundary
|
||||
/// from the pinned result slot. Index order matches `grad_mag_*_ratio`
|
||||
/// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
|
||||
@@ -2325,136 +2330,46 @@ impl GpuDqnTrainer {
|
||||
// NO graph invalidation — lr_dev_ptr is captured, value is read at replay time.
|
||||
}
|
||||
|
||||
/// Recompute adaptive atom positions from learned spacing parameters.
|
||||
/// Call once per epoch (positions are slow-moving).
|
||||
/// Plan 1 Task 9: GPU-driven atom position recompute (all 4 branches, single launch).
|
||||
///
|
||||
/// Per-branch v-range bounds come from the ISV signal bus (slots
|
||||
/// `V_CENTER_{branch}_INDEX` and `V_HALF_{branch}_INDEX`). The bootstrap
|
||||
/// written at construction reproduces `[config.v_min, config.v_max]` for
|
||||
/// every branch before the first `update_eval_v_range` call, so epoch 1
|
||||
/// behaviour matches the pre-ISV path byte-for-byte.
|
||||
pub(crate) fn recompute_atom_positions(&self) -> Result<(), MLError> {
|
||||
/// Dispatches `atoms_update` kernel with grid=(4,1,1), block=(256,1,1).
|
||||
/// Each block handles one branch; reads ISV v-range slots [23..31) and
|
||||
/// spacing_raw params; writes atom_positions_buf in-place. Replaces the
|
||||
/// CPU-side loop over `adaptive_atom_positions` per branch. Must be called
|
||||
/// once per epoch boundary and after each `step_atom_positions` SGD update.
|
||||
pub(crate) fn launch_atoms_update(&self) -> Result<(), MLError> {
|
||||
let na = self.config.num_atoms;
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let shmem = (na * 2) * std::mem::size_of::<f32>();
|
||||
let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr();
|
||||
let isv_dev_ptr = self.isv_signals_dev_ptr;
|
||||
let positions_out = self.atom_positions_buf.raw_ptr();
|
||||
let isv_dev_ptr = self.isv_signals_dev_ptr;
|
||||
|
||||
for branch in 0..4_usize {
|
||||
let spacing_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 52 + branch);
|
||||
let out_offset = (branch * na * std::mem::size_of::<f32>()) as u64;
|
||||
let out_ptr = atom_positions_buf_ptr + out_offset;
|
||||
let branch_i32 = branch as i32;
|
||||
let sp_b0 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 52);
|
||||
let sp_b1 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 53);
|
||||
let sp_b2 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 54);
|
||||
let sp_b3 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 55);
|
||||
let na_i32 = na as i32;
|
||||
|
||||
// ISV v-range diag: log host-side ISV read used by the kernel.
|
||||
// The kernel reads ISV via isv_dev_ptr (device mirror); the
|
||||
// pinned mirror tracks the same writes from update_eval_v_range.
|
||||
let (diag_center, diag_half) = if !self.isv_signals_pinned.is_null() {
|
||||
unsafe {
|
||||
let c = *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch);
|
||||
let h = *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch + 1);
|
||||
(c, h)
|
||||
}
|
||||
} else {
|
||||
(f32::NAN, f32::NAN)
|
||||
};
|
||||
tracing::info!(
|
||||
target: "isv_vrange_diag",
|
||||
branch, diag_center, diag_half,
|
||||
"recompute_atom_positions ISV read"
|
||||
);
|
||||
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.adaptive_atom_kernel)
|
||||
.arg(&spacing_ptr)
|
||||
.arg(&out_ptr)
|
||||
.arg(&(na as i32))
|
||||
.arg(&branch_i32)
|
||||
.arg(&isv_dev_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: shmem as u32,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("adaptive_atom_positions branch {branch}: {e}")))?;
|
||||
}
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.atoms_update_kernel)
|
||||
.arg(&sp_b0)
|
||||
.arg(&sp_b1)
|
||||
.arg(&sp_b2)
|
||||
.arg(&sp_b3)
|
||||
.arg(&positions_out)
|
||||
.arg(&na_i32)
|
||||
.arg(&isv_dev_ptr)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (4, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: shmem as u32,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("atoms_update launch: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Warm-start C51 atom positions from empirical reward quantiles.
|
||||
///
|
||||
/// Called once per epoch after experience collection, before training steps.
|
||||
/// Shared empirical reward quantiles are clamped to each branch's
|
||||
/// ISV-derived `[centre - half, centre + half]` support before being
|
||||
/// tiled into `atom_positions_buf`. Prior to the first ISV write, the
|
||||
/// bootstrap values (centre=0, half=(config.v_max-config.v_min)/2) make
|
||||
/// the clamp equivalent to the old `config.v_{min,max}` clamp for every
|
||||
/// branch — epoch 1 behaviour is preserved byte-for-byte. Per-branch
|
||||
/// specialisation then flows through the existing `step_atom_positions`
|
||||
/// SGD optimiser.
|
||||
///
|
||||
/// The HtoD transfer is ~1us for 4 * num_atoms floats (typically 4 * 51 = 204).
|
||||
pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), MLError> {
|
||||
let na = self.config.num_atoms;
|
||||
if quantiles.len() != na {
|
||||
return Ok(()); // skip on mismatch
|
||||
}
|
||||
|
||||
// Fallback bounds used when the ISV bus has not yet been written —
|
||||
// keep the pre-spec behaviour (full config range, shared across
|
||||
// branches) as a safety net.
|
||||
let v_min_f = self.config.v_min;
|
||||
let v_max_f = self.config.v_max;
|
||||
let isv_ptr = self.isv_signals_pinned;
|
||||
|
||||
let mut host_data = vec![0.0_f32; 4 * na];
|
||||
for branch in 0..4_usize {
|
||||
let (b_min, b_max) = if isv_ptr.is_null() {
|
||||
(v_min_f, v_max_f)
|
||||
} else {
|
||||
let centre = unsafe { *isv_ptr.add(V_CENTER_DIR_INDEX + 2 * branch) };
|
||||
let half = unsafe { *isv_ptr.add(V_CENTER_DIR_INDEX + 2 * branch + 1) };
|
||||
let b_min = centre - half;
|
||||
let b_max = centre + half;
|
||||
// Guard against degenerate (half=0) reads; fall back to full
|
||||
// config range so warm-start never produces zero-range atoms.
|
||||
if (b_max - b_min).abs() < 1e-6 { (v_min_f, v_max_f) }
|
||||
else { (b_min, b_max) }
|
||||
};
|
||||
tracing::info!(
|
||||
target: "isv_vrange_diag",
|
||||
branch, b_min, b_max,
|
||||
isv_ptr_null = isv_ptr.is_null(),
|
||||
"warm_start clamp bounds"
|
||||
);
|
||||
let slot = &mut host_data[branch * na..(branch + 1) * na];
|
||||
for (atom_idx, q) in quantiles.iter().enumerate() {
|
||||
slot[atom_idx] = q.clamp(b_min, b_max);
|
||||
}
|
||||
}
|
||||
|
||||
self.stream.memcpy_htod(&host_data, &mut self.atom_positions_buf)
|
||||
.map_err(|e| MLError::ModelError(format!("warm_start_atom_positions HtoD: {e}")))?;
|
||||
|
||||
// ISV v-range diag: log what we just wrote per branch (min/max of
|
||||
// atom_positions) — host_data is the source of truth (no need for a
|
||||
// DtoH roundtrip since memcpy_htod doesn't mutate source on failure
|
||||
// and succeeded above).
|
||||
for branch in 0..4_usize {
|
||||
let slot = &host_data[branch * na..(branch + 1) * na];
|
||||
let amin = slot.iter().cloned().fold(f32::INFINITY, f32::min);
|
||||
let amax = slot.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
tracing::info!(
|
||||
target: "isv_vrange_diag",
|
||||
branch, amin, amax,
|
||||
"warm_start atom_positions written"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SGD update on adaptive atom spacing_raw parameters.
|
||||
/// Maximizes atom entropy (gradient ascent) — concentrates atoms where mass exists.
|
||||
/// Uses scale_f32_kernel to decay spacing_raw toward zero: softmax(0,...,0) = uniform.
|
||||
@@ -2487,8 +2402,8 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("atom_position_sgd branch {branch}: {e}")))?;
|
||||
}
|
||||
}
|
||||
// Recompute positions from updated spacing_raw
|
||||
self.recompute_atom_positions()?;
|
||||
// Recompute positions from updated spacing_raw via GPU atoms_update kernel.
|
||||
self.launch_atoms_update()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -7070,6 +6985,14 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("kelly_cap_update load: {e}")))?
|
||||
};
|
||||
|
||||
// Plan 1 Task 9: load atoms_update kernel (cold-path, per-epoch + SGD step).
|
||||
let atoms_update_kernel = {
|
||||
let module = stream.context().load_cubin(ATOMS_UPDATE_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("atoms_update cubin load: {e}")))?;
|
||||
module.load_function("atoms_update")
|
||||
.map_err(|e| MLError::ModelError(format!("atoms_update load: {e}")))?
|
||||
};
|
||||
|
||||
// eval_v_range — pinned device-mapped (zero-copy write from CPU).
|
||||
let eval_v_range_pinned: *mut f32 = unsafe {
|
||||
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
|
||||
@@ -9127,6 +9050,7 @@ impl GpuDqnTrainer {
|
||||
epsilon_update_kernel,
|
||||
gamma_update_kernel,
|
||||
kelly_cap_update_kernel,
|
||||
atoms_update_kernel,
|
||||
grad_component_norms_mag: [0.0_f32; 9],
|
||||
grad_component_norms_dir: [0.0_f32; 9],
|
||||
grad_component_norms_trunk: [0.0_f32; 9],
|
||||
|
||||
@@ -2960,9 +2960,10 @@ impl FusedTrainingCtx {
|
||||
/// Per-sample epsilon from IQL expectile gap.
|
||||
pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() }
|
||||
pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms }
|
||||
/// Warm-start C51 atom positions from empirical reward quantiles (once per epoch).
|
||||
pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), crate::MLError> {
|
||||
self.trainer.warm_start_atom_positions(quantiles)
|
||||
/// Plan 1 Task 9: GPU-driven atom position recompute (all 4 branches, single launch).
|
||||
pub(crate) fn launch_atoms_update(&self) -> anyhow::Result<()> {
|
||||
self.trainer.launch_atoms_update()
|
||||
.map_err(|e| anyhow::anyhow!("launch_atoms_update: {e}"))
|
||||
}
|
||||
pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); }
|
||||
pub(crate) fn adaptive_clip_value(&self) -> f32 { self.trainer.adaptive_clip_value() }
|
||||
@@ -3055,12 +3056,6 @@ impl FusedTrainingCtx {
|
||||
self.trainer.set_lr(lr);
|
||||
}
|
||||
|
||||
/// Recompute adaptive atom positions from learned spacing parameters (once per epoch).
|
||||
pub(crate) fn recompute_atom_positions(&self) -> anyhow::Result<()> {
|
||||
self.trainer.recompute_atom_positions()
|
||||
.map_err(|e| anyhow::anyhow!("recompute_atom_positions: {e}"))
|
||||
}
|
||||
|
||||
/// Read adam_step counter.
|
||||
pub(crate) fn adam_step_val(&self) -> i32 { self.trainer.adam_step_val() }
|
||||
|
||||
|
||||
118
crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs
Normal file
118
crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs
Normal file
@@ -0,0 +1,118 @@
|
||||
//! CPU-side observer for the GPU-driven C51 adaptive atom grid.
|
||||
//!
|
||||
//! The `atoms_update` CUDA kernel (Plan 1 Task 9) reads ISV v-range slots
|
||||
//! [23..31) (per-branch centre + half-width) and spacing_raw parameters to
|
||||
//! compute adaptive atom positions for all 4 branches in a single GPU launch.
|
||||
//! This monitor reads the representative `V_CENTER_DIR_INDEX=23` scalar for
|
||||
//! fire-rate tracking, and exposes all 8 v-range ISV slots in `diagnose()`.
|
||||
//!
|
||||
//! The monitor's `read()` returns `isv[V_CENTER_DIR_INDEX]`. `diagnose()`
|
||||
//! exposes the direction branch v-range (centre + half), all branch v-ranges,
|
||||
//! and fire-rate.
|
||||
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::{
|
||||
V_CENTER_DIR_INDEX, V_HALF_DIR_INDEX,
|
||||
V_CENTER_MAG_INDEX, V_HALF_MAG_INDEX,
|
||||
V_CENTER_ORD_INDEX, V_HALF_ORD_INDEX,
|
||||
V_CENTER_URG_INDEX, V_HALF_URG_INDEX,
|
||||
};
|
||||
use crate::trainers::dqn::adaptive_monitor::{
|
||||
AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
|
||||
};
|
||||
|
||||
pub(crate) struct AtomsMonitor {
|
||||
last_center: f32,
|
||||
fire_rate: FireRateStats,
|
||||
}
|
||||
|
||||
impl AtomsMonitor {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self { last_center: 0.0, fire_rate: FireRateStats::default() }
|
||||
}
|
||||
}
|
||||
|
||||
impl AdaptiveMonitor for AtomsMonitor {
|
||||
fn read(&self, isv: &IsvBus<'_>) -> f32 {
|
||||
isv.read(V_CENTER_DIR_INDEX)
|
||||
}
|
||||
|
||||
fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
|
||||
DiagSnapshot::new()
|
||||
.with("v_center_dir", isv.read(V_CENTER_DIR_INDEX))
|
||||
.with("v_half_dir", isv.read(V_HALF_DIR_INDEX))
|
||||
.with("v_center_mag", isv.read(V_CENTER_MAG_INDEX))
|
||||
.with("v_half_mag", isv.read(V_HALF_MAG_INDEX))
|
||||
.with("v_center_ord", isv.read(V_CENTER_ORD_INDEX))
|
||||
.with("v_half_ord", isv.read(V_HALF_ORD_INDEX))
|
||||
.with("v_center_urg", isv.read(V_CENTER_URG_INDEX))
|
||||
.with("v_half_urg", isv.read(V_HALF_URG_INDEX))
|
||||
.with("fire_rate", self.fire_rate.fire_rate())
|
||||
}
|
||||
|
||||
fn observe(&mut self, current: f32) {
|
||||
let fired = (current - self.last_center).abs() > 1e-6;
|
||||
self.fire_rate.record_fire(fired);
|
||||
self.last_center = current;
|
||||
}
|
||||
|
||||
fn fire_rate(&self) -> &FireRateStats {
|
||||
&self.fire_rate
|
||||
}
|
||||
|
||||
fn name(&self) -> &'static str {
|
||||
"atoms"
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_isv(center_dir: f32, half_dir: f32) -> Vec<f32> {
|
||||
let mut v = vec![0.0f32; 49];
|
||||
v[V_CENTER_DIR_INDEX] = center_dir;
|
||||
v[V_HALF_DIR_INDEX] = half_dir;
|
||||
v[V_CENTER_MAG_INDEX] = 0.1;
|
||||
v[V_HALF_MAG_INDEX] = 0.5;
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_returns_v_center_dir_slot() {
|
||||
let mut buf = make_isv(0.05, 2.5);
|
||||
let bus = IsvBus::new(&mut buf);
|
||||
let mon = AtomsMonitor::new();
|
||||
let v = mon.read(&bus);
|
||||
assert!((v - 0.05).abs() < 1e-7, "expected v_center_dir=0.05, got {v}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnose_exposes_all_fields() {
|
||||
let mut buf = make_isv(0.10, 3.0);
|
||||
let bus = IsvBus::new(&mut buf);
|
||||
let mon = AtomsMonitor::new();
|
||||
let snap = mon.diagnose(&bus);
|
||||
assert!(snap.fields.contains_key("v_center_dir"));
|
||||
assert!(snap.fields.contains_key("v_half_dir"));
|
||||
assert!(snap.fields.contains_key("v_center_mag"));
|
||||
assert!(snap.fields.contains_key("v_half_mag"));
|
||||
assert!(snap.fields.contains_key("v_center_ord"));
|
||||
assert!(snap.fields.contains_key("v_half_ord"));
|
||||
assert!(snap.fields.contains_key("v_center_urg"));
|
||||
assert!(snap.fields.contains_key("v_half_urg"));
|
||||
assert!(snap.fields.contains_key("fire_rate"));
|
||||
let center = snap.fields["v_center_dir"];
|
||||
assert!((center - 0.10).abs() < 1e-6, "expected v_center_dir=0.10, got {center}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn observe_tracks_fire_rate_on_change() {
|
||||
let mut mon = AtomsMonitor::new();
|
||||
mon.observe(0.0);
|
||||
mon.observe(0.0);
|
||||
mon.observe(0.05);
|
||||
let rate = mon.fire_rate().fire_rate();
|
||||
assert!((rate - 2.0 / 3.0).abs() < 1e-6,
|
||||
"expected 2/3 fire rate, got {rate}");
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,7 @@
|
||||
//! read those slots for HEALTH_DIAG emission and `controller_activity`
|
||||
//! smoke-test fire-rate tracking. No CPU-side adaptive computation.
|
||||
|
||||
pub(crate) mod atoms_monitor;
|
||||
pub(crate) mod epsilon_monitor;
|
||||
pub(crate) mod gamma_monitor;
|
||||
pub(crate) mod grad_balancer_monitor;
|
||||
|
||||
@@ -165,22 +165,22 @@ impl StateResetRegistry {
|
||||
RegistryEntry {
|
||||
name: "isv_epoch_idx",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[EPOCH_IDX_INDEX=39] — current epoch index; CPU writes at epoch boundary (follow-up task)",
|
||||
description: "ISV[EPOCH_IDX_INDEX=39] — current epoch index; CPU writes at epoch boundary (Plan 1 Task 13)",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_epsilon_eff",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[EPSILON_EFF_INDEX=41] — effective epsilon; GPU epsilon-adaptive kernel fills (follow-up task)",
|
||||
description: "ISV[EPSILON_EFF_INDEX=41] — effective epsilon; GPU epsilon_update kernel fills (Plan 1 Task 14)",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_tau_eff",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[TAU_EFF_INDEX=42] — effective tau; GPU tau-adaptive kernel fills (follow-up task)",
|
||||
description: "ISV[TAU_EFF_INDEX=42] — effective tau; GPU tau_update kernel fills (Plan 1 Task 13)",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_gamma_eff",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[GAMMA_EFF_INDEX=43] — effective gamma; GPU gamma-adaptive kernel fills (follow-up task)",
|
||||
description: "ISV[GAMMA_EFF_INDEX=43] — effective gamma; GPU gamma_update kernel fills (Plan 1 Task 10)",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_kelly_cap_eff",
|
||||
|
||||
@@ -352,10 +352,10 @@ impl DQNTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// Recompute adaptive atom positions from learned spacing_raw params (slow-moving, once per epoch).
|
||||
// Plan 1 Task 9: GPU-driven atom position recompute (4-branch, single launch).
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
if let Err(e) = fused.recompute_atom_positions() {
|
||||
tracing::warn!(epoch, "recompute_atom_positions failed: {e}");
|
||||
if let Err(e) = fused.launch_atoms_update() {
|
||||
tracing::warn!(epoch, "launch_atoms_update failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -517,38 +517,11 @@ impl DQNTrainer {
|
||||
// v8: TD(λ) is now wired in collect_experiences_gpu — self-bootstrap with
|
||||
// rewards as Q(s') approximation, overwrites n-step when td_lambda > 0.01.
|
||||
|
||||
// C51 atom warm-start: extract empirical reward quantiles from the bitonic
|
||||
// sort pipeline and initialize atom positions for this epoch. The existing
|
||||
// step_atom_positions SGD optimizer refines from here during training steps.
|
||||
if let Some(ref mut collector) = self.gpu_experience_collector {
|
||||
let count = collector.last_experience_count();
|
||||
let num_atoms = self.hyperparams.num_atoms;
|
||||
match collector.compute_reward_quantiles(count, num_atoms) {
|
||||
Ok(quantiles) => {
|
||||
// Per-branch clamp now lives inside
|
||||
// `warm_start_atom_positions` (ISV v-range spec Phase 2c,
|
||||
// commit 9deda5f65). That path reads each branch's
|
||||
// `[centre - half, centre + half]` from the ISV pinned
|
||||
// bus and clamps the quantile array per branch. An
|
||||
// outer static clamp to `config.v_{min,max}` here would
|
||||
// be redundant double-clamping to a wider bound — it
|
||||
// hides which layer owns the support and never
|
||||
// contributes once ISV is populated. Pass the raw
|
||||
// quantiles straight through.
|
||||
if let Some(ref mut fused) = self.fused_ctx {
|
||||
if let Err(e) = fused.warm_start_atom_positions(&quantiles) {
|
||||
warn!("C51 atom warm-start failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
debug!("C51 reward quantile extraction failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Robust PopArt: extract median/IQR from the already-sorted reward buffer.
|
||||
// The bitonic sort from compute_reward_quantiles above leaves reward_sort_keys_buf
|
||||
// sorted, so extract_median_iqr just reads 3 positions (tiny DtoH: 3 floats).
|
||||
// Robust PopArt: extract median/IQR from the reward sort buffer.
|
||||
// extract_median_iqr reads 3 positions from the sorted buffer (tiny DtoH: 3 floats).
|
||||
if self.hyperparams.popart_robust {
|
||||
match collector.extract_median_iqr(count) {
|
||||
Ok((median, iqr)) => {
|
||||
|
||||
@@ -71,6 +71,7 @@
|
||||
| `epsilon_update_kernel.cu` (Plan 1 Task 14) | single-thread kernel: reads ISV[39,40,12]; writes ISV[41] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK |
|
||||
| `gamma_update_kernel.cu` (Plan 1 Task 10) | single-thread kernel: reads ISV[12]; writes ISV[43] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; ISV[43] read by fill_gamma_buf on hot path via read_isv_signal_at (pinned, zero-copy) | OK |
|
||||
| `kelly_cap_update_kernel.cu` (Plan 1 Task 11) | single-thread kernel: reads ISV[12] + portfolio_states[n_envs, PS_STRIDE]; aggregates mean half-Kelly × health-safety; writes ISV[44] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; portfolio_states dev ptr passed from GpuExperienceCollector without CPU roundtrip | OK |
|
||||
| `atoms_update_kernel.cu` (Plan 1 Task 9) | 4-block kernel (one per branch): reads ISV[23..31) + spacing_raw params[52..56); writes atom_positions_buf[4 * num_atoms] via softmax + cumsum | COLD-PATH | Per-epoch boundary + SGD step (every 50 steps); grid=(4,1,1), block=(256,1,1); replaces CPU loop that launched adaptive_atom_positions 4× separately | OK |
|
||||
|
||||
## Summary
|
||||
|
||||
|
||||
@@ -52,6 +52,8 @@
|
||||
| `cuda_pipeline/gamma_update_kernel.cu` | `GpuDqnTrainer::launch_gamma_update` → `FusedTrainingCtx::launch_gamma_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 10 — cold-path (per-epoch) | — |
|
||||
| `trainers/dqn/monitors/kelly_cap_monitor.rs` | Read-only observer for kelly_cap_update kernel output (ISV slot 44); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 11 | — |
|
||||
| `cuda_pipeline/kelly_cap_update_kernel.cu` | `GpuDqnTrainer::launch_kelly_cap_update` → `FusedTrainingCtx::launch_kelly_cap_update` → `training_loop.rs` epoch boundary; takes portfolio_states dev ptr + n_envs | Wired | Plan 1 Task 11 — cold-path (per-epoch) | — |
|
||||
| `trainers/dqn/monitors/atoms_monitor.rs` | Read-only observer for atoms_update kernel inputs (ISV v-range slots 23..31); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 9 | — |
|
||||
| `cuda_pipeline/atoms_update_kernel.cu` | `GpuDqnTrainer::launch_atoms_update` → `FusedTrainingCtx::launch_atoms_update` → `training_loop.rs` epoch boundary; replaces per-branch `recompute_atom_positions` CPU loop | Wired | Plan 1 Task 9 — cold-path (per-epoch + SGD step) | — |
|
||||
|
||||
## CUDA Pipeline — Rust Wrappers
|
||||
|
||||
|
||||
Reference in New Issue
Block a user