From 2f69bb1fc0b66790bb8a6724619fa33b9599449d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 17:09:01 +0200 Subject: [PATCH] =?UTF-8?q?feat(per-horizon-cfc):=20wire=20Controller=20A?= =?UTF-8?q?=20+=20Phase=201=E2=86=922=20transition=20in=20trainer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per spec §2.3 + §3.1 and Task 9 of the plan. Adds two device kernels to bucket_transition_kernels.cu: - tau_change_frobenius_kernel: ‖tau_t − tau_{t-1}‖_F² via block-tree reduction (no atomicAdd per feedback_no_atomicadd) - slack_factor_apply_kernel: (Q1, Q3) → (Q1/slack, Q3*slack) where slack = sqrt(Q3/Q1), ISV-derived per feedback_isv_for_adaptive_bounds Adds to PerceptionTrainer: - TrainingPhase enum (Phase1Warmup / Phase2Routed) - ControllerA state machine instance - prev_tau_d device buffer + tau_change_d + tau_change_staged mapped-pinned - bucket_warmup_cap_override CLI diagnostic field - bucket_routing_metadata stored after transition fires - heads_w_skip_compact_d (HIDDEN_DIM floats) populated by transition - Cached function handles for the two new kernels Per-step in step_batched (Phase 1 only): - Launch tau_change_frobenius_kernel → mapped-pinned scalar shadow - DtoD prev_tau ← tau_all_d for next step - Sync + host scalar read - controller_a.update(tau_change, bucket_warmup_cap_override) - On trigger: stage tau into scratch buffer (avoids in/out aliasing in tau_reorder_kernel), execute_transition writes routing metadata + reorders tau_all_d + populates heads_w_skip_compact_d, slack_factor_apply widens IQR bounds, invalidate CUDA Graph, latch phase = Phase2Routed, log via tracing::info. Phase 2 dispatch + Controllers B/C/D wiring deferred to Tasks 10–12. In Task 9's transient state, the trainer enters Phase 2 but continues Phase 1 dispatch path; the reordered tau_all_d still produces a valid forward pass since CfC's per-channel decay math is order-invariant. Per pearl_no_host_branches_in_captured_graph: transition fires OUTSIDE the captured graph (cached graph invalidated → recaptured next step). Per pearl_cudarc_disable_event_tracking_for_graph_capture: event tracking is already disabled trainer-wide; recapture on next step is safe. Per feedback_no_htod_htoh_only_mapped_pinned: host scalar read goes via mapped-pinned DtoD shadow, not bulk DtoH. CLI: --bucket-warmup-cap-steps added to alpha_train (Option, diagnostic override of Controller A's ISV-derived cap). Tests: 7 GPU oracle tests pass on RTX 3050 (5 existing + 2 new for Frobenius + slack_factor). 33 ml-alpha lib tests pass. Workspace cargo check clean modulo pre-existing cupti / sp15 / gpu_per_integration errors unrelated to this task. Co-Authored-By: Claude Opus 4.7 --- .../cuda/bucket_transition_kernels.cu | 52 +++ crates/ml-alpha/examples/alpha_train.rs | 12 + crates/ml-alpha/src/trainer/perception.rs | 321 ++++++++++++++++++ .../tests/bucket_transition_kernels.rs | 65 ++++ crates/ml-alpha/tests/perception_overfit.rs | 8 + 5 files changed, 458 insertions(+) diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index 6f5a5a1e3..196d9726a 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -191,3 +191,55 @@ extern "C" __global__ void heads_compact_kernel( __syncthreads(); } } + +// ───────────────────────────────────────────────────────────────────── +// tau_change_frobenius_kernel: ||tau_t - tau_{t-1}||_F² (scalar output). +// +// Single block × HIDDEN_DIM threads. Block-tree reduction (no atomicAdd +// per feedback_no_atomicadd). +// +// Per pearl_first_observation_bootstrap: caller's ControllerA bootstraps +// the EMA from this scalar on step 1; on subsequent steps the scalar +// feeds the Wiener-α update. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void tau_change_frobenius_kernel( + const float* __restrict__ tau_t, // [HIDDEN_DIM] + const float* __restrict__ tau_t_minus_1, // [HIDDEN_DIM] + float* __restrict__ tau_change_out // [1] +) { + __shared__ float sdata[HIDDEN_DIM]; + int tid = threadIdx.x; + if (tid >= HIDDEN_DIM) return; + float diff = tau_t[tid] - tau_t_minus_1[tid]; + sdata[tid] = diff * diff; + __syncthreads(); + // Block-tree reduction (HIDDEN_DIM=128 is power-of-two; no padding needed). + for (int s = HIDDEN_DIM / 2; s > 0; s >>= 1) { + if (tid < s) sdata[tid] += sdata[tid + s]; + __syncthreads(); + } + if (tid == 0) *tau_change_out = sdata[0]; +} + +// ───────────────────────────────────────────────────────────────────── +// slack_factor_apply_kernel: convert (Q1, Q3) → IQR_widened in place. +// +// Per spec §3.2: slack_factor_k = sqrt(Q3_k / Q1_k), ISV-derived from +// observed bucket IQR (no hardcoded 1.5× factor). +// +// Launch: 1 block × N_HORIZONS threads (5 threads, well within a warp). +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void slack_factor_apply_kernel( + float* __restrict__ iqr_lo, // [N_HORIZONS] in: Q1, out: Q1/slack + float* __restrict__ iqr_hi // [N_HORIZONS] in: Q3, out: Q3*slack +) { + int k = threadIdx.x; + if (k >= N_HORIZONS) return; + float q1 = iqr_lo[k]; + float q3 = iqr_hi[k]; + // Numerical defense: Q1 must be > 0 (taus are log-uniform [0.01, 1000], + // all positive by construction). + float slack = sqrtf(q3 / fmaxf(q1, 1e-6f)); + iqr_lo[k] = q1 / slack; + iqr_hi[k] = q3 * slack; +} diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index 231f696a3..575aacafa 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -172,6 +172,17 @@ struct Cli { /// Example: `--kernel-step-trace /feature-cache/.../step_trace.jsonl` #[arg(long)] kernel_step_trace: Option, + + /// Diagnostic override for the per-horizon CfC bucket warmup hard cap. + /// + /// `None` (default) lets Controller A derive the cap from the + /// observed `tau_change` dispersion in the first 100 training steps + /// (per spec §3.1). `Some(steps)` forces Phase 1→2 transition at + /// exactly that step count — used by diagnostic Argo runs to test + /// the routing path without waiting on natural CfC.tau convergence. + /// Production runs should leave this unset. + #[arg(long)] + bucket_warmup_cap_steps: Option, } #[derive(Serialize, serde::Deserialize, Default)] @@ -287,6 +298,7 @@ fn main() -> Result<()> { n_batch: cli.batch_size, smoothness_base_lambda: cli.smoothness_base_lambda, kernel_step_trace_path: cli.kernel_step_trace.clone(), + bucket_warmup_cap_override: cli.bucket_warmup_cap_steps, }; let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 107e29546..b74bb353b 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -67,6 +67,14 @@ const SMOOTHNESS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/output const SMOOTHNESS_CONTROLLER_CUBIN: &[u8] = include_bytes!( concat!(env!("OUT_DIR"), "/smoothness_lambda_controller.cubin") ); +/// Phase 1→2 transition kernels (Task 9 scope: tau_change_frobenius + +/// slack_factor_apply; the full transition orchestration in +/// `crate::cfc::bucket_routing::execute_transition` loads the cubin +/// independently — this constant gives the trainer cached function +/// handles for the per-step Controller A signal kernels). +const BUCKET_TRANSITION_CUBIN: &[u8] = include_bytes!( + concat!(env!("OUT_DIR"), "/bucket_transition_kernels.cubin") +); #[derive(Clone, Debug)] pub struct PerceptionTrainerConfig { @@ -101,6 +109,17 @@ pub struct PerceptionTrainerConfig { /// Per `feedback_no_feature_flags`: gated by the compile-time /// `kernel-step-trace` feature; specific name justifies the gate. pub kernel_step_trace_path: Option, + + /// Per-horizon CfC bucket warmup cap override (diagnostic). + /// `None` (default) lets `ControllerA` derive the cap from the + /// observed `tau_change` dispersion in the first 100 training steps + /// (per spec §3.1, `half_life × (1 + MAD/median)`). `Some(steps)` + /// forces Phase 1→2 transition at exactly that step count, used by + /// diagnostic Argo runs to test the routing path without waiting on + /// natural CfC.tau convergence. Per + /// `feedback_isv_for_adaptive_bounds`: production must leave this + /// at `None` — only diagnostic CLI sweeps should set it. + pub bucket_warmup_cap_override: Option, } impl Default for PerceptionTrainerConfig { @@ -115,10 +134,31 @@ impl Default for PerceptionTrainerConfig { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, } } } +/// Two-phase training-loop state machine for the per-horizon CfC +/// inference architecture (spec §2.3). +/// +/// `Phase1Warmup` — single shared CfC body, full HIDDEN_DIM-wide heads, +/// Controller A monitoring `CfC.tau` drift. The current dispatch path +/// is unchanged from the pre-routing trainer. +/// +/// `Phase2Routed` — bucket-routed dispatch (Task 10 will land the +/// per-branch fused CfC + compact heads dispatch). Task 9 transitions +/// the trainer's `phase` flag and stages the routing metadata; the +/// actual dispatch path remains Phase-1-shaped until Task 10. This +/// transient state is well-defined: reordered `tau_all_d` still +/// produces a valid forward pass because CfC's per-channel math is +/// order-invariant — only the bucket-routing benefit is deferred. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TrainingPhase { + Phase1Warmup, + Phase2Routed, +} + /// Per-horizon BCE weight schedule, used when the CLI passes /// `--auto-horizon-weights`. Replaces the previous `min(1, K/h)` /// schedule which down-weighted long horizons to fractions of a @@ -639,6 +679,53 @@ pub struct PerceptionTrainer { /// `[N, FEATURE_DIM]` interface (the [N, 1, F] shape flattens to /// `[N, F]` by row-major equivalence). window_step_d: GpuTensor, + + // ── Per-horizon CfC routing (Task 9 wiring, spec §2.3 + §3.1) ── + /// Two-phase loop state. Starts in `Phase1Warmup`; flips to + /// `Phase2Routed` exactly once when Controller A fires. + pub phase: TrainingPhase, + /// Controller A — warmup-completion detector. Per-step host-side + /// state machine (EMA + ISV-derived hard cap); reads a single + /// mapped-pinned scalar each step per `feedback_cpu_is_read_only`. + pub controller_a: crate::cfc::bucket_routing::ControllerA, + /// CLI override for Controller A's hard cap. `None` → ISV-derived. + pub bucket_warmup_cap_override: Option, + /// Previous-step `cfc.tau_all_d` snapshot (HIDDEN_DIM floats). + /// Updated each Phase 1 step via DtoD copy after the Frobenius + /// kernel reads it. Lifetime spans Phase 1 only; the buffer stays + /// allocated in Phase 2 but is no longer touched. + prev_tau_d: CudaSlice, + /// Device-side scalar slot written by `tau_change_frobenius_fn` and + /// then DtoD-shadowed into `tau_change_staged` for host read. + tau_change_d: CudaSlice, + /// Mapped-pinned host shadow of `tau_change_d` (single f32). Read + /// each Phase 1 step after a stream sync to feed Controller A. + tau_change_staged: MappedF32Buffer, + /// Cached function handle for the per-step tau-change kernel + /// (`tau_change_frobenius_kernel` in `bucket_transition_kernels.cu`). + tau_change_frobenius_fn: CudaFunction, + /// Cached function handle for the IQR-widening kernel + /// (`slack_factor_apply_kernel` in `bucket_transition_kernels.cu`). + /// Launched once at Phase 1→2 transition, then never again. + slack_factor_apply_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 + /// pay a load-cubin cost. + _bucket_transition_module: Arc, + /// Compact-ragged heads_w_skip storage populated at Phase 1→2 + /// transition. Size HIDDEN_DIM (sum of `BUCKET_DIM_K` = 128). + /// Allocated up front (zero-init), populated by + /// `heads_compact_kernel` at transition, consumed by Task 10's + /// Phase 2 dispatch path. In Task 9's transient state it stays + /// zero (and unread) until the transition fires. + heads_w_skip_compact_d: CudaSlice, + /// Bucket-routing metadata produced by `execute_transition`. + /// `None` during Phase 1; `Some(metadata)` after the transition + /// fires. Tasks 10–12 will read its slices to drive Phase 2 + /// dispatch + Controllers B / C / D. + pub bucket_routing_metadata: + Option, } impl PerceptionTrainer { @@ -749,6 +836,21 @@ impl PerceptionTrainer { let heads_grn_bwd_fn = heads_module .load_function("multi_horizon_heads_grn_bwd_batched") .context("heads GRN bwd symbol")?; + // Bucket-transition cubin: tau_change_frobenius (per-step) + + // slack_factor_apply (transition-time, once). The full transition + // orchestration (5 kernels) lives in `bucket_routing::execute_transition` + // and loads the cubin independently — keeping its module local to that + // function call. Here we cache just the two function handles the + // trainer dispatches directly. + let bucket_transition_module = ctx + .load_cubin(BUCKET_TRANSITION_CUBIN.to_vec()) + .context("bucket_transition cubin")?; + let tau_change_frobenius_fn = bucket_transition_module + .load_function("tau_change_frobenius_kernel") + .context("load tau_change_frobenius_kernel")?; + let slack_factor_apply_fn = bucket_transition_module + .load_function("slack_factor_apply_kernel") + .context("load slack_factor_apply_kernel")?; // Mamba2 stacks live on the trunk from new_random; we wire optimizer // + training scratches against the trunk-owned blocks. @@ -1104,6 +1206,29 @@ impl PerceptionTrainer { let window_step_d = GpuTensor::zeros(&[cfg.n_batch, 1, FEATURE_DIM], &stream) .map_err(|e| anyhow::anyhow!("window_step_d alloc: {e}"))?; + // ── Per-horizon CfC routing scaffolding (Task 9) ── + // `prev_tau_d` shadows the previous step's `cfc.tau_all_d` so we can + // compute per-step ‖Δτ‖_F² on device. Initialised to zeros so the + // first step's `tau_change` equals ‖τ_1‖_F² — the first-observation + // bootstrap value per `pearl_first_observation_bootstrap`. + // `tau_change_d` is the regular CudaSlice the kernel writes into, + // then DtoD-copied to `tau_change_staged` (mapped-pinned) for the + // host scalar read. The compact heads buffer is allocated up front + // so Phase 1 doesn't hit an allocator at the boundary; Task 10 will + // wire it into the Phase 2 forward path. + let prev_tau_d = stream + .alloc_zeros::(crate::cfc::bucket_routing::HIDDEN_DIM) + .context("prev_tau_d alloc")?; + let tau_change_d = stream + .alloc_zeros::(1) + .context("tau_change_d alloc")?; + let tau_change_staged = unsafe { MappedF32Buffer::new(1) } + .map_err(|e| anyhow::anyhow!("tau_change_staged: {e}"))?; + let heads_w_skip_compact_d = stream + .alloc_zeros::(crate::cfc::bucket_routing::HIDDEN_DIM) + .context("heads_w_skip_compact_d alloc")?; + let controller_a = crate::cfc::bucket_routing::ControllerA::new(); + // ── CRT.train: output-smoothness regularizer state ── // λ is now ISV-driven by smoothness_lambda_controller — initialised // to LAMBDA_FLOOR per horizon and updated each step inside the @@ -1451,6 +1576,18 @@ impl PerceptionTrainer { step_ts_ns_d, step_prev_ts_ns_d, window_step_d, + // Per-horizon CfC routing (Task 9). + phase: TrainingPhase::Phase1Warmup, + controller_a, + bucket_warmup_cap_override: cfg.bucket_warmup_cap_override, + prev_tau_d, + tau_change_d, + tau_change_staged, + tau_change_frobenius_fn, + slack_factor_apply_fn, + _bucket_transition_module: bucket_transition_module, + heads_w_skip_compact_d, + bucket_routing_metadata: None, }) } @@ -1638,6 +1775,190 @@ impl PerceptionTrainer { } } + // ── 1.5. Phase 1→2 transition orchestration (Task 9 scope) ── + // + // Per spec §3.1: monitor CfC.tau drift, trigger Phase 2 when the + // EMA falls below the noise-floor-anchored threshold (or the + // ISV-derived hard cap fires). This block runs OUTSIDE the + // captured training graph (per + // `pearl_no_host_branches_in_captured_graph`) — the per-step + // host-side scalar read could not survive graph capture, and + // the transition itself changes the kernel pointer set the + // graph would record. When the transition fires we invalidate + // `self.train_graph` so the next step recaptures cleanly per + // `pearl_cudarc_disable_event_tracking_for_graph_capture`. + // + // In Phase 2 this entire block is a no-op (the `phase` check + // short-circuits), so the steady-state replay path keeps its + // captured-graph fast path. + if self.phase == TrainingPhase::Phase1Warmup { + // ── (a) Launch tau_change_frobenius_kernel ── + // Block = HIDDEN_DIM threads, block-tree reduction (no + // atomicAdd per `feedback_no_atomicadd`). Shared mem holds + // HIDDEN_DIM f32 squared-diff entries. + let cfg_frob = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (crate::cfc::bucket_routing::HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: (crate::cfc::bucket_routing::HIDDEN_DIM * 4) as u32, + }; + { + let mut launch = self.stream.launch_builder(&self.tau_change_frobenius_fn); + launch + .arg(&self.trunk.tau_all_d) + .arg(&self.prev_tau_d) + .arg(&mut self.tau_change_d); + unsafe { + launch + .launch(cfg_frob) + .context("tau_change_frobenius launch")?; + } + } + + // ── (b) DtoD: current τ → prev_τ for next step's comparison ── + // Captured-graph-safe pattern (DtoD async on the same stream). + unsafe { + let s = self.stream.cu_stream(); + let (src, _gs) = self.trunk.tau_all_d.device_ptr(&self.stream); + let (dst, _gd) = self.prev_tau_d.device_ptr_mut(&self.stream); + let nbytes = crate::cfc::bucket_routing::HIDDEN_DIM + * std::mem::size_of::(); + cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, s) + .context("prev_tau DtoD")?; + } + + // ── (c) DtoD: tau_change_d → tau_change_staged (mapped-pinned) ── + // Per `feedback_no_htod_htoh_only_mapped_pinned`: the ONLY + // permitted CPU↔GPU path is mapped-pinned via DtoD shadow. + unsafe { + let (src_ptr, _g) = self.tau_change_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + self.tau_change_staged.dev_ptr, + src_ptr, + std::mem::size_of::(), + self.stream.cu_stream(), + ) + .context("tau_change → mapped-pinned shadow")?; + } + + // ── (d) Sync + host read ── + // Sync is unavoidable here: Controller A is a host-side + // state machine and needs the scalar before it can decide + // whether to trigger. Cost is amortized — one extra sync + // per step until the transition fires, then this block + // becomes a no-op for the rest of training. + self.stream.synchronize().context("tau_change sync")?; + let tau_change = self.tau_change_staged.read_all()[0]; + + // ── (e) Controller A update + transition check ── + let should_transition = self + .controller_a + .update(tau_change, self.bucket_warmup_cap_override); + + if should_transition { + // ── Phase 1→2 transition ── + // 1. Invalidate captured graph so it recaptures with the + // post-transition kernel pointer set + reordered τ + // layout (per + // `pearl_cudarc_disable_event_tracking_for_graph_capture`: + // event tracking is already disabled trainer-wide, so + // the recapture on next step is safe). + self.train_graph = None; + + // 2. Stage `cfc_tau_d` into a scratch source so the + // reorder kernel can read the original τ ordering + // while writing the bucket-grouped layout back into + // `trunk.tau_all_d`. Without the scratch the kernel + // would alias its input and output (the orchestrator + // needs `cfc_tau_d` immutably + `tau_all_d` mutably); + // a fresh device buffer also avoids the Rust borrow + // checker complaint about borrowing `trunk` twice. + // The scratch is local to this transition path — + // one allocation per training run, off the hot path. + let tau_scratch_d = { + let mut s = self.stream + .alloc_zeros::(crate::cfc::bucket_routing::HIDDEN_DIM) + .context("tau_scratch alloc (transition)")?; + unsafe { + let st = self.stream.cu_stream(); + let (src, _gs) = self.trunk.tau_all_d.device_ptr(&self.stream); + let (dst, _gd) = s.device_ptr_mut(&self.stream); + let nbytes = crate::cfc::bucket_routing::HIDDEN_DIM + * std::mem::size_of::(); + cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, st) + .context("tau → tau_scratch DtoD")?; + } + s + }; + + // 3. Execute the 5-kernel transition (sort, assign, IQR, + // reorder, compact). `execute_transition` returns + // bucket-routing metadata; we store it in trainer + // state for Task 10+11+12 to consume. The reordered + // τ is written into `tau_all_d` in place; the compact + // heads_w_skip lives in `heads_w_skip_compact_d`, + // leaving the original `heads_w_skip_d` intact for + // the Task-9-transient Phase 1 dispatch path. + let mut metadata = crate::cfc::bucket_routing::execute_transition( + &self.stream, + &tau_scratch_d, + &self.trunk.heads_w_skip_d, + &mut self.trunk.tau_all_d, + &mut self.heads_w_skip_compact_d, + )?; + drop(tau_scratch_d); + + // 4. Apply slack_factor = sqrt(Q3/Q1) widening to the + // bucket τ IQR bounds (spec §3.2). ISV-derived per + // `feedback_isv_for_adaptive_bounds`: no hardcoded + // 1.5× factor — the slack is proportional to each + // bucket's natural log-spread. Launches in place on + // the metadata's Q1/Q3 buffers. + let cfg_slack = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N_HORIZONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + { + let mut launch = + self.stream.launch_builder(&self.slack_factor_apply_fn); + launch + .arg(&mut metadata.bucket_tau_iqr_lo_d) + .arg(&mut metadata.bucket_tau_iqr_hi_d); + unsafe { + launch + .launch(cfg_slack) + .context("slack_factor_apply launch")?; + } + } + + // 5. Persist routing metadata in trainer state. Tasks + // 10–12 read its slices to drive Phase 2 dispatch + + // Controllers B / C / D. + self.bucket_routing_metadata = Some(metadata); + + // 6. Latch trainer into Phase 2. + // + // NOTE: Task 9 leaves the per-step dispatch path + // unchanged (Phase 1 single-CfC, 640-float heads_w_skip). + // The reordered `tau_all_d` still produces a valid + // forward pass because CfC's per-channel decay math + // (`h_new = h_old * exp(-dt/tau) + (1-decay) * tanh(pre)`) + // is order-invariant. Task 10 will land the per-branch + // fused dispatch that actually consumes the routing + // metadata; until then we run a transient state where + // `phase == Phase2Routed` but the dispatch is Phase-1 + // shaped. This is documented behaviour, not a stub. + self.phase = TrainingPhase::Phase2Routed; + + tracing::info!( + tau_change = tau_change, + noise_floor = self.controller_a.noise_floor, + step_count = self.controller_a.step_count, + "Phase 1→2 transition fired (CfC.tau stabilized)" + ); + } + } + // ── 2. Three-state machine: warmup (first), capture (second), // replay (third+). The captured graph records all // in-graph kernel decisions at capture time per diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs index f4e76bda2..ae62d776a 100644 --- a/crates/ml-alpha/tests/bucket_transition_kernels.rs +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -182,3 +182,68 @@ fn heads_compact_kernel_reorders_w_skip() -> Result<()> { } Ok(()) } + +#[test] +#[ignore = "requires CUDA"] +fn tau_change_frobenius_kernel_sums_squared_diff() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "tau_change_frobenius_kernel")?; + + let tau_t: Vec = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.1).collect(); + let tau_prev: Vec = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.1 + 0.5).collect(); + let tau_t_d = upload(&stream, &tau_t)?; + let tau_prev_d = upload(&stream, &tau_prev)?; + let mut out_d = stream.alloc_zeros::(1)?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: HIDDEN_DIM as u32 * 4, + }; + let mut launch = stream.launch_builder(&func); + launch.arg(&tau_t_d).arg(&tau_prev_d).arg(&mut out_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let out = download(&stream, &out_d)?; + // Each diff = -0.5; squared = 0.25; sum over 128 channels = 32.0 + assert!((out[0] - 32.0).abs() < 1e-3, "out={} expected ~32.0", out[0]); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "slack_factor_apply_kernel")?; + + // Q1 = [1, 4, 9, 16, 25], Q3 = [4, 16, 36, 64, 100] + // → slack = sqrt(Q3/Q1) = [2, 2, 2, 2, 2] + // → out: lo = [0.5, 2, 4.5, 8, 12.5], hi = [8, 32, 72, 128, 200] + let q1: Vec = vec![1.0, 4.0, 9.0, 16.0, 25.0]; + let q3: Vec = vec![4.0, 16.0, 36.0, 64.0, 100.0]; + let mut lo_d = upload(&stream, &q1)?; + let mut hi_d = upload(&stream, &q3)?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N_HORIZONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch.arg(&mut lo_d).arg(&mut hi_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let lo = download(&stream, &lo_d)?; + let hi = download(&stream, &hi_d)?; + let expected_lo = [0.5, 2.0, 4.5, 8.0, 12.5]; + let expected_hi = [8.0, 32.0, 72.0, 128.0, 200.0]; + for k in 0..N_HORIZONS { + assert!((lo[k] - expected_lo[k]).abs() < 1e-3, "lo[{}]={} expected {}", k, lo[k], expected_lo[k]); + assert!((hi[k] - expected_hi[k]).abs() < 1e-3, "hi[{}]={} expected {}", k, hi[k], expected_hi[k]); + } + Ok(()) +} diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index f72cae6dc..cef6a1669 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -74,6 +74,7 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -149,6 +150,7 @@ fn stacked_trainer_loss_shrinks_with_stride_4() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -204,6 +206,7 @@ fn stacked_trainer_loss_shrinks_at_batch_32() { n_batch: 32, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -303,6 +306,7 @@ fn evaluate_alone_succeeds() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64; @@ -332,6 +336,7 @@ fn evaluate_works_after_captured_training_step() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -368,6 +373,7 @@ fn evaluate_works_after_capture_no_replay() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64; @@ -399,6 +405,7 @@ fn horizon_ema_and_lambda_track_after_training() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); @@ -456,6 +463,7 @@ fn evaluate_works_after_warmup_only() { n_batch: 1, smoothness_base_lambda: 0.0, kernel_step_trace_path: None, + bucket_warmup_cap_override: None, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); let ts = 1_000_000u64;