diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index 196d9726a..2aebb0681 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -243,3 +243,62 @@ extern "C" __global__ void slack_factor_apply_kernel( iqr_lo[k] = q1 / slack; iqr_hi[k] = q3 * slack; } + +// ───────────────────────────────────────────────────────────────────── +// tau_clamp_kernel — Controller B post-Adam projection. +// +// Per spec §3.2 and pearl_adam_normalizes_loss_weights: hard projection +// AFTER each Adam step on the τ slice, bypassing Adam's m/sqrt(v) +// normalization that would no-op a loss-weight modulation. +// +// Per pearl_isv_for_adaptive_bounds: the (lo, hi) inputs are the IQR- +// widened bounds from `slack_factor_apply_kernel` (ISV-derived from each +// bucket's observed Q3/Q1 ratio at the Phase 1→2 transition). +// +// Launch: 1 block × HIDDEN_DIM threads (single warp on Ampere+; tiny). +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void tau_clamp_kernel( + float* __restrict__ tau_all, // [HIDDEN_DIM] + const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] + const float* __restrict__ iqr_lo, // [N_HORIZONS] + const float* __restrict__ iqr_hi // [N_HORIZONS] +) { + int c = threadIdx.x; + if (c >= HIDDEN_DIM) return; + unsigned char k = bucket_id_per_channel[c]; + float v = tau_all[c]; + float lo = iqr_lo[k]; + float hi = iqr_hi[k]; + tau_all[c] = fminf(fmaxf(v, lo), hi); +} + +// ───────────────────────────────────────────────────────────────────── +// heads_lr_multiplier_scale_kernel — Controller C per-bucket LR multiplier. +// +// Per spec §3.3 and pearl_adam_normalizes_loss_weights: rescales the +// per-step parameter delta (post-Adam − pre-Adam) by a per-horizon LR +// multiplier. Bypasses Adam's m/sqrt(v) normalization because the +// rescaling is applied to the parameter delta itself, not the loss +// weight or the Adam moments. +// +// Layout: `param` is row-major [N_HORIZONS × per_horizon_stride]. Each +// horizon owns a contiguous block of `per_horizon_stride` floats. Element +// `i` belongs to horizon `i / per_horizon_stride`. +// +// Launch: enough threads to cover `total_elements` (one per element). +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void heads_lr_multiplier_scale_kernel( + float* __restrict__ param, // [N_HORIZONS * per_horizon_stride] + const float* __restrict__ param_pre, // pre-Adam snapshot, same shape + const float* __restrict__ lr_mult_per_horizon, // [N_HORIZONS] + int per_horizon_stride, + int total_elements +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_elements) return; + int horizon = i / per_horizon_stride; + if (horizon >= N_HORIZONS) return; // safety bound for stride/total mismatch + float delta = param[i] - param_pre[i]; + float scaled_delta = delta * lr_mult_per_horizon[horizon]; + param[i] = param_pre[i] + scaled_delta; +} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index fcc853504..be8d7ab96 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -708,6 +708,15 @@ pub struct PerceptionTrainer { /// (`slack_factor_apply_kernel` in `bucket_transition_kernels.cu`). /// Launched once at Phase 1→2 transition, then never again. slack_factor_apply_fn: CudaFunction, + /// Cached function handle for `tau_clamp_kernel` — Controller B's + /// post-Adam τ projection (spec §3.2). Launched after every + /// `opt_tau.step` in Phase 2; no-op in Phase 1. + tau_clamp_fn: CudaFunction, + /// Cached function handle for `heads_lr_multiplier_scale_kernel` — + /// Controller C's per-bucket LR multiplier on the heads_w_skip Adam + /// delta (spec §3.3, scope reduced to heads_w_skip in Task 11). + /// Launched after every `opt_heads_w_skip.step` in Phase 2. + heads_lr_multiplier_scale_fn: CudaFunction, /// Module handle keeping the bucket-transition cubin alive for the /// lifetime of the cached function handles above. Loaded once at /// construction so the per-step tau_change_frobenius launch doesn't @@ -726,6 +735,40 @@ pub struct PerceptionTrainer { /// dispatch + Controllers B / C / D. pub bucket_routing_metadata: Option, + + // ── Controller C state (per-bucket LR multiplier; spec §3.3) ── + // + // Scope reduction documented in Task 11: only `heads_w_skip_d` is + // rescaled per-horizon (N_HORIZONS × HIDDEN_DIM = 640 contiguous floats, + // per-horizon stride HIDDEN_DIM). The other heads weight groups + // (w1, w2, w_gate, w_main) keep unit LR multipliers; they are entangled + // in the GRN gate/main computation and rescaling them would require + // additional per-bucket scoping out of Task 11 scope. Revisit if + // Smoke 2 fails per `feedback_no_quickfixes`. + /// Pre-Adam snapshot of `trunk.heads_w_skip_d` for Controller C delta + /// rescaling. DtoD-copied each Phase 2 step BEFORE `opt_heads_w_skip.step`. + /// Shape matches `heads_w_skip_d` (N_HORIZONS × HIDDEN_DIM = 640 floats). + heads_w_skip_pre_adam_d: CudaSlice, + /// Per-horizon LR multiplier vector consumed by + /// `heads_lr_multiplier_scale_kernel`. Mapped-pinned per + /// `feedback_no_htod_htoh_only_mapped_pinned`: host computes + /// `1 / (1 + smoothness_ratio_k)` each Phase 2 step and writes into + /// the staging slice; kernel reads via the device-visible pointer. + /// In Phase 1 stays at the zero-init value (kernel never launched). + lr_mult_per_horizon_staging: MappedF32Buffer, + /// Mapped-pinned shadow of `smoothness_jitter_ema_d` populated each + /// Phase 2 step (DtoD into the device-visible buffer, then host read + /// after the smoothness controller's launch barrier). Provides the + /// per-horizon `jitter_ema_k` signal that feeds Controller C's + /// smoothness ratio. + smoothness_jitter_ema_host_d: MappedF32Buffer, + /// Per-horizon target jitter for Controller C (`target_jitter_k` in + /// spec §3.3). Sentinel-bootstrapped per + /// `pearl_first_observation_bootstrap`: stays at 0 until the first + /// non-zero `raw_per_h` observation in Phase 2; then is written + /// directly with that value. After bootstrap, this acts as the + /// `realised_label_variance_k` proxy. + target_jitter_per_h: [f32; N_HORIZONS], } impl PerceptionTrainer { @@ -851,6 +894,12 @@ impl PerceptionTrainer { let slack_factor_apply_fn = bucket_transition_module .load_function("slack_factor_apply_kernel") .context("load slack_factor_apply_kernel")?; + let tau_clamp_fn = bucket_transition_module + .load_function("tau_clamp_kernel") + .context("load tau_clamp_kernel")?; + let heads_lr_multiplier_scale_fn = bucket_transition_module + .load_function("heads_lr_multiplier_scale_kernel") + .context("load heads_lr_multiplier_scale_kernel")?; // Mamba2 stacks live on the trunk from new_random; we wire optimizer // + training scratches against the trunk-owned blocks. @@ -1227,6 +1276,26 @@ impl PerceptionTrainer { let heads_w_skip_compact_d = stream .alloc_zeros::(crate::cfc::bucket_routing::HIDDEN_DIM) .context("heads_w_skip_compact_d alloc")?; + // ── Controller C scratch buffers (Task 11) ── + // Pre-Adam snapshot of heads_w_skip_d. Same shape as the live + // weight tensor (N_HORIZONS × HIDDEN_DIM). + let heads_w_skip_pre_adam_d = stream + .alloc_zeros::(N_HORIZONS * crate::cfc::bucket_routing::HIDDEN_DIM) + .context("heads_w_skip_pre_adam_d alloc")?; + // Per-horizon LR multiplier staging. Mapped-pinned so the host + // writes from `target_jitter_per_h` + jitter EMA each Phase 2 step + // and the kernel reads the device-visible pointer without DtoH. + let lr_mult_per_horizon_staging = unsafe { MappedF32Buffer::new(N_HORIZONS) } + .map_err(|e| anyhow::anyhow!("lr_mult_per_horizon_staging: {e}"))?; + // Bootstrap to 1.0 (no attenuation) so a Phase 1 kernel that + // accidentally launches still preserves the Adam delta. Phase 2's + // first step overwrites this before the launch. + lr_mult_per_horizon_staging.write_from_slice(&[1.0; N_HORIZONS]); + // Mapped-pinned shadow of the per-horizon jitter EMA. Each Phase 2 + // step DtoD-copies `smoothness_jitter_ema_d` into this buffer (5 + // floats); host reads after the step sync. + let smoothness_jitter_ema_host_d = unsafe { MappedF32Buffer::new(N_HORIZONS) } + .map_err(|e| anyhow::anyhow!("smoothness_jitter_ema_host_d: {e}"))?; let controller_a = crate::cfc::bucket_routing::ControllerA::new(); // ── CRT.train: output-smoothness regularizer state ── @@ -1585,9 +1654,16 @@ impl PerceptionTrainer { tau_change_staged, tau_change_frobenius_fn, slack_factor_apply_fn, + tau_clamp_fn, + heads_lr_multiplier_scale_fn, _bucket_transition_module: bucket_transition_module, heads_w_skip_compact_d, bucket_routing_metadata: None, + // Controller C state (Task 11). + heads_w_skip_pre_adam_d, + lr_mult_per_horizon_staging, + smoothness_jitter_ema_host_d, + target_jitter_per_h: [0.0; N_HORIZONS], // sentinel; first obs bootstraps }) } @@ -2004,6 +2080,59 @@ impl PerceptionTrainer { // ── 3. Sync + read mapped-pinned loss. self.stream.synchronize().context("end-of-step sync")?; let loss = unsafe { std::ptr::read_volatile(self.loss_host_d.host_ptr) }; + + // ── 4. Controller C host-side LR multiplier update (spec §3.3) ── + // + // Now that the end-of-step sync has propagated the captured DtoD + // shadows, both `smoothness_jitter_ema_host_d` and + // `smoothness_loss_per_horizon_host_d` carry this step's per-horizon + // values. The host: + // 1. Sentinel-bootstraps `target_jitter_per_h[k]` from the FIRST + // non-zero raw smoothness loss in Phase 2 (per + // `pearl_first_observation_bootstrap`). + // 2. Computes `smoothness_ratio_k = max(0, jitter_ema - target) / + // max(target, ε_floor)` per `pearl_blend_formulas_must_have_permanent_floor`. + // 3. Writes `lr_mult_k = 1 / (1 + smoothness_ratio_k)` (bounded + // (0, 1]) to the mapped-pinned staging slice. The next + // step's captured-graph rescale kernel reads the updated + // values via the device-visible pointer (mapped-pinned + // coherence — no DtoH). + // + // Phase 1: no-op (the rescale kernel is also gated to Phase 2, + // and the staging keeps its zero-init / bootstrap value). + if self.phase == TrainingPhase::Phase2Routed { + let jitter_ema = self.smoothness_jitter_ema_host_d.read_all(); + let raw_per_h = self.smoothness_loss_per_horizon_host_d.read_all(); + let mut lr_mult = [1.0_f32; N_HORIZONS]; + for k in 0..N_HORIZONS { + // Sentinel-bootstrap target_jitter_per_h from first non-zero + // observation. Both EMA and raw can legitimately read as + // sentinel-zero on the first Phase 2 step before the + // smoothness controller has run; once the first non-zero + // sample arrives we replace directly per the pearl. + if self.target_jitter_per_h[k] == 0.0 && raw_per_h[k] > 0.0 { + self.target_jitter_per_h[k] = raw_per_h[k]; + } + let target = self.target_jitter_per_h[k]; + if target <= 0.0 { + // Still pre-bootstrap: leave lr_mult at 1.0 (no + // attenuation). Once the first observation lands the + // next step picks up a non-zero target. + continue; + } + // ε_floor = target / 16 per `pearl_blend_formulas_must_have_permanent_floor` + // (avoids division blow-up if target drifts toward zero). + let eps_floor = target / 16.0; + let safe_target = target.max(eps_floor); + let excess = (jitter_ema[k] - target).max(0.0); + let smoothness_ratio = excess / safe_target; + lr_mult[k] = 1.0 / (1.0 + smoothness_ratio); + } + // Mapped-pinned write: kernel reads via dev_ptr on the next + // graph replay. Volatile write inside `write_from_slice`. + self.lr_mult_per_horizon_staging.write_from_slice(&lr_mult); + } + Ok(loss) } @@ -3111,6 +3240,33 @@ impl PerceptionTrainer { self.opt_w_rec.step(&mut self.trunk.w_rec_d, &self.grad_w_rec_d)?; self.opt_b.step(&mut self.trunk.b_d, &self.grad_b_d)?; self.opt_tau.step(&mut self.trunk.tau_all_d, &self.grad_tau_d)?; + // ── Controller B: post-Adam τ projection (spec §3.2) ── + // Per `pearl_adam_normalizes_loss_weights`, the only way to actually + // bound `tau_all_d` is a hard projection AFTER Adam updates the + // parameter. Loss-weight modulation would no-op because Adam's + // m/sqrt(v) cancels constant multipliers. Phase 2 only — Phase 1's + // captured graph never enters this branch since the transition + // invalidates the graph (`pearl_no_host_branches_in_captured_graph`). + if self.phase == TrainingPhase::Phase2Routed { + if let Some(metadata) = self.bucket_routing_metadata.as_ref() { + let cfg_tau_clamp = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.tau_clamp_fn); + launch + .arg(&mut self.trunk.tau_all_d) + .arg(&metadata.bucket_id_per_channel_d) + .arg(&metadata.bucket_tau_iqr_lo_d) + .arg(&metadata.bucket_tau_iqr_hi_d); + unsafe { + launch + .launch(cfg_tau_clamp) + .context("tau_clamp_kernel launch (Controller B)")?; + } + } + } self.opt_heads_w1.step(&mut self.trunk.heads_w1_d, &self.grad_heads_w1_d)?; self.opt_heads_b1.step(&mut self.trunk.heads_b1_d, &self.grad_heads_b1_d)?; self.opt_heads_w2.step(&mut self.trunk.heads_w2_d, &self.grad_heads_w2_d)?; @@ -3119,7 +3275,56 @@ impl PerceptionTrainer { self.opt_heads_b_gate.step(&mut self.trunk.heads_b_gate_d, &self.grad_heads_b_gate_d)?; self.opt_heads_w_main.step(&mut self.trunk.heads_w_main_d, &self.grad_heads_w_main_d)?; self.opt_heads_b_main.step(&mut self.trunk.heads_b_main_d, &self.grad_heads_b_main_d)?; + // ── Controller C: pre-Adam snapshot of heads_w_skip_d (spec §3.3) ── + // Captured DtoD copy: heads_w_skip_d → heads_w_skip_pre_adam_d. Used + // by the post-Adam rescale below to compute the parameter delta + // (post-Adam − pre-Adam) without ever reading Adam's internal + // m/v moments. Phase 2 only; the rescale launch below is gated + // identically so the snapshot is also gated to preserve memory + // bandwidth in Phase 1. + if self.phase == TrainingPhase::Phase2Routed { + unsafe { + let (src, _gs) = self.trunk.heads_w_skip_d.device_ptr(&self.stream); + let (dst, _gd) = self.heads_w_skip_pre_adam_d.device_ptr_mut(&self.stream); + let nbytes = N_HORIZONS * HIDDEN_DIM * std::mem::size_of::(); + cudarc::driver::result::memcpy_dtod_async( + dst, src, nbytes, self.stream.cu_stream(), + ) + .context("heads_w_skip pre-Adam DtoD snapshot")?; + } + } self.opt_heads_w_skip.step(&mut self.trunk.heads_w_skip_d, &self.grad_heads_w_skip_d)?; + // ── Controller C: per-horizon LR multiplier rescale (spec §3.3) ── + // Computes `heads_w_skip[i] = pre[i] + (post[i] − pre[i]) * lr_mult[h(i)]` + // for every element i, with `h(i) = i / HIDDEN_DIM`. Scope reduction + // documented in the trainer struct doc: only heads_w_skip is rescaled + // here; revisit if Smoke 2 fails. + if self.phase == TrainingPhase::Phase2Routed { + const TOTAL_HEADS_W_SKIP: i32 = (N_HORIZONS * HIDDEN_DIM) as i32; + const PER_HORIZON_STRIDE: i32 = HIDDEN_DIM as i32; + let block: u32 = 64; + let grid: u32 = ((TOTAL_HEADS_W_SKIP as u32) + block - 1) / block; + let cfg_scale = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }; + let lr_mult_dev_ptr = self.lr_mult_per_horizon_staging.dev_ptr; + let mut launch = self + .stream + .launch_builder(&self.heads_lr_multiplier_scale_fn); + launch + .arg(&mut self.trunk.heads_w_skip_d) + .arg(&self.heads_w_skip_pre_adam_d) + .arg(&lr_mult_dev_ptr) + .arg(&PER_HORIZON_STRIDE) + .arg(&TOTAL_HEADS_W_SKIP); + unsafe { + launch + .launch(cfg_scale) + .context("heads_lr_multiplier_scale_kernel launch (Controller C)")?; + } + } self.opt_heads_b_skip.step(&mut self.trunk.heads_b_skip_d, &self.grad_heads_b_skip_d)?; self.opt_ln_gain.step(&mut self.trunk.ln_b_gain_d, &self.grad_ln_gain_d)?; self.opt_ln_bias.step(&mut self.trunk.ln_b_bias_d, &self.grad_ln_bias_d)?; @@ -3174,6 +3379,20 @@ impl PerceptionTrainer { self.stream.cu_stream(), ).context("smoothness_loss_per_horizon_host_d shadow")?; } + // Controller C feeder shadow: per-horizon jitter EMA → mapped-pinned + // host buffer so step_batched can compute lr_mult_per_horizon between + // steps. Captured-graph DtoD on the same stream; visible after the + // end-of-step sync per the standard pattern. + unsafe { + let (src_ptr, _g) = self.smoothness_jitter_ema_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + self.smoothness_jitter_ema_host_d.dev_ptr, + src_ptr, + N_HORIZONS * std::mem::size_of::(), + self.stream.cu_stream(), + ) + .context("smoothness_jitter_ema_host_d shadow")?; + } // ── Shadow device step counter into mapped-pinned host shadow ── // The drain task reads the host shadow at 500ms cadence to know // how many steps have completed since its last drain. Captured diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs index ae62d776a..5f8974567 100644 --- a/crates/ml-alpha/tests/bucket_transition_kernels.rs +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -212,6 +212,120 @@ fn tau_change_frobenius_kernel_sums_squared_diff() -> Result<()> { Ok(()) } +#[test] +#[ignore = "requires CUDA"] +fn tau_clamp_kernel_bounds_tau_within_bucket_iqr() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "tau_clamp_kernel")?; + + // All channels start with τ = 100 + c * 0.1 (way above any bucket's hi). + // Bucket assignment: channel c → bucket c/25 capped at 4. + // IQR bounds: bucket k has [k*10, k*10 + 5]. + // Expected outcome: every channel clamps to its bucket's hi value. + let tau_host: Vec = + (0..HIDDEN_DIM).map(|c| 100.0 + (c as f32) * 0.1).collect(); + let bucket_id_host: Vec = + (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); + let iqr_lo_host: Vec = + (0..N_HORIZONS).map(|k| (k as f32) * 10.0).collect(); + let iqr_hi_host: Vec = + (0..N_HORIZONS).map(|k| (k as f32) * 10.0 + 5.0).collect(); + + let mut tau_d = upload(&stream, &tau_host)?; + let bucket_id_d = upload(&stream, &bucket_id_host)?; + let iqr_lo_d = upload(&stream, &iqr_lo_host)?; + let iqr_hi_d = upload(&stream, &iqr_hi_host)?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&mut tau_d) + .arg(&bucket_id_d) + .arg(&iqr_lo_d) + .arg(&iqr_hi_d); + unsafe { + launch.launch(cfg)?; + } + stream.synchronize()?; + + let tau_out = download(&stream, &tau_d)?; + for c in 0..HIDDEN_DIM { + let k = (c / 25).min(4); + let expected = (k as f32) * 10.0 + 5.0; // upper bound of bucket + assert!( + (tau_out[c] - expected).abs() < 1e-6, + "c={} k={} got={} expected={}", + c, + k, + tau_out[c], + expected + ); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn heads_lr_multiplier_scale_kernel_rescales_delta_per_horizon() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "heads_lr_multiplier_scale_kernel")?; + + // Per-horizon stride 8: param tensor shape [N_HORIZONS * 8] = 40 floats. + let per_horizon_stride: i32 = 8; + let total: i32 = N_HORIZONS as i32 * per_horizon_stride; + let param_pre: Vec = + (0..total).map(|i| (i as f32) * 0.1).collect(); + // Uniform delta of 1.0 applied "by Adam" — simulated post-Adam value. + let param_post: Vec = param_pre.iter().map(|x| x + 1.0).collect(); + // Distinct multipliers per horizon, including 0 (full attenuation) and 1 + // (no change). After scaling: delta'[k] = lr_mult[k] = 0, 0.25, 0.5, 0.75, 1.0. + let lr_mult: Vec = vec![0.0, 0.25, 0.5, 0.75, 1.0]; + + let mut param_d = upload(&stream, ¶m_post)?; + let param_pre_d = upload(&stream, ¶m_pre)?; + let lr_mult_d = upload(&stream, &lr_mult)?; + + let block: u32 = 64; + let grid: u32 = ((total as u32) + block - 1) / block; + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&mut param_d) + .arg(¶m_pre_d) + .arg(&lr_mult_d) + .arg(&per_horizon_stride) + .arg(&total); + unsafe { + launch.launch(cfg)?; + } + stream.synchronize()?; + + let param_out = download(&stream, ¶m_d)?; + for i in 0..total as usize { + let k = i / per_horizon_stride as usize; + let expected = param_pre[i] + lr_mult[k]; + assert!( + (param_out[i] - expected).abs() < 1e-6, + "i={} k={} got={} expected={}", + i, + k, + param_out[i], + expected + ); + } + Ok(()) +} + #[test] #[ignore = "requires CUDA"] fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> {