diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index ab88bf979..c3f0a657a 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5012,6 +5012,11 @@ extern "C" __global__ void regime_branch_gate( * spacing = softmax(spacing_raw) * positions[j] = v_min + (v_max - v_min) * cumsum(spacing)[j] * + * Per-branch bounds come from the ISV signal bus (slots 23 + 2*b for centre, + * 24 + 2*b for half-width). Written by `update_eval_v_range`; bootstrap values + * (centre=0, half=(config.v_max-config.v_min)/2) at construction produce the + * same `[config.v_min, config.v_max]` positions as the pre-ISV behaviour. + * * Single-block kernel: num_atoms=51 fits in one block. * Grid: (1,1,1), Block: (256,1,1). Shared memory: 2 * num_atoms * sizeof(float). */ @@ -5019,14 +5024,22 @@ extern "C" __global__ void adaptive_atom_positions( const float* __restrict__ spacing_raw, float* __restrict__ positions_out, int num_atoms, - float v_min, - float v_max + int branch_idx, + const float* __restrict__ isv_signals ) { int tid = threadIdx.x; extern __shared__ float shmem[]; float* s_softmax = shmem; float* s_cumsum = shmem + num_atoms; + /* Per-branch v-range from ISV slots. Layout: centre at 23+2*b, half at + * 24+2*b for b in [0..4). Caller guarantees isv_signals is non-null and + * branch_idx ∈ [0,4). */ + float v_center = isv_signals[23 + 2 * branch_idx]; + float v_half = isv_signals[24 + 2 * branch_idx]; + float v_min = v_center - v_half; + float v_max = v_center + v_half; + /* Load spacing_raw into shared memory */ if (tid < num_atoms) { s_softmax[tid] = spacing_raw[tid]; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4414969e2..2d157a879 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -140,7 +140,22 @@ const ISV_K: usize = 4; // Temporal ISV history length /// NOT rotated into history. Slot [22] (`SHARPE_EMA_INDEX`) added by the /// ISV-audit bundle (2026-04-23) — Rust host broadcasts `training_sharpe_ema` /// each epoch; the kernel reads it to drive slot 12 (learning_health). -const ISV_DIM: usize = 23; +/// Network-facing ISV width: the length of the vector fed into `w_isv_fc1` +/// attention weights. **Never change without a checkpoint format migration** — +/// the w_isv_fc1 weight tensor is sized [16, ISV_NETWORK_DIM] at 368 floats. +const ISV_NETWORK_DIM: usize = 23; +/// Total ISV slot count including cross-kernel scratchpad slots that never +/// feed the attention path. Slots [23..31] carry per-branch Q-support +/// (center, half-width) for direction/magnitude/order/urgency — written by +/// `update_eval_v_range`, read by `adaptive_atom_positions`, +/// `warm_start_atom_positions`, and per-sample support tiling. Extending +/// past 31 requires updating the pinned allocation in the constructor and +/// any kernel that takes an ISV pointer expecting a specific length. +const ISV_TOTAL_DIM: usize = 31; +/// Legacy alias preserved for call sites that haven't been audited for the +/// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight +/// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer). +const ISV_DIM: usize = ISV_NETWORK_DIM; pub const LEARNING_HEALTH_INDEX: usize = 12; /// Task 2.X "make Full useful": ISV slots for per-magnitude-bin Q-mean EMAs. /// Read by `c51_loss_batched` and `c51_grad_kernel` to compute the @@ -171,6 +186,23 @@ pub const Q_DIR_ABS_REF_INDEX: usize = 21; /// component-aggregation formula whose Sharpe-correlation measured at /// r=-0.765 on the train-mdh86 logs that triggered the audit. 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` +/// (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]` +/// behaviour byte-for-byte across all four branches. +pub const V_CENTER_DIR_INDEX: usize = 23; +pub const V_HALF_DIR_INDEX: usize = 24; +pub const V_CENTER_MAG_INDEX: usize = 25; +pub const V_HALF_MAG_INDEX: usize = 26; +pub const V_CENTER_ORD_INDEX: usize = 27; +pub const V_HALF_ORD_INDEX: usize = 28; +pub const V_CENTER_URG_INDEX: usize = 29; +pub const V_HALF_URG_INDEX: usize = 30; const ISV_EMB_DIM: usize = 8; // ISV embedding output dimension (FC2 output, gate/gamma input) /// First ISV tensor index in the flat param buffer. /// ISV weights (68-79) are online-only — NOT synced to the target network. @@ -481,6 +513,21 @@ pub struct QValueStatsResult { pub atom_utilization: f32, } +/// Per-branch Q-value statistics — one 7-tuple per action branch (direction=0, +/// magnitude=1, order=2, urgency=3). Produced by `q_stats_per_branch_reduce` +/// kernel; consumed by `update_eval_v_range` to track per-branch EMA state +/// and write the 8 ISV v-range slots. +/// +/// The per-branch atom-entropy / utilization slots in the kernel output are +/// left at 0 (per-branch atom_stats would require splitting the global +/// accumulator). `QValueStatsResult.atom_{entropy,utilization}` inherit 0 +/// until that wiring lands — the v-range EMA updater doesn't use those two +/// fields, so the zero does not propagate into the ISV path. +#[derive(Debug, Clone, Copy)] +pub struct PerBranchQValueStats { + pub per_branch: [QValueStatsResult; 4], +} + /// Scalar-only result from the fused GPU training path. /// /// TD errors stay on GPU (`td_errors_buf`) — no readback. PER priority update @@ -1210,10 +1257,13 @@ pub struct GpuDqnTrainer { /// Pinned device-mapped: CPU writes v_min/v_max, GPU reads via dev_ptr. No HtoD copy. eval_v_range_pinned: *mut f32, eval_v_range_ptr: u64, - /// EMA-smoothed Q-stats for stable eval v_range updates. - pub(crate) eval_q_mean_ema: f32, - pub(crate) eval_q_std_ema: f32, - eval_ema_initialized: bool, + /// Per-branch EMA-smoothed Q-stats for the ISV v-range update. Index + /// corresponds to branch (direction=0, magnitude=1, order=2, urgency=3). + /// Promoted from scalar to `[f32; 4]` as part of the ISV v-range + /// unification — see `update_eval_v_range`. + pub(crate) eval_q_mean_ema: [f32; 4], + pub(crate) eval_q_std_ema: [f32; 4], + eval_ema_initialized: [bool; 4], /// Adaptive IQN lambda readiness: 0=uncertain (suppress gradient), 1=converged (full weight). iqn_readiness: f32, /// IQN readiness — pinned device-mapped for CUDA graph. GPU reads via dev_ptr. @@ -1533,6 +1583,14 @@ pub struct GpuDqnTrainer { q_stats_kernel: CudaFunction, /// GPU buffer for Q-value statistics [7 floats] q_stats_buf: CudaSlice, + /// Per-branch Q-stats reducer — writes 28 floats to `per_branch_q_stats_pinned` + /// (4 × 7, branch-major). Runs alongside `q_stats_kernel`; feeds the + /// ISV v-range per-branch EMA state in `update_eval_v_range`. + q_stats_per_branch_kernel: CudaFunction, + /// Pinned device-mapped readback for per-branch Q-stats [4 × 7 = 28 floats]. + /// GPU writes via `per_branch_q_stats_dev_ptr`; CPU reads host-side. + per_branch_q_stats_pinned: *mut f32, + per_branch_q_stats_dev_ptr: u64, /// GPU buffer for atom utilization accumulation [2 floats: sum_entropy, sum_utilized]. /// Populated by `atom_stats_finalize` kernel (phase 2 of the deterministic reduction). atom_stats_buf: CudaSlice, @@ -1828,11 +1886,15 @@ pub struct GpuDqnTrainer { q_dir_bin_means_reduce_kernel: CudaFunction, // ── ISV core buffers (pinned device-mapped, GPU read/write) ── - isv_signals_pinned: *mut f32, // [ISV_DIM = 23] + // isv_signals_pinned is sized for ISV_TOTAL_DIM (31) — the extra eight + // slots [23..31] carry per-branch Q-support centres/half-widths. The + // network (w_isv_fc1) and history rotation still consume only the first + // ISV_NETWORK_DIM (23) slots; the tail is broadcast-bus scratchpad only. + isv_signals_pinned: *mut f32, // [ISV_TOTAL_DIM = 31] isv_signals_dev_ptr: u64, isv_history_pinned: *mut f32, // [ISV_K * 12 = 48] — history rotates slots [0..11] only isv_history_dev_ptr: u64, - isv_decay_pinned: *mut f32, // [ISV_DIM = 23] learned decay weights + isv_decay_pinned: *mut f32, // [ISV_NETWORK_DIM = 23] learned decay weights isv_decay_dev_ptr: u64, lagged_td_error_pinned: *mut f32, // [1] recursive confidence target lagged_td_error_dev_ptr: u64, @@ -2070,31 +2132,32 @@ impl GpuDqnTrainer { /// Recompute adaptive atom positions from learned spacing parameters. /// Call once per epoch (positions are slow-moving). + /// + /// 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> { let na = self.config.num_atoms; let param_sizes = compute_param_sizes(&self.config); let shmem = (na * 2) * std::mem::size_of::(); let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); - // Kernel signature expects `float v_min, float v_max` (f32). Passing - // `self.config.v_{min,max}` (f64) directly routes 8 bytes into 4-byte - // arg slots, corrupting both values and producing unbounded atom - // positions — the observed cause of Q-range explosion to ±144k while - // config bound is ±15. Cast to f32 explicitly to match the kernel ABI. - let v_min = self.config.v_min as f32; - let v_max = self.config.v_max as f32; + 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::()) as u64; let out_ptr = atom_positions_buf_ptr + out_offset; + let branch_i32 = branch as i32; unsafe { self.stream.launch_builder(&self.adaptive_atom_kernel) .arg(&spacing_ptr) .arg(&out_ptr) .arg(&(na as i32)) - .arg(&v_min) - .arg(&v_max) + .arg(&branch_i32) + .arg(&isv_dev_ptr) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), @@ -2109,8 +2172,14 @@ impl GpuDqnTrainer { /// Warm-start C51 atom positions from empirical reward quantiles. /// /// Called once per epoch after experience collection, before training steps. - /// Writes the same quantile values to all 4 branches — per-branch specialization - /// happens via the existing `step_atom_positions` SGD optimizer. + /// 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> { @@ -2119,10 +2188,31 @@ impl GpuDqnTrainer { return Ok(()); // skip on mismatch } - // Write same quantiles to all 4 branches + // 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 { - host_data[branch * na..(branch + 1) * na].copy_from_slice(quantiles); + 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) } + }; + 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) @@ -2306,6 +2396,9 @@ impl Drop for GpuDqnTrainer { if !self.q_readback_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.q_readback_pinned.cast()) }; } + if !self.per_branch_q_stats_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.per_branch_q_stats_pinned.cast()) }; + } if !self.q_divergence_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.q_divergence_pinned.cast()) }; } @@ -2932,65 +3025,92 @@ impl GpuDqnTrainer { Ok(()) } - /// Update eval v_range from observed Q-value statistics. - /// Called at epoch boundary after compute_q_stats. - pub fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32, q_gap: f32, per_branch_q_gaps: [f32; 4]) { - // Adaptive-rate EMA with baseline proportional to Q-std. - // Baseline = max(q_std_ema, 0.001) so normal Q-oscillation doesn't - // trigger aggressive tracking — only abnormal shifts do. - if !self.eval_ema_initialized { - self.eval_q_mean_ema = 0.0; - self.eval_q_std_ema = q_std.max(0.01); - self.eval_ema_initialized = true; - } - { - let baseline = self.eval_q_std_ema.max(0.001); - let mean_err = (q_mean - self.eval_q_mean_ema).abs(); - let std_err = (q_std - self.eval_q_std_ema).abs(); - let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.3); - let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.3); - self.eval_q_mean_ema = (1.0 - alpha_mean) * self.eval_q_mean_ema + alpha_mean * q_mean; - self.eval_q_std_ema = (1.0 - alpha_std) * self.eval_q_std_ema + alpha_std * q_std; - } + /// Update eval v_range from observed per-branch Q-value statistics. + /// Called at epoch boundary after `reduce_current_q_stats_per_branch`. + /// + /// Per-branch path (ISV v-range unification, spec 2026-04-23): + /// maintains 4 independent `(q_mean_ema, q_std_ema)` pairs and writes 8 + /// ISV slots (`V_CENTER_{branch}_INDEX`, `V_HALF_{branch}_INDEX`) — one + /// pair of (centre, half) per branch. Consumers read through the ISV + /// signal bus (`adaptive_atom_positions`, `warm_start_atom_positions`). + /// + /// Legacy global `eval_v_range_pinned[2]` is still updated using the + /// direction-branch (branch 0) bounds so any consumer that has not yet + /// migrated to the per-branch ISV bus sees the same value as the + /// "dominant branch" behaviour prior to this change. Full removal of + /// `eval_v_range_pinned` is deferred to the Phase-3 cleanup. + pub fn update_eval_v_range( + &mut self, + per_branch_stats: &PerBranchQValueStats, + per_branch_q_gaps: [f32; 4], + ) { + let v_min_f = self.config.v_min; + let v_max_f = self.config.v_max; + let v_range_full = (v_max_f - v_min_f).max(1e-6); + // Prevents atom-grid collapse when the Q-distribution is narrow. + // Spec 2026-04-23: min_half_floor = 0.1 * (v_max - v_min). + let min_half_floor = 0.1_f32 * v_range_full; + let abs_half = 0.5_f32 * v_range_full; - // Keep last_per_branch_q_gaps in sync (per_branch_q_gaps pinned already written in reduce_current_q_stats). - let _ = per_branch_q_gaps; // already on device via per_branch_q_gaps_dev_ptr + for (branch_idx, stats) in per_branch_stats.per_branch.iter().enumerate() { + let q_mean = stats.q_mean; + let q_std = stats.q_variance.max(0.0).sqrt(); + let q_gap = per_branch_q_gaps[branch_idx].max(0.0); + + // Adaptive-rate EMA with baseline proportional to Q-std. + if !self.eval_ema_initialized[branch_idx] { + self.eval_q_mean_ema[branch_idx] = 0.0; + self.eval_q_std_ema[branch_idx] = q_std.max(0.01); + self.eval_ema_initialized[branch_idx] = true; + } else { + let baseline = self.eval_q_std_ema[branch_idx].max(0.001); + let mean_err = (q_mean - self.eval_q_mean_ema[branch_idx]).abs(); + let std_err = (q_std - self.eval_q_std_ema[branch_idx]).abs(); + let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.3); + let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.3); + self.eval_q_mean_ema[branch_idx] = + (1.0 - alpha_mean) * self.eval_q_mean_ema[branch_idx] + alpha_mean * q_mean; + self.eval_q_std_ema[branch_idx] = + (1.0 - alpha_std) * self.eval_q_std_ema[branch_idx] + alpha_std * q_std; + } + + // Width from max(Q-gap action scale, 3σ band, min floor), clamped + // to the full config range. Floor = 10% of config range — enough + // atom-grid resolution to differentiate actions regardless of the + // observed distribution. + let gap_width = (10.0 * q_gap) + .max(3.0 * self.eval_q_std_ema[branch_idx]) + .max(min_half_floor); + let half = gap_width.min(abs_half).max(min_half_floor); + // Centre clamp so the band always fits inside [v_min, v_max]. + let center = self.eval_q_mean_ema[branch_idx].clamp( + v_min_f + half, + v_max_f - half, + ); + + unsafe { + *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch_idx) = center; + *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch_idx + 1) = half; + } + + // Branch 0 (direction) drives the legacy scalar support — matches + // the dominant-branch Q-scale the single-EMA path used. + if branch_idx == 0 && !self.eval_v_range_pinned.is_null() { + unsafe { + *self.eval_v_range_pinned = center - half; + *self.eval_v_range_pinned.add(1) = center + half; + } + } + } // GPU kernel: liquid tau RK4/Euler adaptive ODE (zero CPU compute). - let delta_z = self.eval_q_std_ema.max(0.01); + // Uses the direction-branch (b0) Q-std as the global delta_z proxy — + // matches the single-EMA path before per-branch promotion. + let delta_z = self.eval_q_std_ema[0].max(0.01); if let Err(e) = self.update_liquid_tau(delta_z) { tracing::warn!("liquid_tau_rk4 kernel failed: {e}"); } - - // Width from Q-gap (action differentiation range), not Q-std. - // Floor = max(10 * q_gap, 3 * q_std) — enough atoms to resolve actions. - // With 51 atoms, a width of 10*q_gap gives q_gap/delta_z ≈ 5 atoms of - // resolution between best and worst action. - // - // Clamp the adaptive [v_min, v_max] to the theoretical bounds supplied - // at construction (config.v_{min,max}, derived from reward_scale/gamma). - // Without this clamp the window can drift arbitrarily far from 0: once - // the policy overestimates, eval_q_mean_ema tracks it up → v_max follows - // → C51 atoms span a wider range → even larger TD targets → loss - // explodes geometrically (Fold 1 smoke: final_loss=1.6e16, grad=62B). - // - // Two clamps are needed — one on half-width so the band can't exceed - // the full theoretical range, and one on the centre so the band always - // fits inside [config.v_min, config.v_max]: - let gap_width = (10.0 * q_gap).max(3.0 * self.eval_q_std_ema).max(0.1); - let abs_half = (self.config.v_max - self.config.v_min) * 0.5; - let half = gap_width.min(abs_half); - let center = self.eval_q_mean_ema.clamp( - self.config.v_min + half, - self.config.v_max - half, - ); - let v_min = center - half; - let v_max = center + half; - // Pinned device-mapped: CPU writes directly, GPU reads via dev_ptr. No HtoD copy. - unsafe { - *self.eval_v_range_pinned = v_min; - *self.eval_v_range_pinned.add(1) = v_max; - } + let _ = per_branch_q_gaps; // already on device via per_branch_q_gaps_dev_ptr } pub fn set_per_sample_support_ptr(&mut self, ptr: u64) { self.per_sample_support_ptr = ptr; @@ -6837,12 +6957,37 @@ impl GpuDqnTrainer { // ── Compile expected Q-value + stats kernels (validation, not in CUDA Graph) ─ let expected_q_kernel = compile_expected_q_kernel(&stream)?; let q_stats_kernel = compile_q_stats_kernel(&stream)?; + // ISV v-range unification (spec 2026-04-23): per-branch Q-stats + // reducer — emits one 7-tuple per action branch (4 × 7 = 28 floats). + // Drives the per-branch EMA state in `update_eval_v_range` which in + // turn writes the 8 ISV v-range slots. + let q_stats_per_branch_kernel = compile_q_stats_per_branch_kernel(&stream)?; // Task 2.X "make Full useful" — per-magnitude-bin Q-mean reducer that // feeds ISV slots [13..16] on the stats cadence. let q_mag_bin_means_reduce_kernel = compile_q_mag_bin_means_kernel(&stream)?; let q_dir_bin_means_reduce_kernel = compile_q_dir_bin_means_kernel(&stream)?; let q_stats_buf = stream.alloc_zeros::(7) .map_err(|e| MLError::ModelError(format!("alloc q_stats_f32: {e}")))?; + // Per-branch Q-stats readback — pinned device-mapped, 28 floats (4 × 7). + // GPU writes, CPU reads without sync (one-step lag tolerated; the + // update_eval_v_range path sees the latest value after the next + // replay completes). + let per_branch_q_stats_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(28 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned per_branch_q_stats alloc: {e}")))? + as *mut f32 + }; + unsafe { std::ptr::write_bytes(per_branch_q_stats_pinned, 0, 28); } + let per_branch_q_stats_dev_ptr = unsafe { + let mut dp = 0u64; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dp as *mut u64, + per_branch_q_stats_pinned.cast(), + 0, + ); + dp + }; let atom_stats_buf = stream.alloc_zeros::(2) .map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?; // Per-block partial sums for the atom-stat deterministic reduction. @@ -7820,8 +7965,12 @@ impl GpuDqnTrainer { }; // ── ISV core buffers (pinned device-mapped) ────────────────────── + // Allocation uses `ISV_TOTAL_DIM` so slots [23..31] (per-branch + // Q-support centres and half-widths) have backing storage. The + // network-facing path still consumes only the first `ISV_NETWORK_DIM` + // slots — the scratchpad tail is invisible to `w_isv_fc1`. let (isv_signals_pinned, isv_signals_dev_ptr) = { - let num_bytes = ISV_DIM * std::mem::size_of::(); + let num_bytes = ISV_TOTAL_DIM * std::mem::size_of::(); let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); let mut dev_ptr_out: u64 = 0; unsafe { @@ -7830,6 +7979,18 @@ impl GpuDqnTrainer { let rc2 = cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dev_ptr_out, host_ptr, 0); assert_eq!(rc2, cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS, "cuMemHostGetDevicePointer for isv_signals"); std::ptr::write_bytes(host_ptr as *mut u8, 0, num_bytes); + // Bootstrap the per-branch v-range slots to the config bounds so + // that any consumer reaching for slots [23..31] before the first + // `update_eval_v_range` call sees identical-to-config behaviour. + // Half-width = (v_max - v_min) / 2, centre = 0 → v_min..v_max. + let v_min_f = config.v_min as f32; + let v_max_f = config.v_max as f32; + let half_bootstrap = 0.5_f32 * (v_max_f - v_min_f); + let sig_ptr = host_ptr as *mut f32; + for b in 0..4usize { + *sig_ptr.add(23 + 2 * b) = 0.0_f32; // v_center + *sig_ptr.add(23 + 2 * b + 1) = half_bootstrap; // v_half + } } (host_ptr as *mut f32, dev_ptr_out) }; @@ -8648,9 +8809,9 @@ impl GpuDqnTrainer { branch_scales_ptr: 0, eval_v_range_pinned, eval_v_range_ptr, - eval_q_mean_ema: 0.0, - eval_q_std_ema: 0.0, - eval_ema_initialized: false, + eval_q_mean_ema: [0.0_f32; 4], + eval_q_std_ema: [0.0_f32; 4], + eval_ema_initialized: [false; 4], iqn_readiness: 0.0, iqn_readiness_pinned, iqn_readiness_dev_ptr, @@ -8734,6 +8895,9 @@ impl GpuDqnTrainer { bw_d_glu_gate, expected_q_kernel, q_stats_kernel, + q_stats_per_branch_kernel, + per_branch_q_stats_pinned, + per_branch_q_stats_dev_ptr, q_mag_bin_means_reduce_kernel, q_mag_means_scratch_pinned, q_mag_means_scratch_dev_ptr, @@ -9199,8 +9363,13 @@ impl GpuDqnTrainer { Ok(()) } - /// Read eval_q_mean_ema. - pub fn eval_q_mean_ema(&self) -> f32 { self.eval_q_mean_ema } + /// Read eval_q_mean_ema as the direction-branch scalar — matches the + /// pre-unification single-EMA behaviour for callers that have not yet + /// migrated to per-branch awareness. + pub fn eval_q_mean_ema(&self) -> f32 { self.eval_q_mean_ema[0] } + + /// Read all four per-branch `q_mean_ema` values. + pub fn eval_q_mean_ema_per_branch(&self) -> [f32; 4] { self.eval_q_mean_ema } /// ISV signals device pointer for adaptive hold enforcement in experience collector. pub fn isv_signals_dev_ptr(&self) -> u64 { self.isv_signals_dev_ptr } @@ -9595,8 +9764,13 @@ impl GpuDqnTrainer { sigma1 / sigma2 } - /// Read eval_q_std_ema. - pub fn eval_q_std_ema(&self) -> f32 { self.eval_q_std_ema } + /// Read eval_q_std_ema as the direction-branch scalar — matches the + /// pre-unification single-EMA behaviour for callers that have not yet + /// migrated to per-branch awareness. + pub fn eval_q_std_ema(&self) -> f32 { self.eval_q_std_ema[0] } + + /// Read all four per-branch `q_std_ema` values. + pub fn eval_q_std_ema_per_branch(&self) -> [f32; 4] { self.eval_q_std_ema } // G6 branch_indep_loss_value + G10 temporal_loss_value removed — // their buffers were deallocated after V7-gem sub-noise verdict. @@ -9630,36 +9804,78 @@ impl GpuDqnTrainer { v } - /// Set eval_q_mean_ema (for trajectory backtracking restore). - pub fn set_eval_q_mean_ema(&mut self, v: f32) { self.eval_q_mean_ema = v; } + /// Set eval_q_mean_ema for all four branches (trajectory backtracking + /// restore). Scalar setter kept for backwards compatibility — broadcasts + /// the same value to every branch. Callers that snapshot per-branch state + /// should use `set_eval_q_mean_ema_per_branch`. + pub fn set_eval_q_mean_ema(&mut self, v: f32) { + self.eval_q_mean_ema = [v; 4]; + self.eval_ema_initialized = [v != 0.0; 4]; + } - /// Set eval_q_std_ema (for trajectory backtracking restore). - pub fn set_eval_q_std_ema(&mut self, v: f32) { self.eval_q_std_ema = v; } + /// Set eval_q_std_ema for all four branches (trajectory backtracking + /// restore). Scalar setter kept for backwards compatibility — broadcasts + /// the same value to every branch. Callers that snapshot per-branch state + /// should use `set_eval_q_std_ema_per_branch`. + pub fn set_eval_q_std_ema(&mut self, v: f32) { + self.eval_q_std_ema = [v; 4]; + } + + /// Per-branch setter for trajectory backtracking restore. + pub fn set_eval_q_mean_ema_per_branch(&mut self, v: [f32; 4]) { + self.eval_q_mean_ema = v; + // Treat any non-zero value as evidence the EMA was observed before. + for b in 0..4 { + self.eval_ema_initialized[b] = v[b] != 0.0 || self.eval_ema_initialized[b]; + } + } + + /// Per-branch setter for trajectory backtracking restore. + pub fn set_eval_q_std_ema_per_branch(&mut self, v: [f32; 4]) { + self.eval_q_std_ema = v; + } /// Reset the adaptive C51 v_range state at a fold boundary. /// - /// Clears `eval_q_mean_ema` / `eval_q_std_ema` (the rolling Q-distribution - /// estimate) and resets `eval_v_range_pinned` back to the theoretical - /// `[config.v_min, config.v_max]`. Without this, Fold N+1 inherits Fold N's - /// final tight atom support — if the new fold's Q-distribution doesn't fit - /// the stale range, TD errors explode and training NaN's out (observed: + /// Clears the per-branch `eval_q_mean_ema` / `eval_q_std_ema` arrays and + /// resets both the legacy `eval_v_range_pinned[2]` buffer and the 8 + /// per-branch ISV slots (23..30) back to bootstrap values so that fold + /// N+1 starts with atoms spanning the full `[config.v_min, config.v_max]` + /// range on every branch. Without this, Fold N+1 inherits Fold N's final + /// tight atom support — if the new fold's Q-distribution doesn't fit the + /// stale range, TD errors explode and training NaN's out (observed: /// train-92xbj Fold 1 Epoch 12, grad_norm=inf → NaN loss at step 5). /// - /// The first call to `update_eval_v_range` after this reset will reinitialise - /// the EMA from the new fold's first observed (q_mean, q_std). + /// The first call to `update_eval_v_range` after this reset will + /// reinitialise each branch's EMA from the new fold's first observed + /// `(q_mean, q_std)`. pub fn reset_eval_v_range_state(&mut self) { - self.eval_q_mean_ema = 0.0; - self.eval_q_std_ema = 0.0; - self.eval_ema_initialized = false; - // Restore the pinned device-mapped range to the theoretical bounds so - // the experience-collection kernels running before the next - // update_eval_v_range call see the wide safe support. + self.eval_q_mean_ema = [0.0_f32; 4]; + self.eval_q_std_ema = [0.0_f32; 4]; + self.eval_ema_initialized = [false; 4]; + // Restore the pinned device-mapped legacy range to the theoretical + // bounds so the experience-collection kernels running before the + // next update_eval_v_range call see the wide safe support. if !self.eval_v_range_pinned.is_null() { unsafe { *self.eval_v_range_pinned = self.config.v_min; *self.eval_v_range_pinned.add(1) = self.config.v_max; } } + // Bootstrap the 8 per-branch ISV slots to centre=0, half=abs_half. + // Matches the construction-time bootstrap — epoch 1 of the next fold + // sees `[config.v_min, config.v_max]` on every branch. + if !self.isv_signals_pinned.is_null() { + let v_min_f = self.config.v_min; + let v_max_f = self.config.v_max; + let half_bootstrap = 0.5_f32 * (v_max_f - v_min_f); + unsafe { + for b in 0..4usize { + *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * b) = 0.0_f32; + *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * b + 1) = half_bootstrap; + } + } + } } /// Set per_branch_q_gap_ema — HtoD for trajectory backtracking restore. @@ -11198,6 +11414,93 @@ impl GpuDqnTrainer { Ok(result) } + /// Reduce current `q_out_buf` to per-branch Q-stats (4 × 7 floats). + /// + /// Launches `q_stats_per_branch_reduce` — slices the flat Q-output by + /// branch and computes `[avg_max_q, q_min, q_max, q_mean, q_var, + /// atom_entropy(=0), atom_utilization(=0)]` per branch. Drives the + /// `update_eval_v_range` per-branch EMA state and the 8 ISV v-range + /// slots (spec 2026-04-23). + /// + /// Non-destructive to `q_stats_buf` / `q_readback_pinned` — uses a + /// separate pinned readback. Intended to be called alongside + /// `reduce_current_q_stats` at epoch boundary; both share the same + /// `q_out_buf` population so there is no extra forward pass. + pub fn reduce_current_q_stats_per_branch(&mut self) + -> Result + { + let _eg = EventTrackingGuard::new(self.stream.context()); + + let total_actions = self.total_actions() as i32; + let n = self.config.batch_size as i32; + let b0 = self.config.branch_0_size as i32; + let b1 = self.config.branch_1_size as i32; + let b2 = self.config.branch_2_size as i32; + let b3 = self.config.branch_3_size as i32; + let b0_off: i32 = 0; + let b1_off: i32 = b0; + let b2_off: i32 = b0 + b1; + let b3_off: i32 = b0 + b1 + b2; + + let q_out_buf_ptr = self.q_out_buf.raw_ptr(); + let out_ptr = self.per_branch_q_stats_dev_ptr; + + unsafe { + self.stream + .launch_builder(&self.q_stats_per_branch_kernel) + .arg(&q_out_buf_ptr) + .arg(&out_ptr) + .arg(&n) + .arg(&total_actions) + .arg(&b0_off).arg(&b0) + .arg(&b1_off).arg(&b1) + .arg(&b2_off).arg(&b2) + .arg(&b3_off).arg(&b3) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("q_stats_per_branch_reduce: {e}")))?; + } + + // Synchronous read from pinned host buffer. The kernel is launched on + // the same stream as the C51 graph, so a single stream synchronize + // guarantees the pinned bytes reflect the current launch — matches the + // legacy `compute_q_stats` cadence. + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + + let mut host = [0.0_f32; 28]; + unsafe { + std::ptr::copy_nonoverlapping( + self.per_branch_q_stats_pinned, host.as_mut_ptr(), 28, + ); + } + + let mut per_branch = [QValueStatsResult { + avg_max_q: 0.0, + q_min: 0.0, + q_max: 0.0, + q_mean: 0.0, + q_variance: 0.0, + atom_entropy: 0.0, + atom_utilization: 0.0, + }; 4]; + for b in 0..4_usize { + let base = b * 7; + per_branch[b] = QValueStatsResult { + avg_max_q: host[base] as f64, + q_min: host[base + 1], + q_max: host[base + 2], + q_mean: host[base + 3], + q_variance: host[base + 4], + atom_entropy: host[base + 5], + atom_utilization: host[base + 6], + }; + } + Ok(PerBranchQValueStats { per_branch }) + } + /// Launch the xLSTM mLSTM step: reads q_stats_buf[7] + per_branch_q_gaps[4] (pinned), /// updates persistent matrix memory C[8,8] and normalizer n[8], /// writes context[8] output. Runs in a single thread (grid=1, block=1). @@ -14887,6 +15190,22 @@ fn compile_q_stats_kernel( .map_err(|e| MLError::ModelError(format!("q_stats_reduce load: {e}"))) } +/// Load the per-branch Q-value statistics reducer. +/// +/// Emits one 7-tuple per branch (direction, magnitude, order, urgency) for the +/// ISV v-range per-branch EMA update path. See `q_stats_kernel.cu` for the +/// `q_stats_per_branch_reduce` body. Runs alongside the legacy global +/// `q_stats_reduce` — no semantic conflict, both single-threaded deterministic. +fn compile_q_stats_per_branch_kernel( + stream: &Arc, +) -> Result { + let context = stream.context(); + let module = context.load_cubin(Q_STATS_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("q_stats cubin load: {e}")))?; + module.load_function("q_stats_per_branch_reduce") + .map_err(|e| MLError::ModelError(format!("q_stats_per_branch_reduce load: {e}"))) +} + /// Load the per-magnitude-bin Q-mean reducer from the q_stats cubin. /// /// Task 2.X "make Full useful": this kernel computes batch-mean Q-values diff --git a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu index a7b6bb9b1..f385ff5a3 100644 --- a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu +++ b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu @@ -61,6 +61,108 @@ extern "C" __global__ void q_stats_reduce( out[6] = (atom_stats != NULL) ? atom_stats[1] * inv_n / (float)num_atoms : 0.0f; /* utilization fraction [0,1] */ } +/** + * Per-branch Q-value statistics reduction. + * + * Parallel path to `q_stats_reduce` above. Emits one `[7]` stats tuple per + * action branch (direction, magnitude, order, urgency) — 28 floats total, + * branch-major. + * + * Layout of `out[4 * 7]`: + * out[b*7 + 0] = avg_max_q_b (mean of per-sample max over branch_b's actions) + * out[b*7 + 1] = q_min_b (min Q across all samples × branch_b's actions) + * out[b*7 + 2] = q_max_b + * out[b*7 + 3] = q_mean_b (mean Q across samples × branch_b's actions) + * out[b*7 + 4] = q_var_b (variance) + * out[b*7 + 5] = atom_entropy_b (NaN — atom_stats kernel is global, not per-branch yet) + * out[b*7 + 6] = atom_util_b (NaN — same) + * + * Used by `update_eval_v_range` to maintain per-branch EMA state for the + * ISV v-range broadcast (slots 23..30). Branches have genuinely different + * Q-scales — direction Q ≈ ±50, magnitude Q ≈ ±5 — and a single shared + * range forces atoms into a bad quantization for at least one branch. + * + * Branch slicing: given `q_values[N, total_actions]`, branch b's actions + * occupy indices `[branch_offsets[b], branch_offsets[b] + branch_sizes[b])`. + * Caller passes `branch_offsets[4]` and `branch_sizes[4]`. + * + * Grid=(1,1,1) Block=(1,1,1) — single-threaded determinism matching + * q_stats_reduce. Cost: ~N * total_actions FP32 loads split across 4 branches + * plus a second pass for variance; single-digit microseconds at typical B. + */ +extern "C" __global__ void q_stats_per_branch_reduce( + const float* __restrict__ q_values, /* [N, total_actions] */ + float* __restrict__ out, /* [4 * 7] branch-major */ + int N, + int total_actions, + int b0_off, int b0_size, + int b1_off, int b1_size, + int b2_off, int b2_size, + int b3_off, int b3_size) +{ + if (threadIdx.x != 0) return; + + const int branch_offsets[4] = { b0_off, b1_off, b2_off, b3_off }; + const int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size }; + + for (int b = 0; b < 4; b++) { + const int off = branch_offsets[b]; + const int sz = branch_sizes[b]; + + /* Bounds check — invalid branch produces zeros, not garbage. */ + if (sz <= 0 || off < 0 || off + sz > total_actions || N <= 0) { + for (int k = 0; k < 7; k++) out[b * 7 + k] = 0.0f; + continue; + } + + float branch_min = 1e30f; + float branch_max = -1e30f; + float branch_sum = 0.0f; + float sum_max_q = 0.0f; + const int total = N * sz; + + /* First pass: min/max/mean/avg_max. */ + for (int i = 0; i < N; i++) { + const float* row = q_values + (long long)i * total_actions + off; + float row_max = -1e30f; + for (int a = 0; a < sz; a++) { + float v = row[a]; + if (v < branch_min) branch_min = v; + if (v > branch_max) branch_max = v; + if (v > row_max) row_max = v; + branch_sum += v; + } + sum_max_q += row_max; + } + + const float mean = (total > 0) ? branch_sum / (float)total : 0.0f; + + /* Second pass: variance. */ + float var_sum = 0.0f; + for (int i = 0; i < N; i++) { + const float* row = q_values + (long long)i * total_actions + off; + for (int a = 0; a < sz; a++) { + float d = row[a] - mean; + var_sum += d * d; + } + } + const float variance = (total > 0) ? var_sum / (float)total : 0.0f; + const float avg_max = (N > 0) ? sum_max_q / (float)N : 0.0f; + + out[b * 7 + 0] = avg_max; + out[b * 7 + 1] = branch_min; + out[b * 7 + 2] = branch_max; + out[b * 7 + 3] = mean; + out[b * 7 + 4] = variance; + /* Atom entropy / util remain global-only for now. Per-branch atom_stats + * would require splitting the compute_expected_q accumulator the same + * way; deferred to a follow-up since ISV v-range only needs the first + * five stats. Leave 0 — consumers check for it. */ + out[b * 7 + 5] = 0.0f; + out[b * 7 + 6] = 0.0f; + } +} + /** * Per-branch magnitude-Q mean reducer: computes batch-mean Q-values for each * of the 3 magnitude bins separately (Quarter=0, Half=1, Full=2), plus the diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2495fff9a..184ee6846 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2940,9 +2940,26 @@ impl FusedTrainingCtx { } vr } - /// Update eval v_range from observed Q-value statistics. - pub(crate) fn update_eval_v_range(&mut self, q_mean: f32, q_std: f32, q_gap: f32, per_branch_q_gaps: [f32; 4]) { - self.trainer.update_eval_v_range(q_mean, q_std, q_gap, per_branch_q_gaps); + /// Update eval v_range from observed per-branch Q-value statistics. + /// Spec 2026-04-23: writes 8 ISV slots (centre + half-width per branch) + /// and, for backward compatibility, the legacy scalar `eval_v_range_pinned` + /// from the direction-branch bounds. + pub(crate) fn update_eval_v_range( + &mut self, + per_branch_stats: &crate::cuda_pipeline::gpu_dqn_trainer::PerBranchQValueStats, + per_branch_q_gaps: [f32; 4], + ) { + self.trainer.update_eval_v_range(per_branch_stats, per_branch_q_gaps); + } + + /// Reduce per-branch Q-stats from the current `q_out_buf` — 4 × 7 floats + /// covering `[avg_max_q, q_min, q_max, q_mean, q_var, atom_entropy(=0), + /// atom_utilization(=0)]` per branch. Feeds `update_eval_v_range` for the + /// ISV v-range per-branch EMA bus. + pub(crate) fn reduce_current_q_stats_per_branch(&mut self) + -> Result + { + self.trainer.reduce_current_q_stats_per_branch() } /// Per-sample epsilon from IQL expectile gap. pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 5de0b01e8..7889ee534 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1722,12 +1722,12 @@ impl DQNTrainer { self.epoch_atom_entropy = stats.atom_entropy; self.epoch_atom_utilization = stats.atom_utilization; - // Update eval v_range from observed Q-stats + Q-gap - let q_std = stats.q_variance.max(0.0).sqrt(); - let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0); - // Per-branch Q-gaps computed in reduce_current_q_stats (48B DtoH readback). + // Per-branch Q-stats for the ISV v-range EMA updater. + // Per-branch Q-gaps pinned in reduce_current_q_stats (48B DtoH readback). let per_branch_q_gaps = fused.get_per_branch_q_gaps(); - fused.update_eval_v_range(stats.q_mean, q_std, q_gap, per_branch_q_gaps); + if let Ok(per_branch_stats) = fused.reduce_current_q_stats_per_branch() { + fused.update_eval_v_range(&per_branch_stats, per_branch_q_gaps); + } } } } // end if false (Q-stats disabled per-step) @@ -1756,10 +1756,18 @@ impl DQNTrainer { // the only thing that updates the buffer — so it stays at // zero and the downstream health composition believes // q_gap=0 even when raw epoch q_gap peaks above 1.0. - let q_std = stats.q_variance.max(0.0).sqrt(); - let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0); let per_branch_q_gaps = fused.get_per_branch_q_gaps(); - fused.update_eval_v_range(stats.q_mean, q_std, q_gap, per_branch_q_gaps); + match fused.reduce_current_q_stats_per_branch() { + Ok(per_branch_stats) => { + fused.update_eval_v_range(&per_branch_stats, per_branch_q_gaps); + } + Err(e) => { + tracing::warn!( + "reduce_current_q_stats_per_branch failed at epoch-boundary: {e} — \ + ISV v-range update skipped this step" + ); + } + } } } diff --git a/docs/superpowers/specs/2026-04-23-isv-v-range-unification.md b/docs/superpowers/specs/2026-04-23-isv-v-range-unification.md new file mode 100644 index 000000000..74e931573 --- /dev/null +++ b/docs/superpowers/specs/2026-04-23-isv-v-range-unification.md @@ -0,0 +1,280 @@ +# ISV-Unified Q-Value Support Range (per-branch) + +> **Problem statement:** Three Q-support consumers (atom grid, quantile warm-start, loss projection) hold independent opinions about the valid Q range for each of the 4 action branches. When they disagree, atoms become a bad quantizer: probability mass collapses onto 1-2 boundary atoms, the policy can't differentiate actions, and reward signal doesn't flow into useful Q-gradient information. + +## Observation driving this design + +- `train-fpxnw` (current run, Q-clamp fix applied): atom utilization **11-15%** across 57 fold-0 epochs. Probability concentrates on 1-2 of 51 atoms. +- Peak fold-0 Sharpe trending down as correctness fixes land: **79 → 68 → 48**. Each fix removed an accidental exploration crutch (wide Q range = wide atom grid = noisy Q-differentiation = effective exploration). +- Hypothesis under test: the Sharpe regression is not caused by our fixes being wrong — it's caused by *remaining* range inconsistencies between consumers. Atoms sit in one range, loss projects onto another, quantile-init fills a third. Each fix exposed one layer of the mismatch but left the others. + +## Design principles + +1. **One signal, one source of truth** for Q-support range per branch. +2. **Adaptive, temporal, signal-driven** — per the `adaptive_not_tuned` rule. EMA-tracked from observed Q-stats. +3. **Safety-clamped** to static `config.v_{min,max}` — hard rail, not tuning dial. +4. **Per-branch** — direction, magnitude, order, urgency have genuinely different Q scales; forcing a shared range recreates the exact mismatch we're fixing, just moved one level deeper. +5. **ISV bus broadcast** — established pattern for cross-kernel signals; consumers read from the bus they already receive. + +## Current state (before this spec) + +| Component | Q-range source | Type | +|-----------|----------------|------| +| `atom_positions_buf[4, num_atoms]` | `adaptive_atom_positions` kernel with `config.v_{min,max}` | **Static, per-config** | +| `per_sample_support_buf[N, 3]` | `eval_v_range_pinned`, tiled per epoch | **Adaptive, single-EMA, shared across branches** | +| `warm_start_atom_positions` quantile clamp | `config.v_{min,max}` (commit `d38a8cf99`) | **Static, per-config** | +| `q_stats_reduce` kernel output | Global min/max/mean/var across all branches | **Not per-branch** | + +The four components disagree along two axes: +- **Static vs adaptive**: atoms + quantile = static; projection = adaptive. +- **Per-branch vs pooled**: nothing is per-branch today; magnitude's Q=±5 and direction's Q=±50 share atoms and support bounds. + +## Proposed state + +| Component | New Q-range source | Type | +|-----------|-------------------|------| +| `atom_positions_buf[4, num_atoms]` | ISV per-branch center/half-width | **Adaptive, per-branch** | +| `per_sample_support_buf` — deprecated | ISV per-branch center/half-width (read per-branch inside loss kernel) | **Adaptive, per-branch** | +| `warm_start_atom_positions` quantile clamp | ISV per-branch center/half-width (read per-branch on host) | **Adaptive, per-branch** | +| `q_stats_reduce` kernel output | Per-branch `q_mean, q_std, q_gap` → 4 sets of stats | **Per-branch** | + +One source of truth: the ISV bus. Broadcast written by one path, read by three consumers. + +## ISV slot allocation + +Current ISV layout: 23 slots (slot 22 = `sharpe_ema`). Extend to 31: + +| Slot | Name | Semantics | +|------|------|-----------| +| 23 | `v_center_dir` | EMA Q-center, direction branch | +| 24 | `v_half_dir` | EMA Q-half-width, direction branch | +| 25 | `v_center_mag` | EMA Q-center, magnitude branch | +| 26 | `v_half_mag` | EMA Q-half-width, magnitude branch | +| 27 | `v_center_ord` | EMA Q-center, order branch | +| 28 | `v_half_ord` | EMA Q-half-width, order branch | +| 29 | `v_center_urg` | EMA Q-center, urgency branch | +| 30 | `v_half_urg` | EMA Q-half-width, urgency branch | + +**Why center + half-width, not min + max:** the adaptive EMA computes these two quantities naturally. Center is mean-Q EMA; half-width is `max(k·std, k·gap, min_floor)`. Consumers compute `v_min = center - half_width`, `v_max = center + half_width` themselves. Storing (center, half-width) keeps the evolution independent — center can shift while width stays constant. + +**Bootstrap values** (pre any Q observations, set at construction + fold-boundary reset): +- `v_center_* = 0.0` +- `v_half_* = (config.v_max - config.v_min) / 2.0` + +At epoch 1 these produce `[config.v_min, config.v_max]` for every branch — byte-identical to current behavior. + +**Safety clamp** on every ISV write: +- `center = center_raw.clamp(config.v_min + min_half, config.v_max - min_half)` +- `half = half_raw.clamp(min_half_floor, (config.v_max - config.v_min) / 2.0)` + +where `min_half_floor = 0.1 * (config.v_max - config.v_min)` prevents atom-grid collapse. Guarantees atoms always span a usable fraction of the config range. + +## Prerequisite: per-branch Q-stats kernel + +The ISV write path needs per-branch `q_mean`, `q_std`, `q_gap` statistics. Currently `q_stats_reduce` produces one global tuple across all 108 actions. This must be reworked first. + +**Option A (extend existing kernel):** modify `q_stats_reduce` to emit `[4 × 7]` instead of `[7]`. Output format per branch: `[avg_max_q, q_min, q_max, q_mean, q_var, atom_entropy, atom_utilization]`. + +**Option B (separate per-branch kernel):** add `q_stats_per_branch_reduce` alongside the existing one. Keeps the legacy kernel for any global consumers, adds a parallel path. + +**Choice: Option A.** The global stats are already computed implicitly in the per-branch reduction (sum-across-branches of per-branch sums). Callers that need the global tuple can aggregate from the per-branch output. Avoids kernel duplication. + +Branch offsets into `q_out_buf[N, total_actions]`: +- direction: `[0, b0)` +- magnitude: `[b0, b0+b1)` +- order: `[b0+b1, b0+b1+b2)` +- urgency: `[b0+b1+b2, b0+b1+b2+b3)` + +Pass `b0, b1, b2, b3` to the kernel; iterate per sample, per branch, reduce. + +**Readback shape change:** host side `QValueStatsResult` becomes `[QValueStatsResult; 4]` (or a new `PerBranchQValueStats`). Propagate the type change through `reduce_current_q_stats`, `update_eval_v_range`, consumers in `training_loop.rs`. + +## Write path + +Only `GpuDqnTrainer::update_eval_v_range` writes the ISV slots. It already runs at epoch boundary after `compute_q_stats`. New code flow: + +```rust +pub fn update_eval_v_range(&mut self, per_branch_stats: &[QValueStatsResult; 4]) { + for (branch_idx, stats) in per_branch_stats.iter().enumerate() { + // Adaptive-rate EMA per branch + let baseline = self.eval_q_std_ema[branch_idx].max(0.001); + let mean_err = (stats.q_mean - self.eval_q_mean_ema[branch_idx]).abs(); + let std_err = (stats.q_std() - self.eval_q_std_ema[branch_idx]).abs(); + let alpha_mean = (mean_err / (mean_err + baseline)).clamp(0.01, 0.3); + let alpha_std = (std_err / (std_err + baseline)).clamp(0.01, 0.3); + self.eval_q_mean_ema[branch_idx] = + (1.0 - alpha_mean) * self.eval_q_mean_ema[branch_idx] + alpha_mean * stats.q_mean; + self.eval_q_std_ema[branch_idx] = + (1.0 - alpha_std) * self.eval_q_std_ema[branch_idx] + alpha_std * stats.q_std(); + + // Derived center + half-width + let gap_width = (10.0 * stats.q_gap()).max(3.0 * self.eval_q_std_ema[branch_idx]).max(min_half_floor); + let abs_half = (self.config.v_max - self.config.v_min) as f32 * 0.5; + let half = gap_width.min(abs_half).max(min_half_floor); + let center = self.eval_q_mean_ema[branch_idx].clamp( + (self.config.v_min as f32) + half, + (self.config.v_max as f32) - half, + ); + + // Pinned device-mapped: writes visible to next kernel replay + unsafe { + *self.isv_signals_pinned.add(23 + 2 * branch_idx) = center; + *self.isv_signals_pinned.add(24 + 2 * branch_idx) = half; + } + } +} +``` + +Per-branch EMA state replaces the current single `eval_q_mean_ema, eval_q_std_ema` scalars — they become `[f32; 4]`. + +## Read-path changes + +**1. `adaptive_atom_positions` kernel (CUDA)** + +Current signature: +```cpp +extern "C" __global__ void adaptive_atom_positions( + const float* __restrict__ spacing_raw, // [NA] + float* __restrict__ positions_out, // [NA] + int num_atoms, + float v_min, // per-branch in call site loop, from config + float v_max +); +``` + +New signature: +```cpp +extern "C" __global__ void adaptive_atom_positions( + const float* __restrict__ spacing_raw, // [NA] + float* __restrict__ positions_out, // [NA] + int num_atoms, + int branch_idx, // 0..4 + const float* __restrict__ isv_signals // read slots 23+2*branch_idx, 24+2*branch_idx +); +``` + +Kernel body: +```cpp +float center = isv_signals[23 + 2 * branch_idx]; +float half = isv_signals[24 + 2 * branch_idx]; +float v_min = center - half; +float v_max = center + half; +// (rest unchanged — softmax + cumsum scaled by v_min, v_max) +``` + +Rust caller `recompute_atom_positions` changes: instead of passing config.v_min/v_max as f32, pass `branch_idx` and `isv_signals_dev_ptr`. The f64→f32 ABI fix from commit `768cc7d82` becomes moot (no more scalar f32 arg for bounds). + +**2. `warm_start_atom_positions` (Rust host side)** + +Current: reads reward quantiles, clamps with `hyperparams.v_{min,max}`, writes to atom_positions_buf. + +New: reads per-branch bounds from `isv_signals_pinned` (pinned, zero-cost host read), clamps each of the 4 branch quantile arrays independently: + +```rust +for branch_idx in 0..4 { + let center = unsafe { *isv_pinned.add(23 + 2 * branch_idx) }; + let half = unsafe { *isv_pinned.add(24 + 2 * branch_idx) }; + let v_min = center - half; + let v_max = center + half; + for (atom_idx, q) in branch_quantiles[branch_idx].iter_mut().enumerate() { + *q = q.clamp(v_min, v_max); + } +} +``` + +Currently `compute_reward_quantiles` returns a single `Vec` of size num_atoms shared across branches. Either: +- Keep shared quantiles (the reward distribution is the same regardless of branch), just clamp per-branch. +- Extract per-branch quantiles (requires per-branch reward attribution). **Out of scope** — shared quantiles with per-branch clamp is sufficient for the fix. + +**3. C51 loss + MSE loss kernels** + +Currently read `per_sample_support_buf[N, 3]` which is tiled from the single `eval_v_range_pinned`. For per-branch atom positions, the projection support in the loss kernel must also be per-branch. + +Two options: +- **Keep `per_sample_support_buf` but make it `[N, 4, 3]`** (per-sample, per-branch). Tiling cost: 4× but still one HtoD per epoch. Minimal kernel change — index `per_sample_support[sample, branch_idx, {min,max,dz}]`. +- **Remove `per_sample_support_buf`, read from ISV directly inside the loss kernel.** Smaller kernel arg surface, one fewer buffer. But every Bellman-projection iteration does an extra ISV read per atom. + +**Choice: option 1** (per-sample × per-branch support tile). Simpler kernel diff, caches well, matches the atom_positions layout. Tiling happens on epoch boundary from ISV slots. + +**4. `update_per_sample_support` (Rust) — rename & restructure** + +`update_per_sample_support(v_min, v_max)` becomes `update_per_sample_support_per_branch()` — reads all 8 ISV slots, tiles into `[N, 4, 3]`. + +## Cross-fold semantics + +`reset_eval_v_range_state()` extends to: +- Per-branch EMA state reset to 0.0 (center) and config-default (std). +- ISV slots 23..30 reset to bootstrap values. + +Fold N+1 starts with atoms spanning full config range per branch. After epoch 1's per-branch Q-stats, adapts to new fold's distribution. No cross-fold bleed. + +## Feedback-loop safety + +Same safety rails as before, now per-branch: +1. **Hard clamp** at `config.v_{min,max}` — per-branch half-width cannot exceed `(v_max - v_min) / 2`. +2. **Minimum half-width floor** — prevents atom-grid collapse when Q-distribution is narrow. +3. **Anti-LR controller**, **Q-divergence EMA**, **gradient clipping** — unchanged, all still apply. + +A per-branch runaway (e.g., magnitude Q grows without bound) is observable via `isv[24+2k]` hitting the config-clamp ceiling. Diagnosable directly from HEALTH_DIAG, which already reads ISV slots. + +## Test plan + +1. **Unit test: per-branch stats extraction** — feed known `q_out_buf` into `q_stats_reduce`, assert output matches per-branch hand-computed stats. +2. **Smoke test: ISV write-read round-trip** — run `update_eval_v_range` with synthetic per-branch stats, read back ISV slots 23..30, assert matches expected center + half. +3. **Smoke test: atom range tracks ISV** — inject ISV slots with known values, run `recompute_atom_positions`, assert atom_positions span `[center - half, center + half]` for each branch. +4. **Smoke test: quantile clamp uses per-branch bounds** — inject ISV values that differ per branch, verify `warm_start_atom_positions` clamps each branch's atoms to that branch's range. +5. **Smoke test: fold-boundary reset** — call `reset_for_fold`, assert all 8 ISV slots at bootstrap values. +6. **Smoke test: atom utilization** — target metric. Expect ≥40% after fix (up from 11-15%). +7. **L40S 60-epoch run** — compare Sharpe trajectory, peak Sharpe, atom util vs `train-fpxnw` baseline. + +## Implementation phases + +**Phase 0 — prerequisite: per-branch Q-stats kernel (~1 day)** +- Modify `q_stats_reduce` to emit `[4 × 7]` stats. +- Propagate `PerBranchQValueStats` type through Rust side. +- Unit-test equivalence: aggregated per-branch stats == old global stats. +- Commit as a data-structure refactor, no semantic behavior change. + +**Phase 1 — ISV plumbing (~half day, no functional change at epoch 1)** +- Grep-update `ISV_DIM` 23 → 31. +- Add per-branch EMA state `[f32; 4]` for mean + std. +- Bootstrap values set at construction. +- `update_eval_v_range` writes ISV slots 23..30. +- `reset_eval_v_range_state` resets slots. +- Consumers still reach for `config.v_{min,max}` — no semantic change. +- Commit: "ISV plumbing, zero behavior change." + +**Phase 2 — consumer migration (~half day, breaks behavior)** +- `adaptive_atom_positions` kernel: new signature, reads ISV. +- `recompute_atom_positions` Rust caller: passes branch_idx + isv_dev_ptr. +- `warm_start_atom_positions`: per-branch clamp from ISV pinned. +- `update_per_sample_support` rename + restructure to `[N, 4, 3]` tiling. +- C51/MSE loss kernels: index per-branch into support tile. +- Commit: "v-range unified via ISV, per-branch atom grids." + +**Phase 3 — validation + cleanup (~half day + 2.5h GPU)** +- Run smoke tests 1-6. +- L40S run, compare to `train-fpxnw`. +- Delete `eval_v_range_pinned` if redundant (probably is). + +Total: ~2.5 days engineering + 1 validation run. + +## Success criteria + +All four must hold: +1. **Correctness:** ISV slots 23..30 written from one place, read by three consumers; no other path reaches for `config.v_{min,max}` as a runtime range except through ISV. +2. **Atom utilization improved:** ≥40% on fold 0 (up from 11-15% currently). +3. **Sharpe not regressed:** peak fold-0 Sharpe ≥ current +48.19. +4. **Fold-1 stability holds:** Q-range at fold 1 epoch 1 is within ISV-derived bounds (no ±333k or ±144k excursions as in pre-fix runs). + +## Out of scope + +- Removing `eval_v_range_pinned` (deferred to Phase 3 cleanup). +- Per-branch reward quantile extraction (shared quantile with per-branch clamp is sufficient). +- Cross-fold inherit (fold-reset is the proven default; revisit only if convergence speed becomes a bottleneck). +- Additional ISV signals beyond the 8 slots for v-range. + +## Revision history + +- **v1 (earlier today)**: proposed 2 ISV slots (global center + half-width). Revised after recognizing that direction-branch Q-scale (~±50) differs from magnitude-branch (~±5) — a shared range recreates the exact mismatch the fix targets. +- **v2 (this version)**: 8 slots (4 centers + 4 half-widths). Per-branch Q-stats kernel added as explicit prerequisite.