dc76d205eb
Per spec §3.2 + §3.3 and pearl_adam_normalizes_loss_weights.
Controller B (post-Adam τ projection):
- tau_clamp_kernel added to bucket_transition_kernels.cu
- Each channel's τ clamped to its bucket's IQR_widened range after every
Adam step on tau_all_d. Bypasses Adam's m/sqrt(v) normalization
(loss-weight modulation would have been a no-op per the pearl).
Controller C (per-bucket LR multiplier on heads_w_skip):
- heads_lr_multiplier_scale_kernel added.
- Per-step: snapshot heads_w_skip_d (DtoD) before opt_heads_w_skip.step,
run Adam, then rescale the per-horizon delta by
lr_mult_k = 1.0 / (1.0 + smoothness_ratio_k).
- smoothness_ratio_k = max(0, jitter_ema_k - target_jitter_k) /
max(target_jitter_k, target_jitter_k / 16)
with the ε_floor max-with-floor pattern per
pearl_blend_formulas_must_have_permanent_floor.
- target_jitter_k sentinel-bootstrapped from first non-zero raw
smoothness loss per pearl_first_observation_bootstrap.
- jitter EMA shadowed via mapped-pinned smoothness_jitter_ema_host_d
(DtoD captured in graph; host reads after end-of-step sync); host
writes lr_mult_per_horizon_staging (mapped-pinned); next step's
captured-graph rescale kernel reads via dev_ptr.
Scope reduction (Task 11): Controller C scales ONLY heads_w_skip_d
per-horizon (N_HORIZONS × HIDDEN_DIM = 640 contiguous floats, per-horizon
stride HIDDEN_DIM). Other heads weights (w1, w2, w_gate, w_main) stay
unscaled because they're 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.
Both controllers active Phase 2 only; no-op in Phase 1.
GPU oracle tests: 9 total (7 existing + tau_clamp + lr_multiplier_scale).
All 33 ml-alpha lib tests pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
4849 lines
244 KiB
Rust
4849 lines
244 KiB
Rust
//! PerceptionTrainer — stacked Mamba2 -> CfC -> heads with per-position supervision.
|
||
//!
|
||
//! Per spec 2026-05-16 amendment + 2026-05-17 BPTT-unroll amendment.
|
||
//! Topology per training step:
|
||
//! snap_features × seq_len → Mamba2.forward_train_seq → h_enriched_seq [K, hidden_dim]
|
||
//! For each k=0..K-1 with a valid label:
|
||
//! → cfc_step (h_old=0) → h_new_k
|
||
//! → heads → probs_k
|
||
//! → BCE(probs_k, labels[k]) → loss_k
|
||
//!
|
||
//! Backward chain:
|
||
//! For each supervised k (kernels accumulate via += into shared grad
|
||
//! buffers; CfC + heads weights are shared across positions):
|
||
//! grad_probs_k → heads_backward → grad_h_new_k + accum grad_W_heads/b_heads
|
||
//! → cfc_step_backward → accum grad_W_in/W_rec/b + grad_x_k
|
||
//! grad_x_k stored into grad_h_enriched_seq[k]
|
||
//! Once:
|
||
//! grad_h_enriched_seq → Mamba2.backward_from_h_enriched_seq → full Mamba2 grad set
|
||
//!
|
||
//! Per-step supervision densifies the gradient signal ~K× over the
|
||
//! final-step-only design (which trained flat at chance on real ES
|
||
//! data despite working on synthetic overfit). Mamba2's analytical
|
||
//! scan backward unrolls those per-step gradients through the full
|
||
//! 32-step SSM state evolution.
|
||
//!
|
||
//! Optimizers (unchanged):
|
||
//! - 5 CfC AdamWs (W_in, W_rec, b, heads_w, heads_b)
|
||
//! - 1 Mamba2AdamW for all 9 Mamba2 parameter tensors
|
||
//!
|
||
//! Architectural note: CfC runs with h_old=0 each step (no inter-step
|
||
//! recurrence in v1). With h_old=0, the cfc_step degenerates to a
|
||
//! per-cell-tau scaled tanh-FC layer feeding the heads. Inter-step CfC
|
||
//! state (where h_old carries between calls) is a v2 extension.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use anyhow::{Context, Result};
|
||
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
|
||
use cudarc::driver::{
|
||
CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
|
||
LaunchConfig, PushKernelArg,
|
||
};
|
||
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
|
||
use ml_core::device::MlDevice;
|
||
use rand::{Rng, SeedableRng};
|
||
use rand_chacha::ChaCha8Rng;
|
||
|
||
use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM};
|
||
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
|
||
use crate::mamba2_block::{
|
||
Mamba2AdamW, Mamba2AdamWConfig, Mamba2BackwardGradsBuffers, Mamba2BackwardScratch,
|
||
Mamba2BlockForwardScratch, Mamba2BlockStepScratch,
|
||
};
|
||
use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer, MappedI64Buffer};
|
||
use crate::trainer::optim::AdamW;
|
||
|
||
const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
|
||
const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin"));
|
||
const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
|
||
const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin"));
|
||
const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin"));
|
||
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
|
||
const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin"));
|
||
const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin"));
|
||
const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin"));
|
||
const SMOOTHNESS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/output_smoothness.cubin"));
|
||
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 {
|
||
pub seq_len: usize,
|
||
pub mamba2_state_dim: usize,
|
||
pub lr_cfc: f32,
|
||
pub lr_mamba2: f32,
|
||
pub seed: u64,
|
||
/// Per-horizon BCE loss weights — multiplied into the loss + grad
|
||
/// contribution at each (position, horizon) pair. Down-weights
|
||
/// horizons whose forward labels overlap across K positions (h ≫ K
|
||
/// → near-identical labels would otherwise inflate gradient
|
||
/// pressure). All-ones gives the original uniform-mean BCE.
|
||
pub horizon_weights: [f32; N_HORIZONS],
|
||
/// Number of sequences processed per optimizer step. Sets the
|
||
/// effective batch size — kernels switch to batched variants
|
||
/// (cfc_step_batched, multi_horizon_heads_batched, etc.).
|
||
/// `n_batch=1` matches the unbatched behavior exactly.
|
||
pub n_batch: usize,
|
||
/// CRT.train intervention A: output-smoothness regularizer base λ.
|
||
/// Amplitude knob for the ISV-driven `smoothness_lambda_controller`,
|
||
/// which derives per-horizon λ[h] each step from the post-step
|
||
/// jitter EMA (anchored on h30). Default 0.0 = controller produces
|
||
/// LAMBDA_FLOOR uniformly (kernel still launches but effective
|
||
/// regularization is negligible). Set to e.g. 0.01 to enable.
|
||
/// See `docs/superpowers/specs/2026-05-20-crt-train-output-smoothness-design.md`.
|
||
pub smoothness_base_lambda: f32,
|
||
|
||
/// Optional path for the kernel-step-trace JSONL drain. None →
|
||
/// ring not allocated, drain task not spawned. Some(path) → ring
|
||
/// allocated + drain writes JSONL records to `path` (truncated).
|
||
/// 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 {
|
||
fn default() -> Self {
|
||
Self {
|
||
seq_len: 32,
|
||
mamba2_state_dim: 16,
|
||
lr_cfc: 3e-3,
|
||
lr_mamba2: 1e-3,
|
||
seed: 0x4242,
|
||
horizon_weights: [1.0; N_HORIZONS],
|
||
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
|
||
/// percent of the loss (h6000 weight = 0.0053 at K=32) and effectively
|
||
/// excluded them from training.
|
||
///
|
||
/// The 3-fold ISV CV on commit `0171c8c0e` showed the prior schedule
|
||
/// suppressing h6000 — the deployment-relevant horizon — to the point
|
||
/// where the ISV controller's lambda was operating on rounding error.
|
||
/// h6000 ended at 0.685-0.728 across folds; mean_auc CV at 0.717±.
|
||
///
|
||
/// New schedule: UNIFORM `[1.0; N]`. All horizons contribute equally
|
||
/// to the BCE loss; let the ISV controller (boosting per-horizon
|
||
/// gradient via lambda into the trunk) do the per-horizon re-balancing
|
||
/// in a single canonical place. Callers who want a custom schedule
|
||
/// can still pass `--horizon-weights w0,w1,w2,w3,w4`.
|
||
///
|
||
/// `seq_len` and `horizons` are kept on the signature for backward
|
||
/// compatibility with the CLI but are no longer consulted; the
|
||
/// returned schedule is shape-determined, not data-determined.
|
||
pub fn auto_horizon_weights(_seq_len: usize, _horizons: &[usize; N_HORIZONS]) -> [f32; N_HORIZONS] {
|
||
[1.0f32; N_HORIZONS]
|
||
}
|
||
|
||
pub struct PerceptionTrainer {
|
||
cfg: PerceptionTrainerConfig,
|
||
stream: Arc<CudaStream>,
|
||
|
||
// Modules + cached function handles
|
||
_snap_module: Arc<CudaModule>,
|
||
_step_module: Arc<CudaModule>,
|
||
_heads_module: Arc<CudaModule>,
|
||
_bce_module: Arc<CudaModule>,
|
||
// Forward kernels are owned by `self.trunk` (snap_batched_fn,
|
||
// step_batched_fn, heads_grn_fwd_fn, transpose_3d_fn). The trainer
|
||
// keeps only the training-only handles below.
|
||
bce_fn: CudaFunction,
|
||
step_bwd_batched_fn: CudaFunction,
|
||
heads_grn_bwd_fn: CudaFunction,
|
||
|
||
// X3: Mamba2 stack 1 weights moved to `self.trunk.mamba2_stack_1`.
|
||
// Access via `self.trunk.mamba2_l1()` / `mamba2_l1_mut()`. Gradient
|
||
// buffers + AdamW state remain at trainer level.
|
||
pub mamba2_adamw: Mamba2AdamW,
|
||
// X5: Mamba2 stack 2 weights moved to `self.trunk.mamba2_stack_2`.
|
||
// Access via `self.trunk.mamba2_l2()` / `mamba2_l2_mut()`.
|
||
pub mamba2_l2_adamw: Mamba2AdamW,
|
||
mamba2_l2_fwd_scratch: Mamba2BlockForwardScratch,
|
||
mamba2_l2_bwd_scratch: Mamba2BackwardScratch,
|
||
mamba2_l2_grads_buffers: Mamba2BackwardGradsBuffers,
|
||
// ── LN_a (between Mamba2 stacks; Phase 2B) ──
|
||
// The existing LN sits between m2 and CfC (now called LN_b). LN_a
|
||
// is a separate instance with its own learnable gain/bias.
|
||
// X4: LN_a weights moved to `self.trunk.ln_a_gain_d` / `ln_a_bias_d`.
|
||
ln_a_stats_d: CudaSlice<f32>,
|
||
/// `[B, K, HIDDEN_DIM]` — LN_a forward output, fed as input to m2.
|
||
ln_a_out_d: GpuTensor,
|
||
grad_ln_a_gain_per_row_d: CudaSlice<f32>,
|
||
grad_ln_a_bias_per_row_d: CudaSlice<f32>,
|
||
grad_ln_a_gain_d: CudaSlice<f32>,
|
||
grad_ln_a_bias_d: CudaSlice<f32>,
|
||
/// `[B, K, HIDDEN_DIM]` — LN_a backward grad-x: gradient w.r.t.
|
||
/// m1's `h_enriched_seq` output (consumed by m1.backward).
|
||
grad_ln_a_in_d: GpuTensor,
|
||
pub opt_ln_a_gain: AdamW,
|
||
pub opt_ln_a_bias: AdamW,
|
||
/// Pre-allocated forward intermediates for Mamba2 (input projection
|
||
/// output x, A/B projections, h_s2 residual, h_enriched_seq scan
|
||
/// output). Construct once at trainer init; reused every step.
|
||
mamba2_fwd_scratch: Mamba2BlockForwardScratch,
|
||
/// Pre-allocated scratch for the Mamba2 seq backward — the four
|
||
/// largest per-channel buffers (~6 MB each at B=8, K=96).
|
||
mamba2_bwd_scratch: Mamba2BackwardScratch,
|
||
/// Pre-allocated outputs for Mamba2 seq backward (cuBLAS-projection
|
||
/// dw/db tensors + reduction-result tensors + d_x intermediates).
|
||
/// Read by Mamba2AdamW::step_from_buffers — no temporary
|
||
/// Mamba2BackwardGrads wrapper allocated per step.
|
||
mamba2_grads_buffers: Mamba2BackwardGradsBuffers,
|
||
/// Pre-allocated input window for snap_features → Mamba2 fwd.
|
||
/// [B, K, FEATURE_DIM] — overwritten each step by the batched
|
||
/// snap_feature kernel.
|
||
window_tensor_d: GpuTensor,
|
||
/// Pre-allocated transpose of Mamba2's h_enriched_seq into [K, B, H]
|
||
/// layout for contiguous per-K slot access in the trainer loop.
|
||
h_enriched_seq_t_d: GpuTensor,
|
||
/// Pre-allocated per-K gradient accumulator in [K, B, H] layout.
|
||
/// Written by the reverse-order backward K loop; transposed back
|
||
/// to [B, K, H] for Mamba2 backward consumption.
|
||
grad_h_enriched_seq_t_d: GpuTensor,
|
||
/// Pre-allocated [B, K, H] grad input to Mamba2 backward.
|
||
grad_h_enriched_seq_d: GpuTensor,
|
||
|
||
// CfC + heads weights + their AdamWs (6 groups — tau is trained now).
|
||
// X8: CfC weights (w_in, w_rec, b, tau) moved to
|
||
// `self.trunk.w_in_d` / `w_rec_d` / `b_d` / `tau_all_d` (now v2-shaped
|
||
// via CfcConfig.cfc_n_in = HIDDEN_DIM).
|
||
// ── TFT GRN heads (Phase 1.7) ──
|
||
// Per-horizon Gated Residual Network: 2-layer GELU MLP body
|
||
// (eta_2 → eta_1) + GLU gate + skip-projection. Gives per-horizon
|
||
// "linear vs deeper-transform" specialisation, matching the
|
||
// regime-conditional alpha pattern (pearl_snapshot_alpha_is_regime_conditional).
|
||
// X9: GRN head weights moved to `self.trunk.heads_w1_d` / `heads_b1_d`
|
||
// / ... / `heads_b_skip_d`. Gradient buffers + AdamW state stay at
|
||
// trainer level.
|
||
pub opt_w_in: AdamW,
|
||
pub opt_w_rec: AdamW,
|
||
pub opt_b: AdamW,
|
||
pub opt_tau: AdamW,
|
||
pub opt_heads_w1: AdamW,
|
||
pub opt_heads_b1: AdamW,
|
||
pub opt_heads_w2: AdamW,
|
||
pub opt_heads_b2: AdamW,
|
||
pub opt_heads_w_gate: AdamW,
|
||
pub opt_heads_b_gate: AdamW,
|
||
pub opt_heads_w_main: AdamW,
|
||
pub opt_heads_b_main: AdamW,
|
||
pub opt_heads_w_skip: AdamW,
|
||
pub opt_heads_b_skip: AdamW,
|
||
|
||
// Pre-allocated BATCHED snap_feature staging — sized for max B*K
|
||
// snapshots. All staging buffers are MAPPED-PINNED (the only
|
||
// permitted CPU↔GPU path per `feedback_no_htod_htoh_only_mapped_pinned.md`).
|
||
// Per training step: host writes through staging.host_slice_mut(),
|
||
// then one DtoD copy per array → device buffer, then one batched
|
||
// snap_feature_assemble launch instead of B*K single-snapshot launches.
|
||
bk_capacity: usize,
|
||
stg_bid_px_all: MappedF32Buffer, // [B*K, 10]
|
||
stg_bid_sz_all: MappedF32Buffer, // [B*K, 10]
|
||
stg_ask_px_all: MappedF32Buffer, // [B*K, 10]
|
||
stg_ask_sz_all: MappedF32Buffer, // [B*K, 10]
|
||
stg_regime_all: MappedF32Buffer, // [B*K, 6]
|
||
stg_prev_mid: MappedF32Buffer, // [B*K]
|
||
stg_trade_signed_vol: MappedF32Buffer, // [B*K]
|
||
stg_trade_count: MappedI32Buffer, // [B*K]
|
||
stg_ts_ns: MappedI64Buffer, // [B*K]
|
||
stg_prev_ts_ns: MappedI64Buffer, // [B*K]
|
||
bid_px_all_d: CudaSlice<f32>,
|
||
bid_sz_all_d: CudaSlice<f32>,
|
||
ask_px_all_d: CudaSlice<f32>,
|
||
ask_sz_all_d: CudaSlice<f32>,
|
||
/// `prev_bid_sz_all_d` / `prev_ask_sz_all_d` — zero-init device
|
||
/// buffers passed to the snap kernel. The loader doesn't compute
|
||
/// per-snapshot prev-OFI yet, so these stay zero (same as the
|
||
/// single-snapshot path).
|
||
prev_bid_sz_all_d: CudaSlice<f32>,
|
||
prev_ask_sz_all_d: CudaSlice<f32>,
|
||
regime_all_d: CudaSlice<f32>,
|
||
prev_mid_all_d: CudaSlice<f32>,
|
||
trade_signed_vol_all_d: CudaSlice<f32>,
|
||
trade_count_all_d: CudaSlice<i32>,
|
||
ts_ns_all_d: CudaSlice<i64>,
|
||
prev_ts_ns_all_d: CudaSlice<i64>,
|
||
|
||
// Per-K device-resident scratch — pre-allocated once at construction
|
||
// so the K loop in step() does ZERO device allocations.
|
||
/// `h_new_per_k_d[k * n_hid .. (k+1) * n_hid]` holds CfC h_new at
|
||
/// position k. Used as h_old at position k+1 (recurrence carry) and
|
||
/// re-read in backward for cfc_step_bwd's h_old argument.
|
||
h_new_per_k_d: CudaSlice<f32>,
|
||
/// `probs_per_k_d[k * N_HORIZONS .. (k+1) * N_HORIZONS]` — heads
|
||
/// output at position k. Single fused BCE call consumes all K rows.
|
||
probs_per_k_d: CudaSlice<f32>,
|
||
/// Labels for the sequence, uploaded once per step(); NaN slots are
|
||
/// natively masked by the BCE kernel.
|
||
labels_per_k_d: CudaSlice<f32>,
|
||
/// BCE-emitted per-position grad_probs; consumed by heads_bwd.
|
||
grad_probs_per_k_d: CudaSlice<f32>,
|
||
/// Scalar BCE loss output (mean over valid label entries).
|
||
loss_d: CudaSlice<f32>,
|
||
/// Mapped-pinned host shadow of `loss_d`. After the post-step sync,
|
||
/// host_ptr holds the freshly-computed loss — no extra dtoh sync.
|
||
loss_host_d: MappedF32Buffer,
|
||
/// Per-horizon UNWEIGHTED mean BCE — the ISV signal an EMA layer
|
||
/// tracks to compute dynamic per-horizon weights. Refreshed by
|
||
/// the BCE kernel on every training step. Shape: [N_HORIZONS].
|
||
loss_per_horizon_d: CudaSlice<f32>,
|
||
/// Per-horizon smoothness multiplier λ[h]. Initialised to a uniform
|
||
/// `LAMBDA_FLOOR` at construction; updated each step by
|
||
/// `smoothness_lambda_controller` from the post-step jitter EMA.
|
||
/// Read by `output_smoothness_loss_and_grad` on the NEXT step
|
||
/// (one-step closed-loop delay inside the captured graph).
|
||
smoothness_lambda_d: CudaSlice<f32>,
|
||
/// Scalar total smoothness loss output (λ-weighted Σ_h). Output of
|
||
/// `output_smoothness_loss_and_grad`. Mapped-pinned host shadow
|
||
/// `smoothness_loss_host_d` provides the post-step value with no
|
||
/// extra dtoh sync (telemetry only).
|
||
smoothness_loss_d: CudaSlice<f32>,
|
||
smoothness_loss_host_d: MappedF32Buffer,
|
||
/// Per-horizon raw mean-squared-diff (pre-λ). Telemetry feed for
|
||
/// the epoch summary writer. Shape [N_HORIZONS].
|
||
smoothness_loss_per_horizon_d: CudaSlice<f32>,
|
||
smoothness_loss_per_horizon_host_d: MappedF32Buffer,
|
||
/// Cached handle for `output_smoothness_loss_and_grad`.
|
||
smoothness_fn: CudaFunction,
|
||
_smoothness_module: Arc<CudaModule>,
|
||
/// Per-horizon EMA of `raw_per_h` from output_smoothness. Updated
|
||
/// by `smoothness_lambda_controller` each step. Drives the
|
||
/// controller's target derivation.
|
||
smoothness_jitter_ema_d: CudaSlice<f32>,
|
||
/// First-observation sentinel (per pearl_first_observation_bootstrap).
|
||
/// 0 → next step bootstraps EMA from raw_per_h; 1 → Wiener-α update.
|
||
smoothness_jitter_first_obs_d: CudaSlice<i32>,
|
||
/// Cached handle for `smoothness_lambda_controller`.
|
||
smoothness_controller_fn: CudaFunction,
|
||
_smoothness_controller_module: Arc<CudaModule>,
|
||
// ── GPU log ring (feature: kernel-step-trace) ──
|
||
// When the feature is compiled AND `cfg.kernel_step_trace_path` is
|
||
// Some(path), the trainer allocates a mapped-pinned ring + a device
|
||
// step counter, passes their pointers into logging-capable kernels
|
||
// (currently `smoothness_lambda_controller`), and spawns a tokio drain
|
||
// task that polls the ring at 500 ms cadence and writes JSONL records
|
||
// to `path`. When the path is None, every Option below stays None —
|
||
// no allocations, no tick kernel, no drain task, zero overhead even
|
||
// with the feature compiled in.
|
||
// See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_ring: Option<std::sync::Arc<crate::gpu_log::LogRing>>,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_step_counter_d: Option<CudaSlice<i32>>,
|
||
/// Mapped-pinned host shadow of `log_step_counter_d`. Updated via
|
||
/// memcpy_dtod_async inside the captured region; read by the drain
|
||
/// task each tick.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_step_counter_host_shadow: Option<std::sync::Arc<crate::pinned_mem::MappedRecordBuffer<i32>>>,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
gpu_log_tick_fn: Option<CudaFunction>,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
_gpu_log_module: Option<Arc<CudaModule>>,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_drain_stop: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_drain_handle: Option<std::thread::JoinHandle<()>>,
|
||
// X6: LN_b weights (formerly ln_gain_d / ln_bias_d) moved to
|
||
// `self.trunk.ln_b_gain_d` / `ln_b_bias_d`. LN_b is the LayerNorm
|
||
// AFTER Mamba2 stack 2 — stabilises the trunk-to-CfC distribution.
|
||
/// Per-row LN stats `[B*K, 2]` (mean, inv_std) saved by fwd, used
|
||
/// by bwd.
|
||
ln_stats_d: CudaSlice<f32>,
|
||
/// LN forward output `[B*K, HIDDEN]` — fed to the transpose that
|
||
/// produces `h_enriched_seq_t_d`. Replaces direct reads of
|
||
/// Mamba2's `h_enriched_seq` in the K-loop path.
|
||
ln_out_d: CudaSlice<f32>,
|
||
/// Per-row gain-grad scratch `[B*K, HIDDEN]`, reduced into
|
||
/// `grad_ln_gain_d`.
|
||
grad_ln_gain_per_row_d: CudaSlice<f32>,
|
||
/// Per-row bias-grad scratch `[B*K, HIDDEN]`, reduced into
|
||
/// `grad_ln_bias_d`.
|
||
grad_ln_bias_per_row_d: CudaSlice<f32>,
|
||
/// Final LN gain gradient `[HIDDEN]` — Adam consumes this.
|
||
grad_ln_gain_d: CudaSlice<f32>,
|
||
/// Final LN bias gradient `[HIDDEN]` — Adam consumes this.
|
||
grad_ln_bias_d: CudaSlice<f32>,
|
||
/// LN backward grad-x: the gradient w.r.t. Mamba2's `h_enriched_seq`
|
||
/// output. Replaces direct use of `grad_h_enriched_seq_d` when
|
||
/// feeding Mamba2 backward.
|
||
grad_ln_in_d: GpuTensor,
|
||
/// Cached LN kernel handles + module lifetime owner.
|
||
ln_bwd_fn: CudaFunction,
|
||
ln_reduce_fn: CudaFunction,
|
||
_ln_module: Arc<CudaModule>,
|
||
/// AdamW state for the two LN params (gain + bias).
|
||
opt_ln_gain: AdamW,
|
||
opt_ln_bias: AdamW,
|
||
|
||
// ── TFT VSN (Phase 2D) ──
|
||
// Per-position softmax-normalised feature gating at trunk entry:
|
||
// gates = softmax(W_vsn @ x + b_vsn); y = x * gates
|
||
// Lets the model learn which snap_features matter per regime.
|
||
/// `[FEATURE_DIM, FEATURE_DIM]` linear layer for gate logits.
|
||
// X2 migration: VSN weights now live on `self.trunk` (CfcTrunk's
|
||
// pub vsn_w_d / vsn_b_d). PerceptionTrainer wraps a trunk to host
|
||
// the v2 inference graph; the trunk's V1 fields stay unused for now
|
||
// (X8 consolidates against PerceptionTrainer's CfC weights). Gradient
|
||
// buffers + AdamW state for VSN remain at trainer level.
|
||
pub trunk: crate::cfc::trunk::CfcTrunk,
|
||
/// VSN forward output `[B, K, FEATURE_DIM]` — replaces window_tensor_d
|
||
/// as Mamba2's input. window_tensor_d still holds the RAW features
|
||
/// (consumed by VSN bwd as `x`).
|
||
vsn_out_d: GpuTensor,
|
||
/// Saved softmax-normalised gates `[B*K, FEATURE_DIM]` (per-K row).
|
||
vsn_gates_d: CudaSlice<f32>,
|
||
/// VSN bwd writes grad_x here (gradient w.r.t. snap_features). Not
|
||
/// propagated further (snap_features are a non-trainable transform
|
||
/// of raw MBP-10 data); allocated only so the kernel has a buffer.
|
||
vsn_grad_x_d: CudaSlice<f32>,
|
||
grad_vsn_w_d: CudaSlice<f32>,
|
||
grad_vsn_b_d: CudaSlice<f32>,
|
||
pub opt_vsn_w: AdamW,
|
||
pub opt_vsn_b: AdamW,
|
||
vsn_bwd_fn: CudaFunction,
|
||
_vsn_module: Arc<CudaModule>,
|
||
|
||
// ── Attention pool (Phase 3) ──
|
||
// Single-head attention pool over LN_b output. Replaces CfC's
|
||
// zero-initialised h_old at k=0 with a learned content-addressable
|
||
// summary over all K positions. Single learned param: Q [HIDDEN_DIM].
|
||
// X7: attention-pool query weight moved to `self.trunk.attn_q_d`.
|
||
/// Per-sample pooled context `[B, HIDDEN_DIM]` — fed as h_old at k=0.
|
||
attn_context_d: CudaSlice<f32>,
|
||
/// Saved post-softmax attention weights `[B, K]` for bwd.
|
||
attn_weights_d: CudaSlice<f32>,
|
||
grad_attn_q_d: CudaSlice<f32>,
|
||
pub opt_attn_q: AdamW,
|
||
attn_bwd_fn: CudaFunction,
|
||
_attn_module: Arc<CudaModule>,
|
||
|
||
// ── Kendall σ-weighted BCE (axis A, ISV-driven) ──
|
||
// Per-horizon `log σ_h` is computed in closed form by the
|
||
// horizon_ema_and_lambda kernel from loss_ema (Kendall equilibrium
|
||
// σ_h = sqrt(mean_bce_h)). No Adam state — replaces the prior
|
||
// Adam-learned variant which fought m/√v normalization
|
||
// (pearl_adam_normalizes_loss_weights). Single source of truth
|
||
// shared with λ_h: both anchor on loss_ema.
|
||
pub log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS] — ISV output
|
||
pub grad_log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS] — BCE kernel writes here; ignored downstream
|
||
|
||
// ── K-loop parallelization (Phase B) ──
|
||
// Per-batch grad scratch buffers for cfc_step_backward_batched.
|
||
// Zeroed once per training step; the K-loop's 64 bwd calls accumulate
|
||
// into these via +=. After the K-loop, reduce_axis0 sums across B
|
||
// into the final grad buffers (overwrite semantics). AdamW reads the
|
||
// final grads (must run AFTER the reducer).
|
||
cfc_grad_w_in_scratch_d: CudaSlice<f32>, // [B, n_hid, n_in]
|
||
cfc_grad_w_rec_scratch_d: CudaSlice<f32>, // [B, n_hid, n_hid]
|
||
cfc_grad_b_scratch_d: CudaSlice<f32>, // [B, n_hid]
|
||
cfc_grad_tau_scratch_d: CudaSlice<f32>, // [B, n_hid]
|
||
// GRN per-batch grad scratch (Phase B commit 2).
|
||
grn_grad_w1_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID, HIDDEN]
|
||
grn_grad_b1_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID]
|
||
grn_grad_w2_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID, HEAD_MID]
|
||
grn_grad_b2_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID]
|
||
grn_grad_w_gate_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID]
|
||
grn_grad_b_gate_scratch_d: CudaSlice<f32>, // [B, 5]
|
||
grn_grad_w_main_scratch_d: CudaSlice<f32>, // [B, 5, HEAD_MID]
|
||
grn_grad_b_main_scratch_d: CudaSlice<f32>, // [B, 5]
|
||
grn_grad_w_skip_scratch_d: CudaSlice<f32>, // [B, 5, HIDDEN]
|
||
grn_grad_b_skip_scratch_d: CudaSlice<f32>, // [B, 5]
|
||
// VSN per-row grad scratch (Phase B commit 3). n_rows = B * K.
|
||
vsn_grad_w_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM, FEATURE_DIM]
|
||
vsn_grad_b_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM]
|
||
// Attention pool per-batch grad scratch (Phase B commit 4).
|
||
attn_grad_q_scratch_d: CudaSlice<f32>, // [B, HIDDEN_DIM]
|
||
/// Cross-batch reducer kernel: `[B, N] → [N]` via block tree-reduce.
|
||
/// Used for every per-batch grad scratch in the refactored bwd path.
|
||
reduce_axis0_fn: CudaFunction,
|
||
/// Owns the reduce_axis0 cubin lifetime.
|
||
_reduce_module: Arc<CudaModule>,
|
||
/// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised
|
||
/// (sentinel); the horizon_ema_and_lambda kernel detects the
|
||
/// sentinel on step 1 and replaces directly per
|
||
/// `pearl_first_observation_bootstrap.md`.
|
||
loss_ema_d: CudaSlice<f32>,
|
||
/// Per-horizon multiplicative gradient scaler. Updated each step
|
||
/// by horizon_ema_and_lambda. Asymmetric clamp [1.0, 2.0].
|
||
/// Consumed by heads_grn_bwd as `s_lambda` multiplier on the
|
||
/// trunk grad_h contribution.
|
||
lambda_d: CudaSlice<f32>,
|
||
/// EMA(max |z_h|) state for adaptive Z_SCALE in horizon_ema_and_lambda.
|
||
/// Single scalar; tracks the typical inter-horizon spread so
|
||
/// Z_SCALE_ISV maps the observed max-z to LAMBDA_CEILING.
|
||
z_max_ema_d: CudaSlice<f32>,
|
||
/// Cached function handle for the horizon_ema_and_lambda kernel.
|
||
horizon_lambda_fn: CudaFunction,
|
||
/// Module that owns `horizon_lambda_fn`; kept alive so the function
|
||
/// handle stays valid for the trainer's lifetime.
|
||
_horizon_lambda_module: Arc<CudaModule>,
|
||
/// Valid-label count for BCE normalisation.
|
||
valid_d: CudaSlice<i32>,
|
||
/// Per-horizon BCE loss weights, device-resident. Filled from
|
||
/// `cfg.horizon_weights` at construction; immutable thereafter.
|
||
loss_weights_d: CudaSlice<f32>,
|
||
/// Carry of grad_h_old across the reverse-order K backward loop.
|
||
/// Initialised to zero at step start; cfc_step_bwd writes its own
|
||
/// grad_h_old here, which heads_bwd of the previous position adds
|
||
/// in via the new `grad_h_carry` kernel arg.
|
||
grad_h_carry_d: CudaSlice<f32>,
|
||
/// Single-position grad_h_new buffer fed to cfc_step_bwd (heads_bwd
|
||
/// writes to it, fused with the carry).
|
||
grad_h_new_d: CudaSlice<f32>,
|
||
/// Single-position zero h_old buffer (used only at k=0 — first
|
||
/// position of the sequence has no preceding state).
|
||
zero_h_d: CudaSlice<f32>,
|
||
|
||
// Backward-grad accumulators (one per trainable group). Pre-zeroed
|
||
// at the top of step(); cfc_step_bwd / heads_bwd use += semantics.
|
||
grad_w_in_d: CudaSlice<f32>,
|
||
grad_w_rec_d: CudaSlice<f32>,
|
||
grad_b_d: CudaSlice<f32>,
|
||
grad_tau_d: CudaSlice<f32>,
|
||
grad_heads_w1_d: CudaSlice<f32>,
|
||
grad_heads_b1_d: CudaSlice<f32>,
|
||
grad_heads_w2_d: CudaSlice<f32>,
|
||
grad_heads_b2_d: CudaSlice<f32>,
|
||
grad_heads_w_gate_d: CudaSlice<f32>,
|
||
grad_heads_b_gate_d: CudaSlice<f32>,
|
||
grad_heads_w_main_d: CudaSlice<f32>,
|
||
grad_heads_b_main_d: CudaSlice<f32>,
|
||
grad_heads_w_skip_d: CudaSlice<f32>,
|
||
grad_heads_b_skip_d: CudaSlice<f32>,
|
||
// GRN forward intermediates saved for backward, per K position.
|
||
// Stored as [K, B, N_HORIZONS, HEAD_MID] in row-major; per-K slot
|
||
// is [B, N_HORIZONS, HEAD_MID] contiguous.
|
||
z1_per_k_d: CudaSlice<f32>,
|
||
a1_per_k_d: CudaSlice<f32>,
|
||
z2_per_k_d: CudaSlice<f32>,
|
||
// Per-K horizon-scalar intermediates: [K, B, N_HORIZONS].
|
||
gate_logit_per_k_d: CudaSlice<f32>,
|
||
main_per_k_d: CudaSlice<f32>,
|
||
logit_per_k_d: CudaSlice<f32>,
|
||
|
||
// Mapped-pinned staging for label uploads (the only remaining
|
||
// single-call host→device path; staging covers all K*B labels per
|
||
// training step in one DtoD copy).
|
||
stg_labels: MappedF32Buffer,
|
||
|
||
/// Captured CUDA Graph for the training step. First step_batched
|
||
/// call runs uncaptured (warmup); second call captures; third+
|
||
/// replays. Eliminates ~155 individual kernel launches per step.
|
||
train_graph: Option<CudaGraph>,
|
||
/// Captured CUDA Graph for the inference (`forward_only`) path. Same
|
||
/// three-state machine as `train_graph`: first call eager (warmup),
|
||
/// second call captures, third+ replays. Mirrors the forward chain
|
||
/// in `dispatch_evaluate_forward_kernels` (no labels / BCE / sync).
|
||
/// Per `pearl_no_host_branches_in_captured_graph` +
|
||
/// `pearl_cudarc_disable_event_tracking_for_graph_capture`: event
|
||
/// tracking is disabled for the trainer's lifetime at construction;
|
||
/// the dispatch helper performs only kernel launches and DtoD
|
||
/// memcpys on pre-bound device pointers.
|
||
forward_graph: Option<CudaGraph>,
|
||
/// First-call warmup flag for the forward (inference) path. Distinct
|
||
/// from `cublas_warmed` because the inference dispatch hits a
|
||
/// different cuBLAS call shape than the training step (no backward,
|
||
/// no Adam, no reduce-axis0). The first forward_only call runs
|
||
/// eager; the second captures; the third+ replays.
|
||
forward_warmed: bool,
|
||
/// cuBLAS warmup done — required before stream capture can safely
|
||
/// include cuBLAS calls.
|
||
cublas_warmed: bool,
|
||
|
||
// ── CRT Phase A0.5: incremental forward_step state ────────────────
|
||
/// Single-step scratch for Mamba2 L1 (in_dim=FEATURE_DIM). Holds the
|
||
/// persistent SSM register state `x_state` that advances by one
|
||
/// snapshot per `forward_step` call. Construct once at trainer init;
|
||
/// `reset_step_state()` zeros `x_state` for session-gap resets.
|
||
step_scratch_l1: Mamba2BlockStepScratch,
|
||
/// Single-step scratch for Mamba2 L2 (in_dim=HIDDEN_DIM). Same role
|
||
/// as `step_scratch_l1` but for the second SSM stack.
|
||
step_scratch_l2: Mamba2BlockStepScratch,
|
||
/// Persistent CfC hidden state for `forward_step` (`[1, HIDDEN_DIM]`).
|
||
/// Carries CfC's recurrent state across single-event calls — the
|
||
/// step path drops the attention-pool seed (which requires a
|
||
/// K-window) and instead lets CfC's natural decay (`exp(-dt/tau)`)
|
||
/// govern long-range context. Zero-initialised; `reset_step_state()`
|
||
/// rezeroes it.
|
||
cfc_h_state_step_d: CudaSlice<f32>,
|
||
/// Scratch buffer for one forward_step's CfC h_new output
|
||
/// (`[1, HIDDEN_DIM]`). Copied back into `cfc_h_state_step_d` at
|
||
/// end of step so the state thread is single-buffer-managed.
|
||
cfc_h_new_step_d: CudaSlice<f32>,
|
||
/// Single-row VSN output `[1, FEATURE_DIM]` for forward_step.
|
||
/// Distinct from `vsn_out_d` (which is [B, K, FEATURE_DIM] for the
|
||
/// supervised path) to keep the two paths independent.
|
||
vsn_step_out_d: GpuTensor,
|
||
/// Single-row gates scratch `[1, FEATURE_DIM]` for the VSN forward
|
||
/// kernel (variable_selection_fwd writes per-row gates). Unused
|
||
/// downstream in inference but required by the kernel signature.
|
||
vsn_step_gates_d: CudaSlice<f32>,
|
||
/// Single-row LN_a output `[1, HIDDEN_DIM]` for forward_step.
|
||
ln_a_step_out_d: CudaSlice<f32>,
|
||
/// Single-row LN_a stats `[1, 2]` for forward_step (mean + inv_std;
|
||
/// unused downstream but required by layer_norm_fwd).
|
||
ln_a_step_stats_d: CudaSlice<f32>,
|
||
/// Single-row LN_b output `[1, HIDDEN_DIM]` for forward_step (CfC input).
|
||
ln_b_step_out_d: CudaSlice<f32>,
|
||
/// Single-row LN_b stats `[1, 2]` for forward_step.
|
||
ln_b_step_stats_d: CudaSlice<f32>,
|
||
/// Per-horizon intermediate scratches `[1, N_HORIZONS, HEAD_MID_DIM]`
|
||
/// / `[1, N_HORIZONS]` — required by `multi_horizon_heads_grn_fwd_batched`
|
||
/// signature. Unused at inference but the kernel writes them.
|
||
z1_step_d: CudaSlice<f32>,
|
||
a1_step_d: CudaSlice<f32>,
|
||
z2_step_d: CudaSlice<f32>,
|
||
gate_logit_step_d: CudaSlice<f32>,
|
||
main_step_d: CudaSlice<f32>,
|
||
logit_step_d: CudaSlice<f32>,
|
||
/// Single-snapshot staging buffers for forward_step's host→device
|
||
/// path. Sized for one snapshot (10 levels × 4 px+sz fields, 6
|
||
/// regime values, etc.). Separate from the existing K-window
|
||
/// staging (`stg_bid_px_all` etc.) so the two paths don't share
|
||
/// buffer state.
|
||
stg_step_bid_px: MappedF32Buffer,
|
||
stg_step_bid_sz: MappedF32Buffer,
|
||
stg_step_ask_px: MappedF32Buffer,
|
||
stg_step_ask_sz: MappedF32Buffer,
|
||
stg_step_regime: MappedF32Buffer,
|
||
stg_step_prev_mid: MappedF32Buffer,
|
||
stg_step_trade_signed_vol: MappedF32Buffer,
|
||
stg_step_trade_count: MappedI32Buffer,
|
||
stg_step_ts_ns: MappedI64Buffer,
|
||
stg_step_prev_ts_ns: MappedI64Buffer,
|
||
/// Single-snapshot device shadows for the step path. Mirrors
|
||
/// `bid_px_all_d` etc. but sized for 1 snapshot.
|
||
step_bid_px_d: CudaSlice<f32>,
|
||
step_bid_sz_d: CudaSlice<f32>,
|
||
step_ask_px_d: CudaSlice<f32>,
|
||
step_ask_sz_d: CudaSlice<f32>,
|
||
step_prev_bid_sz_d: CudaSlice<f32>,
|
||
step_prev_ask_sz_d: CudaSlice<f32>,
|
||
step_regime_d: CudaSlice<f32>,
|
||
step_prev_mid_d: CudaSlice<f32>,
|
||
step_trade_signed_vol_d: CudaSlice<f32>,
|
||
step_trade_count_d: CudaSlice<i32>,
|
||
step_ts_ns_d: CudaSlice<i64>,
|
||
step_prev_ts_ns_d: CudaSlice<i64>,
|
||
/// Single-row feature output `[1, 1, FEATURE_DIM]` that
|
||
/// snap_feature_assemble_batched writes into for the step path.
|
||
/// Wrapped as a 3-D GpuTensor so it can feed VSN's existing
|
||
/// `[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,
|
||
/// 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
|
||
/// 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>,
|
||
|
||
// ── 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<f32>,
|
||
/// 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 {
|
||
pub fn new(dev: &MlDevice, cfg: &PerceptionTrainerConfig) -> Result<Self> {
|
||
anyhow::ensure!(cfg.seq_len >= 2, "Mamba2 requires seq_len >= 2");
|
||
|
||
// Seed GPU-side weight init (Mamba2 stacks initialise their weights via
|
||
// `OwnedGpuLinear::xavier`, which delegates to
|
||
// `ml_core::cuda_autograd::init::generate_uniform`). Without the
|
||
// thread-local guard, that helper seeds from SystemTime+thread-id and
|
||
// produces non-deterministic weights — every process gets fresh values.
|
||
// CfC/VSN/heads below use an explicit `ChaCha8Rng::seed_from_u64(cfg.seed)`
|
||
// chain already; the guard closes the remaining init paths so the
|
||
// whole trainer is reproducible from `cfg.seed`.
|
||
//
|
||
// Guard is dropped at the end of `new`, restoring default time-based
|
||
// seeding for any later xavier callers outside the trainer (e.g.,
|
||
// dqn / ppo crates that also use OwnedGpuLinear::xavier).
|
||
let _seed_guard = ml_core::cuda_autograd::init::scoped_init_seed(cfg.seed);
|
||
|
||
// X2: construct the wrapped CfcTrunk that owns the v2 inference
|
||
// graph weights. Trunk's V1 fields (CfC + V1 heads + projection)
|
||
// stay unused at this commit; subsequent migrations (X3..X9) move
|
||
// PerceptionTrainer's per-layer weights into the trunk one group
|
||
// at a time, gated by the perception_forward_golden bit-equivalence
|
||
// test.
|
||
let mut trunk = crate::cfc::trunk::CfcTrunk::new_random(
|
||
dev,
|
||
&crate::cfc::trunk::CfcConfig {
|
||
n_in: FEATURE_DIM,
|
||
n_hid: HIDDEN_DIM,
|
||
cfc_n_in: HIDDEN_DIM,
|
||
mamba2_state_dim: cfg.mamba2_state_dim,
|
||
n_batch: cfg.n_batch,
|
||
seq_len: cfg.seq_len,
|
||
},
|
||
cfg.seed,
|
||
)
|
||
.context("PerceptionTrainer: construct wrapped CfcTrunk")?;
|
||
|
||
let stream = dev.cuda_stream().context("trainer stream")?.clone();
|
||
let ctx = dev.cuda_context().context("trainer ctx")?;
|
||
// Disable cudarc's automatic per-allocation read/write event
|
||
// tracking BEFORE allocating any device memory. With tracking
|
||
// enabled, each `device_ptr_mut()` call returns a SyncOnDrop
|
||
// guard that records a CudaEvent on the stream when dropped.
|
||
// Inside a stream-capture region those `event.record(stream)`
|
||
// calls turn the events into "captured events" — which then
|
||
// cannot be waited on from outside the capture, breaking the
|
||
// non-captured eval path with CUDA_ERROR_INVALID_VALUE at the
|
||
// next `cuStreamWaitEvent`. All trainer work runs on a single
|
||
// stream (`self.stream`) and is therefore stream-ordered with
|
||
// no cross-stream sync needed, so event tracking is pure
|
||
// overhead for us. Disabled for the lifetime of this trainer.
|
||
unsafe { ctx.disable_event_tracking(); }
|
||
let _ = ctx.check_err();
|
||
let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("snap cubin")?;
|
||
let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("step cubin")?;
|
||
let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("heads cubin")?;
|
||
let bce_module = ctx.load_cubin(BCE_CUBIN.to_vec()).context("bce cubin")?;
|
||
let horizon_lambda_module = ctx
|
||
.load_cubin(HORIZON_LAMBDA_CUBIN.to_vec())
|
||
.context("horizon_lambda cubin")?;
|
||
let horizon_lambda_fn = horizon_lambda_module
|
||
.load_function("horizon_ema_and_lambda")
|
||
.context("horizon_ema_and_lambda symbol")?;
|
||
let smoothness_module = ctx
|
||
.load_cubin(SMOOTHNESS_CUBIN.to_vec())
|
||
.context("load output_smoothness cubin")?;
|
||
let smoothness_fn = smoothness_module
|
||
.load_function("output_smoothness_loss_and_grad")
|
||
.context("load output_smoothness_loss_and_grad")?;
|
||
let smoothness_controller_module = ctx
|
||
.load_cubin(SMOOTHNESS_CONTROLLER_CUBIN.to_vec())
|
||
.context("load smoothness_lambda_controller cubin")?;
|
||
let smoothness_controller_fn = smoothness_controller_module
|
||
.load_function("smoothness_lambda_controller")
|
||
.context("load smoothness_lambda_controller")?;
|
||
let ln_module = ctx
|
||
.load_cubin(LAYER_NORM_CUBIN.to_vec())
|
||
.context("layer_norm cubin")?;
|
||
let ln_bwd_fn = ln_module
|
||
.load_function("layer_norm_bwd")
|
||
.context("layer_norm_bwd symbol")?;
|
||
let ln_reduce_fn = ln_module
|
||
.load_function("layer_norm_reduce_param_grads")
|
||
.context("layer_norm_reduce_param_grads symbol")?;
|
||
let vsn_module = ctx
|
||
.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec())
|
||
.context("variable_selection cubin")?;
|
||
let vsn_bwd_fn = vsn_module
|
||
.load_function("variable_selection_bwd")
|
||
.context("variable_selection_bwd symbol")?;
|
||
let attn_module = ctx
|
||
.load_cubin(ATTENTION_POOL_CUBIN.to_vec())
|
||
.context("attention_pool cubin")?;
|
||
let attn_bwd_fn = attn_module
|
||
.load_function("attention_pool_bwd")
|
||
.context("attention_pool_bwd symbol")?;
|
||
let reduce_module = ctx
|
||
.load_cubin(REDUCE_AXIS0_CUBIN.to_vec())
|
||
.context("reduce_axis0 cubin")?;
|
||
let reduce_axis0_fn = reduce_module
|
||
.load_function("reduce_axis0")
|
||
.context("reduce_axis0 symbol")?;
|
||
let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?;
|
||
let step_bwd_batched_fn = step_module.load_function("cfc_step_backward_batched")?;
|
||
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")?;
|
||
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.
|
||
let mamba2_adamw = Mamba2AdamW::new(
|
||
trunk.mamba2_l1(),
|
||
Mamba2AdamWConfig {
|
||
lr: cfg.lr_mamba2,
|
||
grad_clip_max_norm: Some(1.0),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.context("Mamba2AdamW::new")?;
|
||
let mamba2_fwd_scratch = Mamba2BlockForwardScratch::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BlockForwardScratch::new")?;
|
||
let mamba2_bwd_scratch = Mamba2BackwardScratch::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BackwardScratch::new")?;
|
||
let mamba2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BackwardGradsBuffers::new")?;
|
||
|
||
// ── Phase 2B: SECOND Mamba2 stack — same shape, on the trunk. ──
|
||
let mamba2_l2_adamw = Mamba2AdamW::new(
|
||
trunk.mamba2_l2(),
|
||
Mamba2AdamWConfig {
|
||
lr: cfg.lr_mamba2,
|
||
grad_clip_max_norm: Some(1.0),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.context("Mamba2AdamW::new (l2)")?;
|
||
let mamba2_l2_fwd_scratch = Mamba2BlockForwardScratch::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BlockForwardScratch::new (l2)")?;
|
||
let mamba2_l2_bwd_scratch = Mamba2BackwardScratch::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BackwardScratch::new (l2)")?;
|
||
let mamba2_l2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BackwardGradsBuffers::new (l2)")?;
|
||
let window_tensor_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, FEATURE_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("window_tensor_d alloc: {e}"))?;
|
||
let h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("h_enriched_seq_t_d alloc: {e}"))?;
|
||
let grad_h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_t_d alloc: {e}"))?;
|
||
let grad_h_enriched_seq_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_d alloc: {e}"))?;
|
||
|
||
// CfC weights (input = h_enriched [HIDDEN_DIM], output = [HIDDEN_DIM])
|
||
let mut r = ChaCha8Rng::seed_from_u64(cfg.seed);
|
||
let n_in = HIDDEN_DIM;
|
||
let n_hid = HIDDEN_DIM;
|
||
let scale_in = (1.0_f32 / n_in as f32).sqrt();
|
||
let scale_rec = (1.0_f32 / n_hid as f32).sqrt();
|
||
let w_in: Vec<f32> = (0..n_hid * n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
|
||
let w_rec: Vec<f32> = (0..n_hid * n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
|
||
let b: Vec<f32> = vec![0.0; n_hid];
|
||
let tau: Vec<f32> = (0..n_hid)
|
||
.map(|_| {
|
||
let u: f32 = r.gen_range(0.0..1.0);
|
||
(-4.6 + u * 11.5).exp()
|
||
})
|
||
.collect();
|
||
// ── TFT GRN heads init ──
|
||
// Xavier-uniform per-tensor: scale = sqrt(1 / fan_in).
|
||
let scale_h_in = (1.0_f32 / n_hid as f32).sqrt(); // HIDDEN → HEAD_MID
|
||
let scale_mid_mid = (1.0_f32 / HEAD_MID_DIM as f32).sqrt(); // HEAD_MID → HEAD_MID / → 1
|
||
let scale_h_skip = (1.0_f32 / n_hid as f32).sqrt(); // HIDDEN → 1
|
||
let heads_w1: Vec<f32> = (0..N_HORIZONS * HEAD_MID_DIM * n_hid)
|
||
.map(|_| r.gen_range(-scale_h_in..scale_h_in)).collect();
|
||
let heads_b1: Vec<f32> = vec![0.0; N_HORIZONS * HEAD_MID_DIM];
|
||
let heads_w2: Vec<f32> = (0..N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM)
|
||
.map(|_| r.gen_range(-scale_mid_mid..scale_mid_mid)).collect();
|
||
let heads_b2: Vec<f32> = vec![0.0; N_HORIZONS * HEAD_MID_DIM];
|
||
let heads_w_gate: Vec<f32> = (0..N_HORIZONS * HEAD_MID_DIM)
|
||
.map(|_| r.gen_range(-scale_mid_mid..scale_mid_mid)).collect();
|
||
let heads_b_gate: Vec<f32> = vec![0.0; N_HORIZONS];
|
||
let heads_w_main: Vec<f32> = (0..N_HORIZONS * HEAD_MID_DIM)
|
||
.map(|_| r.gen_range(-scale_mid_mid..scale_mid_mid)).collect();
|
||
let heads_b_main: Vec<f32> = vec![0.0; N_HORIZONS];
|
||
let heads_w_skip: Vec<f32> = (0..N_HORIZONS * n_hid)
|
||
.map(|_| r.gen_range(-scale_h_skip..scale_h_skip)).collect();
|
||
let heads_b_skip: Vec<f32> = vec![0.0; N_HORIZONS];
|
||
|
||
// X8: upload trainer's PRNG-derived CfC weights into the trunk's
|
||
// (now v2-shaped) CfC slots.
|
||
stream.memcpy_htod(&w_in, &mut trunk.w_in_d).context("trunk.w_in_d upload")?;
|
||
stream.memcpy_htod(&w_rec, &mut trunk.w_rec_d).context("trunk.w_rec_d upload")?;
|
||
stream.memcpy_htod(&b, &mut trunk.b_d).context("trunk.b_d upload")?;
|
||
stream.memcpy_htod(&tau, &mut trunk.tau_all_d).context("trunk.tau_all_d upload")?;
|
||
// X9: upload trainer's GRN heads into the trunk's slots.
|
||
stream.memcpy_htod(&heads_w1, &mut trunk.heads_w1_d).context("trunk.heads_w1 upload")?;
|
||
stream.memcpy_htod(&heads_b1, &mut trunk.heads_b1_d).context("trunk.heads_b1 upload")?;
|
||
stream.memcpy_htod(&heads_w2, &mut trunk.heads_w2_d).context("trunk.heads_w2 upload")?;
|
||
stream.memcpy_htod(&heads_b2, &mut trunk.heads_b2_d).context("trunk.heads_b2 upload")?;
|
||
stream.memcpy_htod(&heads_w_gate, &mut trunk.heads_w_gate_d).context("trunk.heads_w_gate upload")?;
|
||
stream.memcpy_htod(&heads_b_gate, &mut trunk.heads_b_gate_d).context("trunk.heads_b_gate upload")?;
|
||
stream.memcpy_htod(&heads_w_main, &mut trunk.heads_w_main_d).context("trunk.heads_w_main upload")?;
|
||
stream.memcpy_htod(&heads_b_main, &mut trunk.heads_b_main_d).context("trunk.heads_b_main upload")?;
|
||
stream.memcpy_htod(&heads_w_skip, &mut trunk.heads_w_skip_d).context("trunk.heads_w_skip upload")?;
|
||
stream.memcpy_htod(&heads_b_skip, &mut trunk.heads_b_skip_d).context("trunk.heads_b_skip upload")?;
|
||
|
||
let opt_w_in = AdamW::new(dev, n_hid * n_in, cfg.lr_cfc)?;
|
||
let opt_w_rec = AdamW::new(dev, n_hid * n_hid, cfg.lr_cfc)?;
|
||
let mut opt_b = AdamW::new(dev, n_hid, cfg.lr_cfc)?;
|
||
opt_b.wd = 0.0;
|
||
// tau is bounded log-uniformly initialised; small LR + no weight
|
||
// decay so the per-cell decay constants drift smoothly.
|
||
let mut opt_tau = AdamW::new(dev, n_hid, cfg.lr_cfc * 0.1)?;
|
||
opt_tau.wd = 0.0;
|
||
|
||
// ── K-loop parallelization (Phase B): cfc per-batch grad scratch ──
|
||
let cfc_grad_w_in_scratch_d = stream
|
||
.alloc_zeros::<f32>(cfg.n_batch * n_hid * n_in)
|
||
.context("cfc_grad_w_in_scratch_d alloc")?;
|
||
let cfc_grad_w_rec_scratch_d = stream
|
||
.alloc_zeros::<f32>(cfg.n_batch * n_hid * n_hid)
|
||
.context("cfc_grad_w_rec_scratch_d alloc")?;
|
||
let cfc_grad_b_scratch_d = stream
|
||
.alloc_zeros::<f32>(cfg.n_batch * n_hid)
|
||
.context("cfc_grad_b_scratch_d alloc")?;
|
||
let cfc_grad_tau_scratch_d = stream
|
||
.alloc_zeros::<f32>(cfg.n_batch * n_hid)
|
||
.context("cfc_grad_tau_scratch_d alloc")?;
|
||
// GRN heads: 10 AdamW (5 weight + 5 bias). Biases get wd=0 per
|
||
// the existing pattern; weights keep AdamW default wd.
|
||
let opt_heads_w1 = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM * n_hid, cfg.lr_cfc)?;
|
||
let mut opt_heads_b1 = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM, cfg.lr_cfc)?;
|
||
opt_heads_b1.wd = 0.0;
|
||
let opt_heads_w2 = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM, cfg.lr_cfc)?;
|
||
let mut opt_heads_b2 = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM, cfg.lr_cfc)?;
|
||
opt_heads_b2.wd = 0.0;
|
||
let opt_heads_w_gate = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM, cfg.lr_cfc)?;
|
||
let mut opt_heads_b_gate = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc)?;
|
||
opt_heads_b_gate.wd = 0.0;
|
||
let opt_heads_w_main = AdamW::new(dev, N_HORIZONS * HEAD_MID_DIM, cfg.lr_cfc)?;
|
||
let mut opt_heads_b_main = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc)?;
|
||
opt_heads_b_main.wd = 0.0;
|
||
let opt_heads_w_skip = AdamW::new(dev, N_HORIZONS * n_hid, cfg.lr_cfc)?;
|
||
let mut opt_heads_b_skip = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc)?;
|
||
opt_heads_b_skip.wd = 0.0;
|
||
|
||
// Phase B: GRN per-batch grad scratch.
|
||
let grn_grad_w1_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM)?;
|
||
let grn_grad_b1_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?;
|
||
let grn_grad_w2_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM)?;
|
||
let grn_grad_b2_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?;
|
||
let grn_grad_w_gate_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?;
|
||
let grn_grad_b_gate_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS)?;
|
||
let grn_grad_w_main_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?;
|
||
let grn_grad_b_main_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS)?;
|
||
let grn_grad_w_skip_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HIDDEN_DIM)?;
|
||
let grn_grad_b_skip_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS)?;
|
||
|
||
// LayerNorm parameters: gain initialised to 1.0, bias to 0.0.
|
||
// Adam wd=0 — LN params are scale/shift, never penalised.
|
||
// Two instances: LN_a between m1 and m2, LN_b between m2 and CfC.
|
||
let ln_gain_init: Vec<f32> = vec![1.0; HIDDEN_DIM];
|
||
let ln_bias_init: Vec<f32> = vec![0.0; HIDDEN_DIM];
|
||
// X6: upload LN_b weights into the trunk's slots.
|
||
stream
|
||
.memcpy_htod(&ln_gain_init, &mut trunk.ln_b_gain_d)
|
||
.context("trunk.ln_b_gain_d upload")?;
|
||
stream
|
||
.memcpy_htod(&ln_bias_init, &mut trunk.ln_b_bias_d)
|
||
.context("trunk.ln_b_bias_d upload")?;
|
||
let mut opt_ln_gain = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||
opt_ln_gain.wd = 0.0;
|
||
let mut opt_ln_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||
opt_ln_bias.wd = 0.0;
|
||
// X4: upload LN_a weights into the trunk's slots.
|
||
stream
|
||
.memcpy_htod(&ln_gain_init, &mut trunk.ln_a_gain_d)
|
||
.context("trunk.ln_a_gain_d upload")?;
|
||
stream
|
||
.memcpy_htod(&ln_bias_init, &mut trunk.ln_a_bias_d)
|
||
.context("trunk.ln_a_bias_d upload")?;
|
||
let mut opt_ln_a_gain = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||
opt_ln_a_gain.wd = 0.0;
|
||
let mut opt_ln_a_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||
opt_ln_a_bias.wd = 0.0;
|
||
// LN_a + grad scratch buffers (constructed pre-move-of-stream).
|
||
let ln_a_stats_d = stream.alloc_zeros::<f32>(cfg.n_batch * cfg.seq_len * 2)?;
|
||
let ln_a_out_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("ln_a_out_d alloc: {e}"))?;
|
||
let grad_ln_a_gain_per_row_d = stream.alloc_zeros::<f32>(cfg.n_batch * cfg.seq_len * HIDDEN_DIM)?;
|
||
let grad_ln_a_bias_per_row_d = stream.alloc_zeros::<f32>(cfg.n_batch * cfg.seq_len * HIDDEN_DIM)?;
|
||
let grad_ln_a_gain_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
|
||
let grad_ln_a_bias_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
|
||
let grad_ln_a_in_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("grad_ln_a_in_d alloc: {e}"))?;
|
||
|
||
// ── VSN init (Phase 2D) ──
|
||
// W_vsn ∈ [FEATURE_DIM, FEATURE_DIM] init near zero so initial
|
||
// gates ≈ softmax(0) = uniform 1/FEATURE_DIM. The model can then
|
||
// learn departures from uniform. Scale = sqrt(1/FEATURE_DIM).
|
||
//
|
||
// X2 migration: VSN weights now live on `self.trunk` (the wrapped
|
||
// CfcTrunk). We still draw values from the trainer's PRNG chain
|
||
// here to keep the chain state — and therefore every downstream
|
||
// weight (attn_q, etc.) — bit-identical to the pre-X2 layout. The
|
||
// draws are uploaded directly into the trunk's zero-initialised
|
||
// VSN slots via memcpy_htod, replacing the prior trainer-owned
|
||
// vsn_w_d / vsn_b_d allocations.
|
||
let vsn_scale = (1.0_f32 / FEATURE_DIM as f32).sqrt();
|
||
let vsn_w_init: Vec<f32> = (0..FEATURE_DIM * FEATURE_DIM)
|
||
.map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect();
|
||
let vsn_b_init: Vec<f32> = vec![0.0; FEATURE_DIM];
|
||
// Upload the trainer-driven VSN draws into the trunk's
|
||
// zero-initialised VSN slots from X1. From here on,
|
||
// `self.trunk.vsn_w_d` / `self.trunk.vsn_b_d` are the canonical
|
||
// VSN weights for both forward and backward paths.
|
||
stream
|
||
.memcpy_htod(&vsn_w_init, &mut trunk.vsn_w_d)
|
||
.context("trunk.vsn_w_d upload")?;
|
||
stream
|
||
.memcpy_htod(&vsn_b_init, &mut trunk.vsn_b_d)
|
||
.context("trunk.vsn_b_d upload")?;
|
||
let opt_vsn_w = AdamW::new(dev, FEATURE_DIM * FEATURE_DIM, cfg.lr_cfc)?;
|
||
let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?;
|
||
opt_vsn_b.wd = 0.0;
|
||
// Phase B: VSN per-row grad scratch (n_rows = B * K).
|
||
let vsn_grad_w_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * cfg.seq_len * FEATURE_DIM * FEATURE_DIM)?;
|
||
let vsn_grad_b_scratch_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * cfg.seq_len * FEATURE_DIM)?;
|
||
|
||
// ── Attention pool init (Phase 3) ──
|
||
// Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K
|
||
// → context ≈ mean of LN_b output. Model learns content
|
||
// addressing from a near-uniform starting point.
|
||
let attn_q_scale = (1.0_f32 / HIDDEN_DIM as f32).sqrt();
|
||
let attn_q_init: Vec<f32> = (0..HIDDEN_DIM)
|
||
.map(|_| r.gen_range(-attn_q_scale..attn_q_scale)).collect();
|
||
// X7: upload attention-pool query into the trunk's slot.
|
||
stream
|
||
.memcpy_htod(&attn_q_init, &mut trunk.attn_q_d)
|
||
.context("trunk.attn_q_d upload")?;
|
||
let attn_context_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)?;
|
||
let attn_weights_d = stream.alloc_zeros::<f32>(cfg.n_batch * cfg.seq_len)?;
|
||
let grad_attn_q_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
|
||
let opt_attn_q = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||
// Phase B: attn pool per-batch grad scratch.
|
||
let attn_grad_q_scratch_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)?;
|
||
|
||
// Kendall σ — per-horizon learnable `log σ_h` (axis A only;
|
||
// C/D/E reverted after the 2026-05-18 A/B sweep). Init at 0
|
||
// → weight = 1/(2·exp(0)) = 0.5 per horizon at step 0; AdamW
|
||
// updates at 1/4 the head LR per the spec.
|
||
let log_sigma_h_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;
|
||
let grad_log_sigma_h_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;
|
||
// ISV-driven: z_max_ema tracks EMA(max |z_h|) across horizons to set
|
||
// adaptive Z_SCALE for λ. Sentinel = 0 ⇒ first-obs bootstrap.
|
||
let z_max_ema_d = stream.alloc_zeros::<f32>(1)?;
|
||
|
||
// V1 (2026-05-18): per-horizon Q_h trainer state (C24/C25) removed.
|
||
// A/B sweep at commit 83546b5c3 falsified the approach (mean_auc
|
||
// -0.019 vs baseline; h6000 essentially tied). v2 design with
|
||
// horizon-token K-prepend + inverted attention + regime-MoE +
|
||
// Kendall σ-BCE + L2 anchor replaces it. See
|
||
// docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md
|
||
// and the V1-V13 plan.
|
||
|
||
// ── CRT Phase A0.5: forward_step (incremental SSM) state ──
|
||
// Single-snapshot scratch for the event-rate inference path.
|
||
// Mamba2 step kernels write back the SSM register state in-place,
|
||
// so these scratches accumulate context across forward_step calls
|
||
// without re-running over a K-window each time.
|
||
let step_scratch_l1 = Mamba2BlockStepScratch::new(
|
||
&stream, cfg.n_batch, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BlockStepScratch::new (l1)")?;
|
||
let step_scratch_l2 = Mamba2BlockStepScratch::new(
|
||
&stream, cfg.n_batch, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||
).context("Mamba2BlockStepScratch::new (l2)")?;
|
||
let cfc_h_state_step_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)
|
||
.context("cfc_h_state_step_d alloc")?;
|
||
let cfc_h_new_step_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)
|
||
.context("cfc_h_new_step_d alloc")?;
|
||
let vsn_step_out_d = GpuTensor::zeros(&[cfg.n_batch, FEATURE_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("vsn_step_out_d alloc: {e}"))?;
|
||
let vsn_step_gates_d = stream.alloc_zeros::<f32>(cfg.n_batch * FEATURE_DIM)
|
||
.context("vsn_step_gates_d alloc")?;
|
||
let ln_a_step_out_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)
|
||
.context("ln_a_step_out_d alloc")?;
|
||
let ln_a_step_stats_d = stream.alloc_zeros::<f32>(cfg.n_batch * 2)
|
||
.context("ln_a_step_stats_d alloc")?;
|
||
let ln_b_step_out_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)
|
||
.context("ln_b_step_out_d alloc")?;
|
||
let ln_b_step_stats_d = stream.alloc_zeros::<f32>(cfg.n_batch * 2)
|
||
.context("ln_b_step_stats_d alloc")?;
|
||
// probs_step_d / probs_step_host removed: forward_step_into now
|
||
// writes probs directly into a caller-supplied device buffer
|
||
// (typically `LobSimCuda::alpha_probs_d_mut`), eliminating the
|
||
// per-event DtoD-to-host + stream.synchronize that previously
|
||
// round-tripped the 5 probs through pinned host memory just to
|
||
// immediately re-upload them in `broadcast_alpha`.
|
||
let z1_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM).context("z1_step_d alloc")?;
|
||
let a1_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM).context("a1_step_d alloc")?;
|
||
let z2_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS * HEAD_MID_DIM).context("z2_step_d alloc")?;
|
||
let gate_logit_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS).context("gate_logit_step_d alloc")?;
|
||
let main_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS).context("main_step_d alloc")?;
|
||
let logit_step_d = stream.alloc_zeros::<f32>(
|
||
cfg.n_batch * N_HORIZONS).context("logit_step_d alloc")?;
|
||
let stg_step_bid_px = unsafe { MappedF32Buffer::new(cfg.n_batch * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_bid_px: {e}"))?;
|
||
let stg_step_bid_sz = unsafe { MappedF32Buffer::new(cfg.n_batch * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_bid_sz: {e}"))?;
|
||
let stg_step_ask_px = unsafe { MappedF32Buffer::new(cfg.n_batch * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_ask_px: {e}"))?;
|
||
let stg_step_ask_sz = unsafe { MappedF32Buffer::new(cfg.n_batch * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_ask_sz: {e}"))?;
|
||
let stg_step_regime = unsafe { MappedF32Buffer::new(cfg.n_batch * REGIME_DIM) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_regime: {e}"))?;
|
||
let stg_step_prev_mid = unsafe { MappedF32Buffer::new(cfg.n_batch) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_prev_mid: {e}"))?;
|
||
let stg_step_trade_signed_vol = unsafe { MappedF32Buffer::new(cfg.n_batch) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_trade_signed_vol: {e}"))?;
|
||
let stg_step_trade_count = unsafe { MappedI32Buffer::new(cfg.n_batch) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_trade_count: {e}"))?;
|
||
let stg_step_ts_ns = unsafe { MappedI64Buffer::new(cfg.n_batch) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_ts_ns: {e}"))?;
|
||
let stg_step_prev_ts_ns = unsafe { MappedI64Buffer::new(cfg.n_batch) }
|
||
.map_err(|e| anyhow::anyhow!("stg_step_prev_ts_ns: {e}"))?;
|
||
let step_bid_px_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_bid_sz_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_ask_px_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_ask_sz_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_prev_bid_sz_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_prev_ask_sz_d = stream.alloc_zeros::<f32>(cfg.n_batch * 10)?;
|
||
let step_regime_d = stream.alloc_zeros::<f32>(cfg.n_batch * REGIME_DIM)?;
|
||
let step_prev_mid_d = stream.alloc_zeros::<f32>(cfg.n_batch)?;
|
||
let step_trade_signed_vol_d = stream.alloc_zeros::<f32>(cfg.n_batch)?;
|
||
let step_trade_count_d = stream.alloc_zeros::<i32>(cfg.n_batch)?;
|
||
let step_ts_ns_d = stream.alloc_zeros::<i64>(cfg.n_batch)?;
|
||
let step_prev_ts_ns_d = stream.alloc_zeros::<i64>(cfg.n_batch)?;
|
||
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")?;
|
||
// ── 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::<f32>(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 ──
|
||
// λ is now ISV-driven by smoothness_lambda_controller — initialised
|
||
// to LAMBDA_FLOOR per horizon and updated each step inside the
|
||
// captured graph. The base_lambda amplitude knob enters the
|
||
// controller kernel as a scalar arg at launch time.
|
||
const LAMBDA_FLOOR_INIT: f32 = 1.0e-4;
|
||
let smoothness_lambda_host: Vec<f32> = vec![LAMBDA_FLOOR_INIT; N_HORIZONS];
|
||
let smoothness_lambda_d = {
|
||
// Upload via mapped-pinned (per `feedback_no_htod_htoh_only_mapped_pinned`).
|
||
let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) }
|
||
.map_err(|e| anyhow::anyhow!("smoothness λ staging: {e}"))?;
|
||
staging.write_from_slice(&smoothness_lambda_host);
|
||
let mut dst = stream.alloc_zeros::<f32>(N_HORIZONS)
|
||
.context("smoothness_lambda_d alloc")?;
|
||
let nbytes = N_HORIZONS * std::mem::size_of::<f32>();
|
||
unsafe {
|
||
let (dst_ptr, _g) = dst.device_ptr_mut(&stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||
).context("smoothness λ DtoD")?;
|
||
}
|
||
dst
|
||
};
|
||
let smoothness_loss_d = stream.alloc_zeros::<f32>(1)
|
||
.context("smoothness_loss_d alloc")?;
|
||
let smoothness_loss_host_d = unsafe { MappedF32Buffer::new(1) }
|
||
.map_err(|e| anyhow::anyhow!("smoothness_loss_host_d: {e}"))?;
|
||
let smoothness_loss_per_horizon_d = stream.alloc_zeros::<f32>(N_HORIZONS)
|
||
.context("smoothness_loss_per_horizon_d alloc")?;
|
||
let smoothness_loss_per_horizon_host_d = unsafe { MappedF32Buffer::new(N_HORIZONS) }
|
||
.map_err(|e| anyhow::anyhow!("smoothness_loss_per_horizon_host_d: {e}"))?;
|
||
// Controller state buffers.
|
||
let smoothness_jitter_ema_d = stream.alloc_zeros::<f32>(N_HORIZONS)
|
||
.context("smoothness_jitter_ema_d alloc")?;
|
||
let smoothness_jitter_first_obs_d = stream.alloc_zeros::<i32>(1)
|
||
.context("smoothness_jitter_first_obs_d alloc")?;
|
||
|
||
// ── GPU log ring allocation (feature: kernel-step-trace) ──
|
||
// Per `feedback_no_htod_htoh_only_mapped_pinned`: ring is mapped-
|
||
// pinned; step counter shadow is mapped-pinned; the on-stream
|
||
// DtoD memcpy of the device step counter into the host shadow
|
||
// is captured in the graph, so the drainer sees the latest
|
||
// counter without an extra sync.
|
||
//
|
||
// Runtime activation: only allocate + spawn when
|
||
// `cfg.kernel_step_trace_path` is Some(path). Otherwise every
|
||
// field below stays None and the trainer carries zero overhead.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
let (
|
||
log_ring,
|
||
log_step_counter_d,
|
||
log_step_counter_host_shadow,
|
||
gpu_log_tick_fn,
|
||
gpu_log_module,
|
||
log_drain_stop,
|
||
log_drain_handle,
|
||
) = if let Some(trace_path) = cfg.kernel_step_trace_path.clone() {
|
||
let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?);
|
||
let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?;
|
||
let log_step_counter_host_shadow = {
|
||
let buf = unsafe { crate::pinned_mem::MappedRecordBuffer::<i32>::new(1) }
|
||
.map_err(|e| anyhow::anyhow!("log step counter shadow: {e}"))?;
|
||
buf.write_record(0, 0);
|
||
std::sync::Arc::new(buf)
|
||
};
|
||
let (gpu_log_tick_fn, gpu_log_module) = {
|
||
const GPU_LOG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin"));
|
||
let module = ctx.load_cubin(GPU_LOG_CUBIN.to_vec())
|
||
.context("load gpu_log_ring cubin")?;
|
||
let f = module.load_function("gpu_log_tick").context("load gpu_log_tick")?;
|
||
(f, module)
|
||
};
|
||
let log_drain_stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||
// The drain runs on a plain `std::thread`, so spawning is
|
||
// unconditional — it works the same whether the trainer is
|
||
// invoked from `#[tokio::main]` or a synchronous `fn main()`
|
||
// (e.g. `alpha_train`). No runtime-handle gate required.
|
||
let log_drain_handle = Some(crate::gpu_log::spawn_drain_task_to_file(
|
||
log_ring.clone(),
|
||
log_step_counter_host_shadow.clone(),
|
||
log_drain_stop.clone(),
|
||
trace_path.clone(),
|
||
)?);
|
||
(
|
||
Some(log_ring),
|
||
Some(log_step_counter_d),
|
||
Some(log_step_counter_host_shadow),
|
||
Some(gpu_log_tick_fn),
|
||
Some(gpu_log_module),
|
||
Some(log_drain_stop),
|
||
log_drain_handle,
|
||
)
|
||
} else {
|
||
(None, None, None, None, None, None, None)
|
||
};
|
||
|
||
let k = cfg.seq_len;
|
||
Ok(Self {
|
||
cfg: cfg.clone(),
|
||
h_new_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * n_hid)?,
|
||
probs_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
labels_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
grad_probs_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
loss_weights_d: upload(&stream, &cfg.horizon_weights)?,
|
||
loss_d: stream.alloc_zeros::<f32>(1)?,
|
||
loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?,
|
||
loss_per_horizon_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
smoothness_lambda_d,
|
||
smoothness_loss_d,
|
||
smoothness_loss_host_d,
|
||
smoothness_loss_per_horizon_d,
|
||
smoothness_loss_per_horizon_host_d,
|
||
smoothness_fn,
|
||
_smoothness_module: smoothness_module,
|
||
smoothness_jitter_ema_d,
|
||
smoothness_jitter_first_obs_d,
|
||
smoothness_controller_fn,
|
||
_smoothness_controller_module: smoothness_controller_module,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_ring,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_step_counter_d,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_step_counter_host_shadow,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
gpu_log_tick_fn,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
_gpu_log_module: gpu_log_module,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_drain_stop,
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
log_drain_handle, // already Option<_>; field type unchanged
|
||
// LayerNorm state. `ln_gain_d` / `ln_bias_d` constructed
|
||
// above via `upload` (gain=1, bias=0). LN_HIDDEN matches
|
||
// HIDDEN_DIM (128) by construction — checked at kernel
|
||
// launch time. Per-row scratch is [B*K, HIDDEN]; the
|
||
// single param-grad reducer collapses it to [HIDDEN].
|
||
ln_stats_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 2)?,
|
||
ln_out_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
|
||
grad_ln_gain_per_row_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
|
||
grad_ln_bias_per_row_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
|
||
grad_ln_gain_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
|
||
grad_ln_bias_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
|
||
grad_ln_in_d: GpuTensor::zeros(&[cfg.n_batch, k, HIDDEN_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("grad_ln_in_d alloc: {e}"))?,
|
||
ln_bwd_fn,
|
||
ln_reduce_fn,
|
||
_ln_module: ln_module,
|
||
opt_ln_gain,
|
||
opt_ln_bias,
|
||
// VSN (Phase 2D) — params + per-K gates + scratch.
|
||
trunk,
|
||
vsn_out_d: GpuTensor::zeros(&[cfg.n_batch, k, FEATURE_DIM], &stream)
|
||
.map_err(|e| anyhow::anyhow!("vsn_out_d alloc: {e}"))?,
|
||
vsn_gates_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * FEATURE_DIM)?,
|
||
vsn_grad_x_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * FEATURE_DIM)?,
|
||
grad_vsn_w_d: stream.alloc_zeros::<f32>(FEATURE_DIM * FEATURE_DIM)?,
|
||
grad_vsn_b_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
|
||
opt_vsn_w,
|
||
opt_vsn_b,
|
||
vsn_bwd_fn,
|
||
_vsn_module: vsn_module,
|
||
// Attention pool (Phase 3).
|
||
attn_context_d,
|
||
attn_weights_d,
|
||
grad_attn_q_d,
|
||
opt_attn_q,
|
||
attn_bwd_fn,
|
||
_attn_module: attn_module,
|
||
log_sigma_h_d,
|
||
grad_log_sigma_h_d,
|
||
z_max_ema_d,
|
||
// Phase B: cfc per-batch grad scratch + reducer.
|
||
cfc_grad_w_in_scratch_d,
|
||
cfc_grad_w_rec_scratch_d,
|
||
cfc_grad_b_scratch_d,
|
||
cfc_grad_tau_scratch_d,
|
||
grn_grad_w1_scratch_d,
|
||
grn_grad_b1_scratch_d,
|
||
grn_grad_w2_scratch_d,
|
||
grn_grad_b2_scratch_d,
|
||
grn_grad_w_gate_scratch_d,
|
||
grn_grad_b_gate_scratch_d,
|
||
grn_grad_w_main_scratch_d,
|
||
grn_grad_b_main_scratch_d,
|
||
grn_grad_w_skip_scratch_d,
|
||
grn_grad_b_skip_scratch_d,
|
||
vsn_grad_w_scratch_d,
|
||
vsn_grad_b_scratch_d,
|
||
attn_grad_q_scratch_d,
|
||
reduce_axis0_fn,
|
||
_reduce_module: reduce_module,
|
||
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
lambda_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
horizon_lambda_fn,
|
||
_horizon_lambda_module: horizon_lambda_module,
|
||
valid_d: stream.alloc_zeros::<i32>(1)?,
|
||
grad_h_carry_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
|
||
grad_h_new_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
|
||
zero_h_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
|
||
grad_w_in_d: stream.alloc_zeros::<f32>(n_hid * n_in)?,
|
||
grad_w_rec_d: stream.alloc_zeros::<f32>(n_hid * n_hid)?,
|
||
grad_b_d: stream.alloc_zeros::<f32>(n_hid)?,
|
||
grad_tau_d: stream.alloc_zeros::<f32>(n_hid)?,
|
||
grad_heads_w1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * n_hid)?,
|
||
grad_heads_b1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)?,
|
||
grad_heads_w2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM)?,
|
||
grad_heads_b2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)?,
|
||
grad_heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)?,
|
||
grad_heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
grad_heads_w_main_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)?,
|
||
grad_heads_b_main_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
grad_heads_w_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS * n_hid)?,
|
||
grad_heads_b_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||
// GRN per-K intermediate buffers (sized for [K, B, 5, HEAD_MID]
|
||
// and [K, B, 5]) — the K-loop reads/writes per-K slot offsets.
|
||
z1_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?,
|
||
a1_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?,
|
||
z2_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS * HEAD_MID_DIM)?,
|
||
gate_logit_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
main_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
logit_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
|
||
stg_labels: unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?,
|
||
train_graph: None,
|
||
forward_graph: None,
|
||
forward_warmed: false,
|
||
cublas_warmed: false,
|
||
|
||
// Batched snap_feature staging — sized for max B*K snapshots.
|
||
bk_capacity: cfg.n_batch * k,
|
||
stg_bid_px_all: unsafe { MappedF32Buffer::new(cfg.n_batch * k * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_bid_px_all: {e}"))?,
|
||
stg_bid_sz_all: unsafe { MappedF32Buffer::new(cfg.n_batch * k * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_bid_sz_all: {e}"))?,
|
||
stg_ask_px_all: unsafe { MappedF32Buffer::new(cfg.n_batch * k * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_ask_px_all: {e}"))?,
|
||
stg_ask_sz_all: unsafe { MappedF32Buffer::new(cfg.n_batch * k * 10) }
|
||
.map_err(|e| anyhow::anyhow!("stg_ask_sz_all: {e}"))?,
|
||
stg_regime_all: unsafe { MappedF32Buffer::new(cfg.n_batch * k * REGIME_DIM) }
|
||
.map_err(|e| anyhow::anyhow!("stg_regime_all: {e}"))?,
|
||
stg_prev_mid: unsafe { MappedF32Buffer::new(cfg.n_batch * k) }
|
||
.map_err(|e| anyhow::anyhow!("stg_prev_mid: {e}"))?,
|
||
stg_trade_signed_vol: unsafe { MappedF32Buffer::new(cfg.n_batch * k) }
|
||
.map_err(|e| anyhow::anyhow!("stg_trade_signed_vol: {e}"))?,
|
||
stg_trade_count: unsafe { MappedI32Buffer::new(cfg.n_batch * k) }
|
||
.map_err(|e| anyhow::anyhow!("stg_trade_count: {e}"))?,
|
||
stg_ts_ns: unsafe { MappedI64Buffer::new(cfg.n_batch * k) }
|
||
.map_err(|e| anyhow::anyhow!("stg_ts_ns: {e}"))?,
|
||
stg_prev_ts_ns: unsafe { MappedI64Buffer::new(cfg.n_batch * k) }
|
||
.map_err(|e| anyhow::anyhow!("stg_prev_ts_ns: {e}"))?,
|
||
bid_px_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
bid_sz_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
ask_px_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
ask_sz_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
prev_bid_sz_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
prev_ask_sz_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 10)?,
|
||
regime_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * REGIME_DIM)?,
|
||
prev_mid_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k)?,
|
||
trade_signed_vol_all_d: stream.alloc_zeros::<f32>(cfg.n_batch * k)?,
|
||
trade_count_all_d: stream.alloc_zeros::<i32>(cfg.n_batch * k)?,
|
||
ts_ns_all_d: stream.alloc_zeros::<i64>(cfg.n_batch * k)?,
|
||
prev_ts_ns_all_d: stream.alloc_zeros::<i64>(cfg.n_batch * k)?,
|
||
|
||
stream,
|
||
_snap_module: snap_module,
|
||
_step_module: step_module,
|
||
_heads_module: heads_module,
|
||
_bce_module: bce_module,
|
||
bce_fn,
|
||
step_bwd_batched_fn,
|
||
heads_grn_bwd_fn,
|
||
mamba2_adamw,
|
||
mamba2_fwd_scratch,
|
||
mamba2_bwd_scratch,
|
||
mamba2_grads_buffers,
|
||
mamba2_l2_adamw,
|
||
mamba2_l2_fwd_scratch,
|
||
mamba2_l2_bwd_scratch,
|
||
mamba2_l2_grads_buffers,
|
||
// LN_a between Mamba2 stacks.
|
||
ln_a_stats_d,
|
||
ln_a_out_d,
|
||
grad_ln_a_gain_per_row_d,
|
||
grad_ln_a_bias_per_row_d,
|
||
grad_ln_a_gain_d,
|
||
grad_ln_a_bias_d,
|
||
grad_ln_a_in_d,
|
||
opt_ln_a_gain,
|
||
opt_ln_a_bias,
|
||
window_tensor_d,
|
||
h_enriched_seq_t_d,
|
||
grad_h_enriched_seq_t_d,
|
||
grad_h_enriched_seq_d,
|
||
opt_w_in,
|
||
opt_w_rec,
|
||
opt_b,
|
||
opt_tau,
|
||
opt_heads_w1,
|
||
opt_heads_b1,
|
||
opt_heads_w2,
|
||
opt_heads_b2,
|
||
opt_heads_w_gate,
|
||
opt_heads_b_gate,
|
||
opt_heads_w_main,
|
||
opt_heads_b_main,
|
||
opt_heads_w_skip,
|
||
opt_heads_b_skip,
|
||
// CRT Phase A0.5: forward_step state.
|
||
step_scratch_l1,
|
||
step_scratch_l2,
|
||
cfc_h_state_step_d,
|
||
cfc_h_new_step_d,
|
||
vsn_step_out_d,
|
||
vsn_step_gates_d,
|
||
ln_a_step_out_d,
|
||
ln_a_step_stats_d,
|
||
ln_b_step_out_d,
|
||
ln_b_step_stats_d,
|
||
z1_step_d,
|
||
a1_step_d,
|
||
z2_step_d,
|
||
gate_logit_step_d,
|
||
main_step_d,
|
||
logit_step_d,
|
||
stg_step_bid_px,
|
||
stg_step_bid_sz,
|
||
stg_step_ask_px,
|
||
stg_step_ask_sz,
|
||
stg_step_regime,
|
||
stg_step_prev_mid,
|
||
stg_step_trade_signed_vol,
|
||
stg_step_trade_count,
|
||
stg_step_ts_ns,
|
||
stg_step_prev_ts_ns,
|
||
step_bid_px_d,
|
||
step_bid_sz_d,
|
||
step_ask_px_d,
|
||
step_ask_sz_d,
|
||
step_prev_bid_sz_d,
|
||
step_prev_ask_sz_d,
|
||
step_regime_d,
|
||
step_prev_mid_d,
|
||
step_trade_signed_vol_d,
|
||
step_trade_count_d,
|
||
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,
|
||
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
|
||
})
|
||
}
|
||
|
||
/// Set learning rate for all CfC + heads + tau AdamW optimizers.
|
||
/// Mamba2's AdamW carries its own lr config — use [`set_lr_mamba2`].
|
||
pub fn set_lr_cfc(&mut self, lr: f32) {
|
||
self.opt_w_in.lr = lr;
|
||
self.opt_w_rec.lr = lr;
|
||
self.opt_b.lr = lr;
|
||
self.opt_tau.lr = lr * 0.1;
|
||
self.opt_heads_w1.lr = lr;
|
||
self.opt_heads_b1.lr = lr;
|
||
self.opt_heads_w2.lr = lr;
|
||
self.opt_heads_b2.lr = lr;
|
||
self.opt_heads_w_gate.lr = lr;
|
||
self.opt_heads_b_gate.lr = lr;
|
||
self.opt_heads_w_main.lr = lr;
|
||
self.opt_heads_b_main.lr = lr;
|
||
self.opt_heads_w_skip.lr = lr;
|
||
self.opt_heads_b_skip.lr = lr;
|
||
}
|
||
|
||
/// Set Mamba2 AdamW learning rate. Applies to BOTH stack-1 and
|
||
/// stack-2 (Phase 2B) — they share the same lr_mamba2 config.
|
||
pub fn set_lr_mamba2(&mut self, lr: f32) {
|
||
self.mamba2_adamw.config.lr = lr;
|
||
self.mamba2_l2_adamw.config.lr = lr;
|
||
}
|
||
|
||
/// One training step on a single sequence — thin wrapper around
|
||
/// [`step_batched`] with batch size 1. Preserves the existing
|
||
/// API for single-sequence callers (test smokes, etc.).
|
||
pub fn step(
|
||
&mut self,
|
||
snapshots: &[Mbp10RawInput],
|
||
labels_per_position: &[[f32; N_HORIZONS]],
|
||
) -> Result<f32> {
|
||
anyhow::ensure!(
|
||
self.cfg.n_batch == 1,
|
||
"single-sequence step() requires cfg.n_batch == 1; got {}. \
|
||
Use step_batched() for batched training.",
|
||
self.cfg.n_batch
|
||
);
|
||
self.step_batched(&[snapshots], &[labels_per_position])
|
||
}
|
||
|
||
/// Batched training step. Processes `cfg.n_batch` sequences per
|
||
/// optimizer update — uses the batched kernel variants (cfc_step_batched,
|
||
/// multi_horizon_heads_batched, etc.). CfC is RECURRENT across
|
||
/// positions; h_old at step k for sample b is h_new at step k-1 for
|
||
/// the same b (zero at k=0).
|
||
///
|
||
/// `snapshots_batch[b]` is sample b's seq_len-snapshot window.
|
||
/// `labels_batch[b]` is sample b's per-position label rows. NaN
|
||
/// labels are natively masked by the BCE kernel.
|
||
///
|
||
/// Returns mean BCE over the valid (b, position, horizon) triples
|
||
/// in the batch, weighted by `cfg.horizon_weights`.
|
||
/// Test/diagnostic accessor for the per-horizon BCE EMA buffer.
|
||
/// Forces a stream sync + mapped-pinned readback — NOT for the
|
||
/// hot training path. Used by smoke tests to verify the
|
||
/// `horizon_ema_and_lambda` kernel is tracking correctly.
|
||
pub fn loss_ema_snapshot(&self) -> Result<[f32; N_HORIZONS]> {
|
||
let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) }
|
||
.map_err(|e| anyhow::anyhow!("loss_ema staging: {e}"))?;
|
||
unsafe {
|
||
let (src_ptr, _g) = self.loss_ema_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr, src_ptr,
|
||
N_HORIZONS * std::mem::size_of::<f32>(),
|
||
self.stream.cu_stream(),
|
||
).context("loss_ema dtod")?;
|
||
}
|
||
self.stream.synchronize().context("loss_ema sync")?;
|
||
let host = staging.read_all();
|
||
let mut out = [0.0_f32; N_HORIZONS];
|
||
out.copy_from_slice(&host[..N_HORIZONS]);
|
||
Ok(out)
|
||
}
|
||
|
||
/// Test/diagnostic accessor for the per-horizon lambda buffer.
|
||
/// Same constraints as [`Self::loss_ema_snapshot`].
|
||
pub fn lambda_snapshot(&self) -> Result<[f32; N_HORIZONS]> {
|
||
let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) }
|
||
.map_err(|e| anyhow::anyhow!("lambda staging: {e}"))?;
|
||
unsafe {
|
||
let (src_ptr, _g) = self.lambda_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
staging.dev_ptr, src_ptr,
|
||
N_HORIZONS * std::mem::size_of::<f32>(),
|
||
self.stream.cu_stream(),
|
||
).context("lambda dtod")?;
|
||
}
|
||
self.stream.synchronize().context("lambda sync")?;
|
||
let host = staging.read_all();
|
||
let mut out = [0.0_f32; N_HORIZONS];
|
||
out.copy_from_slice(&host[..N_HORIZONS]);
|
||
Ok(out)
|
||
}
|
||
|
||
pub fn step_batched(
|
||
&mut self,
|
||
snapshots_batch: &[&[Mbp10RawInput]],
|
||
labels_batch: &[&[[f32; N_HORIZONS]]],
|
||
) -> Result<f32> {
|
||
let b_sz = self.cfg.n_batch;
|
||
let k_seq = self.cfg.seq_len;
|
||
anyhow::ensure!(
|
||
snapshots_batch.len() == b_sz && labels_batch.len() == b_sz,
|
||
"expected B={} sequences/labels; got snap={} lbl={}",
|
||
b_sz, snapshots_batch.len(), labels_batch.len()
|
||
);
|
||
for (b, snap) in snapshots_batch.iter().enumerate() {
|
||
anyhow::ensure!(
|
||
snap.len() == k_seq,
|
||
"batch[{}] snapshots.len()={} != seq_len={}",
|
||
b, snap.len(), k_seq
|
||
);
|
||
}
|
||
for (b, lbl) in labels_batch.iter().enumerate() {
|
||
anyhow::ensure!(
|
||
lbl.len() == k_seq,
|
||
"batch[{}] labels.len()={} != seq_len={}",
|
||
b, lbl.len(), k_seq
|
||
);
|
||
}
|
||
|
||
debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]);
|
||
let total_snaps = b_sz * k_seq;
|
||
debug_assert!(total_snaps <= self.bk_capacity);
|
||
|
||
// ── 1. Fill ALL mapped-pinned staging on host BEFORE stream
|
||
// ops. Replays read whatever host last wrote.
|
||
{
|
||
let bid_px_h = self.stg_bid_px_all.host_slice_mut();
|
||
let bid_sz_h = self.stg_bid_sz_all.host_slice_mut();
|
||
let ask_px_h = self.stg_ask_px_all.host_slice_mut();
|
||
let ask_sz_h = self.stg_ask_sz_all.host_slice_mut();
|
||
let regime_h = self.stg_regime_all.host_slice_mut();
|
||
for (b_idx, snapshots) in snapshots_batch.iter().enumerate() {
|
||
for (k, snap) in snapshots.iter().enumerate() {
|
||
let n = b_idx * k_seq + k;
|
||
let base10 = n * 10;
|
||
let base6 = n * REGIME_DIM;
|
||
for i in 0..10 {
|
||
bid_px_h[base10 + i] = snap.bid_px[i];
|
||
bid_sz_h[base10 + i] = snap.bid_sz[i];
|
||
ask_px_h[base10 + i] = snap.ask_px[i];
|
||
ask_sz_h[base10 + i] = snap.ask_sz[i];
|
||
}
|
||
for i in 0..REGIME_DIM {
|
||
regime_h[base6 + i] = snap.regime[i];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
{
|
||
let prev_mid_h = self.stg_prev_mid.host_slice_mut();
|
||
let tsv_h = self.stg_trade_signed_vol.host_slice_mut();
|
||
let tc_h = self.stg_trade_count.host_slice_mut();
|
||
let ts_ns_h = self.stg_ts_ns.host_slice_mut();
|
||
let prev_ts_ns_h = self.stg_prev_ts_ns.host_slice_mut();
|
||
for (b_idx, snapshots) in snapshots_batch.iter().enumerate() {
|
||
for (k, snap) in snapshots.iter().enumerate() {
|
||
let n = b_idx * k_seq + k;
|
||
prev_mid_h[n] = snap.prev_mid;
|
||
tsv_h[n] = snap.trade_signed_vol;
|
||
tc_h[n] = snap.trade_count as i32;
|
||
ts_ns_h[n] = snap.ts_ns as i64;
|
||
prev_ts_ns_h[n] = snap.prev_ts_ns as i64;
|
||
}
|
||
}
|
||
}
|
||
// Labels staging fill (host-only). Moved from inside dispatch
|
||
// to here so it happens BEFORE the captured region.
|
||
{
|
||
let dst = self.stg_labels.host_slice_mut();
|
||
for k in 0..k_seq {
|
||
for b_idx in 0..b_sz {
|
||
for h in 0..N_HORIZONS {
|
||
dst[(k * b_sz + b_idx) * N_HORIZONS + h] =
|
||
labels_batch[b_idx][k][h];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── 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
|
||
// `pearl_no_host_branches_in_captured_graph`.
|
||
if self.train_graph.is_some() {
|
||
self.train_graph
|
||
.as_ref()
|
||
.unwrap()
|
||
.launch()
|
||
.context("train graph launch")?;
|
||
} else if !self.cublas_warmed {
|
||
self.dispatch_train_step(b_sz, k_seq, total_snaps)
|
||
.context("train warmup dispatch")?;
|
||
self.cublas_warmed = true;
|
||
} else {
|
||
// Event tracking was disabled at trainer construction; the
|
||
// trainer's CudaSlices have no read/write events, so neither
|
||
// launch_builder nor device_ptr_mut() will try to insert
|
||
// cuStreamWaitEvent / event.record() calls inside the
|
||
// captured region. Capture safely.
|
||
let begin = self.stream.begin_capture(
|
||
CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
|
||
);
|
||
if let Err(e) = begin {
|
||
return Err(anyhow::anyhow!("train begin_capture: {e}"));
|
||
}
|
||
|
||
let dispatch_result = self.dispatch_train_step(b_sz, k_seq, total_snaps);
|
||
|
||
let graph_result = self.stream.end_capture(
|
||
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||
);
|
||
|
||
dispatch_result.context("train graph dispatch (during capture)")?;
|
||
let graph = graph_result
|
||
.context("train graph end_capture")?
|
||
.ok_or_else(|| anyhow::anyhow!(
|
||
"train end_capture returned None — no work captured"
|
||
))?;
|
||
self.train_graph = Some(graph);
|
||
}
|
||
|
||
// ── 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)
|
||
}
|
||
|
||
/// Kernel-dispatch portion of a training step — captured into the
|
||
/// CUDA Graph on the second `step_batched` call.
|
||
fn dispatch_train_step(
|
||
&mut self,
|
||
b_sz: usize,
|
||
k_seq: usize,
|
||
total_snaps: usize,
|
||
) -> Result<()> {
|
||
// ── GPU log ring tick (feature: kernel-step-trace) ──
|
||
// Runs FIRST inside the captured graph so every downstream
|
||
// logging-capable kernel sees the new step counter. Single-
|
||
// thread/single-block kernel: no host malloc, no host branch
|
||
// (per `pearl_cudarc_disable_event_tracking_for_graph_capture`
|
||
// and `pearl_no_host_branches_in_captured_graph`).
|
||
//
|
||
// Activated only when the runtime trace path is Some — i.e.
|
||
// when the trainer was constructed with
|
||
// `cfg.kernel_step_trace_path = Some(_)`. With path = None all
|
||
// the trace fields are None and this block is skipped, so the
|
||
// tick kernel never launches and zero overhead is incurred.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
if let (Some(tick_fn), Some(step_counter_d)) =
|
||
(self.gpu_log_tick_fn.as_ref(), self.log_step_counter_d.as_mut())
|
||
{
|
||
let tick_cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(tick_fn);
|
||
launch.arg(step_counter_d);
|
||
unsafe { launch.launch(tick_cfg).context("gpu_log_tick launch")?; }
|
||
}
|
||
|
||
// DtoD: staging (mapped-pinned, device-visible) → device buffers.
|
||
let n10 = total_snaps * 10;
|
||
let n6 = total_snaps * REGIME_DIM;
|
||
let n1 = total_snaps;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
// f32 arrays
|
||
let (d, _g) = self.bid_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_px_all.dev_ptr, n10 * 4, s)
|
||
.context("bid_px_all dtod")?;
|
||
let (d, _g) = self.bid_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("bid_sz_all dtod")?;
|
||
let (d, _g) = self.ask_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_px_all.dev_ptr, n10 * 4, s)
|
||
.context("ask_px_all dtod")?;
|
||
let (d, _g) = self.ask_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("ask_sz_all dtod")?;
|
||
let (d, _g) = self.regime_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_regime_all.dev_ptr, n6 * 4, s)
|
||
.context("regime_all dtod")?;
|
||
let (d, _g) = self.prev_mid_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_mid.dev_ptr, n1 * 4, s)
|
||
.context("prev_mid_all dtod")?;
|
||
let (d, _g) = self.trade_signed_vol_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_signed_vol.dev_ptr, n1 * 4, s)
|
||
.context("tsv_all dtod")?;
|
||
// i32 / i64 arrays
|
||
let (d, _g) = self.trade_count_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_count.dev_ptr, n1 * 4, s)
|
||
.context("trade_count_all dtod")?;
|
||
let (d, _g) = self.ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("ts_ns_all dtod")?;
|
||
let (d, _g) = self.prev_ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("prev_ts_ns_all dtod")?;
|
||
}
|
||
|
||
// Single batched snap_feature_assemble_batched launch — writes
|
||
// directly into window_tensor's storage at [b_idx * K + k] * 32.
|
||
let tick_size = ES_TICK_SIZE;
|
||
let n_total_i32 = total_snaps as i32;
|
||
let snap_block: u32 = 128;
|
||
let snap_grid: u32 = (total_snaps as u32).div_ceil(snap_block);
|
||
let snap_cfg = LaunchConfig {
|
||
grid_dim: (snap_grid, 1, 1),
|
||
block_dim: (snap_block, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.snap_batched_fn);
|
||
launch
|
||
.arg(&self.bid_px_all_d).arg(&self.bid_sz_all_d)
|
||
.arg(&self.ask_px_all_d).arg(&self.ask_sz_all_d)
|
||
.arg(&self.prev_bid_sz_all_d).arg(&self.prev_ask_sz_all_d)
|
||
.arg(&self.regime_all_d)
|
||
.arg(&self.prev_mid_all_d).arg(&self.trade_signed_vol_all_d)
|
||
.arg(&self.trade_count_all_d)
|
||
.arg(&self.ts_ns_all_d).arg(&self.prev_ts_ns_all_d)
|
||
.arg(&tick_size).arg(&n_total_i32)
|
||
.arg(self.window_tensor_d.data_mut());
|
||
launch.launch(snap_cfg).context("snap_batched fwd")?;
|
||
}
|
||
// No sync — stream-ordered. A sync would trip
|
||
// STREAM_CAPTURE_ISOLATION inside the captured region.
|
||
|
||
// ── 1b. VSN forward — per-position softmax-normalised feature
|
||
// gating over `window_tensor_d` [B, K, FEATURE_DIM].
|
||
// Writes gated features to `vsn_out_d` and saved
|
||
// post-softmax gates to `vsn_gates_d` for backward.
|
||
// Mamba2 then consumes `vsn_out_d` instead of the raw
|
||
// snap features.
|
||
{
|
||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_vsn = LaunchConfig {
|
||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||
block_dim: (64, 1, 1), // VSN_BLOCK = 64 (1 warp + idle threads)
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.vsn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.vsn_w_d)
|
||
.arg(&self.trunk.vsn_b_d)
|
||
.arg(self.window_tensor_d.cuda_data())
|
||
.arg(&n_rows_vsn)
|
||
.arg(self.vsn_out_d.data_mut())
|
||
.arg(&mut self.vsn_gates_d);
|
||
unsafe { launch.launch(cfg_vsn).context("variable_selection_fwd")?; }
|
||
}
|
||
|
||
// ── 2. Mamba2 stack-1 forward — writes into mamba2_fwd_scratch.
|
||
// Consumes vsn_out_d (gated features). in_dim = FEATURE_DIM (40).
|
||
self.trunk
|
||
.mamba2_l1_mut()
|
||
.forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)
|
||
.context("mamba2 (l1) forward_train_seq_into")?;
|
||
|
||
// ── 2a. LayerNorm A — between m1 and m2 (Phase 2B). Reads
|
||
// m1.h_enriched_seq, writes ln_a_out_d.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_a_gain_d)
|
||
.arg(&self.trunk.ln_a_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(self.ln_a_out_d.data_mut())
|
||
.arg(&mut self.ln_a_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("layer_norm_fwd (LN_a)")?; }
|
||
}
|
||
|
||
// ── 2b. Mamba2 stack-2 forward — consumes ln_a_out_d. in_dim = HIDDEN_DIM.
|
||
self.trunk
|
||
.mamba2_l2_mut()
|
||
.forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)
|
||
.context("mamba2 (l2) forward_train_seq_into")?;
|
||
|
||
// ── 2c. LayerNorm B — between m2 and CfC. Stabilises the
|
||
// trunk-to-CfC distribution. One block per row,
|
||
// block-wide tree-reduce for mean/var. Stats saved for
|
||
// backward in ln_stats_d. Reads m2's h_enriched_seq.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_l2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_b_gain_d)
|
||
.arg(&self.trunk.ln_b_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.ln_out_d)
|
||
.arg(&mut self.ln_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("layer_norm_fwd (LN_b)")?; }
|
||
}
|
||
|
||
// ── 2b. Transpose LN-normalised h_enriched [B, K, H] → [K, B, H]
|
||
// into pre-allocated buffer (source = ln_out_d, NOT raw
|
||
// Mamba2 output).
|
||
{
|
||
let block_n3: u32 = 32;
|
||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||
let cfg_tx = LaunchConfig {
|
||
grid_dim: (b_sz as u32, k_seq as u32, grid_z),
|
||
block_dim: (block_n3, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n1 = b_sz as i32;
|
||
let n2 = k_seq as i32;
|
||
let n3 = HIDDEN_DIM as i32;
|
||
let mut launch = self.stream.launch_builder(&self.trunk.transpose_3d_fn);
|
||
launch
|
||
.arg(&self.ln_out_d)
|
||
.arg(self.h_enriched_seq_t_d.data_mut())
|
||
.arg(&n1).arg(&n2).arg(&n3);
|
||
unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; }
|
||
}
|
||
|
||
// ── 2d. Attention pool forward (Phase 3) — produces
|
||
// attn_context_d [B, HIDDEN_DIM] = learned content-summary
|
||
// over all K LN_b output positions. Replaces CfC's
|
||
// zero-initialised h_old at k=0 with this context vector.
|
||
// Shared mem: K floats (scores) + BLOCK floats (reduce) +
|
||
// HIDDEN_DIM floats (context) = (K + 128 + 128) * 4 bytes.
|
||
//
|
||
// Random-sampling training always re-bootstraps the pooled
|
||
// context at sequence start — sequences are independent draws,
|
||
// so the previous step's h_old/attn context carries no useful
|
||
// information across the boundary.
|
||
let need_attn_pool_bootstrap = true;
|
||
if need_attn_pool_bootstrap {
|
||
let k_i32 = k_seq as i32;
|
||
let n_batch_attn = b_sz as i32;
|
||
let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::<f32>();
|
||
let cfg_attn = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (128, 1, 1), // ATTN_BLOCK
|
||
shared_mem_bytes: shared as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.attn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.attn_q_d)
|
||
.arg(&self.ln_out_d)
|
||
.arg(&n_batch_attn).arg(&k_i32)
|
||
.arg(&mut self.attn_context_d)
|
||
.arg(&mut self.attn_weights_d);
|
||
unsafe { launch.launch(cfg_attn).context("attention_pool_fwd")?; }
|
||
}
|
||
|
||
// ── 3. Labels DtoD: staging (filled in step_batched before
|
||
// captured region) → device.
|
||
let total_labels = k_seq * b_sz * N_HORIZONS;
|
||
unsafe {
|
||
let (dst_ptr, _g) = self.labels_per_k_d.device_ptr_mut(&self.stream);
|
||
let nbytes = total_labels * std::mem::size_of::<f32>();
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr, self.stg_labels.dev_ptr, nbytes, self.stream.cu_stream(),
|
||
).context("labels upload DtoD")?;
|
||
}
|
||
|
||
// CfC per-batch grad scratch (Phase B): zero ONCE per step; K-loop
|
||
// bwd accumulates into these, then reduce_axis0 collapses → final
|
||
// grad buffers (OVERWRITE) after the K-loop. AdamW reads the final
|
||
// grads (must run AFTER the reducer).
|
||
self.stream.memset_zeros(&mut self.cfc_grad_w_in_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_w_in_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.cfc_grad_w_rec_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_w_rec_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.cfc_grad_b_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_b_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.cfc_grad_tau_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_tau_scratch: {e}"))?;
|
||
// Phase B commit 4: attention pool per-batch grad_Q scratch.
|
||
self.stream.memset_zeros(&mut self.attn_grad_q_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero attn_grad_q_scratch: {e}"))?;
|
||
// GRN per-batch grad scratch (Phase B commit 2): zero ONCE per
|
||
// step; K-loop bwd accumulates into these, then reduce_axis0
|
||
// collapses → final grad buffers (OVERWRITE) after the K-loop.
|
||
self.stream.memset_zeros(&mut self.grn_grad_w1_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_w1_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_b1_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_b1_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_w2_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_w2_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_b2_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_b2_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_w_gate_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_w_gate_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_b_gate_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_b_gate_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_w_main_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_w_main_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_b_main_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_b_main_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_w_skip_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_w_skip_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.grn_grad_b_skip_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grn_grad_b_skip_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.zero_h_d)
|
||
.map_err(|e| anyhow::anyhow!("zero zero_h: {e}"))?;
|
||
|
||
// A1: decision_stride removed; CRT runs at event rate, dt = 1.0 per step.
|
||
let dt_s = 1.0_f32;
|
||
let n_in_i = HIDDEN_DIM as i32;
|
||
let n_hid_i = HIDDEN_DIM as i32;
|
||
let block_dim = 128u32;
|
||
let _ = block_dim; // legacy 1D layout helper — superseded by block-per-batch launches.
|
||
// CfC fwd: block-per-batch (Phase B). One block per sample;
|
||
// each block's threads parallelize over n_hid channels.
|
||
// Shared layout: s_x [n_in] + s_h_old [n_hid].
|
||
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_cfc = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: cfc_fwd_smem,
|
||
};
|
||
let cfg_heads = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (N_HORIZONS as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
// CfC bwd (Phase B): block-per-batch + 2 * n_hid floats of shared mem
|
||
// (one sd_pre row + sdecay for the block's single bi).
|
||
// Shared layout: sd_pre [n_hid] + sdecay [n_hid] + s_x [n_in] + s_h_old [n_hid].
|
||
// With n_in == n_hid == HIDDEN_DIM, that's 4 × HIDDEN_DIM floats per block.
|
||
let cfc_bwd_smem = (4 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_cfc_bwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: cfc_bwd_smem,
|
||
};
|
||
// Batched heads backward shared mem: sd_z[B, 5].
|
||
let heads_bwd_smem = (b_sz * N_HORIZONS * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_heads_bwd = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: heads_bwd_smem,
|
||
};
|
||
// Heads forward uses block of 128 too (5 threads compute heads,
|
||
// remaining are idle — same kernel layout as backward).
|
||
let cfg_heads_batched_fwd = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let _ = cfg_heads;
|
||
|
||
// Per-K slot strides for [K, B, H] / [K, B, N_HORIZONS] layout.
|
||
let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::<f32>();
|
||
// GRN per-K intermediate stride: [B, N_HORIZONS, HEAD_MID].
|
||
let kb_nh_mid_bytes = b_sz * N_HORIZONS * HEAD_MID_DIM * std::mem::size_of::<f32>();
|
||
let n_batch_i = b_sz as i32;
|
||
|
||
// GRN forward launch config: block-per-sample, threads tile HEAD_MID.
|
||
let cfg_grn_fwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HEAD_MID_DIM as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
// GRN backward launch config: block-per-batch (Phase B). One block
|
||
// per sample; threads = HIDDEN so Pass 5/6 can partition over the
|
||
// output i-dim for COALESCED writes. Passes 2/3/4 use first
|
||
// HEAD_MID threads only; the rest idle until Pass 5.
|
||
// Per-batch scratch + reducer collapses B → final grad after the K-loop.
|
||
let cfg_grn_bwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let _ = cfg_heads_batched_fwd; // GRN fwd uses cfg_grn_fwd now.
|
||
let _ = cfg_heads_bwd; // GRN bwd uses cfg_grn_bwd now.
|
||
|
||
// ── 4. Stream-ordered forward K loop. CfC is RECURRENT —
|
||
// h_old at step k for sample b IS h_new at step k-1 for
|
||
// the same b. Per-K slot is contiguous [B, H] in the
|
||
// [K, B, H] layout. Pointer-offset trick (single mut
|
||
// borrow + raw u64 arithmetic for slot pointers) keeps
|
||
// the launches stream-ordered with no syncs.
|
||
// Phase 3: attention pool produces attn_context_d which replaces
|
||
// the zero initial h_old at k=0. zero_h_ptr is no longer used in
|
||
// the fwd K-loop but is retained for the bwd k=0 case (cfc bwd
|
||
// needs an h_old slot for the input gradient to read from).
|
||
let _zero_h_ptr_unused = {
|
||
let (p, _g) = self.zero_h_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let attn_context_ptr = {
|
||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let henr_t_base = {
|
||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||
let (probs_base, _g_probs) = self.probs_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z1_base, _g_z1) = self.z1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (a1_base, _g_a1) = self.a1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z2_base, _g_z2) = self.z2_per_k_d.device_ptr_mut(&self.stream);
|
||
let (gate_base, _g_gate) = self.gate_logit_per_k_d.device_ptr_mut(&self.stream);
|
||
let (main_base, _g_main) = self.main_per_k_d.device_ptr_mut(&self.stream);
|
||
let (logit_base, _g_logit) = self.logit_per_k_d.device_ptr_mut(&self.stream);
|
||
|
||
for k in 0..k_seq {
|
||
let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64;
|
||
let h_old_k_ptr = if k == 0 {
|
||
attn_context_ptr
|
||
} else {
|
||
h_per_k_base + ((k - 1) * kb_hid_bytes) as u64
|
||
};
|
||
let x_k_ptr = henr_t_base + (k * kb_hid_bytes) as u64;
|
||
let probs_k_ptr = probs_base + (k * kb_nh_bytes) as u64;
|
||
let z1_k_ptr = z1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let a1_k_ptr = a1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let z2_k_ptr = z2_base + (k * kb_nh_mid_bytes) as u64;
|
||
let gate_k_ptr = gate_base + (k * kb_nh_bytes) as u64;
|
||
let main_k_ptr = main_base + (k * kb_nh_bytes) as u64;
|
||
let logit_k_ptr = logit_base + (k * kb_nh_bytes) as u64;
|
||
|
||
// CfC forward dispatch, two-phase per spec §2.3 + Task 10:
|
||
//
|
||
// Phase 1 (warmup): existing single shared `cfc_step_batched`
|
||
// kernel — block-per-batch over all HIDDEN_DIM channels.
|
||
//
|
||
// Phase 2 (routed): fused per-branch kernel
|
||
// `cfc_step_per_branch_fwd` — grid=(B, N_HORIZONS=5, 1),
|
||
// block=(MAX_BUCKET_DIM=28, 1, 1) uniform predicate, with
|
||
// `bucket_channel_offset_d` / `bucket_dim_k_d` routing each
|
||
// block to its quintile of HIDDEN_DIM. Spec §5.4 point 1.
|
||
//
|
||
// The branch on `self.phase` lives OUTSIDE the captured graph's
|
||
// hot kernel scope: the captured train graph is recaptured at
|
||
// the Phase 1→2 transition (cf. `pearl_no_host_branches_in_captured_graph`).
|
||
// Once Phase 2 has been entered, every subsequent capture replays
|
||
// the routed path; the if-branch here is evaluated by host code
|
||
// building the graph, not inside the device-side kernel stream.
|
||
unsafe {
|
||
match self.phase {
|
||
TrainingPhase::Phase1Warmup => {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&x_k_ptr).arg(&h_old_k_ptr)
|
||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
|
||
.arg(&h_new_k_ptr);
|
||
launch.launch(cfg_cfc).context("cfc fwd k batched")?;
|
||
}
|
||
TrainingPhase::Phase2Routed => {
|
||
// grid=(B, N_HORIZONS, 1), block=(MAX_BUCKET_DIM, 1, 1).
|
||
// Shared mem = 2 × HIDDEN_DIM floats (x_local + h_old_local
|
||
// cooperative staging per pearl_cooperative_staging_eliminates_redundant_reads).
|
||
const N_HORIZONS_U32: u32 = N_HORIZONS as u32;
|
||
const MAX_BUCKET_DIM: u32 = 28;
|
||
let cfg_cfc_per_branch = LaunchConfig {
|
||
grid_dim: (b_sz as u32, N_HORIZONS_U32, 1),
|
||
block_dim: (MAX_BUCKET_DIM, 1, 1),
|
||
shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.cfc_step_per_branch_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&x_k_ptr).arg(&h_old_k_ptr)
|
||
.arg(&dt_s).arg(&n_batch_i)
|
||
.arg(&self.trunk.bucket_channel_offset_d).arg(&self.trunk.bucket_dim_k_d)
|
||
.arg(&h_new_k_ptr);
|
||
launch.launch(cfg_cfc_per_branch).context("cfc_step_per_branch_fwd k")?;
|
||
}
|
||
}
|
||
}
|
||
|
||
// multi_horizon_heads_grn_fwd_batched(10 params + h[B*128] + n_batch
|
||
// → probs[B*5] + z1/a1/z2 [B,5,HEAD_MID] + gate_logit/main/logit [B,5]).
|
||
//
|
||
// Task 10 scope note (Option B): the heads dispatch stays at the
|
||
// existing full-`heads_w_skip_d` GRN kernel for BOTH phases. The
|
||
// Phase 2 compact-w_skip path (`heads_block_diagonal_fwd`) and
|
||
// its 128-float `heads_w_skip_compact_d` storage are populated
|
||
// by the transition kernels (Task 9) but NOT yet consumed here:
|
||
// the GRN kernel computes the full per-horizon logit including
|
||
// skip + gate + main sub-components, and routing only the skip
|
||
// term through the compact path requires integration with
|
||
// GRN's main/gate computations (out of Task 10 scope per
|
||
// `feedback_no_partial_refactor`). A follow-up task wires the
|
||
// compact heads dispatch through GRN; the trunk handle
|
||
// `heads_block_diagonal_fwd_fn` is already loaded for that.
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d)
|
||
.arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d)
|
||
.arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d)
|
||
.arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d)
|
||
.arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d)
|
||
.arg(&h_new_k_ptr).arg(&n_batch_i)
|
||
.arg(&probs_k_ptr)
|
||
.arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr)
|
||
.arg(&gate_k_ptr).arg(&main_k_ptr).arg(&logit_k_ptr);
|
||
launch.launch(cfg_grn_fwd).context("heads GRN fwd k")?;
|
||
}
|
||
}
|
||
drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit));
|
||
|
||
// σ grad scratch: BCE kernel writes overwrite-style; no zeroing needed.
|
||
// The output is unused downstream (σ is ISV-driven via horizon_ema_and_lambda).
|
||
|
||
// ── 5. Fused multi-horizon BCE (Kendall σ-weighted, axis A).
|
||
// The kernel doesn't distinguish position-vs-batch
|
||
// since the per-horizon weight is identified by `i % n_horizons`.
|
||
let n_pos_i = (k_seq * b_sz) as i32;
|
||
let n_h_i = N_HORIZONS as i32;
|
||
let bce_cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
{
|
||
let mut launch = self.stream.launch_builder(&self.bce_fn);
|
||
launch
|
||
.arg(&self.probs_per_k_d).arg(&self.labels_per_k_d)
|
||
.arg(&self.loss_weights_d)
|
||
.arg(&self.log_sigma_h_d)
|
||
.arg(&n_pos_i).arg(&n_h_i)
|
||
.arg(&mut self.loss_d)
|
||
.arg(&mut self.loss_per_horizon_d)
|
||
.arg(&mut self.grad_probs_per_k_d)
|
||
.arg(&mut self.valid_d)
|
||
.arg(&mut self.grad_log_sigma_h_d);
|
||
unsafe { launch.launch(bce_cfg).context("bce launch")?; }
|
||
}
|
||
|
||
// ── 5c. CRT.train output-smoothness regularizer ──
|
||
// λ[h]-weighted Σ (p[k]-p[k-1])² across adjacent positions
|
||
// in each sequence. Accumulates the analytic gradient INTO
|
||
// `self.grad_probs_per_k_d` (which BCE just overwrote);
|
||
// the heads_bwd K-loop below consumes the sum.
|
||
//
|
||
// Per `pearl_no_host_branches_in_captured_graph`: kernel
|
||
// launch only — all scalars (k_i, b_i) and buffer pointers
|
||
// are fixed across captures (seq_len and batch_size are
|
||
// config-immutable). λ is uploaded once at construction;
|
||
// when `cfg.smoothness_base_lambda == 0.0` the kernel
|
||
// produces zero loss + zero grad, behaviorally identical
|
||
// to the pre-intervention path.
|
||
let k_seq_i = k_seq as i32;
|
||
let b_sz_i = b_sz as i32;
|
||
let smooth_cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
{
|
||
let mut launch = self.stream.launch_builder(&self.smoothness_fn);
|
||
launch
|
||
.arg(&self.probs_per_k_d)
|
||
.arg(&self.smoothness_lambda_d)
|
||
.arg(&k_seq_i).arg(&b_sz_i)
|
||
.arg(&mut self.smoothness_loss_d)
|
||
.arg(&mut self.smoothness_loss_per_horizon_d)
|
||
.arg(&mut self.grad_probs_per_k_d);
|
||
unsafe { launch.launch(smooth_cfg).context("output_smoothness launch")?; }
|
||
}
|
||
|
||
// ── 5d. ISV-driven λ controller ──
|
||
// Reads the raw per-horizon mean-sq-diff just emitted by
|
||
// the output_smoothness kernel, updates jitter_ema, derives
|
||
// per-horizon target anchored on h30, and produces λ for
|
||
// the NEXT step's output_smoothness launch. Same captured
|
||
// graph; one-step delay between observation and effect
|
||
// (standard closed-loop pattern).
|
||
let base_lambda = self.cfg.smoothness_base_lambda;
|
||
let smooth_ctrl_cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (N_HORIZONS as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
// GPU log ring pointers. When the feature is OFF, both are 0
|
||
// (null). When the feature is ON but the runtime trace path is
|
||
// None (no `--kernel-step-trace`), the ring + step counter
|
||
// Options are also None — pass null and the kernel's
|
||
// `log_record()` short-circuits at the nullptr guard, behaviorally
|
||
// identical to the pre-log kernel.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self
|
||
.log_ring
|
||
.as_ref()
|
||
.map(|r| r.dev_ptr())
|
||
.unwrap_or(0);
|
||
#[cfg(not(feature = "kernel-step-trace"))]
|
||
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = 0;
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = match self.log_step_counter_d.as_ref() {
|
||
Some(slice) => {
|
||
let (p, _g) = slice.device_ptr(&self.stream);
|
||
p
|
||
}
|
||
None => 0,
|
||
};
|
||
#[cfg(not(feature = "kernel-step-trace"))]
|
||
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = 0;
|
||
{
|
||
let mut launch = self.stream.launch_builder(&self.smoothness_controller_fn);
|
||
launch
|
||
.arg(&self.smoothness_loss_per_horizon_d)
|
||
.arg(&mut self.smoothness_jitter_ema_d)
|
||
.arg(&mut self.smoothness_jitter_first_obs_d)
|
||
.arg(&base_lambda)
|
||
.arg(&mut self.smoothness_lambda_d)
|
||
.arg(&log_ring_ptr)
|
||
.arg(&log_step_ptr);
|
||
unsafe { launch.launch(smooth_ctrl_cfg).context("smoothness_lambda_controller launch")?; }
|
||
}
|
||
|
||
// ── 5a. (v2 reserved) — per-horizon Q_h backward removed at V1.
|
||
// v2 backward path will live here in commit V10.
|
||
|
||
// ── 5b. ISV-driven per-horizon controllers. Single kernel emits
|
||
// both:
|
||
// lambda_d[h] — λ multiplier on trunk grad_h
|
||
// log_sigma_h_d[h] — Kendall σ for next-step BCE loss
|
||
// Both anchor on loss_ema (EMA of unweighted per-horizon
|
||
// BCE). σ in closed form (Kendall equilibrium), λ as
|
||
// z-score-scaled multiplier with adaptive Z_SCALE_ISV
|
||
// (derived from z_max_ema). No Adam involvement in σ
|
||
// (per pearl_adam_normalizes_loss_weights).
|
||
{
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.horizon_lambda_fn);
|
||
launch
|
||
.arg(&self.loss_per_horizon_d)
|
||
.arg(&mut self.loss_ema_d)
|
||
.arg(&mut self.z_max_ema_d)
|
||
.arg(&mut self.lambda_d)
|
||
.arg(&mut self.log_sigma_h_d);
|
||
unsafe { launch.launch(cfg).context("horizon_ema_and_lambda launch")?; }
|
||
}
|
||
|
||
// ── 6. Reverse-order backward K loop using pre-allocated
|
||
// grad_h_enriched_seq_t_d as the per-K slot output.
|
||
self.stream.memset_zeros(&mut self.grad_h_carry_d)
|
||
.map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?;
|
||
|
||
// Phase 3: at k=0 the bwd kernel reads `h_old` = attn_context_d
|
||
// (mirroring the forward pass). zero_h_ptr_bwd retained as a
|
||
// legacy fallback / unused alias.
|
||
let _zero_h_ptr_bwd_unused = {
|
||
let (p, _g) = self.zero_h_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let attn_context_ptr_bwd = {
|
||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let henr_t_base_bwd = {
|
||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let (h_per_k_base_bwd, _g_hpk_bwd) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||
let (probs_base_bwd, _g_probs_bwd) = self.probs_per_k_d.device_ptr_mut(&self.stream);
|
||
let (gprobs_base_bwd, _g_gprobs_bwd) = self.grad_probs_per_k_d.device_ptr_mut(&self.stream);
|
||
let (grad_henr_t_base, _g_ghen_t_mut) = self.grad_h_enriched_seq_t_d.data_mut().device_ptr_mut(&self.stream);
|
||
let (z1_base_bwd, _g_z1_bwd) = self.z1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (a1_base_bwd, _g_a1_bwd) = self.a1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z2_base_bwd, _g_z2_bwd) = self.z2_per_k_d.device_ptr_mut(&self.stream);
|
||
let (gate_base_bwd, _g_gate_bwd) = self.gate_logit_per_k_d.device_ptr_mut(&self.stream);
|
||
let (main_base_bwd, _g_main_bwd) = self.main_per_k_d.device_ptr_mut(&self.stream);
|
||
|
||
for k in (0..k_seq).rev() {
|
||
let h_new_k_ptr = h_per_k_base_bwd + (k * kb_hid_bytes) as u64;
|
||
let x_k_ptr = henr_t_base_bwd + (k * kb_hid_bytes) as u64;
|
||
let probs_k_ptr = probs_base_bwd + (k * kb_nh_bytes) as u64;
|
||
let gprobs_k_ptr = gprobs_base_bwd + (k * kb_nh_bytes) as u64;
|
||
let z1_k_ptr = z1_base_bwd + (k * kb_nh_mid_bytes) as u64;
|
||
let a1_k_ptr = a1_base_bwd + (k * kb_nh_mid_bytes) as u64;
|
||
let z2_k_ptr = z2_base_bwd + (k * kb_nh_mid_bytes) as u64;
|
||
let gate_k_ptr = gate_base_bwd + (k * kb_nh_bytes) as u64;
|
||
let main_k_ptr = main_base_bwd + (k * kb_nh_bytes) as u64;
|
||
let h_old_k_ptr = if k == 0 {
|
||
attn_context_ptr_bwd
|
||
} else {
|
||
h_per_k_base_bwd + ((k - 1) * kb_hid_bytes) as u64
|
||
};
|
||
let grad_henr_k_ptr = grad_henr_t_base + (k * kb_hid_bytes) as u64;
|
||
|
||
// GRN backward: chain rule through sigmoid → (skip + sigmoid(gate)*main)
|
||
// → GLU → linear → GELU → linear → trunk. `lambda_d` scales
|
||
// ONLY the trunk gradient; per-horizon param grads are
|
||
// unscaled (per pearl_adam_normalizes_loss_weights.md the
|
||
// effective lever is the trunk gradient, not loss weight).
|
||
// Phase B: per-batch GRN grad scratch + per-batch grad_h_new.
|
||
// K-loop's 64 invocations accumulate into the scratch via +=;
|
||
// reduce_axis0 collapses → final grad after the K-loop.
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.heads_grn_bwd_fn);
|
||
launch
|
||
.arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_w2_d)
|
||
.arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_w_main_d)
|
||
.arg(&self.trunk.heads_w_skip_d)
|
||
.arg(&probs_k_ptr).arg(&gprobs_k_ptr)
|
||
.arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr)
|
||
.arg(&gate_k_ptr).arg(&main_k_ptr)
|
||
.arg(&h_new_k_ptr)
|
||
.arg(&self.grad_h_carry_d)
|
||
.arg(&self.lambda_d)
|
||
.arg(&n_batch_i)
|
||
.arg(&mut self.grn_grad_w1_scratch_d).arg(&mut self.grn_grad_b1_scratch_d)
|
||
.arg(&mut self.grn_grad_w2_scratch_d).arg(&mut self.grn_grad_b2_scratch_d)
|
||
.arg(&mut self.grn_grad_w_gate_scratch_d).arg(&mut self.grn_grad_b_gate_scratch_d)
|
||
.arg(&mut self.grn_grad_w_main_scratch_d).arg(&mut self.grn_grad_b_main_scratch_d)
|
||
.arg(&mut self.grn_grad_w_skip_scratch_d).arg(&mut self.grn_grad_b_skip_scratch_d)
|
||
.arg(&mut self.grad_h_new_d);
|
||
launch.launch(cfg_grn_bwd).context("heads GRN bwd k")?;
|
||
}
|
||
|
||
// cfc_step_bwd_batched: writes grad_h_carry (= grad_h_old_out for
|
||
// position k-1) and grad_x[B, n_in]. grad_x output buffer IS
|
||
// grad_henr_k_ptr — kernel writes directly into the K-slot.
|
||
// Phase B: per-batch param-grad scratch + per-batch grad_h_old / grad_x.
|
||
// Writes to scratch via += across K iterations; reduce_axis0
|
||
// collapses B → final grad after the K-loop (see Pass 8c).
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&x_k_ptr).arg(&h_old_k_ptr).arg(&self.grad_h_new_d)
|
||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
|
||
.arg(&mut self.cfc_grad_w_in_scratch_d)
|
||
.arg(&mut self.cfc_grad_w_rec_scratch_d)
|
||
.arg(&mut self.cfc_grad_b_scratch_d)
|
||
.arg(&mut self.cfc_grad_tau_scratch_d)
|
||
.arg(&mut self.grad_h_carry_d)
|
||
.arg(&grad_henr_k_ptr);
|
||
launch.launch(cfg_cfc_bwd).context("cfc bwd k batched")?;
|
||
}
|
||
}
|
||
drop((_g_hpk_bwd, _g_probs_bwd, _g_gprobs_bwd, _g_ghen_t_mut,
|
||
_g_z1_bwd, _g_a1_bwd, _g_z2_bwd, _g_gate_bwd, _g_main_bwd));
|
||
|
||
// ── 7. Stream stays async — bwd loop kernels, transposes, and
|
||
// Mamba2 bwd are all sequential on the same stream, no
|
||
// sync needed between them. Single sync at the very end
|
||
// to flush before reading the mapped-pinned loss shadow.
|
||
|
||
// ── 7b. Transpose grad_h_enriched_seq_t_d [K, B, H] → [B, K, H]
|
||
// (pre-allocated grad_h_enriched_seq_d). This is the
|
||
// gradient w.r.t. the LN OUTPUT (= grad_y for LN bwd).
|
||
{
|
||
let block_n3: u32 = 32;
|
||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||
let cfg_tx = LaunchConfig {
|
||
grid_dim: (k_seq as u32, b_sz as u32, grid_z),
|
||
block_dim: (block_n3, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n1 = k_seq as i32;
|
||
let n2 = b_sz as i32;
|
||
let n3 = HIDDEN_DIM as i32;
|
||
let mut launch = self.stream.launch_builder(&self.trunk.transpose_3d_fn);
|
||
launch
|
||
.arg(self.grad_h_enriched_seq_t_d.cuda_data())
|
||
.arg(self.grad_h_enriched_seq_d.data_mut())
|
||
.arg(&n1).arg(&n2).arg(&n3);
|
||
unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; }
|
||
}
|
||
|
||
// ── 7c-pre. Attention pool backward (Phase 3). Consumes:
|
||
// Q = self.attn_q_d [HIDDEN_DIM]
|
||
// ln_out (values)= self.ln_out_d [B, K, HIDDEN_DIM]
|
||
// attn_weights = self.attn_weights_d (saved by fwd) [B, K]
|
||
// grad_context = self.grad_h_carry_d (= grad on initial
|
||
// h_old at k=0, which IS attn_context) [B, HIDDEN_DIM]
|
||
// Writes (BOTH ARE +=):
|
||
// grad_attn_q_d += attn-path contribution to Q
|
||
// grad_h_enriched_seq_d (LN_b output grad) += attn-path
|
||
// contribution to ln_out
|
||
// The pre-zero of grad_attn_q_d at step start makes the += a
|
||
// clean overwrite for the Q grad. grad_h_enriched_seq_d already
|
||
// holds the K-loop's contribution at this point — attn's
|
||
// contribution adds on top.
|
||
// Phase B commit 4: block-per-batch attn bwd writes per-batch
|
||
// grad_Q scratch; reducer collapses → final grad_attn_q_d below.
|
||
{
|
||
let k_i32 = k_seq as i32;
|
||
let n_batch_attn = b_sz as i32;
|
||
let shared = (2 * k_seq + 128) * std::mem::size_of::<f32>();
|
||
let cfg_attn_bwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (128, 1, 1), // ATTN_BLOCK
|
||
shared_mem_bytes: shared as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.attn_bwd_fn);
|
||
launch
|
||
.arg(&self.trunk.attn_q_d)
|
||
.arg(&self.ln_out_d)
|
||
.arg(&self.attn_weights_d)
|
||
.arg(&self.grad_h_carry_d)
|
||
.arg(&n_batch_attn).arg(&k_i32)
|
||
.arg(&mut self.attn_grad_q_scratch_d)
|
||
.arg(self.grad_h_enriched_seq_d.data_mut());
|
||
unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; }
|
||
}
|
||
// Attn pool reducer: collapse [B, HIDDEN_DIM] → [HIDDEN_DIM].
|
||
// reduce_axis0: block tiles 32 cols × 8 batches per launch.
|
||
{
|
||
let n_batch_i = b_sz as i32;
|
||
let n_tail_i = HIDDEN_DIM as i32;
|
||
let cfg_red = LaunchConfig {
|
||
grid_dim: (((HIDDEN_DIM + 31) / 32) as u32, 1, 1),
|
||
block_dim: (32, 8, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||
launch
|
||
.arg(&self.attn_grad_q_scratch_d)
|
||
.arg(&n_batch_i)
|
||
.arg(&n_tail_i)
|
||
.arg(&mut self.grad_attn_q_d);
|
||
unsafe { launch.launch(cfg_red).context("reduce attn_grad_q")?; }
|
||
}
|
||
|
||
// ── 7c. LayerNorm B backward (between m2 and CfC). Consumes:
|
||
// x = m2.h_enriched_seq [B, K, H]
|
||
// gain = self.ln_gain_d [H]
|
||
// stats = self.ln_stats_d (from fwd) [B*K, 2]
|
||
// grad_y = self.grad_h_enriched_seq_d [B, K, H]
|
||
// Produces:
|
||
// grad_x = self.grad_ln_in_d [B, K, H]
|
||
// grad_gain/row = self.grad_ln_gain_per_row_d [B*K, H]
|
||
// grad_bias/row = self.grad_ln_bias_per_row_d [B*K, H]
|
||
// Then two reducer launches collapse the per-row
|
||
// scratches into the final [H] param grad buffers
|
||
// (no atomicAdd per feedback_no_atomicadd.md).
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.ln_bwd_fn);
|
||
launch
|
||
.arg(self.mamba2_l2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_b_gain_d)
|
||
.arg(&self.ln_stats_d)
|
||
.arg(self.grad_h_enriched_seq_d.cuda_data())
|
||
.arg(&n_rows_ln)
|
||
.arg(self.grad_ln_in_d.data_mut())
|
||
.arg(&mut self.grad_ln_gain_per_row_d)
|
||
.arg(&mut self.grad_ln_bias_per_row_d);
|
||
unsafe { launch.launch(cfg_ln).context("layer_norm_bwd (LN_b)")?; }
|
||
}
|
||
// 7c.i — reduce per-row grad_gain → [H].
|
||
{
|
||
let cfg_red = LaunchConfig {
|
||
grid_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
|
||
launch
|
||
.arg(&self.grad_ln_gain_per_row_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.grad_ln_gain_d);
|
||
unsafe { launch.launch(cfg_red).context("layer_norm_reduce gain")?; }
|
||
}
|
||
// 7c.ii — reduce per-row grad_bias → [H].
|
||
{
|
||
let cfg_red = LaunchConfig {
|
||
grid_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
|
||
launch
|
||
.arg(&self.grad_ln_bias_per_row_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.grad_ln_bias_d);
|
||
unsafe { launch.launch(cfg_red).context("layer_norm_reduce bias")?; }
|
||
}
|
||
|
||
// ── 8. Mamba2 stack-2 backward. Consumes:
|
||
// input = ln_a_out_d (m2's forward input)
|
||
// d_h_out = grad_ln_in_d (LN_b bwd output)
|
||
// Writes:
|
||
// mamba2_l2_grads_buffers.d_x_from_in = grad w.r.t. m2's
|
||
// input = grad w.r.t. LN_a output (consumed by LN_a bwd below).
|
||
self.trunk
|
||
.mamba2_l2_mut()
|
||
.backward_from_h_enriched_seq_full_into(
|
||
&self.ln_a_out_d,
|
||
&self.mamba2_l2_fwd_scratch,
|
||
&self.grad_ln_in_d,
|
||
&mut self.mamba2_l2_bwd_scratch,
|
||
&mut self.mamba2_l2_grads_buffers,
|
||
)
|
||
.context("mamba2 (l2) backward_from_h_enriched_seq_full_into")?;
|
||
|
||
// ── 8a. LN_a backward. Consumes:
|
||
// x = m1.h_enriched_seq [B, K, H]
|
||
// gain = ln_a_gain_d [H]
|
||
// stats = ln_a_stats_d (from fwd) [B*K, 2]
|
||
// grad_y = m2.d_x_from_in (reshape) [B, K, H] — read flat
|
||
// Writes:
|
||
// grad_x = grad_ln_a_in_d [B, K, H] (fed to m1.bwd)
|
||
// grad_gain/row = grad_ln_a_gain_per_row_d
|
||
// grad_bias/row = grad_ln_a_bias_per_row_d
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.ln_bwd_fn);
|
||
launch
|
||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_a_gain_d)
|
||
.arg(&self.ln_a_stats_d)
|
||
.arg(self.mamba2_l2_grads_buffers.d_x_from_in.cuda_data())
|
||
.arg(&n_rows_ln)
|
||
.arg(self.grad_ln_a_in_d.data_mut())
|
||
.arg(&mut self.grad_ln_a_gain_per_row_d)
|
||
.arg(&mut self.grad_ln_a_bias_per_row_d);
|
||
unsafe { launch.launch(cfg_ln).context("layer_norm_bwd (LN_a)")?; }
|
||
}
|
||
// 8a.i — reduce per-row grad_a_gain → [H].
|
||
{
|
||
let cfg_red = LaunchConfig {
|
||
grid_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
|
||
launch
|
||
.arg(&self.grad_ln_a_gain_per_row_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.grad_ln_a_gain_d);
|
||
unsafe { launch.launch(cfg_red).context("layer_norm_reduce gain (LN_a)")?; }
|
||
}
|
||
// 8a.ii — reduce per-row grad_a_bias → [H].
|
||
{
|
||
let cfg_red = LaunchConfig {
|
||
grid_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
|
||
launch
|
||
.arg(&self.grad_ln_a_bias_per_row_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.grad_ln_a_bias_d);
|
||
unsafe { launch.launch(cfg_red).context("layer_norm_reduce bias (LN_a)")?; }
|
||
}
|
||
|
||
// ── 8b. Mamba2 stack-1 backward. Consumes:
|
||
// input = vsn_out_d (m1's forward input)
|
||
// d_h_out = grad_ln_a_in_d (LN_a bwd output)
|
||
// Writes:
|
||
// mamba2_grads_buffers.d_x_from_in = grad on VSN output
|
||
// (consumed by VSN bwd below).
|
||
self.trunk
|
||
.mamba2_l1_mut()
|
||
.backward_from_h_enriched_seq_full_into(
|
||
&self.vsn_out_d,
|
||
&self.mamba2_fwd_scratch,
|
||
&self.grad_ln_a_in_d,
|
||
&mut self.mamba2_bwd_scratch,
|
||
&mut self.mamba2_grads_buffers,
|
||
)
|
||
.context("mamba2 (l1) backward_from_h_enriched_seq_full_into")?;
|
||
|
||
// ── 8b. VSN backward. Consumes:
|
||
// grad_y = mamba2_grads_buffers.d_x_from_in [B*K, FEATURE_DIM]
|
||
// x = window_tensor_d [B, K, FEATURE_DIM]
|
||
// gates = vsn_gates_d (saved by fwd) [B*K, FEATURE_DIM]
|
||
// Writes:
|
||
// grad_W_vsn, grad_b_vsn — Adam consumes.
|
||
// vsn_grad_x_d — discarded (snap_features non-trainable).
|
||
// Param grads OVERWRITE (kernel uses += but the prior
|
||
// memsets clear them at step start, so a single VSN bwd
|
||
// per step writes the full accumulation cleanly).
|
||
// Phase B commit 3: VSN per-row grad scratch + reducer.
|
||
self.stream.memset_zeros(&mut self.vsn_grad_w_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero vsn_grad_w_scratch: {e}"))?;
|
||
self.stream.memset_zeros(&mut self.vsn_grad_b_scratch_d)
|
||
.map_err(|e| anyhow::anyhow!("zero vsn_grad_b_scratch: {e}"))?;
|
||
{
|
||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_vsn_bwd = LaunchConfig {
|
||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||
block_dim: (64, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.vsn_bwd_fn);
|
||
launch
|
||
.arg(&self.trunk.vsn_w_d)
|
||
.arg(self.window_tensor_d.cuda_data())
|
||
.arg(&self.vsn_gates_d)
|
||
.arg(self.mamba2_grads_buffers.d_x_from_in.cuda_data())
|
||
.arg(&n_rows_vsn)
|
||
.arg(&mut self.vsn_grad_w_scratch_d)
|
||
.arg(&mut self.vsn_grad_b_scratch_d)
|
||
.arg(&mut self.vsn_grad_x_d);
|
||
unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; }
|
||
}
|
||
// VSN reducer: collapse n_rows (= B * K) → final grad buffers.
|
||
// reduce_axis0 tiles 32 columns per block, 8-way batch stride.
|
||
{
|
||
let n_rows_i = (b_sz * k_seq) as i32;
|
||
let n_tail_w_usz = FEATURE_DIM * FEATURE_DIM;
|
||
let cfg_red_w = LaunchConfig {
|
||
grid_dim: (((n_tail_w_usz + 31) / 32) as u32, 1, 1),
|
||
block_dim: (32, 8, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_tail_w = n_tail_w_usz as i32;
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||
launch
|
||
.arg(&self.vsn_grad_w_scratch_d)
|
||
.arg(&n_rows_i)
|
||
.arg(&n_tail_w)
|
||
.arg(&mut self.grad_vsn_w_d);
|
||
launch.launch(cfg_red_w).context("reduce vsn_grad_w")?;
|
||
}
|
||
let cfg_red_b = LaunchConfig {
|
||
grid_dim: (((FEATURE_DIM + 31) / 32) as u32, 1, 1),
|
||
block_dim: (32, 8, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_tail_b = FEATURE_DIM as i32;
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||
launch
|
||
.arg(&self.vsn_grad_b_scratch_d)
|
||
.arg(&n_rows_i)
|
||
.arg(&n_tail_b)
|
||
.arg(&mut self.grad_vsn_b_d);
|
||
launch.launch(cfg_red_b).context("reduce vsn_grad_b")?;
|
||
}
|
||
}
|
||
|
||
// ── 8c. (Phase B) Reduce CfC per-batch grad scratch → final grad buffers.
|
||
// Block tree-reduce across bi; deterministic sum order.
|
||
// Each launch is grid=(ceil(n_tail/32), 1, 1), block=(32, 8, 1)
|
||
// — 32-wide column tile per block, 8-way B-stride per thread.
|
||
// 4 launches, one per cfc param tensor. AdamW (below) reads
|
||
// the final grad buffers; MUST run after these reducers.
|
||
{
|
||
let n_batch_i = b_sz as i32;
|
||
let reduce_at = |n_tail: usize,
|
||
scratch: &CudaSlice<f32>,
|
||
out: &mut CudaSlice<f32>,
|
||
label: &'static str|
|
||
-> Result<()> {
|
||
let cfg = LaunchConfig {
|
||
grid_dim: (((n_tail + 31) / 32) as u32, 1, 1),
|
||
block_dim: (32, 8, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n_tail_i = n_tail as i32;
|
||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||
launch
|
||
.arg(scratch)
|
||
.arg(&n_batch_i)
|
||
.arg(&n_tail_i)
|
||
.arg(out);
|
||
unsafe { launch.launch(cfg).context(label)?; }
|
||
Ok(())
|
||
};
|
||
reduce_at(HIDDEN_DIM * HIDDEN_DIM, &self.cfc_grad_w_in_scratch_d,
|
||
&mut self.grad_w_in_d, "reduce cfc_grad_w_in")?;
|
||
reduce_at(HIDDEN_DIM * HIDDEN_DIM, &self.cfc_grad_w_rec_scratch_d,
|
||
&mut self.grad_w_rec_d, "reduce cfc_grad_w_rec")?;
|
||
reduce_at(HIDDEN_DIM, &self.cfc_grad_b_scratch_d,
|
||
&mut self.grad_b_d, "reduce cfc_grad_b")?;
|
||
reduce_at(HIDDEN_DIM, &self.cfc_grad_tau_scratch_d,
|
||
&mut self.grad_tau_d, "reduce cfc_grad_tau")?;
|
||
|
||
// GRN: 10 reducer launches (Phase B commit 2).
|
||
let nh = N_HORIZONS;
|
||
let mid = HEAD_MID_DIM;
|
||
let h = HIDDEN_DIM;
|
||
reduce_at(nh * mid * h, &self.grn_grad_w1_scratch_d,
|
||
&mut self.grad_heads_w1_d, "reduce grn_grad_w1")?;
|
||
reduce_at(nh * mid, &self.grn_grad_b1_scratch_d,
|
||
&mut self.grad_heads_b1_d, "reduce grn_grad_b1")?;
|
||
reduce_at(nh * mid * mid, &self.grn_grad_w2_scratch_d,
|
||
&mut self.grad_heads_w2_d, "reduce grn_grad_w2")?;
|
||
reduce_at(nh * mid, &self.grn_grad_b2_scratch_d,
|
||
&mut self.grad_heads_b2_d, "reduce grn_grad_b2")?;
|
||
reduce_at(nh * mid, &self.grn_grad_w_gate_scratch_d,
|
||
&mut self.grad_heads_w_gate_d, "reduce grn_grad_w_gate")?;
|
||
reduce_at(nh, &self.grn_grad_b_gate_scratch_d,
|
||
&mut self.grad_heads_b_gate_d, "reduce grn_grad_b_gate")?;
|
||
reduce_at(nh * mid, &self.grn_grad_w_main_scratch_d,
|
||
&mut self.grad_heads_w_main_d, "reduce grn_grad_w_main")?;
|
||
reduce_at(nh, &self.grn_grad_b_main_scratch_d,
|
||
&mut self.grad_heads_b_main_d, "reduce grn_grad_b_main")?;
|
||
reduce_at(nh * h, &self.grn_grad_w_skip_scratch_d,
|
||
&mut self.grad_heads_w_skip_d, "reduce grn_grad_w_skip")?;
|
||
reduce_at(nh, &self.grn_grad_b_skip_scratch_d,
|
||
&mut self.grad_heads_b_skip_d, "reduce grn_grad_b_skip")?;
|
||
}
|
||
|
||
// ── 9. Apply AdamW updates on all 17 param groups: CfC×4 +
|
||
// GRN heads×10 + LN×2 + Mamba2 grouped.
|
||
self.opt_w_in.step(&mut self.trunk.w_in_d, &self.grad_w_in_d)?;
|
||
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)?;
|
||
self.opt_heads_b2.step(&mut self.trunk.heads_b2_d, &self.grad_heads_b2_d)?;
|
||
self.opt_heads_w_gate.step(&mut self.trunk.heads_w_gate_d, &self.grad_heads_w_gate_d)?;
|
||
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::<f32>();
|
||
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)?;
|
||
self.opt_ln_a_gain.step(&mut self.trunk.ln_a_gain_d, &self.grad_ln_a_gain_d)?;
|
||
self.opt_ln_a_bias.step(&mut self.trunk.ln_a_bias_d, &self.grad_ln_a_bias_d)?;
|
||
self.opt_vsn_w.step(&mut self.trunk.vsn_w_d, &self.grad_vsn_w_d)?;
|
||
self.opt_vsn_b.step(&mut self.trunk.vsn_b_d, &self.grad_vsn_b_d)?;
|
||
self.opt_attn_q.step(&mut self.trunk.attn_q_d, &self.grad_attn_q_d)?;
|
||
// Kendall σ is ISV-driven (closed form from loss_ema in
|
||
// horizon_ema_and_lambda) — no Adam step. grad_log_sigma_h_d
|
||
// is still written by the BCE kernel but intentionally ignored.
|
||
|
||
// (v2 reserved) — per-horizon AdamW step removed at V1; v2's six
|
||
// optimizer groups (horizon_tokens, Q_inv, w_fuse + b_fuse,
|
||
// moe_gate, experts, log_sigma) will land in commit V10.
|
||
|
||
// GPU-resident grad-clip + AdamW (zero host roundtrips).
|
||
// Replaces step_from_buffers which did 9× memcpy_dtoh per step.
|
||
self.mamba2_adamw
|
||
.step_from_buffers_gpu_clip(self.trunk.mamba2_l1_mut(), &self.mamba2_grads_buffers)
|
||
.context("mamba2 (l1) AdamW step_from_buffers_gpu_clip")?;
|
||
self.mamba2_l2_adamw
|
||
.step_from_buffers_gpu_clip(self.trunk.mamba2_l2_mut(), &self.mamba2_l2_grads_buffers)
|
||
.context("mamba2 (l2) AdamW step_from_buffers_gpu_clip")?;
|
||
|
||
// Queue loss_d → loss_host_d (mapped-pinned). Captured inside
|
||
// graph; sync + read happen in step_batched outside the region.
|
||
unsafe {
|
||
let (src_ptr, _g) = self.loss_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
self.loss_host_d.dev_ptr, src_ptr, 4, self.stream.cu_stream(),
|
||
).context("loss dtod → mapped-pinned")?;
|
||
}
|
||
// CRT.train: shadow smoothness telemetry buffers to mapped-pinned
|
||
// host slots for zero-sync operator reads. Same stream, same
|
||
// capture region — adjacent to the loss_d shadow above.
|
||
unsafe {
|
||
let (smooth_loss_src_ptr, _g) = self.smoothness_loss_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
self.smoothness_loss_host_d.dev_ptr,
|
||
smooth_loss_src_ptr,
|
||
std::mem::size_of::<f32>(),
|
||
self.stream.cu_stream(),
|
||
).context("smoothness_loss_host_d shadow")?;
|
||
}
|
||
unsafe {
|
||
let (smooth_loss_ph_src_ptr, _g) = self.smoothness_loss_per_horizon_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
self.smoothness_loss_per_horizon_host_d.dev_ptr,
|
||
smooth_loss_ph_src_ptr,
|
||
N_HORIZONS * std::mem::size_of::<f32>(),
|
||
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::<f32>(),
|
||
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
|
||
// inside the same graph as the other shadow memcpys, so a single
|
||
// end-of-step sync makes all of them visible. Skipped when the
|
||
// ring isn't allocated (runtime trace path is None).
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
if let (Some(step_counter_d), Some(shadow)) =
|
||
(self.log_step_counter_d.as_ref(), self.log_step_counter_host_shadow.as_ref())
|
||
{
|
||
unsafe {
|
||
let (src_ptr, _g) = step_counter_d.device_ptr(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
shadow.dev_ptr,
|
||
src_ptr,
|
||
std::mem::size_of::<i32>(),
|
||
self.stream.cu_stream(),
|
||
).context("log_step_counter_host_shadow shadow")?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Single-sequence forward-only eval — thin wrapper around
|
||
/// [`evaluate_batched`] for cfg.n_batch == 1.
|
||
/// X13: Persist the trunk's weights to a Checkpoint bincode file.
|
||
/// Delegates to `self.trunk.save_checkpoint` — gradient buffers and
|
||
/// AdamW state are NOT serialized (training-only, not needed for
|
||
/// inference).
|
||
pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> {
|
||
self.trunk.save_checkpoint(path).context("trunk save_checkpoint")
|
||
}
|
||
|
||
/// X11: accessor for the trainer's `cfg` field — read-only view of
|
||
/// the construction-time configuration. Used by BacktestHarness to
|
||
/// size its snapshot window from `seq_len`.
|
||
pub fn config(&self) -> &PerceptionTrainerConfig {
|
||
&self.cfg
|
||
}
|
||
|
||
/// X11 inference entry point: forward-only pass over `snapshots`
|
||
/// (length `cfg.n_batch * cfg.seq_len` per the trainer's batching).
|
||
/// Returns per-horizon probabilities flattened as `[K, B, N_HORIZONS]`
|
||
/// in row-major layout — same shape as `evaluate_batched`'s second
|
||
/// return value but with no loss computed and no labels required.
|
||
///
|
||
/// The backtester calls this per decision: load `PerceptionTrainer`
|
||
/// from a Checkpoint via `from_checkpoint`, maintain a sliding K-window
|
||
/// of recent snapshots, take the last K position's probs as the
|
||
/// decision signal.
|
||
///
|
||
/// P5: the kernel-dispatch chain is captured into a CUDA Graph on
|
||
/// the second call and replayed on every subsequent call. The host
|
||
/// staging-fill stays OUTSIDE the captured region because input
|
||
/// snapshots vary per call; the captured DtoD copies the latest
|
||
/// staged data into the trunk's pre-bound device pointers. Mirrors
|
||
/// the three-state machine already used by `step_batched`.
|
||
pub fn forward_only(&mut self, snapshots: &[Mbp10RawInput]) -> Result<Vec<f32>> {
|
||
let b_sz = self.cfg.n_batch;
|
||
let k_seq = self.cfg.seq_len;
|
||
anyhow::ensure!(
|
||
snapshots.len() == b_sz * k_seq,
|
||
"forward_only: expected {} snapshots (B={} * K={}); got {}",
|
||
b_sz * k_seq, b_sz, k_seq, snapshots.len()
|
||
);
|
||
|
||
// ── 1. Host staging fill (OUTSIDE the captured region — input
|
||
// data varies per call). All stg_* buffers are mapped-pinned;
|
||
// the captured DtoD picks up whatever host last wrote.
|
||
let total_snaps = b_sz * k_seq;
|
||
debug_assert!(total_snaps <= self.bk_capacity);
|
||
{
|
||
let bid_px_h = self.stg_bid_px_all.host_slice_mut();
|
||
let bid_sz_h = self.stg_bid_sz_all.host_slice_mut();
|
||
let ask_px_h = self.stg_ask_px_all.host_slice_mut();
|
||
let ask_sz_h = self.stg_ask_sz_all.host_slice_mut();
|
||
let regime_h = self.stg_regime_all.host_slice_mut();
|
||
for b_idx in 0..b_sz {
|
||
for k in 0..k_seq {
|
||
let n = b_idx * k_seq + k;
|
||
let snap = &snapshots[n];
|
||
let base10 = n * 10;
|
||
let base6 = n * REGIME_DIM;
|
||
for i in 0..10 {
|
||
bid_px_h[base10 + i] = snap.bid_px[i];
|
||
bid_sz_h[base10 + i] = snap.bid_sz[i];
|
||
ask_px_h[base10 + i] = snap.ask_px[i];
|
||
ask_sz_h[base10 + i] = snap.ask_sz[i];
|
||
}
|
||
for i in 0..REGIME_DIM {
|
||
regime_h[base6 + i] = snap.regime[i];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
{
|
||
let prev_mid_h = self.stg_prev_mid.host_slice_mut();
|
||
let tsv_h = self.stg_trade_signed_vol.host_slice_mut();
|
||
let tc_h = self.stg_trade_count.host_slice_mut();
|
||
let ts_ns_h = self.stg_ts_ns.host_slice_mut();
|
||
let prev_ts_ns_h = self.stg_prev_ts_ns.host_slice_mut();
|
||
for b_idx in 0..b_sz {
|
||
for k in 0..k_seq {
|
||
let n = b_idx * k_seq + k;
|
||
let snap = &snapshots[n];
|
||
prev_mid_h[n] = snap.prev_mid;
|
||
tsv_h[n] = snap.trade_signed_vol;
|
||
tc_h[n] = snap.trade_count as i32;
|
||
ts_ns_h[n] = snap.ts_ns as i64;
|
||
prev_ts_ns_h[n] = snap.prev_ts_ns as i64;
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── 2. Three-state machine: warmup (first call), capture
|
||
// (second call), replay (third+). Mirrors step_batched's
|
||
// pattern at lines 1219-1257. Event tracking is disabled
|
||
// for the trainer's lifetime (see new(): line ~529) so
|
||
// the captured region is free of cuStreamWaitEvent calls
|
||
// per pearl_cudarc_disable_event_tracking_for_graph_capture.
|
||
if self.forward_graph.is_some() {
|
||
self.forward_graph
|
||
.as_ref()
|
||
.unwrap()
|
||
.launch()
|
||
.context("forward graph launch")?;
|
||
} else if !self.forward_warmed {
|
||
self.dispatch_forward_kernels(b_sz, k_seq, total_snaps)
|
||
.context("forward warmup dispatch")?;
|
||
self.forward_warmed = true;
|
||
} else {
|
||
let begin = self.stream.begin_capture(
|
||
CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
|
||
);
|
||
if let Err(e) = begin {
|
||
return Err(anyhow::anyhow!("forward begin_capture: {e}"));
|
||
}
|
||
|
||
let dispatch_result =
|
||
self.dispatch_forward_kernels(b_sz, k_seq, total_snaps);
|
||
|
||
let graph_result = self.stream.end_capture(
|
||
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||
);
|
||
|
||
dispatch_result.context("forward graph dispatch (during capture)")?;
|
||
let graph = graph_result
|
||
.context("forward graph end_capture")?
|
||
.ok_or_else(|| anyhow::anyhow!(
|
||
"forward end_capture returned None — no work captured"
|
||
))?;
|
||
self.forward_graph = Some(graph);
|
||
}
|
||
|
||
// ── 3. Sync + dtoh probs (OUTSIDE capture; download requires sync).
|
||
self.stream.synchronize().context("forward_only end-sync")?;
|
||
download(&self.stream, &self.probs_per_k_d)
|
||
}
|
||
|
||
/// Kernel-dispatch portion of an inference forward — captured into the
|
||
/// CUDA Graph on the second `forward_only` call. Mirrors
|
||
/// `evaluate_batched`'s forward kernel chain (VSN → Mamba2 ×2 → LN ×2
|
||
/// → attn-pool → CfC K-loop → heads) but SKIPS labels DtoD, BCE,
|
||
/// and any stream syncs.
|
||
///
|
||
/// All inputs are read from pre-bound device pointers (`self.stg_*`
|
||
/// device addresses + `self.trunk.*` weights); output is written to
|
||
/// `self.probs_per_k_d`. No host branches, no scalar-arg-changes,
|
||
/// no host-mallocs inside per pearl_no_host_branches_in_captured_graph.
|
||
fn dispatch_forward_kernels(
|
||
&mut self,
|
||
b_sz: usize,
|
||
k_seq: usize,
|
||
total_snaps: usize,
|
||
) -> Result<()> {
|
||
// DtoD: staging (mapped-pinned, device-visible) → device buffers.
|
||
let n10 = total_snaps * 10;
|
||
let n6 = total_snaps * REGIME_DIM;
|
||
let n1 = total_snaps;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (d, _g) = self.bid_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_px_all.dev_ptr, n10 * 4, s)
|
||
.context("fwd bid_px_all dtod")?;
|
||
let (d, _g) = self.bid_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("fwd bid_sz_all dtod")?;
|
||
let (d, _g) = self.ask_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_px_all.dev_ptr, n10 * 4, s)
|
||
.context("fwd ask_px_all dtod")?;
|
||
let (d, _g) = self.ask_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("fwd ask_sz_all dtod")?;
|
||
let (d, _g) = self.regime_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_regime_all.dev_ptr, n6 * 4, s)
|
||
.context("fwd regime_all dtod")?;
|
||
let (d, _g) = self.prev_mid_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_mid.dev_ptr, n1 * 4, s)
|
||
.context("fwd prev_mid_all dtod")?;
|
||
let (d, _g) = self.trade_signed_vol_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_signed_vol.dev_ptr, n1 * 4, s)
|
||
.context("fwd tsv_all dtod")?;
|
||
let (d, _g) = self.trade_count_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_count.dev_ptr, n1 * 4, s)
|
||
.context("fwd trade_count_all dtod")?;
|
||
let (d, _g) = self.ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("fwd ts_ns_all dtod")?;
|
||
let (d, _g) = self.prev_ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("fwd prev_ts_ns_all dtod")?;
|
||
}
|
||
|
||
// snap_feature_assemble_batched: raw MBP-10 SoA → window_tensor [B,K,FEATURE_DIM].
|
||
let tick_size = ES_TICK_SIZE;
|
||
let n_total_i32 = total_snaps as i32;
|
||
let snap_block: u32 = 128;
|
||
let snap_grid: u32 = (total_snaps as u32).div_ceil(snap_block);
|
||
let snap_cfg = LaunchConfig {
|
||
grid_dim: (snap_grid, 1, 1),
|
||
block_dim: (snap_block, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.snap_batched_fn);
|
||
launch
|
||
.arg(&self.bid_px_all_d).arg(&self.bid_sz_all_d)
|
||
.arg(&self.ask_px_all_d).arg(&self.ask_sz_all_d)
|
||
.arg(&self.prev_bid_sz_all_d).arg(&self.prev_ask_sz_all_d)
|
||
.arg(&self.regime_all_d)
|
||
.arg(&self.prev_mid_all_d).arg(&self.trade_signed_vol_all_d)
|
||
.arg(&self.trade_count_all_d)
|
||
.arg(&self.ts_ns_all_d).arg(&self.prev_ts_ns_all_d)
|
||
.arg(&tick_size).arg(&n_total_i32)
|
||
.arg(self.window_tensor_d.data_mut());
|
||
launch.launch(snap_cfg).context("fwd snap_batched")?;
|
||
}
|
||
// No sync — stream-ordered. A sync inside the captured region
|
||
// would trip STREAM_CAPTURE_ISOLATION.
|
||
|
||
// VSN fwd.
|
||
{
|
||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_vsn = LaunchConfig {
|
||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||
block_dim: (64, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.vsn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.vsn_w_d)
|
||
.arg(&self.trunk.vsn_b_d)
|
||
.arg(self.window_tensor_d.cuda_data())
|
||
.arg(&n_rows_vsn)
|
||
.arg(self.vsn_out_d.data_mut())
|
||
.arg(&mut self.vsn_gates_d);
|
||
unsafe { launch.launch(cfg_vsn).context("fwd variable_selection_fwd")?; }
|
||
}
|
||
|
||
// Mamba2 stack-1 fwd.
|
||
self.trunk
|
||
.mamba2_l1_mut()
|
||
.forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)
|
||
.context("fwd mamba2 (l1) fwd_into")?;
|
||
|
||
// LN_a fwd.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_a_gain_d)
|
||
.arg(&self.trunk.ln_a_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(self.ln_a_out_d.data_mut())
|
||
.arg(&mut self.ln_a_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("fwd layer_norm_fwd (LN_a)")?; }
|
||
}
|
||
|
||
// Mamba2 stack-2 fwd.
|
||
self.trunk
|
||
.mamba2_l2_mut()
|
||
.forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)
|
||
.context("fwd mamba2 (l2) fwd_into")?;
|
||
|
||
// LN_b fwd.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_l2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_b_gain_d)
|
||
.arg(&self.trunk.ln_b_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.ln_out_d)
|
||
.arg(&mut self.ln_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("fwd layer_norm_fwd (LN_b)")?; }
|
||
}
|
||
|
||
// Transpose LN [B, K, H] → [K, B, H].
|
||
{
|
||
let block_n3: u32 = 32;
|
||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||
let cfg_tx = LaunchConfig {
|
||
grid_dim: (b_sz as u32, k_seq as u32, grid_z),
|
||
block_dim: (block_n3, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n1 = b_sz as i32;
|
||
let n2 = k_seq as i32;
|
||
let n3 = HIDDEN_DIM as i32;
|
||
let mut launch = self.stream.launch_builder(&self.trunk.transpose_3d_fn);
|
||
launch
|
||
.arg(&self.ln_out_d)
|
||
.arg(self.h_enriched_seq_t_d.data_mut())
|
||
.arg(&n1).arg(&n2).arg(&n3);
|
||
unsafe { launch.launch(cfg_tx).context("fwd transpose h_enriched")?; }
|
||
}
|
||
|
||
// Attention pool fwd.
|
||
{
|
||
let k_i32 = k_seq as i32;
|
||
let n_batch_attn = b_sz as i32;
|
||
let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::<f32>();
|
||
let cfg_attn = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: shared as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.attn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.attn_q_d)
|
||
.arg(&self.ln_out_d)
|
||
.arg(&n_batch_attn).arg(&k_i32)
|
||
.arg(&mut self.attn_context_d)
|
||
.arg(&mut self.attn_weights_d);
|
||
unsafe { launch.launch(cfg_attn).context("fwd attention_pool_fwd")?; }
|
||
}
|
||
|
||
// Per-K CfC + heads. Pointer arithmetic on stable base addresses;
|
||
// no host branches inside the captured region.
|
||
// A1: decision_stride removed; CRT runs at event rate, dt = 1.0 per step.
|
||
let dt_s = 1.0_f32;
|
||
let n_in_i = HIDDEN_DIM as i32;
|
||
let n_hid_i = HIDDEN_DIM as i32;
|
||
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_cfc = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: cfc_fwd_smem,
|
||
};
|
||
let cfg_grn_fwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HEAD_MID_DIM as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::<f32>();
|
||
let kb_nh_mid_bytes = b_sz * N_HORIZONS * HEAD_MID_DIM * std::mem::size_of::<f32>();
|
||
let n_batch_i = b_sz as i32;
|
||
// h_old at k=0 is attn_context_d (learned pooled summary).
|
||
let attn_context_ptr = {
|
||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let henr_t_base = {
|
||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||
let (probs_base, _g_probs) = self.probs_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z1_base, _g_z1) = self.z1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (a1_base, _g_a1) = self.a1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z2_base, _g_z2) = self.z2_per_k_d.device_ptr_mut(&self.stream);
|
||
let (gate_base, _g_gate) = self.gate_logit_per_k_d.device_ptr_mut(&self.stream);
|
||
let (main_base, _g_main) = self.main_per_k_d.device_ptr_mut(&self.stream);
|
||
let (logit_base, _g_logit) = self.logit_per_k_d.device_ptr_mut(&self.stream);
|
||
for k in 0..k_seq {
|
||
let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64;
|
||
let h_old_k_ptr = if k == 0 {
|
||
attn_context_ptr
|
||
} else {
|
||
h_per_k_base + ((k - 1) * kb_hid_bytes) as u64
|
||
};
|
||
let x_k_ptr = henr_t_base + (k * kb_hid_bytes) as u64;
|
||
let probs_k_ptr = probs_base + (k * kb_nh_bytes) as u64;
|
||
let z1_k_ptr = z1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let a1_k_ptr = a1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let z2_k_ptr = z2_base + (k * kb_nh_mid_bytes) as u64;
|
||
let gate_k_ptr = gate_base + (k * kb_nh_bytes) as u64;
|
||
let main_k_ptr = main_base + (k * kb_nh_bytes) as u64;
|
||
let logit_k_ptr = logit_base + (k * kb_nh_bytes) as u64;
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&x_k_ptr).arg(&h_old_k_ptr)
|
||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
|
||
.arg(&h_new_k_ptr);
|
||
launch.launch(cfg_cfc).context("fwd cfc")?;
|
||
}
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d)
|
||
.arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d)
|
||
.arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d)
|
||
.arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d)
|
||
.arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d)
|
||
.arg(&h_new_k_ptr).arg(&n_batch_i)
|
||
.arg(&probs_k_ptr)
|
||
.arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr)
|
||
.arg(&gate_k_ptr).arg(&main_k_ptr).arg(&logit_k_ptr);
|
||
launch.launch(cfg_grn_fwd).context("fwd heads GRN")?;
|
||
}
|
||
}
|
||
drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit));
|
||
Ok(())
|
||
}
|
||
|
||
/// CRT Phase A0.5: single-event incremental forward pass.
|
||
///
|
||
/// Advances the encoder's recurrent state (Mamba2 SSM register state
|
||
/// in `step_scratch_l1` / `_l2`, plus CfC hidden state in
|
||
/// `cfc_h_state_step_d`) by exactly one snapshot and returns the
|
||
/// per-horizon alpha probabilities for the current event.
|
||
///
|
||
/// Per-call work scales as O(hidden_dim × state_dim) — independent
|
||
/// of K — whereas `forward_only` runs a K-window scan (~K× the
|
||
/// arithmetic). This is the structural change Phase A's ≤ 2×
|
||
/// wall-time target needs (see Task A0 cost-investigation memo).
|
||
///
|
||
/// Architectural divergence from `forward_only`:
|
||
/// - No K-window staging: only 1 snapshot is uploaded per call.
|
||
/// - No attention pool: the attention pool requires a K-window of
|
||
/// LN_b outputs to produce a learned `h_old` for CfC at k=0.
|
||
/// `forward_step` instead carries CfC state across calls, so the
|
||
/// k=0 `h_old` comes from the previous step's h_new (zero on first
|
||
/// call / after `reset_step_state()`).
|
||
/// - No K-loop in CfC + heads: exactly one CfC step + one head
|
||
/// projection per call. The K-history is captured implicitly in
|
||
/// the persistent SSM and CfC states.
|
||
///
|
||
/// Sequence semantics: calling forward_step_into N times on snapshots
|
||
/// (s_0 … s_{N-1}) starting from a fresh `reset_step_state()` is
|
||
/// equivalent (within stable-SSM dampening) to `forward_only` on the
|
||
/// same K-window snapshots once the persistent state has accumulated
|
||
/// enough context. The bit-equivalence test (forward_step_golden.rs)
|
||
/// validates this convergence empirically. Per Mamba2's stable gating
|
||
/// (`sigmoid(a) < 1`), state dampens at ~0.5^N — after ~K events the
|
||
/// state is dominated by the recent K snapshots and agrees with
|
||
/// forward_only on the same K within float-rounding tolerance.
|
||
///
|
||
/// The output `alpha_probs_dst` MUST be a device slice of length
|
||
/// `N_HORIZONS` (B=1 enforced). The GRN heads kernel writes
|
||
/// probabilities directly into it — no host staging, no `synchronize`
|
||
/// on the hot path. Per feedback_cpu_is_read_only and
|
||
/// feedback_no_htod_htoh_only_mapped_pinned: probs stay on device
|
||
/// throughout the per-event call chain.
|
||
pub fn forward_step_into(
|
||
&mut self,
|
||
snapshot: &Mbp10RawInput,
|
||
alpha_probs_dst: &mut CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
let b_sz = self.cfg.n_batch;
|
||
anyhow::ensure!(
|
||
b_sz == 1,
|
||
"forward_step_into currently supports n_batch == 1 only (got {})",
|
||
b_sz
|
||
);
|
||
anyhow::ensure!(
|
||
alpha_probs_dst.len() >= N_HORIZONS,
|
||
"forward_step_into: alpha_probs_dst len {} < N_HORIZONS {}",
|
||
alpha_probs_dst.len(), N_HORIZONS
|
||
);
|
||
|
||
// ── 1. Host staging fill (1 snapshot). All buffers are
|
||
// mapped-pinned; the snap_feature kernel reads them after
|
||
// a DtoD copy below.
|
||
{
|
||
let bid_px_h = self.stg_step_bid_px.host_slice_mut();
|
||
let bid_sz_h = self.stg_step_bid_sz.host_slice_mut();
|
||
let ask_px_h = self.stg_step_ask_px.host_slice_mut();
|
||
let ask_sz_h = self.stg_step_ask_sz.host_slice_mut();
|
||
let regime_h = self.stg_step_regime.host_slice_mut();
|
||
for i in 0..10 {
|
||
bid_px_h[i] = snapshot.bid_px[i];
|
||
bid_sz_h[i] = snapshot.bid_sz[i];
|
||
ask_px_h[i] = snapshot.ask_px[i];
|
||
ask_sz_h[i] = snapshot.ask_sz[i];
|
||
}
|
||
for i in 0..REGIME_DIM {
|
||
regime_h[i] = snapshot.regime[i];
|
||
}
|
||
}
|
||
{
|
||
let prev_mid_h = self.stg_step_prev_mid.host_slice_mut();
|
||
let tsv_h = self.stg_step_trade_signed_vol.host_slice_mut();
|
||
let tc_h = self.stg_step_trade_count.host_slice_mut();
|
||
let ts_ns_h = self.stg_step_ts_ns.host_slice_mut();
|
||
let prev_ts_ns_h = self.stg_step_prev_ts_ns.host_slice_mut();
|
||
prev_mid_h[0] = snapshot.prev_mid;
|
||
tsv_h[0] = snapshot.trade_signed_vol;
|
||
tc_h[0] = snapshot.trade_count as i32;
|
||
ts_ns_h[0] = snapshot.ts_ns as i64;
|
||
prev_ts_ns_h[0] = snapshot.prev_ts_ns as i64;
|
||
}
|
||
|
||
// ── 2. DtoD copies: staging (mapped-pinned, device-visible) → device.
|
||
// Per-call cost is tiny (one snapshot = ~280 bytes total).
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let n10 = 10 * 4;
|
||
let n6 = REGIME_DIM * 4;
|
||
let n1f = 4;
|
||
let n1i = 4;
|
||
let n1l = 8;
|
||
let (d, _g) = self.step_bid_px_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_bid_px.dev_ptr, n10, s)
|
||
.context("step bid_px dtod")?;
|
||
let (d, _g) = self.step_bid_sz_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_bid_sz.dev_ptr, n10, s)
|
||
.context("step bid_sz dtod")?;
|
||
let (d, _g) = self.step_ask_px_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_ask_px.dev_ptr, n10, s)
|
||
.context("step ask_px dtod")?;
|
||
let (d, _g) = self.step_ask_sz_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_ask_sz.dev_ptr, n10, s)
|
||
.context("step ask_sz dtod")?;
|
||
let (d, _g) = self.step_regime_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_regime.dev_ptr, n6, s)
|
||
.context("step regime dtod")?;
|
||
let (d, _g) = self.step_prev_mid_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_prev_mid.dev_ptr, n1f, s)
|
||
.context("step prev_mid dtod")?;
|
||
let (d, _g) = self.step_trade_signed_vol_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_trade_signed_vol.dev_ptr, n1f, s)
|
||
.context("step tsv dtod")?;
|
||
let (d, _g) = self.step_trade_count_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_trade_count.dev_ptr, n1i, s)
|
||
.context("step trade_count dtod")?;
|
||
let (d, _g) = self.step_ts_ns_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_ts_ns.dev_ptr, n1l, s)
|
||
.context("step ts_ns dtod")?;
|
||
let (d, _g) = self.step_prev_ts_ns_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_step_prev_ts_ns.dev_ptr, n1l, s)
|
||
.context("step prev_ts_ns dtod")?;
|
||
}
|
||
|
||
// ── 3. snap_feature_assemble: 1 snapshot → 1 row of FEATURE_DIM.
|
||
let tick_size = ES_TICK_SIZE;
|
||
let n_total_i32: i32 = b_sz as i32;
|
||
let snap_block: u32 = 128;
|
||
let snap_grid: u32 = (b_sz as u32).div_ceil(snap_block).max(1);
|
||
let snap_cfg = LaunchConfig {
|
||
grid_dim: (snap_grid, 1, 1),
|
||
block_dim: (snap_block, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.snap_batched_fn);
|
||
launch
|
||
.arg(&self.step_bid_px_d).arg(&self.step_bid_sz_d)
|
||
.arg(&self.step_ask_px_d).arg(&self.step_ask_sz_d)
|
||
.arg(&self.step_prev_bid_sz_d).arg(&self.step_prev_ask_sz_d)
|
||
.arg(&self.step_regime_d)
|
||
.arg(&self.step_prev_mid_d).arg(&self.step_trade_signed_vol_d)
|
||
.arg(&self.step_trade_count_d)
|
||
.arg(&self.step_ts_ns_d).arg(&self.step_prev_ts_ns_d)
|
||
.arg(&tick_size).arg(&n_total_i32)
|
||
.arg(self.window_step_d.data_mut());
|
||
launch.launch(snap_cfg).context("step snap_feature_assemble")?;
|
||
}
|
||
|
||
// ── 4. VSN fwd: [B, FEATURE_DIM] → [B, FEATURE_DIM] (gated).
|
||
{
|
||
let n_rows_vsn: i32 = b_sz as i32;
|
||
let cfg_vsn = LaunchConfig {
|
||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||
block_dim: (64, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.vsn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.vsn_w_d)
|
||
.arg(&self.trunk.vsn_b_d)
|
||
.arg(self.window_step_d.cuda_data())
|
||
.arg(&n_rows_vsn)
|
||
.arg(self.vsn_step_out_d.data_mut())
|
||
.arg(&mut self.vsn_step_gates_d);
|
||
unsafe { launch.launch(cfg_vsn).context("step variable_selection_fwd")?; }
|
||
}
|
||
|
||
// ── 5. Mamba2 L1 step — advances step_scratch_l1.x_state in-place.
|
||
self.trunk
|
||
.mamba2_l1_mut()
|
||
.step_into(&self.vsn_step_out_d, &mut self.step_scratch_l1)
|
||
.context("step mamba2 (l1) step_into")?;
|
||
|
||
// ── 6. LN_a fwd on 1 row.
|
||
{
|
||
let n_rows_ln: i32 = b_sz as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.step_scratch_l1.h_out.cuda_data())
|
||
.arg(&self.trunk.ln_a_gain_d)
|
||
.arg(&self.trunk.ln_a_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.ln_a_step_out_d)
|
||
.arg(&mut self.ln_a_step_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("step LN_a fwd")?; }
|
||
}
|
||
|
||
// ── 7. Mamba2 L2 step — input is ln_a_step_out wrapped as GpuTensor.
|
||
// Reusing the existing ln_a_step_out_d buffer; the L2 step
|
||
// kernel will read [B, HIDDEN_DIM] (single-row).
|
||
{
|
||
// Build a temporary view of ln_a_step_out_d as a [B, HIDDEN_DIM]
|
||
// GpuTensor for step_into's input shape check.
|
||
let ln_a_view = GpuTensor::new(
|
||
self.ln_a_step_out_d.clone(),
|
||
vec![b_sz, HIDDEN_DIM],
|
||
).map_err(|e| anyhow::anyhow!("ln_a_view wrap: {e}"))?;
|
||
self.trunk
|
||
.mamba2_l2_mut()
|
||
.step_into(&ln_a_view, &mut self.step_scratch_l2)
|
||
.context("step mamba2 (l2) step_into")?;
|
||
}
|
||
|
||
// ── 8. LN_b fwd on 1 row.
|
||
{
|
||
let n_rows_ln: i32 = b_sz as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.step_scratch_l2.h_out.cuda_data())
|
||
.arg(&self.trunk.ln_b_gain_d)
|
||
.arg(&self.trunk.ln_b_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.ln_b_step_out_d)
|
||
.arg(&mut self.ln_b_step_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("step LN_b fwd")?; }
|
||
}
|
||
|
||
// ── 9. CfC single step. h_old is the persistent cfc_h_state_step;
|
||
// h_new is written to cfc_h_new_step (then copied back into
|
||
// cfc_h_state_step).
|
||
// A1: decision_stride removed; event-rate dt = 1.0 per CfC step.
|
||
let dt_s: f32 = 1.0;
|
||
let n_in_i: i32 = HIDDEN_DIM as i32;
|
||
let n_hid_i: i32 = HIDDEN_DIM as i32;
|
||
let n_batch_i: i32 = b_sz as i32;
|
||
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_cfc = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: cfc_fwd_smem,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&self.ln_b_step_out_d).arg(&self.cfc_h_state_step_d)
|
||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
|
||
.arg(&mut self.cfc_h_new_step_d);
|
||
launch.launch(cfg_cfc).context("step cfc")?;
|
||
}
|
||
|
||
// Carry-forward: h_new → h_state (in-place via DtoD).
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (src, _gs) = self.cfc_h_new_step_d.device_ptr(&self.stream);
|
||
let (dst, _gd) = self.cfc_h_state_step_d.device_ptr_mut(&self.stream);
|
||
let nbytes = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
cudarc::driver::result::memcpy_dtod_async(dst, src, nbytes, s)
|
||
.context("cfc h_state carry-forward dtod")?;
|
||
}
|
||
|
||
// ── 10. Heads GRN fwd on the new h.
|
||
let cfg_grn_fwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HEAD_MID_DIM as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d)
|
||
.arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d)
|
||
.arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d)
|
||
.arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d)
|
||
.arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d)
|
||
.arg(&self.cfc_h_new_step_d).arg(&n_batch_i)
|
||
// Probs go directly to the caller-supplied device buffer —
|
||
// no host staging, no synchronize. The next consumer (the
|
||
// sim's decision kernels) reads from this same device slice
|
||
// and is launched on the same stream so it is stream-ordered.
|
||
.arg(alpha_probs_dst)
|
||
.arg(&mut self.z1_step_d).arg(&mut self.a1_step_d).arg(&mut self.z2_step_d)
|
||
.arg(&mut self.gate_logit_step_d).arg(&mut self.main_step_d).arg(&mut self.logit_step_d);
|
||
launch.launch(cfg_grn_fwd).context("step heads GRN")?;
|
||
}
|
||
// No DtoD-to-host, no stream.synchronize, no host read. The probs
|
||
// remain on-device for the consumer (decision kernels). The CPU
|
||
// never touches them on the hot path per
|
||
// pearl/feedback_cpu_is_read_only.
|
||
Ok(())
|
||
}
|
||
|
||
/// Test-only convenience: drive `forward_step_into` against a private
|
||
/// device buffer and DtoH-sync the probs into a host array. NOT part
|
||
/// of the production hot path — production callers MUST use
|
||
/// `forward_step_into(snapshot, &mut alpha_probs_dst)` and keep the
|
||
/// probs on-device. This helper exists for integration tests that
|
||
/// previously called the old `forward_step(snap) -> [f32; N]` API and
|
||
/// want to assert on host values; it allocates a one-shot device
|
||
/// buffer, runs forward_step_into into it, then DtoVs once.
|
||
///
|
||
/// Per-call DtoV stops the stream; do NOT use on the per-event hot
|
||
/// path. The function name carries `_returning` as a visible flag for
|
||
/// reviewers that a host roundtrip happens here.
|
||
pub fn forward_step_into_returning(
|
||
&mut self,
|
||
snapshot: &Mbp10RawInput,
|
||
) -> Result<[f32; N_HORIZONS]> {
|
||
let b_sz = self.cfg.n_batch;
|
||
anyhow::ensure!(
|
||
b_sz == 1,
|
||
"forward_step_into_returning supports n_batch == 1 only (got {})",
|
||
b_sz
|
||
);
|
||
let mut probs_d = self.stream
|
||
.alloc_zeros::<f32>(N_HORIZONS)
|
||
.map_err(|e| anyhow::anyhow!("test probs_d alloc: {e}"))?;
|
||
self.forward_step_into(snapshot, &mut probs_d)?;
|
||
let mut host = vec![0.0_f32; N_HORIZONS];
|
||
self.stream
|
||
.memcpy_dtoh(&probs_d, host.as_mut_slice())
|
||
.map_err(|e| anyhow::anyhow!("test probs DtoH: {e}"))?;
|
||
let mut out = [0.0_f32; N_HORIZONS];
|
||
for h in 0..N_HORIZONS {
|
||
out[h] = host[h];
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
/// CRT Phase A0.5: zero the forward_step persistent state — SSM
|
||
/// register state in both Mamba2 step scratches plus the CfC hidden
|
||
/// state — and drop accumulated context. Used by the harness on
|
||
/// session-gap detection so per-day boundaries don't leak state.
|
||
///
|
||
/// Exposed but NOT yet wired into BacktestHarness — session-gap
|
||
/// integration is a downstream task. After A1 deletes the stride
|
||
/// gate, A2 / A3 can pick the call site that fits the harness's
|
||
/// session-gap detection (already present for the max-hold path).
|
||
pub fn reset_step_state(&mut self) -> Result<()> {
|
||
self.step_scratch_l1.reset_state(&self.stream)
|
||
.context("reset step_scratch_l1")?;
|
||
self.step_scratch_l2.reset_state(&self.stream)
|
||
.context("reset step_scratch_l2")?;
|
||
self.stream.memset_zeros(&mut self.cfc_h_state_step_d)
|
||
.map_err(|e| anyhow::anyhow!("reset cfc_h_state_step_d: {e}"))?;
|
||
Ok(())
|
||
}
|
||
|
||
/// CRT Phase A0.5 corrective (concern #1 from a0e81fbdf): bit-identical
|
||
/// seeding from forward_only.
|
||
///
|
||
/// After this call, the step-path persistent state (
|
||
/// step_scratch_l1.x_state,
|
||
/// step_scratch_l2.x_state,
|
||
/// cfc_h_state_step_d
|
||
/// ) reproduces the internal state forward_only had at the end of
|
||
/// processing `window` (K snapshots). Subsequent `forward_step_into`
|
||
/// calls then continue the same trajectory bit-identically to where
|
||
/// forward_only left off, modulo:
|
||
/// - The cuBLAS W_in/W_a/W_b GEMM bit-identity for the K-batched seq
|
||
/// path (one GEMM over [B*K, *]) is what we capture; the scan_fwd_step
|
||
/// replay below reads those projections and produces bit-identical
|
||
/// x_state by construction (same arithmetic, same per-step order).
|
||
/// - The CfC h_state is a DtoD copy of h_new_per_k_d at position K-1
|
||
/// — also bit-identical.
|
||
///
|
||
/// Implementation:
|
||
/// 1. Run `forward_only(window)` — populates mamba2_fwd_scratch.a_proj/b_proj
|
||
/// + mamba2_l2_fwd_scratch.a_proj/b_proj (via batched GEMM over K rows)
|
||
/// + h_new_per_k_d (the CfC K-loop output at every position).
|
||
/// 2. Zero step scratches so the scan_fwd_step replay starts from x=0,
|
||
/// matching scan_fwd_seq's `float x[32] = {0}` initialisation.
|
||
/// 3. For each k=0..K, launch `step_advance_from_seq_row` on L1 reading
|
||
/// a_proj/b_proj at row k. After K calls, step_scratch_l1.x_state
|
||
/// matches scan_fwd_seq's terminal register state for L1.
|
||
/// 4. Same for L2.
|
||
/// 5. DtoD copy h_new_per_k_d at slot K-1 (i.e., the cfc state that
|
||
/// would feed snapshot K if there were one) into cfc_h_state_step_d.
|
||
///
|
||
/// Per memo §4.5 option (a) — extracts forward_only's terminal state
|
||
/// rather than re-introducing the attention pool on the step path.
|
||
pub fn seed_step_state_from_forward_only(
|
||
&mut self,
|
||
window: &[Mbp10RawInput],
|
||
) -> Result<()> {
|
||
let b_sz = self.cfg.n_batch;
|
||
let k_seq = self.cfg.seq_len;
|
||
anyhow::ensure!(
|
||
b_sz == 1,
|
||
"seed_step_state_from_forward_only currently supports n_batch == 1 only (got {})",
|
||
b_sz
|
||
);
|
||
anyhow::ensure!(
|
||
window.len() == b_sz * k_seq,
|
||
"seed window: expected {} snapshots (B={} * K={}); got {}",
|
||
b_sz * k_seq, b_sz, k_seq, window.len()
|
||
);
|
||
|
||
// Step 1: forward_only populates a_proj/b_proj scratches + h_new_per_k_d.
|
||
// We discard the returned probs vec; we only need the side-effect on
|
||
// the trainer's internal scratches.
|
||
let _probs = self.forward_only(window)?;
|
||
|
||
// Step 2: zero step scratches so the scan_fwd_step replay starts
|
||
// from x[s]=0 — matches `mamba2_alpha_scan_fwd_seq`'s register init.
|
||
self.step_scratch_l1.reset_state(&self.stream)
|
||
.context("seed: reset step_scratch_l1 x_state")?;
|
||
self.step_scratch_l2.reset_state(&self.stream)
|
||
.context("seed: reset step_scratch_l2 x_state")?;
|
||
|
||
// Step 3+4: replay K scan_fwd_step calls for L1 then L2. The
|
||
// a_proj/b_proj scratches are sized [B*K, state_dim] in the seq
|
||
// path — row k (with B=1) sits at byte offset k * state_dim * 4.
|
||
let state_dim_l1 = self.mamba2_fwd_scratch.state_dim;
|
||
let row_bytes_l1 = state_dim_l1 * std::mem::size_of::<f32>();
|
||
let (a_l1_base, _g_a_l1) =
|
||
self.mamba2_fwd_scratch.a_proj.cuda_data().device_ptr(&self.stream);
|
||
let (b_l1_base, _g_b_l1) =
|
||
self.mamba2_fwd_scratch.b_proj.cuda_data().device_ptr(&self.stream);
|
||
for k in 0..k_seq {
|
||
let a_row = a_l1_base + (k * row_bytes_l1) as u64;
|
||
let b_row = b_l1_base + (k * row_bytes_l1) as u64;
|
||
self.trunk
|
||
.mamba2_l1()
|
||
.step_advance_from_seq_row(a_row, b_row, &mut self.step_scratch_l1)
|
||
.with_context(|| format!("seed: L1 step at k={k}"))?;
|
||
}
|
||
drop((_g_a_l1, _g_b_l1));
|
||
|
||
let state_dim_l2 = self.mamba2_l2_fwd_scratch.state_dim;
|
||
let row_bytes_l2 = state_dim_l2 * std::mem::size_of::<f32>();
|
||
let (a_l2_base, _g_a_l2) =
|
||
self.mamba2_l2_fwd_scratch.a_proj.cuda_data().device_ptr(&self.stream);
|
||
let (b_l2_base, _g_b_l2) =
|
||
self.mamba2_l2_fwd_scratch.b_proj.cuda_data().device_ptr(&self.stream);
|
||
for k in 0..k_seq {
|
||
let a_row = a_l2_base + (k * row_bytes_l2) as u64;
|
||
let b_row = b_l2_base + (k * row_bytes_l2) as u64;
|
||
self.trunk
|
||
.mamba2_l2()
|
||
.step_advance_from_seq_row(a_row, b_row, &mut self.step_scratch_l2)
|
||
.with_context(|| format!("seed: L2 step at k={k}"))?;
|
||
}
|
||
drop((_g_a_l2, _g_b_l2));
|
||
|
||
// Step 5: copy h_new_per_k_d at slot K-1 into cfc_h_state_step_d.
|
||
// forward_only stored CfC's per-position h_new in h_new_per_k_d
|
||
// laid out as [K, B, HIDDEN_DIM]. Slot K-1 is the LAST cfc h_new —
|
||
// i.e., the value that would feed CfC position K if there were one.
|
||
// That's exactly the h_old forward_step_into expects on its next call.
|
||
let row_bytes_cfc = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (src_base, _gs) = self.h_new_per_k_d.device_ptr(&self.stream);
|
||
let src_at_last = src_base + ((k_seq - 1) * row_bytes_cfc) as u64;
|
||
let (dst, _gd) = self.cfc_h_state_step_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(dst, src_at_last, row_bytes_cfc, s)
|
||
.context("seed: copy h_new_per_k[K-1] → cfc_h_state_step")?;
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
// ── Test-only readback helpers ────────────────────────────────────────
|
||
//
|
||
// These perform one-shot DtoH copies acceptable in test code.
|
||
// NEVER call on the hot path.
|
||
|
||
/// Read `step_scratch_l1.x_state` back to host. Test-only; DtoH+sync.
|
||
pub fn read_step_l1_x_state(&self) -> Result<Vec<f32>> {
|
||
self.step_scratch_l1.read_x_state(&self.stream)
|
||
}
|
||
|
||
/// Read `step_scratch_l2.x_state` back to host. Test-only; DtoH+sync.
|
||
pub fn read_step_l2_x_state(&self) -> Result<Vec<f32>> {
|
||
self.step_scratch_l2.read_x_state(&self.stream)
|
||
}
|
||
|
||
/// Read `cfc_h_state_step_d` back to host. Test-only; DtoH+sync.
|
||
pub fn read_cfc_h_state_step(&self) -> Result<Vec<f32>> {
|
||
let n = self.cfc_h_state_step_d.len();
|
||
let mut host = vec![0.0_f32; n];
|
||
self.stream
|
||
.memcpy_dtoh(&self.cfc_h_state_step_d, &mut host)
|
||
.map_err(|e| anyhow::anyhow!("read_cfc_h_state_step dtoh: {e}"))?;
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| anyhow::anyhow!("read_cfc_h_state_step sync: {e}"))?;
|
||
Ok(host)
|
||
}
|
||
|
||
/// Read the K-1 slot of `h_new_per_k_d` — the CfC h_new at the
|
||
/// final position of the last `forward_only` call. Layout is
|
||
/// `[K, B, HIDDEN_DIM]`; slot K-1 = `[(K-1)*B*HIDDEN_DIM ..]`.
|
||
/// Test-only; DtoH+sync via mapped-pinned staging.
|
||
pub fn read_h_new_per_k_last(&self) -> Result<Vec<f32>> {
|
||
let b_sz = self.cfg.n_batch;
|
||
let k_seq = self.cfg.seq_len;
|
||
let n_hid = HIDDEN_DIM;
|
||
let slot_floats = b_sz * n_hid;
|
||
let slot_bytes = slot_floats * std::mem::size_of::<f32>();
|
||
let staging = unsafe { crate::pinned_mem::MappedF32Buffer::new(slot_floats) }
|
||
.map_err(|e| anyhow::anyhow!("h_new_per_k_last staging alloc: {e}"))?;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (base, _g) = self.h_new_per_k_d.device_ptr(&self.stream);
|
||
let src = base + ((k_seq - 1) * slot_bytes) as u64;
|
||
cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src, slot_bytes, s)
|
||
.map_err(|e| anyhow::anyhow!("h_new_per_k_last DtoD: {e}"))?;
|
||
}
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| anyhow::anyhow!("h_new_per_k_last sync: {e}"))?;
|
||
Ok(staging.read_all())
|
||
}
|
||
|
||
/// X11 checkpoint-loaded constructor: instantiates a PerceptionTrainer
|
||
/// from a Checkpoint file, ready for `forward_only` inference. The
|
||
/// optimizer state + gradient buffers ARE allocated (training-only
|
||
/// state); they're unused at inference but allocating them keeps
|
||
/// the struct invariant simple. For pure-inference deployments where
|
||
/// the memory matters, a leaner "InferenceOnly" variant can be added
|
||
/// later.
|
||
pub fn from_checkpoint(
|
||
dev: &MlDevice,
|
||
cfg: &PerceptionTrainerConfig,
|
||
path: &std::path::Path,
|
||
) -> Result<Self> {
|
||
// Construct a fresh trainer (with random init) to wire up all
|
||
// grad buffers + kernel handles + scratches, then overwrite the
|
||
// trunk's weights from the checkpoint via load_checkpoint.
|
||
let mut trainer = Self::new(dev, cfg)?;
|
||
let trunk_cfg = crate::cfc::trunk::CfcConfig {
|
||
n_in: FEATURE_DIM,
|
||
n_hid: HIDDEN_DIM,
|
||
cfc_n_in: HIDDEN_DIM,
|
||
mamba2_state_dim: cfg.mamba2_state_dim,
|
||
n_batch: cfg.n_batch,
|
||
seq_len: cfg.seq_len,
|
||
};
|
||
let new_trunk = crate::cfc::trunk::CfcTrunk::load_checkpoint(dev, &trunk_cfg, path)
|
||
.context("PerceptionTrainer::from_checkpoint: trunk load")?;
|
||
trainer.trunk = new_trunk;
|
||
Ok(trainer)
|
||
}
|
||
|
||
pub fn evaluate(
|
||
&mut self,
|
||
snapshots: &[Mbp10RawInput],
|
||
labels_per_position: &[[f32; N_HORIZONS]],
|
||
) -> Result<(f32, Vec<f32>)> {
|
||
anyhow::ensure!(
|
||
self.cfg.n_batch == 1,
|
||
"single-sequence evaluate() requires cfg.n_batch == 1; got {}",
|
||
self.cfg.n_batch
|
||
);
|
||
self.evaluate_batched(&[snapshots], &[labels_per_position])
|
||
}
|
||
|
||
/// Forward-only batched evaluation pass. Mirrors the forward chain
|
||
/// of [`step_batched`] but SKIPS backward + grad accumulation +
|
||
/// AdamW. Validation MUST use this — calling step() on val data
|
||
/// leaks the val gradient into trained weights.
|
||
///
|
||
/// Returns (weighted mean BCE, probs[K * B * N_HORIZONS] in
|
||
/// [K, B, N_HORIZONS] row-major layout).
|
||
pub fn evaluate_batched(
|
||
&mut self,
|
||
snapshots_batch: &[&[Mbp10RawInput]],
|
||
labels_batch: &[&[[f32; N_HORIZONS]]],
|
||
) -> Result<(f32, Vec<f32>)> {
|
||
let b_sz = self.cfg.n_batch;
|
||
let k_seq = self.cfg.seq_len;
|
||
anyhow::ensure!(
|
||
snapshots_batch.len() == b_sz && labels_batch.len() == b_sz,
|
||
"expected B={} snapshots+labels; got snap={} lbl={}",
|
||
b_sz, snapshots_batch.len(), labels_batch.len()
|
||
);
|
||
|
||
// Use pre-allocated window_tensor_d — no per-step alloc.
|
||
debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]);
|
||
let total_snaps = b_sz * k_seq;
|
||
debug_assert!(total_snaps <= self.bk_capacity);
|
||
{
|
||
let bid_px_h = self.stg_bid_px_all.host_slice_mut();
|
||
let bid_sz_h = self.stg_bid_sz_all.host_slice_mut();
|
||
let ask_px_h = self.stg_ask_px_all.host_slice_mut();
|
||
let ask_sz_h = self.stg_ask_sz_all.host_slice_mut();
|
||
let regime_h = self.stg_regime_all.host_slice_mut();
|
||
for (b_idx, snapshots) in snapshots_batch.iter().enumerate() {
|
||
for (k, snap) in snapshots.iter().enumerate() {
|
||
let n = b_idx * k_seq + k;
|
||
let base10 = n * 10;
|
||
let base6 = n * REGIME_DIM;
|
||
for i in 0..10 {
|
||
bid_px_h[base10 + i] = snap.bid_px[i];
|
||
bid_sz_h[base10 + i] = snap.bid_sz[i];
|
||
ask_px_h[base10 + i] = snap.ask_px[i];
|
||
ask_sz_h[base10 + i] = snap.ask_sz[i];
|
||
}
|
||
for i in 0..REGIME_DIM {
|
||
regime_h[base6 + i] = snap.regime[i];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
{
|
||
let prev_mid_h = self.stg_prev_mid.host_slice_mut();
|
||
let tsv_h = self.stg_trade_signed_vol.host_slice_mut();
|
||
let tc_h = self.stg_trade_count.host_slice_mut();
|
||
let ts_ns_h = self.stg_ts_ns.host_slice_mut();
|
||
let prev_ts_ns_h = self.stg_prev_ts_ns.host_slice_mut();
|
||
for (b_idx, snapshots) in snapshots_batch.iter().enumerate() {
|
||
for (k, snap) in snapshots.iter().enumerate() {
|
||
let n = b_idx * k_seq + k;
|
||
prev_mid_h[n] = snap.prev_mid;
|
||
tsv_h[n] = snap.trade_signed_vol;
|
||
tc_h[n] = snap.trade_count as i32;
|
||
ts_ns_h[n] = snap.ts_ns as i64;
|
||
prev_ts_ns_h[n] = snap.prev_ts_ns as i64;
|
||
}
|
||
}
|
||
}
|
||
let n10 = total_snaps * 10;
|
||
let n6 = total_snaps * REGIME_DIM;
|
||
let n1 = total_snaps;
|
||
unsafe {
|
||
let s = self.stream.cu_stream();
|
||
let (d, _g) = self.bid_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_px_all.dev_ptr, n10 * 4, s)
|
||
.context("eval bid_px_all dtod")?;
|
||
let (d, _g) = self.bid_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_bid_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("eval bid_sz_all dtod")?;
|
||
let (d, _g) = self.ask_px_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_px_all.dev_ptr, n10 * 4, s)
|
||
.context("eval ask_px_all dtod")?;
|
||
let (d, _g) = self.ask_sz_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ask_sz_all.dev_ptr, n10 * 4, s)
|
||
.context("eval ask_sz_all dtod")?;
|
||
let (d, _g) = self.regime_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_regime_all.dev_ptr, n6 * 4, s)
|
||
.context("eval regime_all dtod")?;
|
||
let (d, _g) = self.prev_mid_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_mid.dev_ptr, n1 * 4, s)
|
||
.context("eval prev_mid_all dtod")?;
|
||
let (d, _g) = self.trade_signed_vol_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_signed_vol.dev_ptr, n1 * 4, s)
|
||
.context("eval tsv_all dtod")?;
|
||
let (d, _g) = self.trade_count_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_trade_count.dev_ptr, n1 * 4, s)
|
||
.context("eval trade_count_all dtod")?;
|
||
let (d, _g) = self.ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("eval ts_ns_all dtod")?;
|
||
let (d, _g) = self.prev_ts_ns_all_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(d, self.stg_prev_ts_ns.dev_ptr, n1 * 8, s)
|
||
.context("eval prev_ts_ns_all dtod")?;
|
||
}
|
||
let tick_size = ES_TICK_SIZE;
|
||
let n_total_i32 = total_snaps as i32;
|
||
let snap_block: u32 = 128;
|
||
let snap_grid: u32 = (total_snaps as u32).div_ceil(snap_block);
|
||
let snap_cfg = LaunchConfig {
|
||
grid_dim: (snap_grid, 1, 1),
|
||
block_dim: (snap_block, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.snap_batched_fn);
|
||
launch
|
||
.arg(&self.bid_px_all_d).arg(&self.bid_sz_all_d)
|
||
.arg(&self.ask_px_all_d).arg(&self.ask_sz_all_d)
|
||
.arg(&self.prev_bid_sz_all_d).arg(&self.prev_ask_sz_all_d)
|
||
.arg(&self.regime_all_d)
|
||
.arg(&self.prev_mid_all_d).arg(&self.trade_signed_vol_all_d)
|
||
.arg(&self.trade_count_all_d)
|
||
.arg(&self.ts_ns_all_d).arg(&self.prev_ts_ns_all_d)
|
||
.arg(&tick_size).arg(&n_total_i32)
|
||
.arg(self.window_tensor_d.data_mut());
|
||
launch.launch(snap_cfg).context("eval snap_batched fwd")?;
|
||
}
|
||
self.stream.synchronize().context("eval snap_batched sync")?;
|
||
|
||
// VSN fwd — same as training path. Eval consumes vsn_out_d.
|
||
{
|
||
let n_rows_vsn: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_vsn = LaunchConfig {
|
||
grid_dim: (n_rows_vsn as u32, 1, 1),
|
||
block_dim: (64, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.vsn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.vsn_w_d)
|
||
.arg(&self.trunk.vsn_b_d)
|
||
.arg(self.window_tensor_d.cuda_data())
|
||
.arg(&n_rows_vsn)
|
||
.arg(self.vsn_out_d.data_mut())
|
||
.arg(&mut self.vsn_gates_d);
|
||
unsafe { launch.launch(cfg_vsn).context("eval variable_selection_fwd")?; }
|
||
}
|
||
|
||
// Mamba2 stack-1 fwd → m1.h_enriched_seq.
|
||
self.trunk
|
||
.mamba2_l1_mut()
|
||
.forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)
|
||
.context("eval mamba2 (l1) fwd_into")?;
|
||
|
||
// LN_a forward → ln_a_out_d.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_a_gain_d)
|
||
.arg(&self.trunk.ln_a_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(self.ln_a_out_d.data_mut())
|
||
.arg(&mut self.ln_a_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("eval layer_norm_fwd (LN_a)")?; }
|
||
}
|
||
|
||
// Mamba2 stack-2 fwd → m2.h_enriched_seq.
|
||
self.trunk
|
||
.mamba2_l2_mut()
|
||
.forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)
|
||
.context("eval mamba2 (l2) fwd_into")?;
|
||
|
||
// LN_b forward over m2.h_enriched_seq — same as training path so
|
||
// eval reads the same distribution downstream layers were
|
||
// trained on.
|
||
{
|
||
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
|
||
let cfg_ln = LaunchConfig {
|
||
grid_dim: (n_rows_ln as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.ln_fwd_fn);
|
||
launch
|
||
.arg(self.mamba2_l2_fwd_scratch.h_enriched_seq.cuda_data())
|
||
.arg(&self.trunk.ln_b_gain_d)
|
||
.arg(&self.trunk.ln_b_bias_d)
|
||
.arg(&n_rows_ln)
|
||
.arg(&mut self.ln_out_d)
|
||
.arg(&mut self.ln_stats_d);
|
||
unsafe { launch.launch(cfg_ln).context("eval layer_norm_fwd (LN_b)")?; }
|
||
}
|
||
|
||
// Transpose LN-normalised h_enriched [B, K, H] → [K, B, H].
|
||
{
|
||
let block_n3: u32 = 32;
|
||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||
let cfg_tx = LaunchConfig {
|
||
grid_dim: (b_sz as u32, k_seq as u32, grid_z),
|
||
block_dim: (block_n3, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let n1 = b_sz as i32;
|
||
let n2 = k_seq as i32;
|
||
let n3 = HIDDEN_DIM as i32;
|
||
let mut launch = self.stream.launch_builder(&self.trunk.transpose_3d_fn);
|
||
launch
|
||
.arg(&self.ln_out_d)
|
||
.arg(self.h_enriched_seq_t_d.data_mut())
|
||
.arg(&n1).arg(&n2).arg(&n3);
|
||
unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; }
|
||
}
|
||
|
||
// Attention pool fwd — same as training, populates attn_context_d
|
||
// for use as k=0 h_old in the eval K-loop below.
|
||
{
|
||
let k_i32 = k_seq as i32;
|
||
let n_batch_attn = b_sz as i32;
|
||
let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::<f32>();
|
||
let cfg_attn = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (128, 1, 1),
|
||
shared_mem_bytes: shared as u32,
|
||
};
|
||
let mut launch = self.stream.launch_builder(&self.trunk.attn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.attn_q_d)
|
||
.arg(&self.ln_out_d)
|
||
.arg(&n_batch_attn).arg(&k_i32)
|
||
.arg(&mut self.attn_context_d)
|
||
.arg(&mut self.attn_weights_d);
|
||
unsafe { launch.launch(cfg_attn).context("eval attention_pool_fwd")?; }
|
||
}
|
||
|
||
// Upload labels [K, B, N_HORIZONS].
|
||
let total_labels = k_seq * b_sz * N_HORIZONS;
|
||
{
|
||
let dst = self.stg_labels.host_slice_mut();
|
||
for k in 0..k_seq {
|
||
for b_idx in 0..b_sz {
|
||
for h in 0..N_HORIZONS {
|
||
dst[(k * b_sz + b_idx) * N_HORIZONS + h] = labels_batch[b_idx][k][h];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
unsafe {
|
||
let (dst_ptr, _g) = self.labels_per_k_d.device_ptr_mut(&self.stream);
|
||
cudarc::driver::result::memcpy_dtod_async(
|
||
dst_ptr, self.stg_labels.dev_ptr,
|
||
total_labels * std::mem::size_of::<f32>(), self.stream.cu_stream(),
|
||
).context("eval labels upload")?;
|
||
}
|
||
self.stream.memset_zeros(&mut self.zero_h_d).map_err(|e| anyhow::anyhow!("zero h: {e}"))?;
|
||
|
||
// Per-K forward (recurrent CfC + heads, batched).
|
||
// A1: decision_stride removed; event-rate dt = 1.0 per CfC step.
|
||
let dt_s = 1.0_f32;
|
||
let n_in_i = HIDDEN_DIM as i32;
|
||
let n_hid_i = HIDDEN_DIM as i32;
|
||
// Phase B fix-up: block-per-batch (Phase B), same launch contract as
|
||
// training's K-loop. BEFORE this fix the eval path used the legacy
|
||
// grid=(1,1,1) config which only processed batch 0 — batches
|
||
// 1..B-1 got GARBAGE h_new. Caught by t6z89-vs-txftz cluster A/B
|
||
// (val_loss=0.83, mean_auc=chance) on 2026-05-17.
|
||
// Shared layout: s_x [n_in] + s_h_old [n_hid].
|
||
let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||
let cfg_cfc = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||
shared_mem_bytes: cfc_fwd_smem,
|
||
};
|
||
let cfg_grn_fwd = LaunchConfig {
|
||
grid_dim: (b_sz as u32, 1, 1),
|
||
block_dim: (HEAD_MID_DIM as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
|
||
let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::<f32>();
|
||
let kb_nh_mid_bytes = b_sz * N_HORIZONS * HEAD_MID_DIM * std::mem::size_of::<f32>();
|
||
let n_batch_i = b_sz as i32;
|
||
let _zero_h_ptr_unused_eval = {
|
||
let (p, _g) = self.zero_h_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
// Phase 3: eval uses attn_context_d as h_old at k=0, mirroring training.
|
||
let attn_context_ptr = {
|
||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let henr_t_base = {
|
||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||
p
|
||
};
|
||
let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||
let (probs_base, _g_probs) = self.probs_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z1_base, _g_z1) = self.z1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (a1_base, _g_a1) = self.a1_per_k_d.device_ptr_mut(&self.stream);
|
||
let (z2_base, _g_z2) = self.z2_per_k_d.device_ptr_mut(&self.stream);
|
||
let (gate_base, _g_gate) = self.gate_logit_per_k_d.device_ptr_mut(&self.stream);
|
||
let (main_base, _g_main) = self.main_per_k_d.device_ptr_mut(&self.stream);
|
||
let (logit_base, _g_logit) = self.logit_per_k_d.device_ptr_mut(&self.stream);
|
||
for k in 0..k_seq {
|
||
let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64;
|
||
let h_old_k_ptr = if k == 0 {
|
||
attn_context_ptr
|
||
} else {
|
||
h_per_k_base + ((k - 1) * kb_hid_bytes) as u64
|
||
};
|
||
let x_k_ptr = henr_t_base + (k * kb_hid_bytes) as u64;
|
||
let probs_k_ptr = probs_base + (k * kb_nh_bytes) as u64;
|
||
let z1_k_ptr = z1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let a1_k_ptr = a1_base + (k * kb_nh_mid_bytes) as u64;
|
||
let z2_k_ptr = z2_base + (k * kb_nh_mid_bytes) as u64;
|
||
let gate_k_ptr = gate_base + (k * kb_nh_bytes) as u64;
|
||
let main_k_ptr = main_base + (k * kb_nh_bytes) as u64;
|
||
let logit_k_ptr = logit_base + (k * kb_nh_bytes) as u64;
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
|
||
launch
|
||
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
|
||
.arg(&x_k_ptr).arg(&h_old_k_ptr)
|
||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
|
||
.arg(&h_new_k_ptr);
|
||
launch.launch(cfg_cfc).context("eval cfc fwd")?;
|
||
}
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn);
|
||
launch
|
||
.arg(&self.trunk.heads_w1_d).arg(&self.trunk.heads_b1_d)
|
||
.arg(&self.trunk.heads_w2_d).arg(&self.trunk.heads_b2_d)
|
||
.arg(&self.trunk.heads_w_gate_d).arg(&self.trunk.heads_b_gate_d)
|
||
.arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d)
|
||
.arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d)
|
||
.arg(&h_new_k_ptr).arg(&n_batch_i)
|
||
.arg(&probs_k_ptr)
|
||
.arg(&z1_k_ptr).arg(&a1_k_ptr).arg(&z2_k_ptr)
|
||
.arg(&gate_k_ptr).arg(&main_k_ptr).arg(&logit_k_ptr);
|
||
launch.launch(cfg_grn_fwd).context("eval heads GRN fwd")?;
|
||
}
|
||
}
|
||
drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit));
|
||
|
||
// BCE for loss reporting (uses same per-horizon weights as train).
|
||
let n_pos_i = (k_seq * b_sz) as i32;
|
||
let n_h_i = N_HORIZONS as i32;
|
||
let bce_cfg = LaunchConfig {
|
||
grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
let mut launch = self.stream.launch_builder(&self.bce_fn);
|
||
launch
|
||
.arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d)
|
||
.arg(&self.log_sigma_h_d)
|
||
.arg(&n_pos_i).arg(&n_h_i)
|
||
.arg(&mut self.loss_d)
|
||
.arg(&mut self.loss_per_horizon_d)
|
||
.arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d)
|
||
.arg(&mut self.grad_log_sigma_h_d);
|
||
launch.launch(bce_cfg).context("eval bce")?;
|
||
}
|
||
self.stream.synchronize().context("eval sync")?;
|
||
|
||
let loss = download(&self.stream, &self.loss_d)?[0];
|
||
let probs = download(&self.stream, &self.probs_per_k_d)?;
|
||
Ok((loss, probs))
|
||
}
|
||
|
||
/// Read per-position probabilities for the most-recent forward
|
||
/// (used by validation / inference). Returns [K, N_HORIZONS] row-major.
|
||
pub fn last_probs_per_position(&self) -> Result<Vec<f32>> {
|
||
download(&self.stream, &self.probs_per_k_d)
|
||
}
|
||
|
||
/// Last-step per-horizon raw mean-squared-diff (pre-λ). Reads from
|
||
/// the mapped-pinned host shadow populated by `step_batched`. Returns
|
||
/// all zeros when the smoothness regularizer is disabled
|
||
/// (`cfg.smoothness_base_lambda == 0.0`).
|
||
pub fn last_smoothness_loss_per_horizon(&self) -> [f32; N_HORIZONS] {
|
||
let raw = self.smoothness_loss_per_horizon_host_d.read_all();
|
||
let mut out = [0.0f32; N_HORIZONS];
|
||
for (i, v) in raw.into_iter().take(N_HORIZONS).enumerate() {
|
||
out[i] = v;
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Convenience for callers that only need the LAST position's probs
|
||
/// (legacy API — final-position AUC scoring).
|
||
pub fn last_probs(&self) -> Result<[f32; N_HORIZONS]> {
|
||
let all = self.last_probs_per_position()?;
|
||
let last = self.cfg.seq_len - 1;
|
||
let mut out = [0f32; N_HORIZONS];
|
||
out.copy_from_slice(&all[last * N_HORIZONS..(last + 1) * N_HORIZONS]);
|
||
Ok(out)
|
||
}
|
||
}
|
||
|
||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||
let n = host.len();
|
||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("stacked upload staging: {e}"))?;
|
||
staging.write_from_slice(host);
|
||
let mut dst = stream.alloc_zeros::<f32>(n).context("stacked upload alloc")?;
|
||
if n > 0 {
|
||
let nbytes = n * std::mem::size_of::<f32>();
|
||
unsafe {
|
||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
|
||
.context("stacked upload DtoD")?;
|
||
}
|
||
}
|
||
Ok(dst)
|
||
}
|
||
|
||
|
||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||
let n = src.len();
|
||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||
.map_err(|e| anyhow::anyhow!("stacked download staging: {e}"))?;
|
||
let nbytes = n * std::mem::size_of::<f32>();
|
||
unsafe {
|
||
let (src_ptr, _g) = src.device_ptr(stream);
|
||
cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
|
||
.context("stacked download DtoD")?;
|
||
}
|
||
stream.synchronize().context("stacked download sync")?;
|
||
Ok(staging.read_all())
|
||
}
|
||
|
||
// ── Drop: shut down the GPU log drain thread (feature-gated) ──
|
||
// Setting the stop flag lets the drainer exit cleanly on its next 500ms
|
||
// tick; the drain thread observes the flag at the top of each iteration,
|
||
// flushes the BufWriter, and exits. We then `join()` synchronously so
|
||
// the trace file is fully flushed by the time the trainer is dropped.
|
||
// The bounded 500 ms tick guarantees the join completes promptly.
|
||
#[cfg(feature = "kernel-step-trace")]
|
||
impl Drop for PerceptionTrainer {
|
||
fn drop(&mut self) {
|
||
use std::sync::atomic::Ordering;
|
||
if let Some(stop) = self.log_drain_stop.as_ref() {
|
||
stop.store(true, Ordering::Relaxed);
|
||
}
|
||
if let Some(h) = self.log_drain_handle.take() {
|
||
// Drain thread observes the stop flag on its next tick (≤500 ms),
|
||
// flushes the BufWriter, and exits. `join()` blocks until the
|
||
// thread terminates — guarantees the JSONL tail lands on disk.
|
||
if let Err(e) = h.join() {
|
||
tracing::warn!(target: "gpu_log",
|
||
error = ?e,
|
||
"gpu_log_drain thread join failed");
|
||
}
|
||
}
|
||
}
|
||
}
|