feat(per-horizon-cfc): wire Controller A + Phase 1→2 transition in trainer
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<u64>,
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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;
|
||||
}
|
||||
|
||||
@@ -172,6 +172,17 @@ struct Cli {
|
||||
/// Example: `--kernel-step-trace /feature-cache/.../step_trace.jsonl`
|
||||
#[arg(long)]
|
||||
kernel_step_trace: Option<std::path::PathBuf>,
|
||||
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
#[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")?;
|
||||
|
||||
|
||||
@@ -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<std::path::PathBuf>,
|
||||
|
||||
/// 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<u64>,
|
||||
}
|
||||
|
||||
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<u64>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<CudaModule>,
|
||||
/// 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<f32>,
|
||||
/// 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<crate::cfc::bucket_routing::BucketRoutingMetadata>,
|
||||
}
|
||||
|
||||
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::<f32>(crate::cfc::bucket_routing::HIDDEN_DIM)
|
||||
.context("prev_tau_d alloc")?;
|
||||
let tau_change_d = stream
|
||||
.alloc_zeros::<f32>(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::<f32>(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::<f32>();
|
||||
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::<f32>(),
|
||||
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::<f32>(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::<f32>();
|
||||
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
|
||||
|
||||
@@ -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<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.1).collect();
|
||||
let tau_prev: Vec<f32> = (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::<f32>(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<f32> = vec![1.0, 4.0, 9.0, 16.0, 25.0];
|
||||
let q3: Vec<f32> = 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(())
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user