diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 4beb147a8..b15ea79e5 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -160,6 +160,7 @@ const KERNELS: &[&str] = &[ "multi_head_policy_backward", // Phase 2A-B (2026-06-03): backward through mixture + per-head softmax + gating softmax. Two kernels: `_backward_pi` and `_backward_gate`. Per-batch grad scratch; caller reduces via reduce_axis0. Spec ADDENDUM 2026-06-03 §R.6 "multi_head_policy_aux_prior", // Phase 2A-B (2026-06-03): per-head auxiliary KL prior gradient — additive into grad_pi_logits_k. β read from ISV slot 764. Spec ADDENDUM 2026-06-03 §R.4 "multi_head_policy_aggregate_diag", // Phase 2A-D (2026-06-03): device-aggregated gate-and-head diagnostics. Reduces forward outputs (gate_probs, pi_probs_k) over B → 25 ISV slots (gate_probs_mean[K] / gate_argmax_mass[K] / gate_entropy_mean / per_head_entropy_mean[K]) for the multi-head specialization verdict signal. Tree-reduce, no atomicAdd. + "rl_gate_lr_multiplier_controller", // Phase 2A-D fix B1.3 (2026-06-03): adaptive gate-LR multiplier controller. Reads gate_entropy_mean (slot 781) via EMA → escalates +0.3 %/step when entropy_ema > 0.85·log(K) (under-learning), decays −5 %/step when < 0.20·log(K) (collapse risk), leaves alone in healthy band. Single-thread launch (1,1,1)/(1,1,1). ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_gate_lr_multiplier_controller.cu b/crates/ml-alpha/cuda/rl_gate_lr_multiplier_controller.cu new file mode 100644 index 000000000..9b435267b --- /dev/null +++ b/crates/ml-alpha/cuda/rl_gate_lr_multiplier_controller.cu @@ -0,0 +1,121 @@ +// rl_gate_lr_multiplier_controller.cu — adaptive gate-LR multiplier (B1.3). +// +// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md +// Phase 2A-D fix B1.3 dispatch (2026-06-03). +// +// Reads (from ISV): +// RL_POLICY_GATE_ENTROPY_MEAN_INDEX (781) — produced per-step +// by `multi_head_policy_aggregate_diag.cu` +// (Phase 2A-D); mean over batches +// of `−Σ_k p[b,k]·log p[b,k]`. +// RL_POLICY_GATE_ENTROPY_EMA_INDEX (791) — EMA state; this kernel +// owns it. +// RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX (792) — monotonic 0→1 latch. +// RL_POLICY_GATE_LR_MULTIPLIER_INDEX (790) — output, also the +// previous-step value. +// RL_POLICY_NUM_HEADS_INDEX (761) — K (active head count), +// drives regime-adaptive +// entropy thresholds. +// +// Writes (to ISV): +// slot 791 — updated EMA +// slot 792 — bootstrap-done flag (latched to 1.0 on first observation) +// slot 790 — updated multiplier (clamped to [MIN_FLOOR, MAX_CEIL]) +// +// Control law (per `pearl_dead_signal_resurrection_discipline`): +// • EMA Wiener-α = 0.02 (~50-step horizon) — smooths per-step noise. +// • Escalate +0.3 %/step (compounds to ≈+50 % over 137 steps) when +// `entropy_ema > 0.85·log(K)` — gate is "under-learning" / stuck +// near uniform. +// • Decay −5 %/step (emergency back-off) when +// `entropy_ema < 0.20·log(K)` — gate is at risk of collapse to +// single-head routing. +// • Healthy band [0.20·log(K), 0.85·log(K)]: leave multiplier alone +// (current value is working — no need to hunt). +// • Clamp to [1.0, 50.0] per `pearl_bootstrap_must_respect_clamp_range` +// so the bootstrap 5.0 lies strictly inside the range. +// +// Per `pearl_bootstrap_must_respect_clamp_range`: bootstrap 5.0 ∈ [1.0, +// 50.0]; controller never snaps to a clamp boundary on its first signal. +// +// Per `pearl_welford_trade_count_is_step_not_trade`: the first-observation +// EMA initialisation is gated by a dedicated `bootstrap_done` flag at +// slot 792; the controller cannot rely on `prev == bootstrap` because 5.0 +// is also a valid steady-state value the controller may visit after +// escalating from a lower point. +// +// Per `feedback_no_atomicadd`: single-thread launch (1×1×1), no atomics. +// Per `feedback_cpu_is_read_only`: pure device-side; reads + writes are +// all ISV slots. +// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs. + +#include +#include + +// ISV slot indices (must match crates/ml-alpha/src/rl/isv_slots.rs). +#define RL_POLICY_NUM_HEADS_INDEX 761 +#define RL_POLICY_GATE_ENTROPY_MEAN_INDEX 781 +#define RL_POLICY_GATE_LR_MULTIPLIER_INDEX 790 +#define RL_POLICY_GATE_ENTROPY_EMA_INDEX 791 +#define RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX 792 + +// Tuning constants (B1.3 design — anchored on the B1.2 dose-response). +// All values are dimensionless ratios or rate scalars. +#define GATE_LR_CTRL_EMA_ALPHA 0.02f // Wiener-α — ~50-step horizon +#define GATE_LR_CTRL_OVER_FRAC 0.85f // multiplier of log(K) — under-learning threshold +#define GATE_LR_CTRL_COLLAPSE_FRAC 0.20f // multiplier of log(K) — collapse threshold +#define GATE_LR_CTRL_ESCALATE_RATE 1.003f // +0.3 %/step (≈+50 % in 137 steps) +#define GATE_LR_CTRL_DECAY_RATE 0.95f // −5 %/step (emergency back-off) +#define GATE_LR_CTRL_MIN_FLOOR 1.0f // never below 1× (effectively off) +#define GATE_LR_CTRL_MAX_CEIL 50.0f // saturation cap (B1.2 plateau ≈20×) + +extern "C" __global__ void rl_gate_lr_multiplier_controller(float* isv) { + // Single-thread kernel — launched (1,1,1)/(1,1,1). + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + const float entropy_now = isv[RL_POLICY_GATE_ENTROPY_MEAN_INDEX]; + const float boot_done = isv[RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX]; + const float k_active = isv[RL_POLICY_NUM_HEADS_INDEX]; + + // ── EMA update (first-observation bootstrap pattern) ───────────── + float entropy_ema; + if (boot_done < 0.5f) { + // First call: replace EMA with current observation; latch flag. + // Necessary because slot 791 sentinel-zeroes at trainer init and + // blending zero with the first real measurement (~log(K) ≈ 1.10 + // for K=3) would lag the controller by ~50 steps before it ever + // sees a representative value. + entropy_ema = entropy_now; + isv[RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX] = 1.0f; + } else { + entropy_ema = (1.0f - GATE_LR_CTRL_EMA_ALPHA) * isv[RL_POLICY_GATE_ENTROPY_EMA_INDEX] + + GATE_LR_CTRL_EMA_ALPHA * entropy_now; + } + isv[RL_POLICY_GATE_ENTROPY_EMA_INDEX] = entropy_ema; + + // ── Regime-adaptive thresholds (scale with K) ──────────────────── + // K ≥ 1 by construction (multi-head policy bootstraps NUM_HEADS to 3 + // or higher in the same flag-on block); guard with fmaxf so logf is + // well-defined even if a future caller writes K=0 (would degenerate + // to single-head and the controller becomes a no-op via the healthy + // band → multiplier left alone, which is the right behaviour). + const float log_k = logf(fmaxf(k_active, 1.0f)); + const float over_threshold = GATE_LR_CTRL_OVER_FRAC * log_k; + const float collapse_threshold = GATE_LR_CTRL_COLLAPSE_FRAC * log_k; + + // ── Multiplier update ───────────────────────────────────────────── + float multiplier = isv[RL_POLICY_GATE_LR_MULTIPLIER_INDEX]; + if (entropy_ema > over_threshold) { + // Under-learning: slowly compound up. + multiplier *= GATE_LR_CTRL_ESCALATE_RATE; + } else if (entropy_ema < collapse_threshold) { + // Collapse risk: fast back-off. + multiplier *= GATE_LR_CTRL_DECAY_RATE; + } + // Healthy band — leave multiplier alone. + + // Clamp; bootstrap 5.0 sits strictly inside [1.0, 50.0]. + multiplier = fmaxf(GATE_LR_CTRL_MIN_FLOOR, + fminf(multiplier, GATE_LR_CTRL_MAX_CEIL)); + isv[RL_POLICY_GATE_LR_MULTIPLIER_INDEX] = multiplier; +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index e46a8c1ca..69d675262 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -1895,15 +1895,63 @@ pub const RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX: usize = 782; /// structurally weak near uniform g_k ≈ 1/K because all heads contribute /// similarly, so W_gate moves too slowly to escape the uniform basin. /// Scaling only the gate Adams (NOT the head Adams) accelerates routing -/// learning without destabilising head training. Range [0.5, 50.0] for -/// the ISV controller — outside that we expect either degenerate noise -/// (too high) or no effect (too low). Slot defaults to 0.0 sentinel when -/// the multi-head flag is off; the bootstrap write at slot 790 happens -/// inside the same `if self.use_multi_head_policy { ... }` block as -/// slots 761-764, so flag-off bit-equality is preserved (0² contributes -/// 0 to the `isv_state` checksum). +/// learning without destabilising head training. Range [1.0, 50.0] for +/// the ISV controller (B1.3) — outside that we expect either degenerate +/// noise (too high) or no effect (too low). Slot defaults to 0.0 +/// sentinel when the multi-head flag is off; the bootstrap write at slot +/// 790 happens inside the same `if self.use_multi_head_policy { ... }` +/// block as slots 761-764, so flag-off bit-equality is preserved (0² +/// contributes 0 to the `isv_state` checksum). +/// +/// B1.3 (2026-06-03): now driven by `rl_gate_lr_multiplier_controller.cu` +/// — escalates +0.3%/step when `gate_entropy_ema > 0.85·log(K)` +/// (under-learning), decays −5%/step when `< 0.20·log(K)` (collapse risk), +/// leaves alone in the healthy band. Bootstrap value 5.0 is the +/// empirically-validated B1.2 starting point; the controller adapts from +/// there. Bootstrap MUST lie in `[MIN_FLOOR=1.0, MAX_CEIL=50.0]` per +/// `pearl_bootstrap_must_respect_clamp_range` — 5.0 satisfies that. pub const RL_POLICY_GATE_LR_MULTIPLIER_INDEX: usize = 790; +/// EMA of `RL_POLICY_GATE_ENTROPY_MEAN_INDEX` (slot 781) maintained by +/// `rl_gate_lr_multiplier_controller.cu` with α = 0.02 (~50-step horizon). +/// Bootstrap-on-first-observation pattern (per +/// `pearl_welford_trade_count_is_step_not_trade`): the first time the +/// controller runs with `RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX` +/// at 0, it sets `entropy_ema = entropy_now` and flips the bootstrap-done +/// flag to 1.0 (monotonic, never cleared) so the EMA never blends with +/// the sentinel-zero starting value. +/// +/// Range [0, log(K)]; log(3) ≈ 1.0986 for default K=3. This EMA is the +/// controller's "is the gate currently stuck / specialising / collapsing" +/// signal; it smooths out per-step noise in the underlying entropy +/// measurement before the controller compares against the +/// regime-adaptive thresholds (0.85·log(K) over-threshold, 0.20·log(K) +/// collapse-threshold). +/// +/// Flag-off path: slot stays sentinel-zero (controller launch is +/// gated on `use_multi_head_policy`), so 0² preserves the +/// `isv_state` checksum bit-equality with pre-B1.3 HEAD 6f3639bfb. +pub const RL_POLICY_GATE_ENTROPY_EMA_INDEX: usize = 791; + +/// Bootstrap-done flag for `rl_gate_lr_multiplier_controller.cu`. +/// Monotonic 0 → 1 — set ONCE the first time the controller runs (when +/// the gate-entropy aggregator has produced its first slot-781 value); +/// NEVER cleared, even across `reset_session_state` boundaries. +/// +/// The flag exists because the EMA at slot 791 needs a first-observation +/// initialisation (set EMA to the current measurement, not blend it with +/// the zero starting value — see +/// `pearl_welford_trade_count_is_step_not_trade` for the canonical +/// failure mode this prevents). A standalone flag (rather than the +/// `prev == bootstrap` test used by older controllers) is required here +/// because slot 790's bootstrap value (5.0) is also a valid steady-state +/// value the controller may land on after escalating from a lower point. +/// +/// Flag-off path: slot stays sentinel-zero (controller launch is +/// gated on `use_multi_head_policy`), so 0² preserves the +/// `isv_state` checksum bit-equality with pre-B1.3 HEAD 6f3639bfb. +pub const RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX: usize = 792; + /// Last RL-allocated slot index (exclusive). /// Pre-risk-stack: 662. Post-Fix-B: 685. Post-v9 (warmup boundary): 696. /// Post-regime-observer (F1.1): 716. Post-warmed-flag addendum: 717. @@ -1921,4 +1969,5 @@ pub const RL_POLICY_GATE_LR_MULTIPLIER_INDEX: usize = 790; /// Post-multi-head-policy Phase 2A-A (4 policy-gating slots): 764. /// Post-multi-head-policy Phase 2A-D (25 aggregate diag slots): 789. /// Post-Phase 2A-D fix B1 (gate LR multiplier slot 790): 790. -pub const RL_SLOTS_END: usize = 791; +/// Post-Phase 2A-D fix B1.3 (gate-LR controller state slots 791-792): 792. +pub const RL_SLOTS_END: usize = 793; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 415130405..cea4569ce 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -299,6 +299,8 @@ const RL_KELLY_FRACTION_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_kelly_fraction_controller.cubin")); const RL_SURFER_SCAFFOLD_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_surfer_scaffold_controller.cubin")); +const RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_gate_lr_multiplier_controller.cubin")); const RL_EVAL_WARMUP_DECAY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_eval_warmup_decay.cubin")); // Regime observer (F1.2 + F1.3) — flat-count block-reduce + state-machine @@ -888,6 +890,13 @@ pub struct IntegratedTrainer { rl_kelly_fraction_controller_fn: CudaFunction, _rl_surfer_scaffold_controller_module: Arc, rl_surfer_scaffold_controller_fn: CudaFunction, + // Phase 2A-D fix B1.3 (2026-06-03): adaptive gate-LR multiplier + // controller. Reads gate_entropy_mean (slot 781) via EMA → drives + // RL_POLICY_GATE_LR_MULTIPLIER_INDEX (slot 790) up when the gate is + // stuck near uniform, down when at risk of collapse. Launched per + // step in the flag-on path inside `step_with_lobsim_reward_and_train`. + _rl_gate_lr_multiplier_controller_module: Arc, + rl_gate_lr_multiplier_controller_fn: CudaFunction, // v9 (2026-05-31): defensive eval-boundary calibration kernel. _rl_eval_warmup_decay_module: Arc, rl_eval_warmup_decay_fn: CudaFunction, @@ -1975,6 +1984,13 @@ impl IntegratedTrainer { let rl_surfer_scaffold_controller_fn = rl_surfer_scaffold_controller_module .load_function("rl_surfer_scaffold_controller") .context("load rl_surfer_scaffold_controller")?; + // Phase 2A-D fix B1.3 — adaptive gate-LR multiplier controller. + let rl_gate_lr_multiplier_controller_module = ctx + .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec()) + .context("load rl_gate_lr_multiplier_controller cubin")?; + let rl_gate_lr_multiplier_controller_fn = rl_gate_lr_multiplier_controller_module + .load_function("rl_gate_lr_multiplier_controller") + .context("load rl_gate_lr_multiplier_controller")?; // v9 — defensive eval-boundary calibration kernel. let rl_eval_warmup_decay_module = ctx .load_cubin(RL_EVAL_WARMUP_DECAY_CUBIN.to_vec()) @@ -3161,6 +3177,8 @@ impl IntegratedTrainer { rl_kelly_fraction_controller_fn, _rl_surfer_scaffold_controller_module: rl_surfer_scaffold_controller_module, rl_surfer_scaffold_controller_fn, + _rl_gate_lr_multiplier_controller_module: rl_gate_lr_multiplier_controller_module, + rl_gate_lr_multiplier_controller_fn, _rl_eval_warmup_decay_module: rl_eval_warmup_decay_module, rl_eval_warmup_decay_fn, _rl_regime_flat_count_module: rl_regime_flat_count_module, @@ -4171,12 +4189,24 @@ impl IntegratedTrainer { // bootstrap 5.0 — addresses gate-not-learning observed in // 3-seed mid-smoke at HEAD 5f10bcde3 (gate_entropy stayed // pinned at MAX, argmax_mass = 1.0). See slot doc-comment. - let mhp_isv_constants: [(usize, f32); 5] = [ - (crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX, 3.0_f32), - (crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, 0.549_f32), - (crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, 0.959_f32), - (crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX, 0.05_f32), - (crate::rl::isv_slots::RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0_f32), + // + // Phase 2A-D fix B1.3 (2026-06-03): slot 791 (entropy EMA) + // and slot 792 (bootstrap-done flag) bootstrap to 0.0 — the + // controller's first-observation pattern flips the flag to + // 1.0 on its first launch and replaces the EMA with the + // current `gate_entropy_mean` measurement, so the 0.0 + // sentinel is never blended into the EMA. Per + // `pearl_bootstrap_must_respect_clamp_range`, slot 790's + // bootstrap 5.0 lies strictly inside the controller's + // [1.0, 50.0] clamp range. + let mhp_isv_constants: [(usize, f32); 7] = [ + (crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX, 3.0_f32), + (crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, 0.549_f32), + (crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, 0.959_f32), + (crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX, 0.05_f32), + (crate::rl::isv_slots::RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0_f32), + (crate::rl::isv_slots::RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.0_f32), + (crate::rl::isv_slots::RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 0.0_f32), ]; for (slot, value) in mhp_isv_constants.iter() { let slot_i32 = *slot as i32; @@ -4528,6 +4558,29 @@ impl IntegratedTrainer { Ok(()) } + /// Adaptive gate-LR multiplier controller (Phase 2A-D fix B1.3, + /// 2026-06-03). Single-thread kernel; reads `gate_entropy_mean` + /// (slot 781) via an EMA → escalates `RL_POLICY_GATE_LR_MULTIPLIER_INDEX` + /// (slot 790) when the gate is stuck near-uniform (under-learning), + /// decays it when the gate is at risk of collapsing to a single head. + /// Must run AFTER `MultiHeadPolicy::emit_diag_stats` (which writes + /// slot 781) and BEFORE the next step's `lr_pi × gate_lr_mult` read at + /// the top of `step_*_body` (which uses the result for the gate Adams). + pub fn launch_rl_gate_lr_multiplier_controller(&self) -> Result<()> { + let mut args = RawArgs::new(); + args.push_ptr(self.isv_dev_ptr); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.rl_gate_lr_multiplier_controller_fn.cu_function(), + (1, 1, 1), (1, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("rl_gate_lr_multiplier_controller: {:?}", e))?; + } + Ok(()) + } + /// Adaptive risk management — Layer 4 Kelly fraction controller (spec 2026-05-30). /// Standalone launcher; production calls go through `rl_fused_controllers`. pub fn launch_rl_kelly_fraction_controller(&self) -> Result<()> { @@ -8136,6 +8189,20 @@ impl IntegratedTrainer { self.launch_rl_surfer_scaffold_controller() .context("launch_rl_surfer_scaffold_controller")?; + // Phase 2A-D fix B1.3 (2026-06-03) — adaptive gate-LR multiplier + // controller. Flag-gated on `use_multi_head_policy`: the captured + // graph baked at construction time mirrors the flag state, so + // flag-off runs do NOT include this launch and bit-equality vs + // pre-B1.3 HEAD 6f3639bfb is preserved. Must run AFTER the + // forward-path aggregator (`emit_diag_stats`) wrote + // `gate_entropy_mean` to slot 781 for this step, and BEFORE the + // next step's `lr_gate = lr_pi × ISV[790]` host read at the top + // of `step_*_body`. + if self.use_multi_head_policy { + self.launch_rl_gate_lr_multiplier_controller() + .context("launch_rl_gate_lr_multiplier_controller")?; + } + // PopArt: normalize rewards in-place (replaces apply_reward_scale). // Welford-EMA on batch mean+variance → whitened rewards → ISV // stats for V-correct. Grid=(1), Block=(min(B, 256)). Shared mem @@ -11838,7 +11905,10 @@ impl IntegratedTrainer { use crate::rl::isv_slots::{ RL_POLICY_AUX_PRIOR_BETA_INDEX, RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX, + RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, + RL_POLICY_GATE_ENTROPY_EMA_INDEX, RL_POLICY_GATE_ENTROPY_MEAN_INDEX, + RL_POLICY_GATE_LR_MULTIPLIER_INDEX, RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX, RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, @@ -11870,6 +11940,16 @@ impl IntegratedTrainer { "gate_argmax_mass": gate_argmax_mass, "gate_entropy_mean": isv[RL_POLICY_GATE_ENTROPY_MEAN_INDEX], "per_head_entropy_mean": per_head_entropy_mean, + // Phase 2A-D fix B1.3 (2026-06-03) — adaptive + // gate-LR multiplier controller state. Slot 790 + // is the controller output (consumed by the + // host-side LR-set at the top of the next step's + // `step_*_body`); slot 791 is the EMA over slot + // 781; slot 792 is the first-observation + // bootstrap-done flag (monotonic 0 → 1). + "gate_lr_multiplier": isv[RL_POLICY_GATE_LR_MULTIPLIER_INDEX], + "gate_entropy_ema": isv[RL_POLICY_GATE_ENTROPY_EMA_INDEX], + "gate_controller_bootstrap_done": isv[RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX], }), ); } diff --git a/crates/ml-alpha/tests/multi_head_policy_invariants.rs b/crates/ml-alpha/tests/multi_head_policy_invariants.rs index 27454c346..8447b5174 100644 --- a/crates/ml-alpha/tests/multi_head_policy_invariants.rs +++ b/crates/ml-alpha/tests/multi_head_policy_invariants.rs @@ -40,8 +40,8 @@ //! Run with: //! `cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture` -use anyhow::Result; -use cudarc::driver::{CudaSlice, CudaStream}; +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use ml_alpha::cfc::snap_features::REGIME_DIM; use ml_alpha::heads::HIDDEN_DIM; use ml_alpha::pinned_mem::MappedF32Buffer; @@ -49,8 +49,12 @@ use ml_alpha::rl::common::N_ACTIONS; use ml_alpha::rl::isv_slots::{ RL_POLICY_AUX_PRIOR_BETA_INDEX, RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX, + RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, + RL_POLICY_GATE_ENTROPY_EMA_INDEX, RL_POLICY_GATE_ENTROPY_MEAN_INDEX, + RL_POLICY_GATE_LR_MULTIPLIER_INDEX, RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX, + RL_POLICY_NUM_HEADS_INDEX, RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX, RL_SLOTS_END, }; @@ -59,6 +63,14 @@ use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub}; use ml_core::device::MlDevice; use std::sync::Arc; +/// Phase 2A-D fix B1.3 — `rl_gate_lr_multiplier_controller.cu` cubin, +/// loaded directly so the four controller invariants exercise the kernel +/// in isolation (no trainer dependency). +const RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/rl_gate_lr_multiplier_controller.cubin" +)); + /// Relative tolerance for sum-to-1 invariants. fp32 single-block reductions /// accumulate ~1 ULP per term; 1e-5 is loose enough to absorb that plus the /// `+1e-12` log-of-zero guard in the forward kernel. @@ -1357,3 +1369,326 @@ fn aggregate_diag_entropy_pins_top_and_bottom() -> Result<()> { ); Ok(()) } + +// ───────────────────────────────────────────────────────────────────── +// Phase 2A-D fix B1.3 — gate-LR multiplier controller invariants. +// +// The controller reads `gate_entropy_mean` (slot 781) via an EMA → +// escalates / decays / leaves alone `RL_POLICY_GATE_LR_MULTIPLIER_INDEX` +// (slot 790) based on whether the EMA sits above 0.85·log(K), +// below 0.20·log(K), or inside the healthy band. +// +// These four tests exercise the kernel in isolation by allocating a +// fresh `MappedF32Buffer`, seeding the relevant slots, and launching +// the standalone cubin. No `MultiHeadPolicy` instance is required. +// ───────────────────────────────────────────────────────────────────── + +/// Allocate a zero-initialised ISV buffer covering all RL slots. +fn isv_zero() -> Result { + let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) } + .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?; + for slot in 0..RL_SLOTS_END { + isv.write_record(slot, 0.0_f32); + } + Ok(isv) +} + +/// Launch the standalone `rl_gate_lr_multiplier_controller` once with a +/// pre-populated ISV. Single-thread (1,1,1)/(1,1,1) — mirrors the trainer's +/// `launch_rl_gate_lr_multiplier_controller`. Synchronises before returning +/// so subsequent `isv.read_record` calls see kernel writes through the +/// mapped pinned host pointer. +fn run_gate_lr_controller_once( + stream: &Arc, + func: &cudarc::driver::CudaFunction, + isv: &MappedF32Buffer, +) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(func); + launch.arg(&isv.dev_ptr); + unsafe { + launch.launch(cfg).context("gate-LR controller launch")?; + } + stream.synchronize().context("gate-LR controller sync")?; + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_lr_controller_bootstrap_initializes_ema_to_first_observation() -> Result<()> { + // Phase 2A-D fix B1.3: on the first launch with `bootstrap_done = 0`, + // the EMA must be set to the current observation (NOT blended with the + // sentinel-zero starting value) and the bootstrap-done flag must + // latch to 1.0. This is the canonical first-observation pattern from + // `pearl_welford_trade_count_is_step_not_trade`. + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let ctx = dev.cuda_context().context("controller ctx")?; + let module = ctx + .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec()) + .context("load rl_gate_lr_multiplier_controller cubin")?; + let func = module + .load_function("rl_gate_lr_multiplier_controller") + .context("load rl_gate_lr_multiplier_controller fn")?; + + let isv = isv_zero()?; + // Seed canonical bootstrap state matching trainer construction: + // - K = 3 + // - multiplier bootstrap = 5.0 + // - entropy_now = 0.5 (inside the healthy band for K=3: + // [0.20·log(3) ≈ 0.220, 0.85·log(3) ≈ 0.934], so the multiplier + // stays put — isolates the EMA-bootstrap effect) + // - entropy_ema = 0 (will be overwritten by first-observation) + // - bootstrap_done = 0 (un-bootstrapped) + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0); + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.5); + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.0); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 0.0); + + run_gate_lr_controller_once(&stream, &func, &isv)?; + + let ema_after = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX); + let boot_after = isv.read_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX); + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + assert!( + (ema_after - 0.5).abs() < 1e-5, + "first-observation EMA = {ema_after}, expected 0.5 (must replace, not blend with 0.0)" + ); + assert!( + (boot_after - 1.0).abs() < 1e-6, + "bootstrap_done after first call = {boot_after}, expected 1.0 (monotonic latch)" + ); + // Healthy-band leaves multiplier alone (5.0 unchanged within 1 ULP). + assert!( + (mult_after - 5.0).abs() < 1e-5, + "healthy-band multiplier = {mult_after}, expected 5.0 (left alone)" + ); + + // Second call should EMA-blend (α = 0.02): with entropy_now still 0.5 and + // ema = 0.5 the result is still 0.5 (idempotent fixed point) — verifies + // we're now in the post-bootstrap branch, not re-bootstrapping. + run_gate_lr_controller_once(&stream, &func, &isv)?; + let ema_after_2 = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX); + let boot_after_2 = isv.read_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX); + assert!( + (ema_after_2 - 0.5).abs() < 1e-5, + "second-call EMA = {ema_after_2}, expected 0.5 (fixed point — confirms post-bootstrap branch)" + ); + assert!( + (boot_after_2 - 1.0).abs() < 1e-6, + "bootstrap_done after second call = {boot_after_2}, expected 1.0 (latch persists)" + ); + eprintln!("PASS — first-observation bootstrap: EMA snaps to obs, flag latches to 1.0"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_lr_controller_escalates_when_entropy_above_threshold() -> Result<()> { + // Phase 2A-D fix B1.3: when `entropy_ema > 0.85·log(K)` the controller + // must compound the multiplier upward at +0.3 %/step (≈+50 % over 137 + // steps for K=3). Drive entropy_now = 1.05 (above 0.85·log(3) ≈ 0.934), + // already past the bootstrap branch, and run 100 steps. Multiplier + // must strictly increase. + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let ctx = dev.cuda_context().context("controller ctx")?; + let module = ctx + .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec()) + .context("load rl_gate_lr_multiplier_controller cubin")?; + let func = module + .load_function("rl_gate_lr_multiplier_controller") + .context("load rl_gate_lr_multiplier_controller fn")?; + + let isv = isv_zero()?; + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + // 1.05 > 0.85·log(3) ≈ 0.934 → escalate branch + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05); + // Pre-bootstrap the EMA so we skip the first-observation pin (1.05 is + // already above threshold so the EMA stays above threshold while the + // controller compounds). + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0); + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0); + + let mult_before = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + let n_steps = 100; + for _ in 0..n_steps { + run_gate_lr_controller_once(&stream, &func, &isv)?; + } + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + + // Analytical expectation: 5.0 × 1.003^100 ≈ 5.0 × 1.3493 ≈ 6.747. + // Allow generous slack to absorb fp32 compounding noise. + let expected = 5.0_f32 * 1.003_f32.powi(n_steps as i32); + assert!( + mult_after > 6.0, + "after {n_steps} escalate steps: mult {mult_before} → {mult_after}, expected > 6.0 (analytical {expected:.4})" + ); + assert!( + (mult_after - expected).abs() < 0.1, + "after {n_steps} escalate steps: mult = {mult_after}, expected ≈ {expected:.4} (within ±0.1)" + ); + // EMA must still be above the over-threshold (otherwise the test + // doesn't actually exercise the escalate branch end-to-end). + let ema_after = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX); + let over_thr = 0.85_f32 * 3.0_f32.ln(); + assert!( + ema_after > over_thr, + "EMA after {n_steps} steps = {ema_after}, expected > over_threshold {over_thr} (test ran in escalate branch)" + ); + eprintln!("PASS — escalate branch: {mult_before} → {mult_after} (expected ≈ {expected:.4}) over {n_steps} steps"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_lr_controller_decays_when_entropy_below_collapse() -> Result<()> { + // Phase 2A-D fix B1.3: when `entropy_ema < 0.20·log(K)` the controller + // must back off the multiplier fast (×0.95/step). Drive + // entropy_now = 0.1 (below 0.20·log(3) ≈ 0.220), pre-bootstrap the EMA + // to the same value, and run 100 steps. Multiplier should decay + // rapidly toward — and likely hit — the MIN_FLOOR of 1.0. + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let ctx = dev.cuda_context().context("controller ctx")?; + let module = ctx + .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec()) + .context("load rl_gate_lr_multiplier_controller cubin")?; + let func = module + .load_function("rl_gate_lr_multiplier_controller") + .context("load rl_gate_lr_multiplier_controller fn")?; + + let isv = isv_zero()?; + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + // 0.1 < 0.20·log(3) ≈ 0.220 → decay branch + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.1); + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.1); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0); + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0); + + let mult_before = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + let n_steps = 100; + for _ in 0..n_steps { + run_gate_lr_controller_once(&stream, &func, &isv)?; + } + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + + // Analytical: 5.0 × 0.95^100 ≈ 5.0 × 5.92e-3 ≈ 0.0296 — pre-clamp would + // be well below MIN_FLOOR. After clamping we expect exactly 1.0. + assert!( + mult_after < 1.5, + "after {n_steps} decay steps: mult {mult_before} → {mult_after}, expected < 1.5" + ); + assert!( + (mult_after - 1.0).abs() < 1e-5, + "after {n_steps} decay steps: mult = {mult_after}, expected ≈ 1.0 (clamped to MIN_FLOOR)" + ); + eprintln!("PASS — decay branch: {mult_before} → {mult_after} (clamped to MIN_FLOOR=1.0) over {n_steps} steps"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn gate_lr_controller_clamps_at_min_floor_and_max_ceiling() -> Result<()> { + // Phase 2A-D fix B1.3: the controller MUST clamp to [1.0, 50.0] + // per `pearl_bootstrap_must_respect_clamp_range`. Push the multiplier + // hard against each boundary by pre-setting an extreme value and + // letting the kernel apply the active branch for many steps. + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let ctx = dev.cuda_context().context("controller ctx")?; + let module = ctx + .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec()) + .context("load rl_gate_lr_multiplier_controller cubin")?; + let func = module + .load_function("rl_gate_lr_multiplier_controller") + .context("load rl_gate_lr_multiplier_controller fn")?; + + // ── Case 1: pre-set above MAX_CEIL with escalate branch active ── + // Pre-setting 100.0 (well above MAX_CEIL=50.0) and entropy above + // the over-threshold should clamp DOWN to 50.0 on first kernel call. + { + let isv = isv_zero()?; + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05); + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0); + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 100.0); + + run_gate_lr_controller_once(&stream, &func, &isv)?; + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + assert!( + mult_after <= 50.0 + 1e-5, + "above-ceil clamp: pre = 100.0 → post = {mult_after}, expected ≤ 50.0" + ); + // 100.0 × 1.003 = 100.3, clamped to 50.0. + assert!( + (mult_after - 50.0).abs() < 1e-5, + "above-ceil clamp: post = {mult_after}, expected exact 50.0" + ); + } + + // ── Case 2: many escalate steps from inside the band shouldn't exceed MAX_CEIL ── + { + let isv = isv_zero()?; + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05); + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0); + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0); + + // 5 × 1.003^800 ≈ 5 × 10.94 ≈ 54.7 — pre-clamp value would exceed + // ceiling; the controller must hold it at 50.0. + for _ in 0..800 { + run_gate_lr_controller_once(&stream, &func, &isv)?; + } + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + assert!( + mult_after <= 50.0 + 1e-5, + "long-escalate ceiling: 800 steps from 5.0 → {mult_after}, expected ≤ 50.0" + ); + assert!( + (mult_after - 50.0).abs() < 1e-5, + "long-escalate ceiling: 800 steps from 5.0 → {mult_after}, expected ≈ 50.0" + ); + } + + // ── Case 3: pre-set below MIN_FLOOR with decay branch active ── + { + let isv = isv_zero()?; + isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0); + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.1); + isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.1); + isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0); + // Set BELOW MIN_FLOOR — the clamp must raise it back to 1.0. + isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 0.5); + + run_gate_lr_controller_once(&stream, &func, &isv)?; + let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX); + assert!( + mult_after >= 1.0 - 1e-5, + "below-floor clamp: pre = 0.5 → post = {mult_after}, expected ≥ 1.0" + ); + // 0.5 × 0.95 = 0.475, clamped UP to 1.0. + assert!( + (mult_after - 1.0).abs() < 1e-5, + "below-floor clamp: post = {mult_after}, expected exact 1.0" + ); + } + eprintln!("PASS — multiplier clamps at both boundaries: [1.0, 50.0]"); + Ok(()) +}