diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 931568813..29e719737 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -112,6 +112,7 @@ const KERNELS: &[&str] = &[ "rl_hindsight_track", // HER Phase 1: per-step mid-price ring + peak tracking for backward hindsight "rl_hindsight_inject", // HER Phase 2: backward inject — synthetic replay push on done if peak >> actual "rl_hindsight_forward", // HER Phase 3: forward continuation — evaluates closed trades after lookahead + "rl_lr_from_mapped_pinned", // Mega-graph: LR controller reads loss from mapped-pinned dev_ptrs (eliminates host loss readback between reward + training graphs) "rl_increment_step", // device-resident step counter bump (ISV[548] += 1.0); graph-safe prereq — removes scalar current_step from all downstream kernel args "rl_fused_controllers", // fused kernel: 10 RL ISV controllers in one launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda) — saves 9 kernel launch overheads (~40-80μs/step) "rl_popart_normalize", // PopArt: Welford-EMA reward normalization (replaces apply_reward_scale) diff --git a/crates/ml-alpha/cuda/adamw_step.cu b/crates/ml-alpha/cuda/adamw_step.cu index 68ea650f1..63ef16b4a 100644 --- a/crates/ml-alpha/cuda/adamw_step.cu +++ b/crates/ml-alpha/cuda/adamw_step.cu @@ -44,3 +44,41 @@ extern "C" __global__ void adamw_step( theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]); } + +// Mega-graph variant: reads LR from an ISV device pointer at a given +// slot index instead of taking a scalar argument. This allows the LR +// to vary across graph replays (the ISV slot is modified in-place by +// the rl_lr_from_mapped_pinned controller kernel which is captured +// earlier in the same graph). All other args are identical to adamw_step. +extern "C" __global__ void adamw_step_isv_lr( + float* __restrict__ theta, + const float* __restrict__ grad, + float* __restrict__ m, + float* __restrict__ v, + int n_params, + const float* __restrict__ isv, + int lr_slot, + float beta1, + float beta2, + float eps, + float wd, + int step +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_params) return; + + const float lr = isv[lr_slot]; + + const float g = grad[i]; + const float m_n = beta1 * m[i] + (1.0f - beta1) * g; + const float v_n = beta2 * v[i] + (1.0f - beta2) * g * g; + m[i] = m_n; + v[i] = v_n; + + const float bc1 = 1.0f - powf(beta1, (float) step); + const float bc2 = 1.0f - powf(beta2, (float) step); + const float m_hat = m_n / fmaxf(bc1, 1e-12f); + const float v_hat = v_n / fmaxf(bc2, 1e-12f); + + theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]); +} diff --git a/crates/ml-alpha/cuda/grad_h_accumulate.cu b/crates/ml-alpha/cuda/grad_h_accumulate.cu index 871d14f63..613bf0c12 100644 --- a/crates/ml-alpha/cuda/grad_h_accumulate.cu +++ b/crates/ml-alpha/cuda/grad_h_accumulate.cu @@ -36,3 +36,20 @@ extern "C" __global__ void grad_h_accumulate_scaled( if (i >= n) return; grad_h_encoder[i] += lambda * grad_h_head[i]; } + +// Mega-graph variant: reads lambda from ISV device pointer at given slot. +// This allows the lambda to vary across graph replays (ISV is modified +// in-place by controller kernels captured earlier in the same graph). +extern "C" __global__ void grad_h_accumulate_scaled_isv( + const float* __restrict__ grad_h_head, + const float* __restrict__ isv, + int lambda_slot, + float batch_inv, // 1.0 / b_size + int n, + float* __restrict__ grad_h_encoder +) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float lambda = isv[lambda_slot] * batch_inv; + grad_h_encoder[i] += lambda * grad_h_head[i]; +} diff --git a/crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu b/crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu new file mode 100644 index 000000000..e0832874c --- /dev/null +++ b/crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu @@ -0,0 +1,134 @@ +// rl_lr_from_mapped_pinned.cu — GPU-side LR controller that reads loss +// observations from mapped-pinned device pointers instead of host- +// supplied scalar arguments. This eliminates the host-side loss readback +// between the reward graph and training graph, enabling mega-graph +// capture of the entire per-step pipeline. +// +// The mapped-pinned buffers (ss_q_loss_mapped, ss_pi_loss_mapped, +// ss_v_loss_sum_mapped) have stable device pointers. The GPU writes +// loss values during backward kernels; this kernel reads them from the +// same physical memory via the device-side pointer. The reads are +// coherent because this kernel launches AFTER the backward kernels on +// the same stream (stream ordering guarantees visibility). +// +// Internally delegates to the same plateau_decay_head logic as +// rl_lr_controller.cu. The only difference is the loss observation +// source: pointer dereference instead of scalar argument. + +#define RL_LR_BCE_INDEX 412 +#define RL_LR_Q_INDEX 413 +#define RL_LR_PI_INDEX 414 +#define RL_LR_V_INDEX 415 +#define RL_LR_AUX_INDEX 416 + +#define RL_LR_BOOTSTRAP_INDEX 462 +#define RL_LR_MIN_INDEX 463 +#define RL_LR_MAX_INDEX 464 +#define RL_LR_LOSS_EMA_ALPHA_INDEX 465 +#define RL_LR_DECAY_FACTOR_INDEX 466 + +#define RL_IMPROVEMENT_THRESHOLD_INDEX 455 +#define RL_PLATEAU_PATIENCE_INDEX 456 +#define RL_LR_WARMUP_STEPS_INDEX 461 + +__device__ __forceinline__ void plateau_decay_head( + float* isv, + int lr_idx, + float observed_loss, + int loss_ema_slot, + int best_slot, + int counter_slot, + int warmup_slot +) { + const float lr_prev = isv[lr_idx]; + const float lr_bootstrap = isv[RL_LR_BOOTSTRAP_INDEX]; + + if (lr_prev == 0.0f) { + isv[lr_idx] = lr_bootstrap; + return; + } + + if (observed_loss == 0.0f) return; + if (loss_ema_slot < 0) return; + + const float loss_ema_prev = isv[loss_ema_slot]; + const float loss_ema_alpha = isv[RL_LR_LOSS_EMA_ALPHA_INDEX]; + float loss_ema_new; + if (loss_ema_prev == 0.0f) { + loss_ema_new = observed_loss; + } else { + loss_ema_new = (1.0f - loss_ema_alpha) * loss_ema_prev + + loss_ema_alpha * observed_loss; + } + isv[loss_ema_slot] = loss_ema_new; + + const float warmup_prev = isv[warmup_slot]; + const float warmup_target = isv[RL_LR_WARMUP_STEPS_INDEX]; + if (warmup_prev < warmup_target) { + isv[best_slot] = loss_ema_new; + isv[counter_slot] = 0.0f; + isv[warmup_slot] = warmup_prev + 1.0f; + return; + } + + const float improvement_threshold = isv[RL_IMPROVEMENT_THRESHOLD_INDEX]; + const float plateau_patience = isv[RL_PLATEAU_PATIENCE_INDEX]; + const float best_prev = isv[best_slot]; + + if (loss_ema_new < best_prev * improvement_threshold) { + isv[best_slot] = loss_ema_new; + isv[counter_slot] = 0.0f; + return; + } + + const float counter_next = isv[counter_slot] + 1.0f; + if (counter_next >= plateau_patience) { + const float lr_min = isv[RL_LR_MIN_INDEX]; + const float decay_factor = isv[RL_LR_DECAY_FACTOR_INDEX]; + const float lr_new = fmaxf(lr_min, lr_prev * decay_factor); + isv[lr_idx] = lr_new; + isv[counter_slot] = 0.0f; + } else { + isv[counter_slot] = counter_next; + } +} + +extern "C" __global__ void rl_lr_from_mapped_pinned( + float* __restrict__ isv, + const float* __restrict__ q_loss_ptr, // mapped-pinned dev_ptr (1 float) + const float* __restrict__ pi_loss_ptr, // mapped-pinned dev_ptr (1 float) + const float* __restrict__ v_loss_sum_ptr, // mapped-pinned dev_ptr (1 float) + int b_size, // batch size for V loss normalization + int q_loss_ema_slot, + int q_best_slot, + int q_counter_slot, + int q_warmup_slot, + int pi_loss_ema_slot, + int pi_best_slot, + int pi_counter_slot, + int pi_warmup_slot, + int v_loss_ema_slot, + int v_best_slot, + int v_counter_slot, + int v_warmup_slot +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + // Read losses from mapped-pinned device pointers. These are the + // same physical pages the backward kernels wrote to — coherent + // because this kernel is stream-ordered after them. + const float observed_loss_q = *q_loss_ptr; + const float observed_loss_pi = *pi_loss_ptr; + // V loss is stored as a sum across the batch; normalize to per-sample. + const float observed_loss_v = (b_size > 0) ? (*v_loss_sum_ptr / (float)b_size) : 0.0f; + + // BCE and AUX: perception-owned heads, pass slot=-1 to skip. + plateau_decay_head(isv, RL_LR_BCE_INDEX, 0.0f, -1, -1, -1, -1); + plateau_decay_head(isv, RL_LR_Q_INDEX, observed_loss_q, + q_loss_ema_slot, q_best_slot, q_counter_slot, q_warmup_slot); + plateau_decay_head(isv, RL_LR_PI_INDEX, observed_loss_pi, + pi_loss_ema_slot, pi_best_slot, pi_counter_slot, pi_warmup_slot); + plateau_decay_head(isv, RL_LR_V_INDEX, observed_loss_v, + v_loss_ema_slot, v_best_slot, v_counter_slot, v_warmup_slot); + plateau_decay_head(isv, RL_LR_AUX_INDEX, 0.0f, -1, -1, -1, -1); +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 43bd82041..c769b998a 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -979,6 +979,11 @@ fn main() -> Result<()> { 0 }; + // Mega-graph: capture the entire per-step pipeline into a single + // cuGraphLaunch call. Eliminates ~150 individual kernel launches + // and ~140ms of host-side Rust overhead per step. + trainer.enable_mega_graph(); + let t_start = std::time::Instant::now(); for step in start_step..cli.n_steps { // GPU-resident data loading: sample_and_gather + gather_next + diff --git a/crates/ml-alpha/src/rl/reward.rs b/crates/ml-alpha/src/rl/reward.rs index 5c7b635ab..b2cd54419 100644 --- a/crates/ml-alpha/src/rl/reward.rs +++ b/crates/ml-alpha/src/rl/reward.rs @@ -46,9 +46,57 @@ use anyhow::Result; use cudarc::driver::CudaSlice; +use cudarc::driver::sys::CUfunction; use crate::cfc::snap_features::Mbp10RawInput; +/// Cached raw device pointers and kernel function handles for the lobsim. +/// Extracted once per step to bypass trait dispatch + CudaSlice borrow +/// overhead in the mega-graph hot path. All pointers are stable (pre- +/// allocated at lobsim construction time, never reallocated). +pub struct LobSimRawPtrs { + // ── Book update (apply_snapshot_from_device) ──────────────────── + pub bid_px: u64, + pub bid_sz: u64, + pub ask_px: u64, + pub ask_sz: u64, + pub books: u64, + pub prev_mid: u64, + pub atr_mid_ema: u64, + pub snapshots_skipped: u64, + pub min_reasonable_px: u64, + pub max_reasonable_px: u64, + pub book_update_fn: CUfunction, + + // ── Fill (step_fill_from_market_targets) ──────────────────────── + pub market_targets: u64, + pub pos: u64, + pub cost_per_lot_per_side: u64, + pub total_fees_per_b: u64, + pub submit_market_fn: CUfunction, + + // ── PnL track (step_pnl_track) ───────────────────────────────── + pub open_trade_state: u64, + pub trade_log: u64, + pub trade_log_head: u64, + pub trail_hwm: u64, + pub zero_vwap_at_open: u64, + pub saturated_vwap_at_open: u64, + pub defensive_exit_clamp: u64, + pub conv_signed_ema: u64, + pub diag_hold_hist: u64, + pub diag_outcome_n: u64, + pub diag_outcome_sum_pnl: u64, + pub diag_outcome_n_wins: u64, + pub pnl_track_fn: CUfunction, + + // ── Dimensions ───────────────────────────────────────────────── + pub n_backtests: i32, + pub pos_bytes: i32, + /// `TRADE_LOG_CAP` from ml-backtesting — passed to `pnl_track` kernel. + pub trade_log_cap: i32, +} + /// Narrow device-oriented backend the integrated RL trainer's /// `step_with_lobsim` invokes for its LOB-simulator interaction. /// Implemented for `LobSimCuda` in `ml-backtesting`. See module-level @@ -126,4 +174,10 @@ pub trait RlLobBackend { ask_px_src: u64, ask_sz_src: u64, ) -> Result<()>; + + /// Extract all raw device pointers and kernel function handles needed + /// for mega-graph capture. Returns a [`LobSimRawPtrs`] struct that the + /// trainer caches to bypass trait dispatch inside the captured section. + /// All pointers are stable for the lobsim's lifetime. + fn raw_ptrs(&self) -> LobSimRawPtrs; } diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 16708341f..55921509e 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -115,6 +115,8 @@ const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); const RL_LR_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_lr_controller.cubin")); +const RL_LR_FROM_MAPPED_PINNED_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_lr_from_mapped_pinned.cubin")); // Phase R1: cubin includes for the 7 RL adaptive controllers // (γ / τ / ε / entropy_coef / n_rollout_steps / per_α / reward_scale). @@ -386,6 +388,9 @@ impl LobPtrs { } } +// LobSimRawPtrs from crate::rl::reward is used via lobsim.raw_ptrs() +// in step_with_lobsim_gpu_body for the inlined lobsim kernel launches. + /// Configuration for [`IntegratedTrainer`]. Wraps a `PerceptionTrainerConfig` /// (the encoder + BCE + aux side) plus RL-specific overrides for the new /// heads. @@ -535,6 +540,13 @@ pub struct IntegratedTrainer { /// only (no host stall), saving ~2ms/step host idle time. train_done_event: CudaEvent, + /// Mega-graph Phase 3: when true, all PER operations (sample, + /// priority_update, tree_rebuild) run on `self.raw_stream` instead + /// of `self.raw_train_stream`, and cross-stream events are skipped. + /// Requires K=1 (single training iteration per env step). Set by + /// the mega-graph path before calling `step_with_lobsim_reward_and_train`. + mega_graph_single_stream: bool, + // ── CUDA Graph capture for the RL step pipeline ───────────────── // Three-state machine: first step = warmup (eager), second = // capture, third+ = replay. Split around apply_snapshot (HtoD @@ -546,9 +558,21 @@ pub struct IntegratedTrainer { training_graph: Option, training_warmup_done: bool, + /// Mega-graph: single CUDA graph capturing the ENTIRE per-step + /// pipeline (gathers + encoder + Q/V/pi + gates + lobsim + reward + + /// PER + training + target). Replaces all 4 separate graphs when + /// `mega_graph_enabled` is true. + mega_graph: Option, + mega_warmup_done: bool, + /// When true, `step_with_lobsim_gpu` uses the mega-graph path + /// instead of the 4 separate graphs. Set after Phase 4 is validated. + mega_graph_enabled: bool, + // ── Kernel handles for grad combine + cross-batch reduce ────────── _grad_h_module: Arc, grad_h_accumulate_fn: CudaFunction, + /// Mega-graph variant: reads lambda from ISV device pointer. + grad_h_accumulate_isv_fn: CudaFunction, _reduce_axis0_module: Arc, reduce_axis0_fn: CudaFunction, @@ -561,6 +585,17 @@ pub struct IntegratedTrainer { /// (sentinel zero → `LR_BOOTSTRAP = 1e-3`). rl_lr_controller_fn: CudaFunction, + /// Mega-graph variant: reads loss from mapped-pinned device pointers + /// instead of host-supplied scalar args. Eliminates the host-side + /// loss readback between the reward and training graphs. + _rl_lr_from_mapped_pinned_module: Arc, + rl_lr_from_mapped_pinned_fn: CudaFunction, + + /// Mega-graph variant: AdamW that reads LR from ISV device pointer + /// instead of a scalar arg. Loaded from the same cubin as adamw_step. + _adamw_isv_lr_module: Arc, + adamw_step_isv_lr_fn: CudaFunction, + // ── 7 RL adaptive controllers (Phase R1) ────────────────────────── // Each controller emits ONE float into its dedicated ISV slot // (γ→400, τ→401, ε→402, entropy_coef→403, n_rollout_steps→404, @@ -1284,6 +1319,9 @@ impl IntegratedTrainer { let grad_h_accumulate_fn = grad_h_module .load_function("grad_h_accumulate_scaled") .context("load grad_h_accumulate_scaled")?; + let grad_h_accumulate_isv_fn = grad_h_module + .load_function("grad_h_accumulate_scaled_isv") + .context("load grad_h_accumulate_scaled_isv")?; let reduce_axis0_module = ctx .load_cubin(REDUCE_AXIS0_CUBIN.to_vec()) .context("load reduce_axis0 cubin")?; @@ -1296,6 +1334,19 @@ impl IntegratedTrainer { let rl_lr_controller_fn = rl_lr_controller_module .load_function("rl_lr_controller") .context("load rl_lr_controller")?; + let rl_lr_from_mapped_pinned_module = ctx + .load_cubin(RL_LR_FROM_MAPPED_PINNED_CUBIN.to_vec()) + .context("load rl_lr_from_mapped_pinned cubin")?; + let rl_lr_from_mapped_pinned_fn = rl_lr_from_mapped_pinned_module + .load_function("rl_lr_from_mapped_pinned") + .context("load rl_lr_from_mapped_pinned")?; + // Mega-graph AdamW variant — loaded from the same cubin as adamw_step. + let adamw_module_for_isv = ctx + .load_cubin(include_bytes!(concat!(env!("OUT_DIR"), "/adamw_step.cubin")).to_vec()) + .context("load adamw_step cubin (isv_lr variant)")?; + let adamw_step_isv_lr_fn = adamw_module_for_isv + .load_function("adamw_step_isv_lr") + .context("load adamw_step_isv_lr")?; // Phase R1: load the 7 RL adaptive controllers. let rl_gamma_controller_module = ctx @@ -2327,18 +2378,27 @@ impl IntegratedTrainer { sample_done_event, replay_done_event, train_done_event, + mega_graph_single_stream: false, prefill_graph: None, postfill_graph: None, reward_graph: None, graph_warmup_done: false, training_graph: None, training_warmup_done: false, + mega_graph: None, + mega_warmup_done: false, + mega_graph_enabled: false, _grad_h_module: grad_h_module, grad_h_accumulate_fn, + grad_h_accumulate_isv_fn, _reduce_axis0_module: reduce_axis0_module, reduce_axis0_fn, _rl_lr_controller_module: rl_lr_controller_module, rl_lr_controller_fn, + _rl_lr_from_mapped_pinned_module: rl_lr_from_mapped_pinned_module, + rl_lr_from_mapped_pinned_fn, + _adamw_isv_lr_module: adamw_module_for_isv, + adamw_step_isv_lr_fn, _rl_gamma_controller_module: rl_gamma_controller_module, rl_gamma_controller_fn, _rl_target_tau_controller_module: rl_target_tau_controller_module, @@ -3888,59 +3948,46 @@ impl IntegratedTrainer { let b_size = self.cfg.perception.n_batch; // ── Step 2: per-head LR controller emit ────────────────────── - // The controller emits the bootstrap target on first call - // (sentinel-zero → 1e-3); subsequent calls Wiener-α blend. - // Uses last_*_loss populated by the deferred readback above - // (one-step delayed). At construction these are 0.0 (sentinel) - // and the controller's cold-start gate (early-return on - // observed_loss == 0) holds LR at LR_BOOTSTRAP for the first - // step. - self.launch_rl_lr_controller( - self.last_q_loss, - self.last_pi_loss, - self.last_v_loss, - ) - .context("rl_lr_controller launch")?; + if self.mega_graph_enabled { + // Mega-graph path: use the mapped-pinned variant that reads + // loss observations from device pointers. The LR controller + // writes into ISV[412..417]; the adamw_step_isv_lr kernel + // reads from those same ISV slots — no host roundtrip. + self.launch_rl_lr_from_mapped_pinned() + .context("rl_lr_from_mapped_pinned launch")?; + } else { + // Standard path: host reads losses, passes as scalar args. + self.launch_rl_lr_controller( + self.last_q_loss, + self.last_pi_loss, + self.last_v_loss, + ) + .context("rl_lr_controller launch")?; - let lambdas = read_loss_lambdas_from_isv(self.isv_host_slice()); - // Mutate each per-head Adam's lr field from the ISV mapped-pinned - // host_ptr. The controller has just emitted into ISV[412..417], - // so this read reflects this step's learning rate. Reading ALL - // five before touching ANY AdamW keeps the field-by-field - // mutations borrow-checker safe. - let lr_bce = self.read_isv_host(RL_LR_BCE_INDEX); - let lr_q = self.read_isv_host(RL_LR_Q_INDEX); - let lr_pi = self.read_isv_host(RL_LR_PI_INDEX); - let lr_v = self.read_isv_host(RL_LR_V_INDEX); - let lr_aux = self.read_isv_host(RL_LR_AUX_INDEX); - // The aux head LR is consumed by the PerceptionTrainer's aux - // optimisers; the perception trainer reads it through its own - // mutator (`set_lr_aux`). The integrated trainer only needs to - // propagate the value — its own Adam instances are Q / π / V. - // BCE LR currently rides the existing `lr_cfc` on the perception - // side (the BCE head's grad path is interlocked with the CfC - // backward through the K-loop). Phase E.3+ separates the BCE - // optimiser; until then we keep the BCE slot in ISV updated for - // forward-looking diagnostics + the controller's first-obs - // bootstrap, and we propagate `lr_aux` via the perception - // trainer's existing mutator hook. - let _ = lr_bce; - self.perception.set_lr_aux(lr_aux); - self.dqn_w_adam.lr = lr_q; - self.dqn_b_adam.lr = lr_q; - self.policy_w_adam.lr = lr_pi; - self.policy_b_adam.lr = lr_pi; - self.value_w_adam.lr = lr_v; - self.value_b_adam.lr = lr_v; - // SP20 P3 FRD head — LR from dedicated ISV slot (498-503 block). - let lr_frd = self.read_isv_host(crate::rl::isv_slots::RL_FRD_LR_INDEX); - self.frd_w1_adam.lr = lr_frd; - self.frd_b1_adam.lr = lr_frd; - self.frd_w2_adam.lr = lr_frd; - self.frd_b2_adam.lr = lr_frd; - // Outcome head shares the FRD LR slot (both are lightweight aux heads). - self.outcome_w_adam.lr = lr_frd; - self.outcome_b_adam.lr = lr_frd; + // Mutate each per-head Adam's lr field from the ISV + // mapped-pinned host_ptr. Only needed in non-mega mode + // because the standard adamw_step takes lr as a scalar arg. + let lr_bce = self.read_isv_host(RL_LR_BCE_INDEX); + let lr_q = self.read_isv_host(RL_LR_Q_INDEX); + let lr_pi = self.read_isv_host(RL_LR_PI_INDEX); + let lr_v = self.read_isv_host(RL_LR_V_INDEX); + let lr_aux = self.read_isv_host(RL_LR_AUX_INDEX); + let _ = lr_bce; + self.perception.set_lr_aux(lr_aux); + self.dqn_w_adam.lr = lr_q; + self.dqn_b_adam.lr = lr_q; + self.policy_w_adam.lr = lr_pi; + self.policy_b_adam.lr = lr_pi; + self.value_w_adam.lr = lr_v; + self.value_b_adam.lr = lr_v; + let lr_frd = self.read_isv_host(crate::rl::isv_slots::RL_FRD_LR_INDEX); + self.frd_w1_adam.lr = lr_frd; + self.frd_b1_adam.lr = lr_frd; + self.frd_w2_adam.lr = lr_frd; + self.frd_b2_adam.lr = lr_frd; + self.outcome_w_adam.lr = lr_frd; + self.outcome_b_adam.lr = lr_frd; + } // Borrow encoder hidden state for forward kernels. let h_t_borrow: &CudaSlice = self.perception.h_t_view(); @@ -3951,7 +3998,7 @@ impl IntegratedTrainer { // capture, third+ = replay. All device pointers are stable // (pre-allocated ss_* fields). Loss DtoDs (async, no sync) // stay OUTSIDE the captured region. - if self.training_graph.is_some() { + if self.training_graph.is_some() && !self.mega_graph_enabled { unsafe { raw_graph_launch( self.training_graph.as_ref().unwrap().cu_graph_exec, @@ -3959,7 +4006,7 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("training graph launch: {:?}", e))?; } } else { - let capturing_training = self.training_warmup_done; + let capturing_training = self.training_warmup_done && !self.mega_graph_enabled; if !self.training_warmup_done { self.training_warmup_done = true; } @@ -4177,24 +4224,48 @@ impl IntegratedTrainer { // Host read deferred to after graph capture region. // ── Step 9: Adam updates on each head's w and b ────────────── - self.dqn_w_adam - .step(&mut self.dqn_head.w_d, &self.ss_q_grad_w_d) - .context("dqn_w_adam.step")?; - self.dqn_b_adam - .step(&mut self.dqn_head.b_d, &self.ss_q_grad_b_d) - .context("dqn_b_adam.step")?; - self.policy_w_adam - .step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d) - .context("policy_w_adam.step")?; - self.policy_b_adam - .step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d) - .context("policy_b_adam.step")?; - self.value_w_adam - .step(&mut self.value_head.w_d, &self.ss_v_grad_w_d) - .context("value_w_adam.step")?; - self.value_b_adam - .step(&mut self.value_head.b_d, &self.ss_v_grad_b_d) - .context("value_b_adam.step")?; + if self.mega_graph_enabled { + // Mega-graph path: Adam reads LR from ISV device pointer. + let isv_fn = self.adamw_step_isv_lr_fn.cu_function(); + let isv = self.isv_dev_ptr; + self.dqn_w_adam + .step_isv_lr(&mut self.dqn_head.w_d, &self.ss_q_grad_w_d, isv_fn, isv, RL_LR_Q_INDEX as i32) + .context("dqn_w_adam.step_isv_lr")?; + self.dqn_b_adam + .step_isv_lr(&mut self.dqn_head.b_d, &self.ss_q_grad_b_d, isv_fn, isv, RL_LR_Q_INDEX as i32) + .context("dqn_b_adam.step_isv_lr")?; + self.policy_w_adam + .step_isv_lr(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d, isv_fn, isv, RL_LR_PI_INDEX as i32) + .context("policy_w_adam.step_isv_lr")?; + self.policy_b_adam + .step_isv_lr(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d, isv_fn, isv, RL_LR_PI_INDEX as i32) + .context("policy_b_adam.step_isv_lr")?; + self.value_w_adam + .step_isv_lr(&mut self.value_head.w_d, &self.ss_v_grad_w_d, isv_fn, isv, RL_LR_V_INDEX as i32) + .context("value_w_adam.step_isv_lr")?; + self.value_b_adam + .step_isv_lr(&mut self.value_head.b_d, &self.ss_v_grad_b_d, isv_fn, isv, RL_LR_V_INDEX as i32) + .context("value_b_adam.step_isv_lr")?; + } else { + self.dqn_w_adam + .step(&mut self.dqn_head.w_d, &self.ss_q_grad_w_d) + .context("dqn_w_adam.step")?; + self.dqn_b_adam + .step(&mut self.dqn_head.b_d, &self.ss_q_grad_b_d) + .context("dqn_b_adam.step")?; + self.policy_w_adam + .step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d) + .context("policy_w_adam.step")?; + self.policy_b_adam + .step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d) + .context("policy_b_adam.step")?; + self.value_w_adam + .step(&mut self.value_head.w_d, &self.ss_v_grad_w_d) + .context("value_w_adam.step")?; + self.value_b_adam + .step(&mut self.value_head.b_d, &self.ss_v_grad_b_d) + .context("value_b_adam.step")?; + } // ── Step 9b: Spectral decouple on Q logits + π logits ──────── // Adds lambda * mean(logits^2) penalty to the Q and π loss accumulators @@ -4276,12 +4347,24 @@ impl IntegratedTrainer { &mut self.ss_outcome_grad_b_d, )?; // Adam step on outcome W and b. - self.outcome_w_adam - .step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d) - .context("outcome_w_adam.step")?; - self.outcome_b_adam - .step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d) - .context("outcome_b_adam.step")?; + if self.mega_graph_enabled { + let isv_fn = self.adamw_step_isv_lr_fn.cu_function(); + let isv = self.isv_dev_ptr; + let frd_slot = crate::rl::isv_slots::RL_FRD_LR_INDEX as i32; + self.outcome_w_adam + .step_isv_lr(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d, isv_fn, isv, frd_slot) + .context("outcome_w_adam.step_isv_lr")?; + self.outcome_b_adam + .step_isv_lr(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d, isv_fn, isv, frd_slot) + .context("outcome_b_adam.step_isv_lr")?; + } else { + self.outcome_w_adam + .step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d) + .context("outcome_w_adam.step")?; + self.outcome_b_adam + .step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d) + .context("outcome_b_adam.step")?; + } } // ── SP20 P3 FRD backward chain (F.4) ───────────────────────── @@ -4330,18 +4413,36 @@ impl IntegratedTrainer { reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_frd_grad_b2_pb_d, b_size, frd_out_dim, &mut self.ss_frd_grad_b2_d)?; // Adam steps — sentinel labels yield zero grads → no // weight movement (only momentum decay). - self.frd_w1_adam - .step(&mut self.frd_head.w1_d, &self.ss_frd_grad_w1_d) - .context("frd_w1_adam.step")?; - self.frd_b1_adam - .step(&mut self.frd_head.b1_d, &self.ss_frd_grad_b1_d) - .context("frd_b1_adam.step")?; - self.frd_w2_adam - .step(&mut self.frd_head.w2_d, &self.ss_frd_grad_w2_d) - .context("frd_w2_adam.step")?; - self.frd_b2_adam - .step(&mut self.frd_head.b2_d, &self.ss_frd_grad_b2_d) - .context("frd_b2_adam.step")?; + if self.mega_graph_enabled { + let isv_fn = self.adamw_step_isv_lr_fn.cu_function(); + let isv = self.isv_dev_ptr; + let frd_slot = crate::rl::isv_slots::RL_FRD_LR_INDEX as i32; + self.frd_w1_adam + .step_isv_lr(&mut self.frd_head.w1_d, &self.ss_frd_grad_w1_d, isv_fn, isv, frd_slot) + .context("frd_w1_adam.step_isv_lr")?; + self.frd_b1_adam + .step_isv_lr(&mut self.frd_head.b1_d, &self.ss_frd_grad_b1_d, isv_fn, isv, frd_slot) + .context("frd_b1_adam.step_isv_lr")?; + self.frd_w2_adam + .step_isv_lr(&mut self.frd_head.w2_d, &self.ss_frd_grad_w2_d, isv_fn, isv, frd_slot) + .context("frd_w2_adam.step_isv_lr")?; + self.frd_b2_adam + .step_isv_lr(&mut self.frd_head.b2_d, &self.ss_frd_grad_b2_d, isv_fn, isv, frd_slot) + .context("frd_b2_adam.step_isv_lr")?; + } else { + self.frd_w1_adam + .step(&mut self.frd_head.w1_d, &self.ss_frd_grad_w1_d) + .context("frd_w1_adam.step")?; + self.frd_b1_adam + .step(&mut self.frd_head.b1_d, &self.ss_frd_grad_b1_d) + .context("frd_b1_adam.step")?; + self.frd_w2_adam + .step(&mut self.frd_head.w2_d, &self.ss_frd_grad_w2_d) + .context("frd_w2_adam.step")?; + self.frd_b2_adam + .step(&mut self.frd_head.b2_d, &self.ss_frd_grad_b2_d) + .context("frd_b2_adam.step")?; + } // Host read deferred to after graph capture region. } @@ -4375,37 +4476,89 @@ impl IntegratedTrainer { // batch-size invariance. Without this, doubling B doubles the // encoder gradient from each RL head. let b_inv = 1.0 / b_size as f32; - accumulate_grad_h( - &self.stream, - &self.grad_h_accumulate_fn, - &self.ss_pi_grad_h_t_d, - lambdas.pi * b_inv, - &mut self.grad_h_t_combined_d, - )?; - accumulate_grad_h( - &self.stream, - &self.grad_h_accumulate_fn, - &self.ss_v_grad_h_t_d, - lambdas.v * b_inv, - &mut self.grad_h_t_combined_d, - )?; - accumulate_grad_h( - &self.stream, - &self.grad_h_accumulate_fn, - &self.ss_frd_grad_h_t_d, - lambdas.frd * b_inv, - &mut self.grad_h_t_combined_d, - )?; - { - let outcome_lambda = - self.read_isv_host(crate::rl::isv_slots::RL_OUTCOME_AUX_LAMBDA_INDEX); + if self.mega_graph_enabled { + // Mega-graph: read lambda from ISV on-device so the values + // update across graph replays (not baked at capture time). + let isv_fn = self.grad_h_accumulate_isv_fn.cu_function(); + let isv = self.isv_dev_ptr; + let n = self.grad_h_t_combined_d.len() as i32; + let raw_s = self.raw_stream; + // Helper closure for the ISV-lambda variant + let launch_isv = |grad_ptr: u64, slot: i32, dst_ptr: u64| -> Result<()> { + let grid_x = (n as u32).div_ceil(256); + let mut args = RawArgs::new(); + args.push_ptr(grad_ptr); + args.push_ptr(isv); + args.push_i32(slot); + args.push_f32(b_inv); + args.push_i32(n); + args.push_ptr(dst_ptr); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + isv_fn, + (grid_x, 1, 1), (256, 1, 1), 0, + raw_s, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("grad_h_accumulate_isv: {:?}", e))?; + } + Ok(()) + }; + let dst = self.grad_h_t_combined_d.raw_ptr(); + launch_isv( + self.ss_pi_grad_h_t_d.raw_ptr(), + crate::rl::isv_slots::RL_LOSS_LAMBDA_PI_INDEX as i32, + dst, + )?; + launch_isv( + self.ss_v_grad_h_t_d.raw_ptr(), + crate::rl::isv_slots::RL_LOSS_LAMBDA_V_INDEX as i32, + dst, + )?; + launch_isv( + self.ss_frd_grad_h_t_d.raw_ptr(), + crate::rl::isv_slots::RL_FRD_LAMBDA_INDEX as i32, + dst, + )?; + launch_isv( + self.ss_outcome_grad_h_t_d.raw_ptr(), + crate::rl::isv_slots::RL_OUTCOME_AUX_LAMBDA_INDEX as i32, + dst, + )?; + } else { + let lambdas = read_loss_lambdas_from_isv(self.isv_host_slice()); accumulate_grad_h( &self.stream, &self.grad_h_accumulate_fn, - &self.ss_outcome_grad_h_t_d, - outcome_lambda * b_inv, + &self.ss_pi_grad_h_t_d, + lambdas.pi * b_inv, &mut self.grad_h_t_combined_d, )?; + accumulate_grad_h( + &self.stream, + &self.grad_h_accumulate_fn, + &self.ss_v_grad_h_t_d, + lambdas.v * b_inv, + &mut self.grad_h_t_combined_d, + )?; + accumulate_grad_h( + &self.stream, + &self.grad_h_accumulate_fn, + &self.ss_frd_grad_h_t_d, + lambdas.frd * b_inv, + &mut self.grad_h_t_combined_d, + )?; + { + let outcome_lambda = + self.read_isv_host(crate::rl::isv_slots::RL_OUTCOME_AUX_LAMBDA_INDEX); + accumulate_grad_h( + &self.stream, + &self.grad_h_accumulate_fn, + &self.ss_outcome_grad_h_t_d, + outcome_lambda * b_inv, + &mut self.grad_h_t_combined_d, + )?; + } } // ── Step 11: encoder backward (Phase E.3a) ─────────────────── @@ -5910,7 +6063,7 @@ impl IntegratedTrainer { // ~20 kernels from extract_realized_pnl_delta through // var_over_abs_mean. Same three-state machine as Graph A/A2 — // captures after A2 is done, replays on subsequent steps. - if self.reward_graph.is_some() { + if self.reward_graph.is_some() && !self.mega_graph_enabled { unsafe { raw_graph_launch( self.reward_graph.as_ref().unwrap().cu_graph_exec, @@ -5918,7 +6071,7 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("reward graph launch: {:?}", e))?; } } else { - let capturing_reward = self.postfill_graph.is_some(); + let capturing_reward = self.postfill_graph.is_some() && !self.mega_graph_enabled; if capturing_reward { self.stream .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) @@ -6620,8 +6773,13 @@ impl IntegratedTrainer { let k_max = self .read_isv_host(crate::rl::isv_slots::RL_K_LOOP_MAX_INDEX) .max(1.0) as usize; - let k_updates = ((n_rollout_steps / k_divisor).round() as usize) - .clamp(1, k_max); + let k_updates = if self.mega_graph_single_stream { + // Mega-graph captures a fixed K=1 topology. The graph's + // kernel sequence must be identical every step. + 1 + } else { + ((n_rollout_steps / k_divisor).round() as usize).clamp(1, k_max) + }; self.last_k_updates = k_updates; // K-loop split (Path B): first iter does the full env-step- @@ -6635,13 +6793,26 @@ impl IntegratedTrainer { // settles at K=1 (advantage_var_ratio drops with batch size), // but the split protects against pathological regimes and // makes the actor/critic separation explicit. + // Mega-graph Phase 3: select PER stream. When + // mega_graph_single_stream is true, all PER ops run on + // self.raw_stream (no cross-stream events needed — everything + // is stream-ordered). K is forced to 1. + let per_stream = if self.mega_graph_single_stream { + self.raw_stream + } else { + self.raw_train_stream + }; + // Event-based cross-stream sync before K-loop: record // push_done_event on self.stream after push_ring + push_flush. // train_stream waits on this event before reading the replay // buffer — no host-blocking cuStreamSynchronize needed. - unsafe { - raw_event_record(self.push_done_event.cu_event(), self.raw_stream) - .map_err(|e| anyhow::anyhow!("push_done_event record: {:?}", e))?; + // Skipped in mega-graph mode (single stream, no cross-stream deps). + if !self.mega_graph_single_stream { + unsafe { + raw_event_record(self.push_done_event.cu_event(), self.raw_stream) + .map_err(|e| anyhow::anyhow!("push_done_event record: {:?}", e))?; + } } let mut stats = IntegratedStepStats::default(); @@ -6651,7 +6822,8 @@ impl IntegratedTrainer { // train_stream is ordered after the tree rebuild from the // previous iteration (same stream), so no cross-stream wait // is needed — train_stream already serialises its own work. - if k_iter == 0 { + // Skipped in mega-graph mode (single stream). + if k_iter == 0 && !self.mega_graph_single_stream { unsafe { raw_stream_wait_event(self.raw_train_stream, self.push_done_event.cu_event()) .map_err(|e| anyhow::anyhow!("train_stream wait push_done: {:?}", e))?; @@ -6687,7 +6859,7 @@ impl IntegratedTrainer { raw_launch( self.rl_per_sample_fn.cu_function(), (b_size as u32, 1, 1), (1, 1, 1), 0, - self.raw_train_stream, + per_stream, &mut ptrs[..args.len()], ).map_err(|e| anyhow::anyhow!("rl_per_sample: {:?}", e))?; } @@ -6695,11 +6867,14 @@ impl IntegratedTrainer { // Event-based sync: record sample_done_event on train_stream // so self.stream can wait on it before reading sampled buffers. - unsafe { - raw_event_record(self.sample_done_event.cu_event(), self.raw_train_stream) - .map_err(|e| anyhow::anyhow!("sample_done_event record: {:?}", e))?; - raw_stream_wait_event(self.raw_stream, self.sample_done_event.cu_event()) - .map_err(|e| anyhow::anyhow!("stream wait sample_done: {:?}", e))?; + // Skipped in mega-graph mode (single stream, ordered by default). + if !self.mega_graph_single_stream { + unsafe { + raw_event_record(self.sample_done_event.cu_event(), self.raw_train_stream) + .map_err(|e| anyhow::anyhow!("sample_done_event record: {:?}", e))?; + raw_stream_wait_event(self.raw_stream, self.sample_done_event.cu_event()) + .map_err(|e| anyhow::anyhow!("stream wait sample_done: {:?}", e))?; + } } if k_iter == 0 { @@ -6715,16 +6890,14 @@ impl IntegratedTrainer { // Record replay_done_event on self.stream after BOTH // step_synthetic (k_iter==0) and dqn_replay_step (k_iter>0). // train_stream waits on this before the PER priority update - // reads td_per_sample_d. Previously k_iter==0 relied on - // step_synthetic's host-blocking stream.synchronize() to make - // self.stream's work globally visible; now that the sync is - // removed (fully async loss readback), the event-based wait - // provides the cross-stream ordering for both paths. - unsafe { - raw_event_record(self.replay_done_event.cu_event(), self.raw_stream) - .map_err(|e| anyhow::anyhow!("replay_done_event record: {:?}", e))?; - raw_stream_wait_event(self.raw_train_stream, self.replay_done_event.cu_event()) - .map_err(|e| anyhow::anyhow!("train_stream wait replay_done (pre-priority): {:?}", e))?; + // reads td_per_sample_d. Skipped in mega-graph mode. + if !self.mega_graph_single_stream { + unsafe { + raw_event_record(self.replay_done_event.cu_event(), self.raw_stream) + .map_err(|e| anyhow::anyhow!("replay_done_event record: {:?}", e))?; + raw_stream_wait_event(self.raw_train_stream, self.replay_done_event.cu_event()) + .map_err(|e| anyhow::anyhow!("train_stream wait replay_done (pre-priority): {:?}", e))?; + } } // PER priority update — writes new priorities from per-sample @@ -6748,7 +6921,7 @@ impl IntegratedTrainer { raw_launch( self.rl_per_update_priority_fn.cu_function(), (1, 1, 1), (b_size as u32, 1, 1), smem, - self.raw_train_stream, + per_stream, &mut ptrs[..args.len()], ).map_err(|e| anyhow::anyhow!("rl_per_update_priority: {:?}", e))?; } @@ -6767,22 +6940,28 @@ impl IntegratedTrainer { raw_launch( self.rl_per_tree_rebuild_fn.cu_function(), (128, 1, 1), (256, 1, 1), 0, - self.raw_train_stream, + per_stream, &mut ptrs[..args.len()], ).map_err(|e| anyhow::anyhow!("rl_per_tree_rebuild: {:?}", e))?; } } } - // Record train_done_event on train_stream after the last K-loop - // iteration's tree rebuild. The NEXT step_with_lobsim's - // stream.wait(&train_done_event) ensures the replay buffer is - // coherent before push_ring/push_flush write new transitions. - // GPU-side event wait replaces the prior host-blocking - // train_stream.synchronize() — saves ~2ms/step. - unsafe { - raw_event_record(self.train_done_event.cu_event(), self.raw_train_stream) - .map_err(|e| anyhow::anyhow!("train_done_event record: {:?}", e))?; + // Record train_done_event after the last K-loop iteration's tree + // rebuild. In mega-graph mode, this goes on self.raw_stream (the + // single stream). In multi-stream mode, on train_stream. + // The NEXT step_with_lobsim's stream.wait(&train_done_event) + // ensures the replay buffer is coherent before push_ring/push_flush. + if !self.mega_graph_single_stream { + unsafe { + raw_event_record(self.train_done_event.cu_event(), self.raw_train_stream) + .map_err(|e| anyhow::anyhow!("train_done_event record: {:?}", e))?; + } + } else { + unsafe { + raw_event_record(self.train_done_event.cu_event(), self.raw_stream) + .map_err(|e| anyhow::anyhow!("train_done_event record (single stream): {:?}", e))?; + } } // Target-net soft update (Phase R5 + R6) — runs once per step @@ -6856,6 +7035,32 @@ impl IntegratedTrainer { /// 7. Pre-snapshot graph pipeline (Q/V/IQN/Pi/gates) /// 8. Post-snapshot pipeline (lobsim fill, reward, controllers) /// 9. PER push/sample, K-loop training, target-net update + /// Enable the mega-graph path. After calling this, subsequent + /// `step_with_lobsim_gpu` calls use a single CUDA graph capturing + /// the entire per-step pipeline. The first step after enabling + /// executes eagerly (warmup), the second captures, and all + /// subsequent steps replay the captured graph. + /// + /// Requires K=1 (the mega-graph captures a fixed K=1 topology). + pub fn enable_mega_graph(&mut self) { + self.mega_graph_enabled = true; + self.mega_graph_single_stream = true; + self.mega_warmup_done = false; + self.mega_graph = None; + // Disable sub-graph state machines — the mega-graph captures + // everything. Drop existing sub-graphs so the state machines + // always take the `else` (eager) branch. Sub-capture is + // prevented by the `!self.mega_graph_enabled` guards added to + // each sub-graph's capture condition. + self.prefill_graph = None; + self.postfill_graph = None; + self.reward_graph = None; + self.training_graph = None; + // Propagate to the perception trainer so its sub-graphs also + // run eagerly (train_graph, forward_graph_no_scatter). + self.perception.mega_graph_enabled = true; + } + pub fn step_with_lobsim_gpu( &mut self, gpu_loader: &mut crate::data::gpu_dataset::GpuDataLoader, @@ -6868,10 +7073,45 @@ impl IntegratedTrainer { anyhow::bail!("step_with_lobsim_gpu: empty batch (n_batch = 0)"); } + // ── Mega-graph fast path: single cuGraphLaunch per step. ────── + if self.mega_graph_enabled { + if let Some(ref mega) = self.mega_graph { + // Replay: 1 cuGraphLaunch, zero host work between kernels. + unsafe { + raw_graph_launch( + mega.cu_graph_exec, + self.raw_stream, + ).map_err(|e| anyhow::anyhow!("mega graph launch: {:?}", e))?; + } + // DiagFrame from mapped-pinned reads (background, outside graph). + self.step_synthetic_read_deferred_stats()?; + return Ok(self.last_step_stats.clone()); + } + // Fall through to warmup / capture below. + } + + // ── Mega-graph warmup / capture state machine ───────────────── + // Step 0: eager (warmup). Step 1: begin_capture → eager → end_capture. + // Step 2+: handled by fast path above (cuGraphLaunch). + let mega_capturing = self.mega_graph_enabled && self.mega_warmup_done && self.mega_graph.is_none(); + if self.mega_graph_enabled && !self.mega_warmup_done { + self.mega_warmup_done = true; + } + if mega_capturing { + self.stream + .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) + .map_err(|e| anyhow::anyhow!("mega begin_capture: {e}"))?; + } + // ── Event-based inter-stream sync (same as step_with_lobsim) ── - unsafe { - raw_stream_wait_event(self.raw_stream, self.train_done_event.cu_event()) - .map_err(|e| anyhow::anyhow!("stream wait train_done: {:?}", e))?; + // Skipped during mega-graph capture (events cannot be recorded + // inside a capture region on a different stream). In mega-graph + // mode, all work is single-stream so the ordering is implicit. + if !self.mega_graph_enabled { + unsafe { + raw_stream_wait_event(self.raw_stream, self.train_done_event.cu_event()) + .map_err(|e| anyhow::anyhow!("stream wait train_done: {:?}", e))?; + } } // ── Step 0: bump device-resident step counter (ISV[548]). ── @@ -7004,7 +7244,21 @@ impl IntegratedTrainer { // SoA's last-snapshot position into self.ts_ns_d. // // We call step_with_lobsim_gpu_body to run the shared pipeline. - self.step_with_lobsim_gpu_body(lobsim, &soa, b_size, seq_len) + let result = self.step_with_lobsim_gpu_body(lobsim, &soa, b_size, seq_len); + + // ── Mega-graph: end capture after all kernels have been recorded. + if mega_capturing { + let graph = self + .stream + .end_capture( + CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .context("mega end_capture")? + .ok_or_else(|| anyhow::anyhow!("mega end_capture returned None"))?; + self.mega_graph = Some(graph); + } + + result } /// Shared body for `step_with_lobsim_gpu`. Runs the graph pipelines, @@ -7029,7 +7283,7 @@ impl IntegratedTrainer { let lob = LobPtrs::new(lobsim); // ── Graph A: pre-snapshot kernel pipeline ────────────────────── - if self.prefill_graph.is_some() { + if self.prefill_graph.is_some() && !self.mega_graph_enabled { unsafe { raw_graph_launch( self.prefill_graph.as_ref().unwrap().cu_graph_exec, @@ -7037,7 +7291,7 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("prefill graph launch: {:?}", e))?; } } else { - let capturing_prefill = self.graph_warmup_done; + let capturing_prefill = self.graph_warmup_done && !self.mega_graph_enabled; if !self.graph_warmup_done { self.graph_warmup_done = true; } @@ -7316,16 +7570,48 @@ impl IntegratedTrainer { // Layout: SoA arrays are [B * K * BOOK_LEVELS] f32. Batch 0's // last snapshot is at position (K-1), byte offset = // (seq_len - 1) * BOOK_LEVELS * sizeof(f32). + // Mega-graph Phase 1: inline apply_snapshot_from_device as + // raw_launch on self.raw_stream. Bypasses the trait method which + // internally launches on the lobsim's own stream (would break + // graph capture). The 4 DtoD copies + 1 book_update kernel are + // now stream-ordered on self.raw_stream like all other kernels. { let book_levels = crate::cfc::snap_features::BOOK_LEVELS; let snap_byte_offset = ((seq_len - 1) * book_levels * std::mem::size_of::()) as u64; - lobsim.apply_snapshot_from_device( - soa.bid_px + snap_byte_offset, - soa.bid_sz + snap_byte_offset, - soa.ask_px + snap_byte_offset, - soa.ask_sz + snap_byte_offset, - ).context("step_with_lobsim_gpu: apply_snapshot_from_device")?; + let lob_raw = lobsim.raw_ptrs(); + let bytes = book_levels * std::mem::size_of::(); + unsafe { + raw_memcpy_dtod_async(lob_raw.bid_px, soa.bid_px + snap_byte_offset, bytes, self.raw_stream) + .map_err(|e| anyhow::anyhow!("apply_snapshot raw: bid_px DtoD: {:?}", e))?; + raw_memcpy_dtod_async(lob_raw.bid_sz, soa.bid_sz + snap_byte_offset, bytes, self.raw_stream) + .map_err(|e| anyhow::anyhow!("apply_snapshot raw: bid_sz DtoD: {:?}", e))?; + raw_memcpy_dtod_async(lob_raw.ask_px, soa.ask_px + snap_byte_offset, bytes, self.raw_stream) + .map_err(|e| anyhow::anyhow!("apply_snapshot raw: ask_px DtoD: {:?}", e))?; + raw_memcpy_dtod_async(lob_raw.ask_sz, soa.ask_sz + snap_byte_offset, bytes, self.raw_stream) + .map_err(|e| anyhow::anyhow!("apply_snapshot raw: ask_sz DtoD: {:?}", e))?; + } + let mut args = RawArgs::new(); + args.push_ptr(lob_raw.bid_px); + args.push_ptr(lob_raw.bid_sz); + args.push_ptr(lob_raw.ask_px); + args.push_ptr(lob_raw.ask_sz); + args.push_ptr(lob_raw.books); + args.push_ptr(lob_raw.prev_mid); + args.push_ptr(lob_raw.atr_mid_ema); + args.push_ptr(lob_raw.snapshots_skipped); + args.push_ptr(lob_raw.min_reasonable_px); + args.push_ptr(lob_raw.max_reasonable_px); + args.push_i32(lob_raw.n_backtests); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + lob_raw.book_update_fn, + (lob_raw.n_backtests as u32, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("apply_snapshot raw: book_update: {:?}", e))?; + } } // // ts_ns: DtoD from the SoA's last-snapshot ts_ns into self.ts_ns_d @@ -7347,7 +7633,7 @@ impl IntegratedTrainer { let b_size_i = b_size as i32; // ── Graph A2: post-snapshot / pre-fill kernel pipeline ───────── - if self.postfill_graph.is_some() { + if self.postfill_graph.is_some() && !self.mega_graph_enabled { unsafe { raw_graph_launch( self.postfill_graph.as_ref().unwrap().cu_graph_exec, @@ -7355,7 +7641,7 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("postfill graph launch: {:?}", e))?; } } else { - let capturing_postfill = self.prefill_graph.is_some(); + let capturing_postfill = self.prefill_graph.is_some() && !self.mega_graph_enabled; if capturing_postfill { self.stream .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) @@ -7535,12 +7821,64 @@ impl IntegratedTrainer { } } // end else (postfill warmup / capture dispatch) - // Fill kernel — ts_ns=0 since we don't have host-side timestamp - // in the GPU data path. The timestamp is diagnostic-only (trade - // log entries); RL training quality is unaffected. - lobsim - .step_fill_from_market_targets(0) - .context("step_with_lobsim_gpu: lobsim.step_fill_from_market_targets")?; + // Mega-graph Phase 1: inline step_fill_from_market_targets + + // step_pnl_track as raw_launch on self.raw_stream. Bypasses the + // trait method which internally launches on the lobsim's own + // stream. ts_ns=0 since we don't have host-side timestamp in the + // GPU data path; diagnostic-only, RL training quality unaffected. + { + let lob_raw = lobsim.raw_ptrs(); + // submit_market_immediate: fill orders from market_targets_d + { + let mut args = RawArgs::new(); + args.push_ptr(lob_raw.books); + args.push_ptr(lob_raw.market_targets); + args.push_ptr(lob_raw.pos); + args.push_ptr(lob_raw.min_reasonable_px); + args.push_ptr(lob_raw.max_reasonable_px); + args.push_ptr(lob_raw.cost_per_lot_per_side); + args.push_ptr(lob_raw.total_fees_per_b); + args.push_i32(lob_raw.n_backtests); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + lob_raw.submit_market_fn, + (lob_raw.n_backtests as u32, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("step_fill raw: submit_market: {:?}", e))?; + } + } + // pnl_track: detect segment_complete, emit TradeRecord + { + let mut args = RawArgs::new(); + args.push_ptr(lob_raw.pos); + args.push_ptr(lob_raw.open_trade_state); + args.push_ptr(lob_raw.trade_log); + args.push_ptr(lob_raw.trade_log_head); + args.push_u64(0_u64); // ts_ns = 0 (diagnostic-only) + args.push_i32(lob_raw.trade_log_cap); + args.push_i32(lob_raw.n_backtests); + args.push_ptr(lob_raw.trail_hwm); + args.push_ptr(lob_raw.zero_vwap_at_open); + args.push_ptr(lob_raw.saturated_vwap_at_open); + args.push_ptr(lob_raw.defensive_exit_clamp); + args.push_ptr(lob_raw.conv_signed_ema); + args.push_ptr(lob_raw.diag_hold_hist); + args.push_ptr(lob_raw.diag_outcome_n); + args.push_ptr(lob_raw.diag_outcome_sum_pnl); + args.push_ptr(lob_raw.diag_outcome_n_wins); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + lob_raw.pnl_track_fn, + (lob_raw.n_backtests as u32, 1, 1), (32, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("step_fill raw: pnl_track: {:?}", e))?; + } + } + } // Activate GPU encoder path for the K-loop's step_synthetic call. // step_synthetic checks this flag and dispatches @@ -7606,7 +7944,42 @@ impl IntegratedTrainer { Ok(()) } - + /// Mega-graph variant of `launch_rl_lr_controller`. Reads loss + /// observations from the mapped-pinned device pointers instead of + /// host-supplied scalars. This kernel can be captured inside a CUDA + /// graph because its inputs are device pointers (stable addresses), + /// not host scalars that would be baked in at capture time. + fn launch_rl_lr_from_mapped_pinned(&self) -> Result<()> { + let b_size = self.cfg.perception.n_batch as i32; + let mut args = RawArgs::new(); + args.push_ptr(self.isv_dev_ptr); + args.push_ptr(self.ss_q_loss_dev_ptr); + args.push_ptr(self.ss_pi_loss_dev_ptr); + args.push_ptr(self.ss_v_loss_sum_dev_ptr); + args.push_i32(b_size); + args.push_i32(crate::rl::isv_slots::RL_LR_Q_LOSS_EMA_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_Q_BEST_LOSS_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_Q_STEPS_SINCE_BEST_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_Q_WARMUP_COUNTER_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_PI_LOSS_EMA_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_PI_BEST_LOSS_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_PI_STEPS_SINCE_BEST_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_PI_WARMUP_COUNTER_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_V_LOSS_EMA_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_V_BEST_LOSS_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_V_STEPS_SINCE_BEST_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_LR_V_WARMUP_COUNTER_INDEX as i32); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.rl_lr_from_mapped_pinned_fn.cu_function(), + (1, 1, 1), (1, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("rl_lr_from_mapped_pinned: {:?}", e))?; + } + Ok(()) + } /// Phase E.2 helper: combine a per-head grad_h_t into the encoder's /// accumulator slot via `grad_h_encoder[i] += λ × grad_h_head[i]`. diff --git a/crates/ml-alpha/src/trainer/optim.rs b/crates/ml-alpha/src/trainer/optim.rs index e31db2939..b0b333130 100644 --- a/crates/ml-alpha/src/trainer/optim.rs +++ b/crates/ml-alpha/src/trainer/optim.rs @@ -122,6 +122,58 @@ impl AdamW { &mut self.v } + /// Mega-graph variant of `step()`. Reads LR from an ISV device + /// pointer at `lr_slot` instead of the host-side `self.lr` field. + /// This allows the LR to vary across graph replays because the ISV + /// is modified in-place by the `rl_lr_from_mapped_pinned` controller + /// kernel captured earlier in the same graph. + /// + /// `isv_lr_fn` is the `adamw_step_isv_lr` kernel function handle. + /// `isv_ptr` is the stable ISV device pointer. + pub fn step_isv_lr( + &mut self, + theta: &mut CudaSlice, + grad: &CudaSlice, + isv_lr_fn: cudarc::driver::sys::CUfunction, + isv_ptr: u64, + lr_slot: i32, + ) -> Result<()> { + let n = theta.len(); + assert_eq!(grad.len(), n, "grad/theta size mismatch"); + assert_eq!(self.m.len(), n, "m size mismatch"); + assert_eq!(self.v.len(), n, "v size mismatch"); + + self.step_count_host += 1; + + let n_params_i = n as i32; + let grid_x = (n as u32).div_ceil(256); + { + let mut args = RawArgs::new(); + args.push_ptr(theta.raw_ptr()); + args.push_ptr(grad.raw_ptr()); + args.push_ptr(self.m.raw_ptr()); + args.push_ptr(self.v.raw_ptr()); + args.push_i32(n_params_i); + args.push_ptr(isv_ptr); + args.push_i32(lr_slot); + args.push_f32(self.beta1); + args.push_f32(self.beta2); + args.push_f32(self.eps); + args.push_f32(self.wd); + args.push_i32(self.step_count_host); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + isv_lr_fn, + (grid_x, 1, 1), (256, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("adamw_step_isv_lr launch: {:?}", e))?; + } + } + Ok(()) + } + /// Return the current step counter value. No GPU sync needed — /// the counter is host-resident. pub fn step_count(&self) -> i32 { diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index f6914812c..0213945e8 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -718,6 +718,13 @@ pub struct PerceptionTrainer { /// include cuBLAS calls. cublas_warmed: bool, + /// Mega-graph mode: when true, all perception sub-graph state + /// machines (train_graph, forward_graph, forward_graph_no_scatter) + /// run eagerly (no sub-capture, no sub-replay). The mega-graph in + /// the integrated trainer captures the entire pipeline including + /// perception kernels. + pub mega_graph_enabled: bool, + /// Event recorded after every training graph launch in /// `step_batched_from_device`. The integrated trainer syncs on /// this event at the start of the NEXT step @@ -2316,6 +2323,7 @@ impl PerceptionTrainer { forward_graph_no_scatter: None, forward_no_scatter_warmed: false, cublas_warmed: false, + mega_graph_enabled: false, training_done_event, // AoS staging — single mapped-pinned buffer for B*K Mbp10RawInput. @@ -3235,7 +3243,7 @@ impl PerceptionTrainer { // 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() { + if self.train_graph.is_some() && !self.mega_graph_enabled { self.train_graph .as_ref() .unwrap() @@ -3245,6 +3253,9 @@ impl PerceptionTrainer { self.dispatch_train_step(b_sz, k_seq, total_snaps) .context("train warmup dispatch")?; self.cublas_warmed = true; + } else if self.mega_graph_enabled { + self.dispatch_train_step(b_sz, k_seq, total_snaps) + .context("train eager dispatch (mega-graph mode)")?; } else { // Event tracking was disabled at trainer construction; the // trainer's CudaSlices have no read/write events, so neither @@ -3886,7 +3897,7 @@ impl PerceptionTrainer { // Three-state machine for the training graph — identical to // step_batched but dispatches `dispatch_train_step_no_scatter` // (skips the AoS→SoA scatter since SoA buffers are pre-filled). - if self.train_graph.is_some() { + if self.train_graph.is_some() && !self.mega_graph_enabled { self.train_graph .as_ref() .unwrap() @@ -3896,6 +3907,11 @@ impl PerceptionTrainer { self.dispatch_train_step_no_scatter(b_sz, k_seq, total_snaps) .context("train warmup dispatch (from_device)")?; self.cublas_warmed = true; + } else if self.mega_graph_enabled { + // Mega-graph mode: always dispatch eagerly — the mega-graph + // in the integrated trainer captures these kernel launches. + self.dispatch_train_step_no_scatter(b_sz, k_seq, total_snaps) + .context("train eager dispatch (mega-graph mode)")?; } else { let begin = self.stream.begin_capture( CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED, @@ -4090,7 +4106,7 @@ impl PerceptionTrainer { // warmup → capture → replay with a new graph that excludes the // scatter. Let's use this approach since it's zero overhead // after the second call. - if self.forward_graph_no_scatter.is_some() { + if self.forward_graph_no_scatter.is_some() && !self.mega_graph_enabled { self.forward_graph_no_scatter .as_ref() .unwrap() @@ -4100,6 +4116,10 @@ impl PerceptionTrainer { self.dispatch_forward_kernels_no_scatter(b_sz, k_seq, total_snaps) .context("forward_no_scatter warmup dispatch")?; self.forward_no_scatter_warmed = true; + } else if self.mega_graph_enabled { + // Mega-graph mode: always dispatch eagerly. + self.dispatch_forward_kernels_no_scatter(b_sz, k_seq, total_snaps) + .context("forward_no_scatter eager dispatch (mega-graph mode)")?; } else { use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode}; let begin = self.stream.begin_capture( diff --git a/crates/ml-backtesting/src/sim/mod.rs b/crates/ml-backtesting/src/sim/mod.rs index 36bb307bf..8fea4b1cb 100644 --- a/crates/ml-backtesting/src/sim/mod.rs +++ b/crates/ml-backtesting/src/sim/mod.rs @@ -1433,6 +1433,43 @@ impl ml_alpha::rl::reward::RlLobBackend for LobSimCuda { ask_sz_src, ) } + + fn raw_ptrs(&self) -> ml_alpha::rl::reward::LobSimRawPtrs { + ml_alpha::rl::reward::LobSimRawPtrs { + bid_px: self.bid_px_d.raw_ptr(), + bid_sz: self.bid_sz_d.raw_ptr(), + ask_px: self.ask_px_d.raw_ptr(), + ask_sz: self.ask_sz_d.raw_ptr(), + books: self.books_d.raw_ptr(), + prev_mid: self.prev_mid_d.raw_ptr(), + atr_mid_ema: self.atr_mid_ema_d.raw_ptr(), + snapshots_skipped: self.snapshots_skipped_d.raw_ptr(), + min_reasonable_px: self.min_reasonable_px_d.raw_ptr(), + max_reasonable_px: self.max_reasonable_px_d.raw_ptr(), + book_update_fn: self.book_update_fn.cu_function(), + market_targets: self.market_targets_d.raw_ptr(), + pos: self.pos_d.raw_ptr(), + cost_per_lot_per_side: self.cost_per_lot_per_side_d.raw_ptr(), + total_fees_per_b: self.total_fees_per_b_d.raw_ptr(), + submit_market_fn: self.submit_market_fn.cu_function(), + open_trade_state: self.open_trade_state_d.raw_ptr(), + trade_log: self.trade_log_d.raw_ptr(), + trade_log_head: self.trade_log_head_d.raw_ptr(), + trail_hwm: self.trail_hwm_d.raw_ptr(), + zero_vwap_at_open: self.zero_vwap_at_open_d.raw_ptr(), + saturated_vwap_at_open: self.saturated_vwap_at_open_d.raw_ptr(), + defensive_exit_clamp: self.defensive_exit_clamp_d.raw_ptr(), + conv_signed_ema: self.conv_signed_ema_d.raw_ptr(), + diag_hold_hist: self.diag_hold_hist_d.raw_ptr(), + diag_outcome_n: self.diag_outcome_n_d.raw_ptr(), + diag_outcome_sum_pnl: self.diag_outcome_sum_pnl_d.raw_ptr(), + diag_outcome_n_wins: self.diag_outcome_n_wins_d.raw_ptr(), + pnl_track_fn: self.pnl_track_fn.cu_function(), + n_backtests: self.n_backtests as i32, + pos_bytes: std::mem::size_of::() as i32, + trade_log_cap: crate::lob::TRADE_LOG_CAP as i32, + } + } } // Re-open the impl block for `LobSimCuda` so subsequent methods (if any