diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index d375b16b6..3cd594823 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -28,10 +28,11 @@ extern "C" __global__ void c51_grad_kernel( const float* __restrict__ liquid_mod, /* [4] pinned device-mapped per-branch modulators */ const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = linear. */ const float* __restrict__ q_mean_ema_ptr, /* [1] pinned device-mapped Q-mean EMA */ - /* Task 2.X adaptive magnitude fix — ISV signal bus [ISV_DIM=15]. Slots - * [13]=q_mag_spread_ema, [14]=q_dir_spread_ema drive an adaptive bin - * weight applied to the magnitude branch (d==1). NULL → identity - * (safe: reduces to pre-fix behaviour). */ + /* Task 2.X adaptive magnitude fix + Task 2.Y adaptive direction fix — + * ISV signal bus [ISV_DIM=22]. Task 2.X slots [13..16] drive an adaptive + * bin weight applied to the magnitude branch (d==1); Task 2.Y slots + * [17..21] drive an adaptive bin weight applied to the direction branch + * (d==0). NULL → identity (safe: reduces to pre-fix behaviour). */ const float* __restrict__ isv_signals) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -114,6 +115,74 @@ extern "C" __global__ void c51_grad_kernel( } } + /* Task 2.Y "make direction branch useful at eval": scale + * direction-branch gradient by the signal-driven bin weight applied + * to branch_ce in c51_loss_batched (d==0). Consistency is required — + * forward CE and backward gradient must share the same scalar + * multiplier so the update step follows the true gradient of the + * (scaled) loss. + * + * Mirror of magnitude block above. See c51_loss_kernel.cu:: + * get_direction_bin_weight for rationale. Component A reads ISV + * slots [17..21] (per-direction Q-mean EMAs + |Q| ref) and fires + * when the taken direction's Q-mean lags max. Component B shares + * ISV slot [12] (learning_health) with the magnitude mechanism. + * dir_bias_signal[k] = {1.0, 0.5, 1.0, 0.0} for Short/Hold/Long/Flat: + * Flat always returns bin_weight=1.0 (no amplification of the bin + * the mechanism is meant to correct AGAINST). Bounded ∈ [1, 2]. */ + if (d == 0 && isv_signals != NULL) { + int a0 = branch_actions[0]; + if (b0_size > 0 && a0 >= 0 && a0 < b0_size && a0 < 4) { + /* Component A: bin-specific bias-vs-argmax (direction) */ + float q_short = isv_signals[17]; + float q_hold = isv_signals[18]; + float q_long = isv_signals[19]; + float q_flat = isv_signals[20]; + float q_dir_abs_ref = isv_signals[21]; + float means[4] = { q_short, q_hold, q_long, q_flat }; + float max_mean = means[0]; + for (int k = 1; k < 4; k++) + max_mean = fmaxf(max_mean, means[k]); + float frac_bin = 0.0f; + if (q_dir_abs_ref > 1e-6f) { + float bias_gap = fmaxf(0.0f, max_mean - means[a0]); + frac_bin = fminf(1.0f, bias_gap / fmaxf(q_dir_abs_ref, 1e-6f)); + } + /* Component B: shared meta-health stress signal */ + float health = fminf(1.0f, fmaxf(0.0f, isv_signals[12])); + float frac_health = 1.0f - health; + /* Composite */ + float collapse_frac = fminf(1.0f, frac_bin + frac_health); + if (collapse_frac > 0.0f) { + /* Architectural shape — NOT a tuning knob. Same + * rationale as get_direction_bin_weight in c51_loss_kernel. */ + float dir_bias_signal = 0.0f; + if (a0 == 0) dir_bias_signal = 1.0f; /* Short */ + else if (a0 == 1) dir_bias_signal = 0.5f; /* Hold */ + else if (a0 == 2) dir_bias_signal = 1.0f; /* Long */ + else if (a0 == 3) dir_bias_signal = 0.0f; /* Flat */ + if (dir_bias_signal > 0.0f) { + /* Cross-branch compression boost — mirror of the + * c51_loss_kernel get_direction_bin_weight formulation. + * When direction |Q|-scale is much tighter than + * magnitude |Q|-scale, amplify more aggressively. + * ISV-driven; self-disables when spreads equalise. */ + float q_abs_ref_mag = isv_signals[16]; + float compression = 0.0f; + if (q_abs_ref_mag > 1e-6f) { + compression = fmaxf(0.0f, + 1.0f - q_dir_abs_ref / fmaxf(q_abs_ref_mag, 1e-6f)); + compression = fminf(1.0f, compression); + } + float compression_boost = 1.0f + compression; /* ∈ [1,2] */ + float bin_weight = 1.0f + + collapse_frac * dir_bias_signal * compression_boost; + d_combined *= bin_weight; + } + } + } + } + d_val_sum += d_combined; float branch_scale = branch_scales[b * 4 + d]; diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index c53a2dcd4..8b380bc8a 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -300,6 +300,130 @@ __device__ __forceinline__ float get_magnitude_bin_weight( return 1.0f + collapse_frac * mag_bias_signal; } +/* Task 2.Y "make direction branch useful at eval" — signal-driven direction + * bin weight. Mirror-identical composite to `get_magnitude_bin_weight` above, + * applied to the direction axis (b0, 4 bins: Short=0, Hold=1, Long=2, Flat=3). + * + * Why this exists: post-Task-2.X diagnostics revealed the eval-mode Quarter=1 + * collapse is NOT magnitude-side but direction-side. Across 3 independent + * smoke runs on HEAD fa8d54661, EVAL_DIR_DIST Hold+Flat fractions were + * {0.000, 0.847, 0.452} — highly run-variant because C51's expected-Q bias + * flattens direction Q-values into a tight band, so strict-argmax picks + * whichever bin nudges above by <1e-3 (non-stationary across seeds). The + * kernel comment at experience_kernels.cu:825-829 names this exact risk: + * + * "Training Boltzmann (C51's expected Q structurally favors Flat → zero + * PnL variance → tightest distribution → highest softmax expected Q; + * argmax amplifies into irrecoverable Flat dominance..." + * + * ── Component A: bin-specific bias-vs-argmax (ISV slots [17..20] + [21]) ── + * max_mean_dir = max(Q_mean[Short,Hold,Long,Flat]) + * bias_gap_dir = max(0, max_mean_dir − Q_mean[a0]) + * frac_bin = min(1, bias_gap_dir / max(Q_abs_ref_dir, eps)) + * + * ── Component B: meta-health stress (ISV slot [12], shared with Task 2.X) ── + * frac_health = 1 − clamp(learning_health, 0, 1) + * + * ── Composite ── + * collapse_frac = min(1, frac_bin + frac_health) + * bin_weight = 1 + collapse_frac * dir_bias_signal(a0) ∈ [1, 2] + * + * dir_bias_signal per direction (architectural shape, NOT a tunable knob): + * Short (a0=0): 1.0 — tradable direction, amplify under-valued + * Hold (a0=1): 0.5 — position-preserving, moderate amplification + * Long (a0=2): 1.0 — tradable direction, amplify under-valued + * Flat (a0=3): 0.0 — no-trade; Flat is already over-estimated by C51's + * low-variance bias, so DO NOT amplify it further. + * bin_weight → 1.0 for Flat samples regardless of + * collapse_frac, correctly preventing the mechanism + * from reinforcing the pathology it is designed to fix. + * + * This layout is the direction-axis analogue of Task 2.X's monotonic + * `(a1 + 1) / b1_size` shape for magnitude — same architectural spirit + * (encode branch structure into the bin weight), different shape because + * direction bins are NOT a monotonic stake-size gradient (Short/Hold/Long/Flat + * are categorical). + * + * Self-disables when: + * * Component A: the taken direction IS the argmax of the per-bin + * Q-means (bias_gap = 0). + * * Component B: learning_health → 1 (training stabilised). + * * Flat samples always return 1.0 regardless of the above. + * + * Returns 1.0 if ISV is unavailable or a0 is out of range. `eps=1e-6f` + * matches the magnitude helper's numerical guard. */ +__device__ __forceinline__ float get_direction_bin_weight( + const float* __restrict__ isv_signals_ptr, + int a0, + int b0_size +) { + if (isv_signals_ptr == NULL) return 1.0f; + if (b0_size <= 0 || a0 < 0 || a0 >= b0_size || a0 >= 4) return 1.0f; + + /* Component A: bin-specific bias-vs-argmax (direction) */ + float q_short = isv_signals_ptr[17]; + float q_hold = isv_signals_ptr[18]; + float q_long = isv_signals_ptr[19]; + float q_flat = isv_signals_ptr[20]; + float q_dir_abs_ref = isv_signals_ptr[21]; + float mean_by_bin[4] = { q_short, q_hold, q_long, q_flat }; + float max_mean = mean_by_bin[0]; + for (int k = 1; k < 4; k++) + max_mean = fmaxf(max_mean, mean_by_bin[k]); + float frac_bin = 0.0f; + if (q_dir_abs_ref > 1e-6f) { + float bias_gap = fmaxf(0.0f, max_mean - mean_by_bin[a0]); + frac_bin = fminf(1.0f, bias_gap / fmaxf(q_dir_abs_ref, 1e-6f)); + } + + /* Component B: meta-health stress signal (shared ISV slot with magnitude) */ + float health = fminf(1.0f, fmaxf(0.0f, isv_signals_ptr[12])); + float frac_health = 1.0f - health; + + /* Composite — bounded ∈ [0, 1]. */ + float collapse_frac = fminf(1.0f, frac_bin + frac_health); + if (collapse_frac <= 0.0f) return 1.0f; + + /* Direction-specific bias signal (architectural shape constant): + * {Short=1.0, Hold=0.5, Long=1.0, Flat=0.0}. Flat returns 1.0 identity + * to prevent the mechanism from amplifying the bin it is meant to + * correct AGAINST. */ + float dir_bias_signal = 0.0f; + if (a0 == 0) dir_bias_signal = 1.0f; /* Short */ + else if (a0 == 1) dir_bias_signal = 0.5f; /* Hold */ + else if (a0 == 2) dir_bias_signal = 1.0f; /* Long */ + else if (a0 == 3) dir_bias_signal = 0.0f; /* Flat — do not amplify */ + + if (dir_bias_signal <= 0.0f) return 1.0f; + + /* Cross-branch compression boost (ISV-driven, no static knobs): when the + * direction-branch |Q|-scale is much tighter than the magnitude-branch + * |Q|-scale, the direction Q-values are in a compressed band that the + * default bin_weight ∈ [1, 2] cannot escape. Boost the amplification in + * proportion to how compressed the direction axis is relative to magnitude. + * + * compression = max(0, 1 - q_dir_abs_ref / max(q_abs_ref_mag, eps)) ∈ [0, 1] + * + * When direction spread matches magnitude spread, compression = 0, boost = 1, + * preserving the existing [1, 2] bound. When direction is much tighter (the + * observed pre-fix pathology — direction 16-73× tighter than magnitude), + * compression → 1 and boost → 2, yielding a bin_weight ∈ [1, 4]. Pure + * ISV-driven; self-disables as direction Q-values differentiate. + * + * The mag-branch ISV abs-ref (slot [16]) is the natural comparator — + * Task 2.X's own scale reference. This re-uses the existing bus without + * introducing a config knob or hard-coded target spread. */ + float q_abs_ref_mag = isv_signals_ptr[16]; + float compression = 0.0f; + if (q_abs_ref_mag > 1e-6f) { + compression = fmaxf(0.0f, 1.0f - q_dir_abs_ref / fmaxf(q_abs_ref_mag, 1e-6f)); + compression = fminf(1.0f, compression); + } + float compression_boost = 1.0f + compression; /* ∈ [1, 2] */ + + return 1.0f + collapse_frac * dir_bias_signal * compression_boost; +} + /* ══════════════════════════════════════════════════════════════════════ * MAIN KERNEL: c51_loss_batched (float arithmetic, BF16 I/O) * ══════════════════════════════════════════════════════════════════════ */ @@ -826,6 +950,20 @@ extern "C" __global__ void c51_loss_batched( branch_ce *= mag_bin_weight; } + /* Task 2.Y adaptive direction fix: scale direction-branch CE by a + * signal-driven bin weight (mirror of the magnitude mechanism + * applied to the direction axis). Identity (bin_weight=1.0) for + * Flat samples always, and for any sample whose direction Q-mean + * is the argmax across {Short, Hold, Long, Flat}. Self-disables + * once the direction head's Q-values differentiate healthily + * (frac_bin → 0) AND training stabilises (frac_health → 0). + * See get_direction_bin_weight. */ + if (d == 0) { + float dir_bin_weight = get_direction_bin_weight( + isv_signals, branch_action[0], branch_sizes[0]); + branch_ce *= dir_bin_weight; + } + total_ce += branch_ce; /* ═══ STEP f: Online-target Q-divergence for adaptive tau ══ diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 2740efe79..9ef950e38 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5536,7 +5536,7 @@ extern "C" __global__ void risk_budget_backward( /* ================================================================== */ extern "C" __global__ void isv_signal_update( - float* __restrict__ isv_signals, /* [ISV_DIM=17] pinned device-mapped */ + float* __restrict__ isv_signals, /* [ISV_DIM=22] pinned device-mapped */ float* __restrict__ isv_history, /* [K*12] pinned — history only rotates slots [0..11] */ float* __restrict__ lagged_td_error, /* [1] pinned — for recursive confidence target */ const float* __restrict__ q_mean_ptr, /* [1] pinned — batch Q-mean */ @@ -5555,7 +5555,13 @@ extern "C" __global__ void isv_signal_update( * the same stream before isv_signal_update. NULL = freeze slots. */ const float* __restrict__ q_mag_means_ptr, /* [mag_size] per-bin Q mean */ const float* __restrict__ q_abs_ref_ptr, /* [1] max(|Q_mean|) across mag bins */ - int mag_size /* typically 3 (Quarter/Half/Full) */ + int mag_size, /* typically 3 (Quarter/Half/Full) */ + /* Task 2.Y "make direction branch useful at eval" — per-direction Q-mean + * array + absolute-scale reference. Produced by q_dir_bin_means_reduce + * on the same stream before isv_signal_update. NULL = freeze slots. */ + const float* __restrict__ q_dir_means_ptr, /* [dir_size] per-bin Q mean */ + const float* __restrict__ q_dir_abs_ref_ptr,/* [1] max(|Q_mean|) across dir bins */ + int dir_size /* typically 4 (Short/Hold/Long/Flat) */ ) { if (threadIdx.x != 0) return; @@ -5668,6 +5674,43 @@ extern "C" __global__ void isv_signal_update( isv_signals[16] = (1.0f - alpha) * isv_signals[16] + alpha * fmaxf(q_abs_ref_ptr[0], 0.0f); } + + /* ── Task 2.Y "make direction branch useful at eval": per-direction Q-mean + * EMAs [17..20] and absolute-Q-scale reference EMA [21] ─────────── + * + * Direct structural mirror of Task 2.X's magnitude fix above, applied to + * branch 0 (direction). Consumed by c51_loss_batched / c51_grad_kernel to + * derive + * max_mean = max(Q_mean(Short), Q_mean(Hold), Q_mean(Long), Q_mean(Flat)) + * bias_gap = max(0, max_mean − Q_mean(a0)) + * collapse_frac = min(1, bias_gap / max(|Q|-scale, eps) + (1 - learning_health)) + * bin_weight = 1 + collapse_frac * dir_bias_signal(a0) + * + * where dir_bias_signal[k] = {1.0, 0.5, 0.0, 1.0} for Short/Hold/Long/Flat + * — architectural shape that amplifies tradable directions (Short/Long) + * and never amplifies Flat (the already-over-estimated low-variance bin + * that C51's expected-Q structurally favours). + * + * EMA tau=0.05 matches the existing alpha=0.05 pattern used for slots + * [13..16] (Task 2.X) and the rest of the ISV bus. + * + * NULL-guarded — missing inputs leave the slots unchanged, preserving + * backward compatibility during staged build-outs. + * + * NOT rotated through isv_history — already EMA-smoothed. isv_forward + * reads only the first 12 slots. */ + if (q_dir_means_ptr != 0) { + int ub = (dir_size < 4) ? dir_size : 4; /* clamp to the 4 reserved slots */ + for (int k = 0; k < ub; k++) { + int slot = 17 + k; + isv_signals[slot] = (1.0f - alpha) * isv_signals[slot] + + alpha * q_dir_means_ptr[k]; + } + } + if (q_dir_abs_ref_ptr != 0) { + isv_signals[21] = (1.0f - alpha) * isv_signals[21] + + alpha * fmaxf(q_dir_abs_ref_ptr[0], 0.0f); + } } /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7e3a7f447..a4c605981 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -105,13 +105,31 @@ const ISV_K: usize = 4; // Temporal ISV history length /// bins. Scale reference for the collapse-fraction normaliser /// (keeps the bin-weight formula scale-invariant as Q-values /// grow over training). +/// [17..20] Task 2.Y "make direction branch useful at eval" — per-direction +/// Q-value EMAs. Layout mirrors [13..15]: +/// [17] = Q_mean_ema(Short), [18] = Q_mean_ema(Hold), +/// [19] = Q_mean_ema(Long), [20] = Q_mean_ema(Flat). +/// Populated by `q_dir_bin_means_reduce` on the stats cadence, +/// smoothed with tau=0.05. The C51 loss / gradient kernels derive +/// an adaptive per-direction bin weight from these EMAs by the +/// same bias-vs-argmax composite applied to the magnitude axis, +/// with the direction-specific architectural shape +/// dir_bias_signal[k] = {1.0, 0.5, 0.0, 1.0} (Short/Hold/Long/Flat) +/// — tradable directions (Short/Long) are amplified; Flat is +/// never amplified (it is the already-over-estimated low-variance +/// bin); Hold sits between (position-preserving but not trading). +/// [21] q_dir_abs_ref_ema — max(|Q_mean_ema[k]|) across the 4 +/// direction bins. Scale reference for the direction-branch +/// collapse-fraction normaliser. /// /// Slots [0..11] populated by `isv_signal_update` kernel and rotated through /// `isv_history` for temporal decay in `isv_forward`. Slot [12] written /// outside `isv_signal_update` and NOT rotated into history. Slots [13..16] /// written by `isv_signal_update` (added with the Task 2.X adaptive magnitude /// fix) but NOT rotated into history — they are already EMA-smoothed. -const ISV_DIM: usize = 17; +/// Slots [17..21] likewise written by `isv_signal_update` (added with the +/// Task 2.Y direction-branch fix) and NOT rotated into history. +const ISV_DIM: usize = 22; 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 @@ -122,6 +140,19 @@ pub const Q_MAG_MEAN_FULL_INDEX: usize = 15; /// Task 2.X "make Full useful": absolute-Q-scale reference used to /// normalise the bias signal into a scale-invariant collapse fraction. pub const Q_ABS_REF_INDEX: usize = 16; +/// Task 2.Y "make direction branch useful at eval": ISV slots for +/// per-direction Q-mean EMAs. Read by `c51_loss_batched` and `c51_grad_kernel` +/// to compute the bias-vs-argmax-direction signal that drives the adaptive +/// direction-branch bin weighting. Layout matches the direction action +/// encoding (Short=0, Hold=1, Long=2, Flat=3). +pub const Q_DIR_MEAN_SHORT_INDEX: usize = 17; +pub const Q_DIR_MEAN_HOLD_INDEX: usize = 18; +pub const Q_DIR_MEAN_LONG_INDEX: usize = 19; +pub const Q_DIR_MEAN_FLAT_INDEX: usize = 20; +/// Task 2.Y: absolute-Q-scale reference for the direction branch (mirror of +/// `Q_ABS_REF_INDEX` for magnitude). Keeps the direction-branch +/// collapse-fraction scale-invariant as Q-values grow over training. +pub const Q_DIR_ABS_REF_INDEX: usize = 21; 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. @@ -1699,6 +1730,17 @@ pub struct GpuDqnTrainer { q_abs_ref_scratch_dev_ptr: u64, q_mag_bin_means_reduce_kernel: CudaFunction, + // Task 2.Y "make direction branch useful at eval" — per-direction-bin + // Q-mean scratch arrays written by `q_dir_bin_means_reduce` (runs with + // q_stats on the stats cadence), read by isv_signal_update to drive + // the EMA slots [17..21]. The means scratch is sized for direction-branch + // capacity (MAX_DIR=4 in the kernel); the abs_ref scratch is a scalar. + q_dir_means_scratch_pinned: *mut f32, + q_dir_means_scratch_dev_ptr: u64, + q_dir_abs_ref_scratch_pinned: *mut f32, + q_dir_abs_ref_scratch_dev_ptr: u64, + q_dir_bin_means_reduce_kernel: CudaFunction, + // ── ISV core buffers (pinned device-mapped, GPU read/write) ── isv_signals_pinned: *mut f32, // [ISV_DIM = 15] isv_signals_dev_ptr: u64, @@ -2195,6 +2237,12 @@ impl Drop for GpuDqnTrainer { if !self.q_abs_ref_scratch_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.q_abs_ref_scratch_pinned.cast()) }; } + if !self.q_dir_means_scratch_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.q_dir_means_scratch_pinned.cast()) }; + } + if !self.q_dir_abs_ref_scratch_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.q_dir_abs_ref_scratch_pinned.cast()) }; + } if !self.isv_signals_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.isv_signals_pinned.cast()) }; } @@ -3113,6 +3161,50 @@ impl GpuDqnTrainer { Ok(()) } + /// Task 2.Y "make direction branch useful at eval" — mirror of + /// `launch_q_mag_bin_means_reduce` applied to the direction branch. + /// Writes: + /// * `q_dir_means_scratch[0..dir_size]` = mean Q over batch for each + /// direction action (Short=0, Hold=1, Long=2, Flat=3). + /// * `q_dir_abs_ref_scratch[0]` = max(|mean Q|) across direction bins. + /// + /// These feed ISV slots [17..21] (EMA) via `update_isv_signals`. The + /// C51 loss / grad kernels then compute the bias-vs-argmax-direction + /// composite signal driving `get_direction_bin_weight`. + /// + /// Runs outside CUDA Graph capture on the stats cadence. Single-thread + /// single-block, ~N * dir_size f32 loads — microseconds at typical B. + pub(crate) fn launch_q_dir_bin_means_reduce( + &self, + batch_size: usize, + ) -> Result<(), MLError> { + let total_actions = self.total_actions() as i32; + let n = batch_size as i32; + /* Direction is the first branch in the factored-action layout — + * offset 0, size = config.branch_0_size (4 for the 4-branch DQN). */ + let dir_off: i32 = 0; + let dir_size = self.config.branch_0_size as i32; + let q_out_buf_ptr = self.q_out_buf.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.q_dir_bin_means_reduce_kernel) + .arg(&q_out_buf_ptr) + .arg(&self.q_dir_means_scratch_dev_ptr) + .arg(&self.q_dir_abs_ref_scratch_dev_ptr) + .arg(&n) + .arg(&total_actions) + .arg(&dir_off) + .arg(&dir_size) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("q_dir_bin_means_reduce: {e}")))?; + } + Ok(()) + } + /// Pure GPU ISV signal update — reads all pinned pointers, updates ISV /// [ISV_DIM=17] + history [K,12]. Includes 4 regime awareness signals /// computed from ADX/CUSUM in the states buffer. Slots [13..16] are the @@ -3125,6 +3217,7 @@ impl GpuDqnTrainer { let k = ISV_K as i32; let state_dim = ml_core::state_layout::STATE_DIM as i32; let mag_size_i32 = self.config.branch_1_size as i32; + let dir_size_i32 = self.config.branch_0_size as i32; unsafe { self.stream.launch_builder(&self.isv_signal_update_kernel) .arg(&self.isv_signals_dev_ptr) @@ -3147,6 +3240,13 @@ impl GpuDqnTrainer { .arg(&self.q_mag_means_scratch_dev_ptr) .arg(&self.q_abs_ref_scratch_dev_ptr) .arg(&mag_size_i32) + // Task 2.Y "make direction branch useful at eval" — + // per-direction-bin Q-mean scratch array + absolute-scale + // reference scalar feed ISV slots [17..20] (Q_mean per bin) + // and [21] (|Q| reference for direction). + .arg(&self.q_dir_means_scratch_dev_ptr) + .arg(&self.q_dir_abs_ref_scratch_dev_ptr) + .arg(&dir_size_i32) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), @@ -6503,6 +6603,7 @@ impl GpuDqnTrainer { // 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}")))?; let atom_stats_buf = stream.alloc_zeros::(2) @@ -7434,6 +7535,53 @@ impl GpuDqnTrainer { (host_ptr as *mut f32, dev_ptr_out) }; + // Task 2.Y "make direction branch useful at eval" — per-direction-bin + // Q-mean pinned array (sized for the 4-branch direction head, + // matching MAX_DIR=4 in q_dir_bin_means_reduce) plus the absolute- + // Q-scale reference scalar. Written by `q_dir_bin_means_reduce` + // (launched with q_stats); read by isv_signal_update to populate + // ISV slots [17..21]. + let q_dir_means_slots: usize = 4; + let (q_dir_means_scratch_pinned, q_dir_means_scratch_dev_ptr) = { + let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut dev_ptr_out: u64 = 0; + unsafe { + let rc = cudarc::driver::sys::cuMemAllocHost_v2( + &mut host_ptr, + q_dir_means_slots * std::mem::size_of::(), + ); + assert_eq!(rc, cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS, "cuMemAllocHost for q_dir_means_scratch"); + 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 q_dir_means_scratch"); + std::ptr::write_bytes(host_ptr as *mut u8, 0, q_dir_means_slots * std::mem::size_of::()); + } + (host_ptr as *mut f32, dev_ptr_out) + }; + + let (q_dir_abs_ref_scratch_pinned, q_dir_abs_ref_scratch_dev_ptr) = { + let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut dev_ptr_out: u64 = 0; + unsafe { + let rc = cudarc::driver::sys::cuMemAllocHost_v2( + &mut host_ptr, + std::mem::size_of::(), + ); + assert_eq!(rc, cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS, "cuMemAllocHost for q_dir_abs_ref_scratch"); + 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 q_dir_abs_ref_scratch"); + *(host_ptr as *mut f32) = 0.0; + } + (host_ptr as *mut f32, dev_ptr_out) + }; + // ── ISV core buffers (pinned device-mapped) ────────────────────── let (isv_signals_pinned, isv_signals_dev_ptr) = { let num_bytes = ISV_DIM * std::mem::size_of::(); @@ -8342,6 +8490,11 @@ impl GpuDqnTrainer { q_mag_means_scratch_dev_ptr, q_abs_ref_scratch_pinned, q_abs_ref_scratch_dev_ptr, + q_dir_bin_means_reduce_kernel, + q_dir_means_scratch_pinned, + q_dir_means_scratch_dev_ptr, + q_dir_abs_ref_scratch_pinned, + q_dir_abs_ref_scratch_dev_ptr, q_stats_buf, atom_stats_buf, atom_stats_block_sums_buf, @@ -8817,6 +8970,31 @@ impl GpuDqnTrainer { } } + /// Task 2.Y "make direction branch useful at eval" diagnostic — read the + /// per-direction-bin Q-mean EMAs and absolute-scale reference from ISV + /// (pinned, zero-copy). Layout mirrors `read_isv_magnitude_bin_q_means`. + /// + /// Returns `(q_short, q_hold, q_long, q_flat, q_dir_abs_ref)`. The C51 + /// loss / gradient kernels consume these via `get_direction_bin_weight`: + /// max_mean = max(q_short, q_hold, q_long, q_flat) + /// bias_gap = max(0, max_mean − q_mean[a0]) + /// collapse_frac = min(1, bias_gap / max(q_dir_abs_ref, 1e-6) + + /// (1 − learning_health)) + /// bin_weight = 1 + collapse_frac * dir_bias_signal(a0) + /// + /// Returns all zeros if ISV is not allocated. + pub fn read_isv_direction_bin_q_means(&self) -> (f32, f32, f32, f32, f32) { + if self.isv_signals_pinned.is_null() { return (0.0, 0.0, 0.0, 0.0, 0.0); } + unsafe { + let s = *self.isv_signals_pinned.add(Q_DIR_MEAN_SHORT_INDEX); + let h = *self.isv_signals_pinned.add(Q_DIR_MEAN_HOLD_INDEX); + let l = *self.isv_signals_pinned.add(Q_DIR_MEAN_LONG_INDEX); + let f = *self.isv_signals_pinned.add(Q_DIR_MEAN_FLAT_INDEX); + let r = *self.isv_signals_pinned.add(Q_DIR_ABS_REF_INDEX); + (s, h, l, f, r.max(0.0)) + } + } + /// B4/G5: Read ISV health index and regime stability from pinned host memory. /// Returns (health [0,1], regime_stability [0,1]). Falls back to (0.5, 0.5) if /// pinned pointer is null. @@ -10562,6 +10740,13 @@ impl GpuDqnTrainer { let mag_means_batch_size = self.config.batch_size; self.launch_q_mag_bin_means_reduce(mag_means_batch_size)?; + // Task 2.Y "make direction branch useful at eval" — per-direction + // Q-mean reducer (mirror of the magnitude reducer one axis up). + // Writes the 4 per-bin Q-means + absolute-scale reference to pinned + // scratch. Runs on the same stream; isv_signal_update reads the + // values directly via pinned device pointers (no sync). + self.launch_q_dir_bin_means_reduce(mag_means_batch_size)?; + // ISV signal update — pure GPU, reads all pinned scalars, updates ISV vector self.update_isv_signals()?; @@ -14253,6 +14438,23 @@ fn compile_q_mag_bin_means_kernel( .map_err(|e| MLError::ModelError(format!("q_mag_bin_means_reduce load: {e}"))) } +/// Load the per-direction-bin Q-mean reducer from the q_stats cubin. +/// +/// Task 2.Y "make direction branch useful at eval": this kernel computes +/// batch-mean Q-values per direction-bin (Short/Hold/Long/Flat) plus an +/// absolute-scale reference, feeding ISV slots [17..21] that drive the +/// adaptive C51 direction bin weight. See q_stats_kernel.cu for the kernel +/// body (q_dir_bin_means_reduce). +fn compile_q_dir_bin_means_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_dir_bin_means cubin load: {e}")))?; + module.load_function("q_dir_bin_means_reduce") + .map_err(|e| MLError::ModelError(format!("q_dir_bin_means_reduce load: {e}"))) +} + // ── Shared memory sizing ──────────────────────────────────────────────────── /// Query the hardware's max shared memory per block via cuDeviceGetAttribute. diff --git a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu index 5c3775848..a7b6bb9b1 100644 --- a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu +++ b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu @@ -154,3 +154,72 @@ extern "C" __global__ void q_mag_bin_means_reduce( } q_abs_ref_out[0] = abs_ref; } + +/** + * Per-branch direction-Q mean reducer: computes batch-mean Q-values for each + * of the 4 direction bins (Short=0, Hold=1, Long=2, Flat=3), plus the + * absolute-Q-scale reference, from `q_values` of shape [N, total_actions] + * sliced at `dir_off..dir_off+dir_size`. Normally `dir_off=0, dir_size=4` + * because direction is the first branch in the factored-action layout. + * + * Task 2.Y "make direction branch useful at eval" adaptive fix — direct + * structural mirror of `q_mag_bin_means_reduce` above. The direction axis + * suffers the same C51 expected-Q low-variance bias documented at + * experience_kernels.cu:825-829: Flat has zero PnL variance → tightest + * distributional Q → systematic over-pricing vs Short/Long (high variance). + * At eval, strict-argmax from near-uniform direction Q-values picks + * Hold/Flat in most states, which then forces mag=0 via the + * experience_kernels.cu:896-897 ABI invariant. Task 2.Y lifts this by + * applying a bin-weighted C51 gradient correction driven by these per-bin + * Q-mean EMAs. + * + * Output layout: + * q_dir_means_out[0] = Q_mean(Short) over N samples, + * q_dir_means_out[1] = Q_mean(Hold), + * q_dir_means_out[2] = Q_mean(Long), + * q_dir_means_out[3] = Q_mean(Flat). + * q_dir_abs_ref_out[0] = max(|Q_mean[k]|) across direction bins. + * + * Pure reduction — no training-path feedback and no atomicAdd. Single-block, + * single-thread to match q_stats_reduce / q_mag_bin_means_reduce determinism. + * + * Launch: grid=(1,1,1), block=(1,1,1). Cost: ~N * dir_size FP32 loads, + * single-digit microseconds at typical B. + */ +extern "C" __global__ void q_dir_bin_means_reduce( + const float* __restrict__ q_values, /* [N, total_actions] */ + float* __restrict__ q_dir_means_out, /* [dir_size] — typically [4] */ + float* __restrict__ q_dir_abs_ref_out, /* [1] absolute-value max of bin means */ + int N, + int total_actions, + int dir_off, int dir_size) +{ + if (threadIdx.x != 0) return; + + /* Hard ceiling — covers the 4-branch direction head (Short/Hold/Long/Flat). + * Sized generously to tolerate any future direction-branch extension. */ + const int MAX_DIR = 4; + float sums[MAX_DIR]; + for (int k = 0; k < MAX_DIR; k++) sums[k] = 0.0f; + + if (N <= 0 || total_actions <= 0 || dir_size <= 0 || dir_size > MAX_DIR + || dir_off < 0 || dir_off + dir_size > total_actions) { + for (int k = 0; k < dir_size && k < MAX_DIR; k++) q_dir_means_out[k] = 0.0f; + q_dir_abs_ref_out[0] = 0.0f; + return; + } + + for (int i = 0; i < N; i++) { + const float* row = q_values + (long long)i * total_actions; + for (int k = 0; k < dir_size; k++) sums[k] += row[dir_off + k]; + } + + float inv_n = 1.0f / (float)N; + float abs_ref = 0.0f; + for (int k = 0; k < dir_size; k++) { + float mean_k = sums[k] * inv_n; + q_dir_means_out[k] = mean_k; + abs_ref = fmaxf(abs_ref, fabsf(mean_k)); + } + q_dir_abs_ref_out[0] = abs_ref; +} diff --git a/crates/ml/src/trainers/dqn/smoke_tests/magnitude_distribution.rs b/crates/ml/src/trainers/dqn/smoke_tests/magnitude_distribution.rs index 7e886e87e..5e234d242 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/magnitude_distribution.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/magnitude_distribution.rs @@ -71,7 +71,8 @@ fn test_magnitude_distribution() -> Result<()> { // (magnitude gets forced to 0 by kernel for Hold/Flat), so a magnitude- // only bin weight cannot shift eval_dist. Surfacing this lets the // operator attribute the failure to direction vs magnitude. - if let Some(dir_dist) = trainer.last_eval_direction_dist() { + let eval_dir_dist = trainer.last_eval_direction_dist(); + if let Some(dir_dist) = eval_dir_dist { println!( "[EVAL_DIR_DIST] Short={:.3} Hold={:.3} Long={:.3} Flat={:.3}", dir_dist[0], dir_dist[1], dir_dist[2], dir_dist[3] @@ -96,6 +97,40 @@ fn test_magnitude_distribution() -> Result<()> { qq, qh, qf, q_abs_ref, max_mean, bias_gap_full, collapse_frac_full ); } + + // Task 2.Y diagnostic — ISV per-direction-bin Q-mean EMAs + |Q|-scale + // reference. Mirror of the magnitude diagnostic above, applied to the + // direction axis. The direction bin weight fires when the argmax of + // (Short/Hold/Long/Flat) differs from the taken direction, AND the + // taken direction is not Flat (Flat is never amplified — it is the + // already-over-estimated low-variance bin the mechanism corrects + // AGAINST). collapse_frac_dir_hold and collapse_frac_dir_flat report + // the bias-vs-argmax fraction for Hold and Flat respectively; these + // are the bins whose over-selection at eval the fix must reduce. + if let Some((qs, qh, ql, qf, q_dir_abs_ref)) = + trainer.last_isv_direction_bin_q_means() + { + let max_mean_dir = qs.max(qh).max(ql).max(qf); + let bias_gap_hold = (max_mean_dir - qh).max(0.0); + let bias_gap_flat = (max_mean_dir - qf).max(0.0); + let collapse_frac_dir_hold = if q_dir_abs_ref > 1e-6 { + (bias_gap_hold / q_dir_abs_ref).min(1.0) + } else { + 0.0 + }; + let collapse_frac_dir_flat = if q_dir_abs_ref > 1e-6 { + (bias_gap_flat / q_dir_abs_ref).min(1.0) + } else { + 0.0 + }; + println!( + "[ISV_DIR_MEANS] q_s={:.4} q_h={:.4} q_l={:.4} q_f={:.4} \ + q_dir_abs_ref={:.4} max_mean_dir={:.4} \ + collapse_frac_dir_hold={:.3} collapse_frac_dir_flat={:.3}", + qs, qh, ql, qf, q_dir_abs_ref, max_mean_dir, + collapse_frac_dir_hold, collapse_frac_dir_flat + ); + } assert!( eh + ef >= 0.30, "H10 regression: eval Half + Full share {:.3} < 0.30 (eq={:.3} eh={:.3} ef={:.3}) \ @@ -130,5 +165,43 @@ fn test_magnitude_distribution() -> Result<()> { escalation path is richer ISV signals / kernel-side arithmetic.", ef, eq, eh, ef ); + + // Task 2.Y "make direction branch useful at eval" — the ISV-adaptive + // C51 direction bin weight (see c51_loss_kernel.cu:: + // get_direction_bin_weight and c51_grad_kernel.cu's d==0 block) must + // prevent the direction-branch strict-argmax from collapsing onto + // Hold or Flat. Pre-fix baseline on HEAD fa8d54661 across 3 runs + // observed Hold+Flat ∈ {0.000, 0.847, 0.452} — hyper-variant because + // C51's expected-Q bias flattens direction Q-values into a tight band + // that argmax picks from arbitrarily. The fix pushes the tradable + // directions (Short/Long) up by an ISV-driven bin weight; Flat is + // never amplified. Self-disables when direction Q-values differentiate + // healthily. + // + // Per-run gate is lenient (≤ 0.60) because direction collapse is + // hyper-variant across seeds — operator should run 3 times and compute + // the median to assess real mechanism engagement (scoping §7 / §8). + if let Some(dir_dist) = eval_dir_dist { + let hold_plus_flat = dir_dist[1] + dir_dist[3]; + assert!( + hold_plus_flat <= 0.60, + "Task 2.Y adaptive direction: eval Hold+Flat share {:.3} > 0.60 \ + (Short={:.3} Hold={:.3} Long={:.3} Flat={:.3}) — ISV direction \ + bin-weight mechanism failed to lift tradable-direction Q-values \ + above the Hold/Flat argmax threshold. Investigate: \ + (a) ISV slots [17..21] populated each stats cadence (check \ + q_dir_bin_means_reduce + isv_signal_update wiring), \ + (b) get_direction_bin_weight fires in C51 loss (branch_ce scale) \ + AND c51_grad_kernel d==0 block (consistent with forward), \ + (c) dir_bias_signal layout {{Short=1, Hold=0.5, Long=1, Flat=0}} — \ + Flat must NEVER be amplified. Note: direction collapse is \ + hyper-variant across seeds; single-run failure on the tail of \ + the distribution is acceptable per scoping §7 if median over 3 \ + independent runs ≤ 0.50. Do NOT fall back to static tuning \ + knobs — per feedback_adaptive_not_tuned.md escalation is richer \ + ISV signals / kernel-side arithmetic.", + hold_plus_flat, dir_dist[0], dir_dist[1], dir_dist[2], dir_dist[3] + ); + } Ok(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 6a111f399..948b4a8dc 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -1546,6 +1546,18 @@ impl DQNTrainer { .map(|ctx| ctx.trainer().read_isv_magnitude_bin_q_means()) } + /// Task 2.Y diagnostic — ISV slots [17..21] (per-direction-bin Q-mean + /// EMAs + absolute-scale reference). Returns + /// `(q_mean_short, q_mean_hold, q_mean_long, q_mean_flat, q_dir_abs_ref)` + /// or `None` if the fused trainer has not been initialised yet. Used by + /// the magnitude_distribution smoke test to log the direction-branch + /// bin-weight engagement state. + pub fn last_isv_direction_bin_q_means(&self) -> Option<(f32, f32, f32, f32, f32)> { + self.fused_ctx + .as_ref() + .map(|ctx| ctx.trainer().read_isv_direction_bin_q_means()) + } + /// Per-epoch magnitude-branch action entropy history. /// Each entry is (epoch_0_indexed, normalized_entropy ∈ [0, 1]). /// Used by exploration_coverage smoke test.