perf(cuda): mega-graph pipeline — 10-12× speedup (6.4 → 68-76 sps)
Captures the ENTIRE per-step training pipeline into ONE cuGraphLaunch, eliminating ~150 individual kernel launches and ~143ms of host-side Rust overhead per step. Phase 1: Lobsim → raw_launch - apply_snapshot_from_device inlined as 4× raw_memcpy_dtod + raw_launch - step_fill_from_market_targets inlined as 2× raw_launch - LobSimRawPtrs trait method caches 30+ device pointers Phase 2: LR controller → GPU kernel - New rl_lr_from_mapped_pinned.cu reads losses from device pointers - New adamw_step_isv_lr kernel reads LR from ISV instead of scalar arg - All 12 Adam .step() calls use step_isv_lr() in mega-graph mode Phase 3: PER → main stream - mega_graph_single_stream flag routes PER to self.raw_stream - Cross-stream events skipped, K forced to 1 Phase 4: Mega-graph capture/replay - Three-state machine: warmup → capture → replay - enable_mega_graph() propagates to perception trainer - Perception sub-graph guards prevent sub-captures inside mega-capture - grad_h_accumulate_scaled_isv reads lambda from ISV on-device Validated: 500 steps, 68-76 sps sustained, no NaN, all losses finite. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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)
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
|
||||
@@ -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];
|
||||
}
|
||||
|
||||
134
crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu
Normal file
134
crates/ml-alpha/cuda/rl_lr_from_mapped_pinned.cu
Normal file
@@ -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);
|
||||
}
|
||||
@@ -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 +
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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<CudaGraph>,
|
||||
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<CudaGraph>,
|
||||
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<CudaModule>,
|
||||
grad_h_accumulate_fn: CudaFunction,
|
||||
/// Mega-graph variant: reads lambda from ISV device pointer.
|
||||
grad_h_accumulate_isv_fn: CudaFunction,
|
||||
_reduce_axis0_module: Arc<CudaModule>,
|
||||
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<CudaModule>,
|
||||
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<CudaModule>,
|
||||
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<f32> = 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::<f32>()) 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::<f32>();
|
||||
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]`.
|
||||
|
||||
@@ -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<f32>,
|
||||
grad: &CudaSlice<f32>,
|
||||
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 {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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::<crate::lob::PosFlat>() as i32,
|
||||
trade_log_cap: crate::lob::TRADE_LOG_CAP as i32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Re-open the impl block for `LobSimCuda` so subsequent methods (if any
|
||||
|
||||
Reference in New Issue
Block a user