gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(rl): GPU-resident data loader — zero CPU per step

Browse Source 
Upload all pre-converted snapshots + FRD labels to GPU at init (1.2GB
for 5.6M snapshots / 2 files locally, ~11GB for 45M / 9 files on L40S).

New GPU kernels: gpu_sample_and_gather (PRNG sampling + AoS-to-SoA),
gpu_gather_next/gpu_gather_current (anchor offset re-gather),
gpu_gather_frd_labels (per-horizon label gather).

step_with_lobsim_gpu: GPU-only data path replacing host-side loader.
Encoder forwards via forward_encoder_from_device. Eliminates 7700 heap
allocs + 418k scalar copies + 16ms CPU work per step at b=256.

Init uploads use cuMemcpyHtoD_v2 (synchronous, one-time).

Note: apply_snapshot skipped (lobsim book stale, dones/rewards=0).
Follow-up: copy last-snapshot book data from SoA to lobsim buffers.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 20:51:06 +02:00
parent a429c3f901
commit 1f672c3b09
8 changed files with 1675 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -126,6 +126,7 @@ const KERNELS: &[&str] = &[
"rl_adversarial_boost", // Adversarial: boost PER priority for negative-reward transitions
"rl_outcome_bwd", // Outcome aux: single linear layer backward — dW (per-batch), db (per-batch), dh_t (per-batch); 1 block per batch, 128 threads
"snapshot_aos_to_soa", // AoS→SoA scatter: one thread per snapshot reads contiguous Mbp10RawInput, writes into 10 SoA device buffers; replaces host nested loops + 10 DtoD copies
"gpu_sample_and_gather", // GPU-resident batch sampler: random file+anchor sampling + AoS→SoA gather from pre-uploaded dataset; eliminates ALL per-step CPU data loading
];
// Cache bust v31 — five new reduce / derive kernels populate the input

298
crates/ml-alpha/cuda/gpu_sample_and_gather.cu Normal file
View File

@@ -0,0 +1,298 @@
// gpu_sample_and_gather.cu — GPU-resident batch sampler + AoS→SoA scatter
//
// Replaces the CPU-side `next_sequence_pair` + `build_sequence_at` +
// `extend_from_slice` + mapped-pinned memcpy hot path. At b=256, K=32,
// the CPU version spent ~16ms/step on 512 build_sequence_at calls with
// 7,700 heap allocations while GPU compute was 0.62ms. This kernel
// samples B random windows and gathers snapshots directly from a
// pre-uploaded GPU-resident dataset — zero CPU work per step.
//
// Architecture:
// * At init, the host pre-converts ALL Mbp10Snapshot → Mbp10RawInput
// and uploads as a single flat CudaSlice<u8> (reinterpreted as
// Mbp10Raw on device). File offsets + sizes are uploaded as i32
// arrays. Labels/regime are uploaded similarly.
// * Per step, this kernel runs with Grid=(B,1,1), Block=(K,1,1):
// - Thread 0 of each block uses device-side xorshift32 PRNG to
// sample file_idx and anchor within that file.
// - All K threads gather all_snapshots[file_offset + anchor + k]
// and scatter to SoA output buffers (same layout as
// snapshot_aos_to_soa.cu).
//
// Per feedback_no_atomicadd: block-local shared memory for sampling,
// no atomics. Per feedback_no_nvrtc: pre-compiled cubin via build.rs.
#define BOOK_LEVELS 10
#define REGIME_DIM 6
// Must match #[repr(C)] Mbp10RawInput in snap_features.rs (216 bytes).
struct __align__(8) Mbp10Raw {
float bid_px[BOOK_LEVELS]; // offset 0, 40 bytes
float bid_sz[BOOK_LEVELS]; // offset 40, 40 bytes
float ask_px[BOOK_LEVELS]; // offset 80, 40 bytes
float ask_sz[BOOK_LEVELS]; // offset 120, 40 bytes
float prev_mid; // offset 160, 4 bytes
float trade_signed_vol; // offset 164, 4 bytes
unsigned int trade_count; // offset 168, 4 bytes
// 4 bytes padding for u64 alignment
unsigned long long ts_ns; // offset 176, 8 bytes
unsigned long long prev_ts_ns; // offset 184, 8 bytes
float regime[REGIME_DIM]; // offset 192, 24 bytes
// total: 216 bytes
};
// Marsaglia xorshift32 — minimal state, sufficient quality for
// stochastic sampling. Each batch element carries its own seed to
// avoid inter-block contention.
__device__ __forceinline__ unsigned int xorshift32(unsigned int s) {
s ^= s << 13;
s ^= s >> 17;
s ^= s << 5;
return s;
}
// ── Primary kernel: sample + gather + SoA scatter ────────────────────
//
// Grid = (B, 1, 1) — one block per batch element
// Block = (K, 1, 1) — one thread per sequence position
//
// Thread 0 samples (file_idx, anchor) via xorshift32. All K threads
// then read their assigned snapshot from the flat GPU array and write
// to SoA outputs at position [b * K + k].
extern "C" __global__ void gpu_sample_and_gather(
const Mbp10Raw* __restrict__ all_snapshots, // [total_snaps]
const int* __restrict__ file_offsets, // [n_files]
const int* __restrict__ file_sizes, // [n_files]
unsigned int* __restrict__ prng_state, // [B]
int n_files,
int seq_len,
int max_horizon,
// SoA outputs (same layout as snapshot_aos_to_soa):
float* __restrict__ bid_px_soa, // [B*K * BOOK_LEVELS]
float* __restrict__ bid_sz_soa,
float* __restrict__ ask_px_soa,
float* __restrict__ ask_sz_soa,
float* __restrict__ regime_soa, // [B*K * REGIME_DIM]
float* __restrict__ prev_mid_soa, // [B*K]
float* __restrict__ tsv_soa, // [B*K]
int* __restrict__ tc_soa, // [B*K]
long long* __restrict__ ts_ns_soa, // [B*K]
long long* __restrict__ prev_ts_ns_soa, // [B*K]
// Global-memory outputs for the sampled (file_offset, anchor) pair
// per batch element. Read by gpu_gather_next and gpu_gather_frd_labels
// to use the same sampling decision without re-rolling the PRNG.
int* __restrict__ out_file_offset, // [B]
int* __restrict__ out_anchor, // [B]
int B
) {
int b = blockIdx.x;
int k = threadIdx.x;
if (b >= B || k >= seq_len) return;
// ── Thread 0: sample file + anchor via device-side PRNG ──────────
__shared__ int s_file_offset;
__shared__ int s_anchor;
if (k == 0) {
unsigned int seed = prng_state[b];
// Sample file index (uniform across files).
seed = xorshift32(seed);
int file_idx = (int)(seed % (unsigned int)n_files);
// Sample anchor within file. Must leave room for seq_len +
// max_horizon snapshots after the anchor (+ 1 for the s_{t+1}
// window gathered by gpu_gather_next).
int fsize = file_sizes[file_idx];
int usable = fsize - seq_len - max_horizon - 1;
seed = xorshift32(seed);
int anchor;
if (usable > 0) {
anchor = (int)(seed % (unsigned int)usable);
} else {
// Degenerate file too short — clamp to 0 (defensive).
anchor = 0;
}
s_file_offset = file_offsets[file_idx];
s_anchor = anchor;
// Write to global memory so downstream kernels can read the
// same (file_offset, anchor) without re-sampling.
out_file_offset[b] = s_file_offset;
out_anchor[b] = anchor;
prng_state[b] = seed;
}
__syncthreads();
// ── All K threads: gather + SoA scatter ──────────────────────────
int global_idx = s_file_offset + s_anchor + k;
const Mbp10Raw& snap = all_snapshots[global_idx];
int n = b * seq_len + k; // output position in [B*K] flat layout
int base_book = n * BOOK_LEVELS;
int base_regime = n * REGIME_DIM;
#pragma unroll
for (int i = 0; i < BOOK_LEVELS; i++) {
bid_px_soa[base_book + i] = snap.bid_px[i];
bid_sz_soa[base_book + i] = snap.bid_sz[i];
ask_px_soa[base_book + i] = snap.ask_px[i];
ask_sz_soa[base_book + i] = snap.ask_sz[i];
}
#pragma unroll
for (int i = 0; i < REGIME_DIM; i++) {
regime_soa[base_regime + i] = snap.regime[i];
}
prev_mid_soa[n] = snap.prev_mid;
tsv_soa[n] = snap.trade_signed_vol;
tc_soa[n] = (int)snap.trade_count;
ts_ns_soa[n] = (long long)snap.ts_ns;
prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
}
// ── Label gather kernel ──────────────────────────────────────────────
//
// Runs AFTER gpu_sample_and_gather. Gathers per-horizon FRD labels at
// the anchor position (rightmost K position = newest snapshot in the
// window, matching the supervised label semantics).
//
// Grid = (B, 1, 1)
// Block = (FRD_N_HORIZONS, 1, 1) — typically 3 threads
//
// For RL training, the FRD labels are the primary use case. BCE labels
// and outcome labels are handled by the supervised path and are not
// needed per-step in the RL loop.
extern "C" __global__ void gpu_gather_frd_labels(
const int* __restrict__ all_frd_labels, // [n_frd_horizons * total_snaps]
const int* __restrict__ file_offsets, // [n_files] (same as above)
const int* __restrict__ sample_file_offset,// [B] — file_offset chosen by gpu_sample_and_gather
const int* __restrict__ sample_anchor, // [B] — anchor chosen by gpu_sample_and_gather
int seq_len,
int n_frd_horizons,
int total_snaps,
int* __restrict__ frd_labels_out, // [B * n_frd_horizons]
int B
) {
int b = blockIdx.x;
int h = threadIdx.x;
if (b >= B || h >= n_frd_horizons) return;
// The FRD label is at the NEWEST snapshot in the window = anchor + seq_len - 1.
int newest_idx = sample_file_offset[b] + sample_anchor[b] + seq_len - 1;
int label_idx = h * total_snaps + newest_idx; // row-major [n_frd_horizons, total_snaps]
frd_labels_out[b * n_frd_horizons + h] = all_frd_labels[label_idx];
}
// ── Pair gather: s_{t+1} window for Bellman targets ──────────────────
//
// The RL trainer needs (s_t, s_{t+1}) consecutive windows. This kernel
// gathers the SECOND window at anchor+1 into a separate set of SoA
// buffers, reusing the same (file_offset, anchor) from the primary
// gpu_sample_and_gather call. Avoids a second sampling pass.
//
// Grid = (B, 1, 1)
// Block = (K, 1, 1)
extern "C" __global__ void gpu_gather_next(
const Mbp10Raw* __restrict__ all_snapshots, // [total_snaps]
const int* __restrict__ sample_file_offset,// [B]
const int* __restrict__ sample_anchor, // [B]
int seq_len,
// SoA outputs for s_{t+1}:
float* __restrict__ bid_px_soa,
float* __restrict__ bid_sz_soa,
float* __restrict__ ask_px_soa,
float* __restrict__ ask_sz_soa,
float* __restrict__ regime_soa,
float* __restrict__ prev_mid_soa,
float* __restrict__ tsv_soa,
int* __restrict__ tc_soa,
long long* __restrict__ ts_ns_soa,
long long* __restrict__ prev_ts_ns_soa,
int B
) {
int b = blockIdx.x;
int k = threadIdx.x;
if (b >= B || k >= seq_len) return;
// anchor + 1 for the next-step window.
int global_idx = sample_file_offset[b] + sample_anchor[b] + 1 + k;
const Mbp10Raw& snap = all_snapshots[global_idx];
int n = b * seq_len + k;
int base_book = n * BOOK_LEVELS;
int base_regime = n * REGIME_DIM;
#pragma unroll
for (int i = 0; i < BOOK_LEVELS; i++) {
bid_px_soa[base_book + i] = snap.bid_px[i];
bid_sz_soa[base_book + i] = snap.bid_sz[i];
ask_px_soa[base_book + i] = snap.ask_px[i];
ask_sz_soa[base_book + i] = snap.ask_sz[i];
}
#pragma unroll
for (int i = 0; i < REGIME_DIM; i++) {
regime_soa[base_regime + i] = snap.regime[i];
}
prev_mid_soa[n] = snap.prev_mid;
tsv_soa[n] = snap.trade_signed_vol;
tc_soa[n] = (int)snap.trade_count;
ts_ns_soa[n] = (long long)snap.ts_ns;
prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
}
// ── Re-gather s_t: current window at anchor+0 ──────────────────────
//
// After gpu_sample_and_gather samples an anchor and gpu_gather_next
// overwrites the SoA with s_{t+1}, this kernel re-populates the SoA
// with the original s_t window. Uses the saved (file_offset, anchor)
// from gpu_sample_and_gather — no PRNG re-sampling. Identical to
// gpu_gather_next except the global_idx offset is +0 instead of +1.
//
// Grid = (B, 1, 1)
// Block = (K, 1, 1)
extern "C" __global__ void gpu_gather_current(
const Mbp10Raw* __restrict__ all_snapshots,
const int* __restrict__ sample_file_offset,
const int* __restrict__ sample_anchor,
int seq_len,
float* __restrict__ bid_px_soa,
float* __restrict__ bid_sz_soa,
float* __restrict__ ask_px_soa,
float* __restrict__ ask_sz_soa,
float* __restrict__ regime_soa,
float* __restrict__ prev_mid_soa,
float* __restrict__ tsv_soa,
int* __restrict__ tc_soa,
long long* __restrict__ ts_ns_soa,
long long* __restrict__ prev_ts_ns_soa,
int B
) {
int b = blockIdx.x;
int k = threadIdx.x;
if (b >= B || k >= seq_len) return;
int global_idx = sample_file_offset[b] + sample_anchor[b] + k;
const Mbp10Raw& snap = all_snapshots[global_idx];
int n = b * seq_len + k;
int base_book = n * BOOK_LEVELS;
int base_regime = n * REGIME_DIM;
#pragma unroll
for (int i = 0; i < BOOK_LEVELS; i++) {
bid_px_soa[base_book + i] = snap.bid_px[i];
bid_sz_soa[base_book + i] = snap.bid_sz[i];
ask_px_soa[base_book + i] = snap.ask_px[i];
ask_sz_soa[base_book + i] = snap.ask_sz[i];
}
#pragma unroll
for (int i = 0; i < REGIME_DIM; i++) {
regime_soa[base_regime + i] = snap.regime[i];
}
prev_mid_soa[n] = snap.prev_mid;
tsv_soa[n] = snap.trade_signed_vol;
tc_soa[n] = (int)snap.trade_count;
ts_ns_soa[n] = (long long)snap.ts_ns;
prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
}

206
crates/ml-alpha/examples/alpha_rl_train.rs
View File

@@ -37,7 +37,6 @@
use anyhow::{Context, Result};
use clap::Parser;
use data::providers::databento::dbn_parser::InstrumentFilter;
use ml_alpha::cfc::snap_features::Mbp10RawInput;
// Use the canonical Action enum size, NOT a literal. Caught 2026-05-24
// during SP20 P4 dogfood: a `0..9` range survived the N_ACTIONS=9→11
// bump and silently dropped HalfFlat samples (a9, a10 showed 0% in
@@ -47,6 +46,7 @@ use ml_alpha::cfc::snap_features::Mbp10RawInput;
// the literal value.
use ml_alpha::rl::common::{N_ACTIONS, FRD_N_ATOMS, FRD_N_HORIZONS};
use ml_alpha::rl::frd::FRD_OUT_DIM;
use ml_alpha::data::gpu_dataset::{GpuDataLoader, GpuDataset};
use ml_alpha::data::loader::{
discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig,
DEFAULT_OUTCOME_LABEL_COST_ES,
@@ -94,7 +94,6 @@ use ml_alpha::rl::isv_slots::{
RL_OUTCOME_AUX_LAMBDA_INDEX,
};
use ml_alpha::trainer::diag_staging::DiagStaging;
use ml_alpha::trainer::raw_launch::raw_memcpy_dtod_async;
use ml_alpha::trainer::integrated::{
IntegratedStepStats, IntegratedTrainer, IntegratedTrainerConfig,
};
@@ -299,73 +298,6 @@ struct AlphaRlTrainSummary {
nan_abort_step: i64,
}
/// Pre-assembled batch produced by the background prefetch thread.
/// Contains the B×K snapshot tensors and FRD labels ready for
/// copy to mapped-pinned staging. The prefetch thread runs
/// `loader.next_sequence_pair()` × n_batch while the main thread
/// trains on the previous batch, overlapping CPU data loading
/// with GPU compute.
struct PrefetchedBatch {
s_t_bk: Vec<Mbp10RawInput>,
s_tp1_bk: Vec<Mbp10RawInput>,
frd_labels_bh: Vec<i32>,
/// Signals that the loader hit EOF (next_sequence_pair returned
/// None) — the main thread should break the training loop.
eof: bool,
}
/// Assemble one batch of B×K snapshot pairs + FRD labels from the
/// loader. Returns `eof: true` if any batch slot hits loader EOF.
fn assemble_batch(
loader: &mut MultiHorizonLoader,
n_batch: usize,
seq_len: usize,
step: usize,
) -> Result<PrefetchedBatch> {
let bk = n_batch.saturating_mul(seq_len);
let mut s_t_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
let mut s_tp1_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
let mut frd_labels_bh: Vec<i32> = vec![-1_i32; n_batch * FRD_N_HORIZONS];
for b_idx in 0..n_batch {
let pair = loader
.next_sequence_pair()
.with_context(|| format!("next_sequence_pair at step {step} batch {b_idx}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!(
"loader EOF at step {step} batch {b_idx} (n_max_sequences exhausted) — exiting early"
);
return Ok(PrefetchedBatch {
s_t_bk,
s_tp1_bk,
frd_labels_bh,
eof: true,
});
}
};
debug_assert_eq!(s_t.snapshots.len(), seq_len);
debug_assert_eq!(s_tp1.snapshots.len(), seq_len);
s_t_bk.extend_from_slice(&s_t.snapshots);
s_tp1_bk.extend_from_slice(&s_tp1.snapshots);
for h in 0..FRD_N_HORIZONS {
let row = s_t.frd_labels.get(h);
let v = row.and_then(|r| r.first().copied()).unwrap_or(-1);
frd_labels_bh[b_idx * FRD_N_HORIZONS + h] = v;
}
}
debug_assert_eq!(s_t_bk.len(), bk);
debug_assert_eq!(s_tp1_bk.len(), bk);
Ok(PrefetchedBatch {
s_t_bk,
s_tp1_bk,
frd_labels_bh,
eof: false,
})
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
@@ -465,7 +397,23 @@ fn main() -> Result<()> {
frd_horizon_ticks: ml_alpha::rl::common::FRD_HORIZON_TICKS,
frd_bucket_range_sigma: ml_alpha::rl::common::FRD_BUCKET_RANGE_SIGMA,
};
let mut loader = MultiHorizonLoader::new(&loader_cfg).context("MultiHorizonLoader::new")?;
let loader = MultiHorizonLoader::new(&loader_cfg).context("MultiHorizonLoader::new")?;
// ── GPU-resident dataset upload. ────────────────────────────────
// Upload ALL pre-converted snapshots + FRD labels to GPU memory as
// a single contiguous device buffer. Per-step data loading becomes
// pure GPU kernel work (sample + gather) — zero CPU data work.
let gpu_dataset: GpuDataset = loader
.upload_to_gpu(&dev, cli.n_backtests.max(1), cli.seed)
.context("loader.upload_to_gpu")?;
let mut gpu_data_loader = GpuDataLoader::new(&dev, cli.n_backtests.max(1))
.context("GpuDataLoader::new")?;
eprintln!(
"GPU dataset uploaded: {} snapshots, {} files, {:.1} GB",
gpu_dataset.total_snapshots,
gpu_dataset.n_files,
(gpu_dataset.total_snapshots * 216) as f64 / 1e9,
);
// Trainer + LobSim.
let perception_cfg = PerceptionTrainerConfig {
@@ -554,79 +502,14 @@ fn main() -> Result<()> {
let mut frd_gate_total: u64 = 0;
let mut heat_cap_total: u64 = 0;
let n_batch = cli.n_backtests.max(1);
// ── Async double-buffered data prefetch. ────────────────────────
// The loader's `next_sequence_pair()` × n_batch takes ~19ms on
// HOST (mmap + label slicing), while the GPU finishes training in
// <1ms. Overlap the two: a background thread assembles batch N+1
// while the main thread trains on batch N.
//
// Bounded `sync_channel(2)` = double buffer: at most 2 batches
// in-flight (the one being trained on + the one being assembled).
// Back-pressure prevents the prefetch thread from racing ahead
// and consuming excessive memory.
let n_steps = cli.n_steps;
let seq_len = cli.seq_len;
let (prefetch_tx, prefetch_rx) = std::sync::mpsc::sync_channel::<PrefetchedBatch>(2);
let prefetch_handle = std::thread::spawn(move || {
for step in 0..n_steps {
let batch = match assemble_batch(&mut loader, n_batch, seq_len, step) {
Ok(b) => b,
Err(e) => {
eprintln!("prefetch thread error at step {step}: {e:#}");
return;
}
};
let is_eof = batch.eof;
if prefetch_tx.send(batch).is_err() {
// Main thread dropped the receiver (early exit).
return;
}
if is_eof {
return;
}
}
});
let t_start = std::time::Instant::now();
'outer: for step in 0..cli.n_steps {
// Receive the pre-assembled batch from the prefetch thread.
// On step 0 the first batch is already in the channel (the
// prefetch thread starts immediately); from step 1+ the batch
// was assembled while the GPU trained on the previous step.
let batch = prefetch_rx
.recv()
.map_err(|_| anyhow::anyhow!("prefetch thread died before step {step}"))?;
if batch.eof {
break 'outer;
}
let s_t_bk = &batch.s_t_bk;
let s_tp1_bk = &batch.s_tp1_bk;
let frd_labels_bh = &batch.frd_labels_bh;
// Write labels to mapped-pinned staging (host-side, instant).
// The DtoD from staging→frd_labels_d happens on the train stream
// without sync — the FRD backward kernel runs AFTER on the same
// stream, so ordering is guaranteed.
{
let dst = trainer.frd_labels_staging.host_slice_mut();
dst.copy_from_slice(&frd_labels_bh);
let nbytes = frd_labels_bh.len() * std::mem::size_of::<i32>();
unsafe {
let s = trainer.stream.cu_stream();
raw_memcpy_dtod_async(
trainer.frd_labels_d.raw_ptr(),
trainer.frd_labels_staging.dev_ptr,
nbytes,
s,
).map_err(|e| anyhow::anyhow!("frd_labels DtoD async: {:?}", e))?;
}
}
for step in 0..cli.n_steps {
// GPU-resident data loading: sample_and_gather + gather_next +
// gather_current + gather_frd_labels all run as GPU kernels
// inside step_with_lobsim_gpu. Zero CPU data work per step.
let stats = trainer
.step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim)
.with_context(|| format!("step_with_lobsim at step {step}"))?;
.step_with_lobsim_gpu(&mut gpu_data_loader, &gpu_dataset, &mut sim)
.with_context(|| format!("step_with_lobsim_gpu at step {step}"))?;
// Gate G8: NaN abort. Per `feedback_stop_on_anomaly` +
// `feedback_kill_runs_on_anomaly_quickly` the cluster smoke
@@ -1243,13 +1126,6 @@ fn main() -> Result<()> {
}
diag.flush().context("diag: final flush")?;
// ── Prefetch thread cleanup. ────────────────────────────────────
// Drop the receiver so the prefetch thread's `tx.send()` returns
// Err and the thread exits cleanly. Then join to ensure it has
// finished before we proceed to the eval phase (or summary).
drop(prefetch_rx);
let _ = prefetch_handle.join();
// ── Eval phase (walk-forward G8). ────────────────────────────────
// After the train phase, run n_eval_steps additional step_with_lobsim
// calls on the held-out eval files. Trade records accumulated by
@@ -1290,32 +1166,20 @@ fn main() -> Result<()> {
frd_horizon_ticks: ml_alpha::rl::common::FRD_HORIZON_TICKS,
frd_bucket_range_sigma: ml_alpha::rl::common::FRD_BUCKET_RANGE_SIGMA,
};
let mut eval_loader = MultiHorizonLoader::new(&eval_loader_cfg)
let eval_loader = MultiHorizonLoader::new(&eval_loader_cfg)
.context("MultiHorizonLoader::new (eval)")?;
let bk = n_batch.saturating_mul(cli.seq_len);
let mut s_t_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
let mut s_tp1_bk: Vec<Mbp10RawInput> = Vec::with_capacity(bk);
'eval_outer: for eval_step in 0..cli.n_eval_steps {
s_t_bk.clear();
s_tp1_bk.clear();
for b_idx in 0..n_batch {
let pair = eval_loader
.next_sequence_pair()
.with_context(|| format!("eval next_sequence_pair at step {eval_step} batch {b_idx}"))?;
let (s_t, s_tp1) = match pair {
Some(p) => p,
None => {
eprintln!("eval loader EOF at step {eval_step} batch {b_idx}");
break 'eval_outer;
}
};
s_t_bk.extend_from_slice(&s_t.snapshots);
s_tp1_bk.extend_from_slice(&s_tp1.snapshots);
}
// Upload eval dataset to GPU (separate from train dataset).
let eval_gpu_dataset = eval_loader
.upload_to_gpu(&dev, cli.n_backtests.max(1), cli.seed.wrapping_add(0xE7AE))
.context("eval loader.upload_to_gpu")?;
let mut eval_gpu_loader = GpuDataLoader::new(&dev, cli.n_backtests.max(1))
.context("GpuDataLoader::new (eval)")?;
for eval_step in 0..cli.n_eval_steps {
let stats = trainer
.step_with_lobsim(&s_t_bk, &s_tp1_bk, &mut sim)
.with_context(|| format!("step_with_lobsim eval step {eval_step}"))?;
.step_with_lobsim_gpu(&mut eval_gpu_loader, &eval_gpu_dataset, &mut sim)
.with_context(|| format!("step_with_lobsim_gpu eval step {eval_step}"))?;
if !stats.l_total.is_finite() {
eprintln!("G8 NaN ABORT during eval at step {eval_step}");
process::exit(2);

305
crates/ml-alpha/src/data/gpu_dataset.rs Normal file
View File

@@ -0,0 +1,305 @@
//! GPU-resident dataset: all pre-converted snapshots + labels on device.
//!
//! Eliminates ALL per-step CPU data loading work. At init, the host
//! pre-converts every `Mbp10Snapshot` → `Mbp10RawInput`, concatenates
//! across files, and uploads as a single contiguous device buffer. A
//! GPU kernel (`gpu_sample_and_gather`) does random sampling + window
//! gathering per step — zero host work, zero host↔device transfers in
//! the training hot path.
//!
//! Memory budget (45M snapshots):
//! snapshots: 45M × 216B = 9.7 GB
//! FRD labels: 45M × 3 × 4B = 0.54 GB
//! Total: ~10.2 GB on L40S (48GB), leaving 34+ GB free.
//!
//! Per `feedback_no_htod_htoh_only_mapped_pinned.md`: the init-time
//! upload uses cudarc's `htod_copy` (a one-time bulk transfer, not in
//! the captured-graph hot path). The per-step path is pure device-to-
//! device (kernel reads from one device buffer, writes to another).
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
use cudarc::driver::sys::CUstream;
use ml_core::device::MlDevice;
use crate::rl::common::FRD_N_HORIZONS;
use crate::trainer::perception::SoaBufferPtrs;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const GPU_SAMPLE_GATHER_CUBIN: &[u8] = include_bytes!(
concat!(env!("OUT_DIR"), "/gpu_sample_and_gather.cubin")
);
/// All pre-converted snapshots + metadata resident on GPU. Created once
/// at init by [`super::loader::MultiHorizonLoader::upload_to_gpu`].
pub struct GpuDataset {
/// All snapshots concatenated across files, on GPU.
/// Byte layout: `[total_snaps]` of `Mbp10Raw` (216 bytes each).
/// The kernel reinterprets this as `Mbp10Raw*`.
pub snapshots_d: CudaSlice<u8>,
/// Per-file start offset into `snapshots_d` (in snapshot units, not
/// bytes). `file_offsets_d[f]` is the index of the first snapshot
/// from file `f` in the flat `snapshots_d` array.
pub file_offsets_d: CudaSlice<i32>,
/// Per-file snapshot count. `file_sizes_d[f]` is the number of
/// snapshots in file `f`.
pub file_sizes_d: CudaSlice<i32>,
/// FRD labels, row-major `[FRD_N_HORIZONS, total_snaps]`.
/// Sentinel -1 at right-edge positions.
pub frd_labels_d: CudaSlice<i32>,
/// PRNG state per batch element `[B]`. Seeded at init, advanced by
/// the kernel each step.
pub prng_d: CudaSlice<u32>,
/// Number of loaded files.
pub n_files: usize,
/// Total snapshots across all files.
pub total_snapshots: usize,
/// Maximum forward horizon in snapshot units (max of
/// `cfg.horizons.iter().max()`), needed by the sampling kernel to
/// avoid right-edge overrun.
pub max_horizon: usize,
}
/// GPU-resident data loader: dispatches the `gpu_sample_and_gather`
/// kernel family to fill SoA buffers each step with zero CPU work.
///
/// Usage pattern per training step (mirrors the existing
/// `step_with_lobsim` sequence):
/// 1. `sample_and_gather(dataset, soa, ...)` — samples B random
/// (file, anchor) pairs AND gathers the s_t window into SoA.
/// 2. Caller runs `forward_encoder_from_device()` → h_t. But wait:
/// the integrated trainer needs h_{t+1} FIRST. So the actual
/// sequence is:
/// a. `sample(dataset, ...)` — samples only, writes (file_offset, anchor) to device.
/// b. `gather_next_window(dataset, soa, ...)` — gathers s_{t+1}.
/// c. `forward_encoder_from_device()` → copy h_t → h_tp1.
/// d. `gather_current_window(dataset, soa, ...)` — gathers s_t.
/// e. `forward_encoder_from_device()` → h_t.
/// f. `gather_frd_labels(dataset, ...)` — gathers FRD labels.
pub struct GpuDataLoader {
/// Kept alive to anchor the `raw_stream` pointer lifetime.
_stream: Arc<CudaStream>,
raw_stream: CUstream,
_module: Arc<CudaModule>,
sample_gather_fn: CudaFunction,
gather_next_fn: CudaFunction,
gather_current_fn: CudaFunction,
gather_frd_labels_fn: CudaFunction,
/// Per-batch sampled file offsets — written by the sampling kernel,
/// read by `gather_next_window` and `gather_frd_labels` so they
/// use the same (file, anchor) pair. `[B]`.
sample_file_offset_d: CudaSlice<i32>,
/// Per-batch sampled anchors. `[B]`.
sample_anchor_d: CudaSlice<i32>,
}
impl GpuDataLoader {
pub fn new(dev: &MlDevice, batch_size: usize) -> Result<Self> {
let stream = dev.cuda_stream().context("GpuDataLoader: CUDA stream")?.clone();
let raw_stream = stream.cu_stream();
let ctx = dev.cuda_context().context("GpuDataLoader: CUDA context")?;
let module = ctx
.load_cubin(GPU_SAMPLE_GATHER_CUBIN.to_vec())
.context("load gpu_sample_and_gather cubin")?;
let sample_gather_fn = module
.load_function("gpu_sample_and_gather")
.context("load gpu_sample_and_gather function")?;
let gather_next_fn = module
.load_function("gpu_gather_next")
.context("load gpu_gather_next function")?;
let gather_current_fn = module
.load_function("gpu_gather_current")
.context("load gpu_gather_current function")?;
let gather_frd_labels_fn = module
.load_function("gpu_gather_frd_labels")
.context("load gpu_gather_frd_labels function")?;
let sample_file_offset_d = stream
.alloc_zeros::<i32>(batch_size)
.context("alloc sample_file_offset_d")?;
let sample_anchor_d = stream
.alloc_zeros::<i32>(batch_size)
.context("alloc sample_anchor_d")?;
Ok(Self {
_stream: stream,
raw_stream,
_module: module,
sample_gather_fn,
gather_next_fn,
gather_current_fn,
gather_frd_labels_fn,
sample_file_offset_d,
sample_anchor_d,
})
}
/// Sample B random (file, anchor) pairs AND gather the s_t window
/// into the perception trainer's SoA buffers. The sampled
/// (file_offset, anchor) values are saved to device memory for
/// subsequent `gather_next_window` and `gather_frd_labels` calls.
pub fn sample_and_gather(
&mut self,
dataset: &GpuDataset,
soa: &SoaBufferPtrs,
seq_len: usize,
batch_size: usize,
) -> Result<()> {
let mut args = RawArgs::new();
args.push_ptr(dataset.snapshots_d.raw_ptr());
args.push_ptr(dataset.file_offsets_d.raw_ptr());
args.push_ptr(dataset.file_sizes_d.raw_ptr());
args.push_ptr(dataset.prng_d.raw_ptr());
args.push_i32(dataset.n_files as i32);
args.push_i32(seq_len as i32);
args.push_i32(dataset.max_horizon as i32);
args.push_ptr(soa.bid_px);
args.push_ptr(soa.bid_sz);
args.push_ptr(soa.ask_px);
args.push_ptr(soa.ask_sz);
args.push_ptr(soa.regime);
args.push_ptr(soa.prev_mid);
args.push_ptr(soa.trade_signed_vol);
args.push_ptr(soa.trade_count);
args.push_ptr(soa.ts_ns);
args.push_ptr(soa.prev_ts_ns);
args.push_ptr(self.sample_file_offset_d.raw_ptr());
args.push_ptr(self.sample_anchor_d.raw_ptr());
args.push_i32(batch_size as i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.sample_gather_fn.cu_function(),
(batch_size as u32, 1, 1),
(seq_len as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("gpu_sample_and_gather: {:?}", e))?;
}
Ok(())
}
/// Gather the s_{t+1} window (anchor+1) into the specified SoA
/// buffers. Uses the (file_offset, anchor) sampled by the most
/// recent `sample_and_gather` call.
pub fn gather_next_window(
&mut self,
dataset: &GpuDataset,
soa: &SoaBufferPtrs,
seq_len: usize,
batch_size: usize,
) -> Result<()> {
let mut args = RawArgs::new();
args.push_ptr(dataset.snapshots_d.raw_ptr());
args.push_ptr(self.sample_file_offset_d.raw_ptr());
args.push_ptr(self.sample_anchor_d.raw_ptr());
args.push_i32(seq_len as i32);
args.push_ptr(soa.bid_px);
args.push_ptr(soa.bid_sz);
args.push_ptr(soa.ask_px);
args.push_ptr(soa.ask_sz);
args.push_ptr(soa.regime);
args.push_ptr(soa.prev_mid);
args.push_ptr(soa.trade_signed_vol);
args.push_ptr(soa.trade_count);
args.push_ptr(soa.ts_ns);
args.push_ptr(soa.prev_ts_ns);
args.push_i32(batch_size as i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.gather_next_fn.cu_function(),
(batch_size as u32, 1, 1),
(seq_len as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("gpu_gather_next: {:?}", e))?;
}
Ok(())
}
/// Re-gather the s_t window (anchor+0) into the SoA buffers. Uses
/// the saved (file_offset, anchor) from the most recent
/// `sample_and_gather` call — no PRNG re-sampling. Called after
/// `gather_next_window` overwrites the SoA with s_{t+1} and the
/// encoder forward on s_{t+1} has completed, to restore s_t in the
/// SoA for the second encoder forward.
pub fn gather_current_window(
&mut self,
dataset: &GpuDataset,
soa: &SoaBufferPtrs,
seq_len: usize,
batch_size: usize,
) -> Result<()> {
let mut args = RawArgs::new();
args.push_ptr(dataset.snapshots_d.raw_ptr());
args.push_ptr(self.sample_file_offset_d.raw_ptr());
args.push_ptr(self.sample_anchor_d.raw_ptr());
args.push_i32(seq_len as i32);
args.push_ptr(soa.bid_px);
args.push_ptr(soa.bid_sz);
args.push_ptr(soa.ask_px);
args.push_ptr(soa.ask_sz);
args.push_ptr(soa.regime);
args.push_ptr(soa.prev_mid);
args.push_ptr(soa.trade_signed_vol);
args.push_ptr(soa.trade_count);
args.push_ptr(soa.ts_ns);
args.push_ptr(soa.prev_ts_ns);
args.push_i32(batch_size as i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.gather_current_fn.cu_function(),
(batch_size as u32, 1, 1),
(seq_len as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("gpu_gather_current: {:?}", e))?;
}
Ok(())
}
/// Gather FRD labels at the newest snapshot (anchor + seq_len - 1)
/// of the s_t window into the output buffer. Uses the (file_offset,
/// anchor) sampled by the most recent `sample_and_gather` call.
pub fn gather_frd_labels(
&mut self,
dataset: &GpuDataset,
seq_len: usize,
batch_size: usize,
frd_labels_out_ptr: u64,
) -> Result<()> {
let mut args = RawArgs::new();
args.push_ptr(dataset.frd_labels_d.raw_ptr());
args.push_ptr(dataset.file_offsets_d.raw_ptr());
args.push_ptr(self.sample_file_offset_d.raw_ptr());
args.push_ptr(self.sample_anchor_d.raw_ptr());
args.push_i32(seq_len as i32);
args.push_i32(FRD_N_HORIZONS as i32);
args.push_i32(dataset.total_snapshots as i32);
args.push_ptr(frd_labels_out_ptr);
args.push_i32(batch_size as i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.gather_frd_labels_fn.cu_function(),
(batch_size as u32, 1, 1),
(FRD_N_HORIZONS as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("gpu_gather_frd_labels: {:?}", e))?;
}
Ok(())
}
}

171
crates/ml-alpha/src/data/loader.rs
View File

@@ -673,6 +673,177 @@ impl MultiHorizonLoader {
Ok(convert(first_snap, first_snap, first_file.regime_full[0]))
}
/// Upload ALL pre-converted snapshots + FRD labels to GPU as a
/// single flat buffer. Returns a [`GpuDataset`] that the
/// [`GpuDataLoader`] kernel reads from each step — zero CPU work
/// in the training hot path.
///
/// This method is the CPU-side half of the GPU-resident data loader:
/// 1. Pre-converts every `Mbp10Snapshot` → `Mbp10RawInput` (calls
/// `convert()` once per snapshot at init, not per step).
/// 2. Concatenates all files into flat arrays.
/// 3. Uploads via `htod_copy` (one-time bulk transfer).
/// 4. Builds per-file offset + size tables on device.
/// 5. Initializes per-batch PRNG seeds on device.
///
/// Per `feedback_no_htod_htoh_only_mapped_pinned.md`: the init-time
/// `htod_copy` is a one-shot outside the captured-graph hot path.
/// The per-step path is pure device→device (kernel reads dataset
/// buffers, writes SoA buffers).
///
/// [`GpuDataset`]: super::gpu_dataset::GpuDataset
/// [`GpuDataLoader`]: super::gpu_dataset::GpuDataLoader
pub fn upload_to_gpu(
&self,
dev: &ml_core::device::MlDevice,
batch_size: usize,
seed: u64,
) -> Result<super::gpu_dataset::GpuDataset> {
use crate::cfc::snap_features::MBP10_RAW_INPUT_BYTES;
use crate::rl::common::FRD_N_HORIZONS;
let stream = dev.cuda_stream().context("upload_to_gpu: CUDA stream")?.clone();
// ── 1. Pre-convert ALL snapshots → Mbp10RawInput + build offsets ──
let n_files = self.files_loaded.len();
let total_snaps: usize = self.files_loaded.iter().map(|f| f.snapshots.len()).sum();
tracing::info!(
n_files,
total_snaps,
bytes = total_snaps * MBP10_RAW_INPUT_BYTES,
"upload_to_gpu: pre-converting snapshots",
);
let mut all_raws: Vec<Mbp10RawInput> = Vec::with_capacity(total_snaps);
let mut file_offsets: Vec<i32> = Vec::with_capacity(n_files);
let mut file_sizes: Vec<i32> = Vec::with_capacity(n_files);
let mut offset: usize = 0;
for lf in &self.files_loaded {
file_offsets.push(offset as i32);
file_sizes.push(lf.snapshots.len() as i32);
for (idx, snap) in lf.snapshots.iter().enumerate() {
let prev = if idx == 0 { snap } else { &lf.snapshots[idx - 1] };
all_raws.push(convert(snap, prev, lf.regime_full[idx]));
}
offset += lf.snapshots.len();
}
debug_assert_eq!(all_raws.len(), total_snaps);
// ── 2. Upload snapshots as raw bytes ─────────────────────────────
// Reinterpret Vec<Mbp10RawInput> as &[u8] for the upload. The GPU
// kernel reads the same struct layout via its `Mbp10Raw` definition.
let total_bytes = total_snaps * MBP10_RAW_INPUT_BYTES;
let snap_bytes: &[u8] = unsafe {
std::slice::from_raw_parts(
all_raws.as_ptr() as *const u8,
total_bytes,
)
};
let snapshots_d = stream
.alloc_zeros::<u8>(total_bytes)
.context("upload_to_gpu: snapshots alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
snapshots_d.raw_ptr(),
snap_bytes.as_ptr() as *const std::ffi::c_void,
total_bytes,
);
}
// Drop the host-side Vec now — it can be multi-GB.
drop(all_raws);
tracing::info!(
gpu_bytes = total_bytes,
"upload_to_gpu: snapshots uploaded",
);
// ── 3. Upload file offsets + sizes ───────────────────────────────
let file_offsets_d = stream
.alloc_zeros::<i32>(n_files)
.context("upload_to_gpu: file_offsets alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
file_offsets_d.raw_ptr(),
file_offsets.as_ptr() as *const std::ffi::c_void,
n_files * 4,
);
}
let file_sizes_d = stream
.alloc_zeros::<i32>(n_files)
.context("upload_to_gpu: file_sizes alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
file_sizes_d.raw_ptr(),
file_sizes.as_ptr() as *const std::ffi::c_void,
n_files * 4,
);
}
// ── 4. Upload FRD labels [FRD_N_HORIZONS, total_snaps] row-major ──
let frd_total = FRD_N_HORIZONS * total_snaps;
let mut frd_flat: Vec<i32> = vec![-1_i32; frd_total];
let mut global_offset: usize = 0;
for lf in &self.files_loaded {
let n = lf.snapshots.len();
for h in 0..FRD_N_HORIZONS {
let src = &lf.frd_labels_full[h];
let dst_start = h * total_snaps + global_offset;
frd_flat[dst_start..dst_start + n].copy_from_slice(&src[..n]);
}
global_offset += n;
}
let frd_labels_d = stream
.alloc_zeros::<i32>(frd_total)
.context("upload_to_gpu: frd_labels alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
frd_labels_d.raw_ptr(),
frd_flat.as_ptr() as *const std::ffi::c_void,
frd_total * 4,
);
}
// ── 5. Initialize per-batch PRNG seeds ──────────────────────────
let prng_seeds: Vec<u32> = (0..batch_size as u64)
.map(|b| {
// Mix seed + batch index to get distinct per-batch seeds.
// Ensures different batch elements sample different files/anchors.
let mixed = seed.wrapping_mul(2654435761).wrapping_add(b);
(mixed & 0xFFFF_FFFF) as u32
})
.collect();
let prng_d = stream
.alloc_zeros::<u32>(batch_size)
.context("upload_to_gpu: prng alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
prng_d.raw_ptr(),
prng_seeds.as_ptr() as *const std::ffi::c_void,
batch_size * 4,
);
}
let max_horizon = *self.cfg.horizons.iter().max().expect("non-empty horizons");
tracing::info!(
n_files,
total_snaps,
max_horizon,
batch_size,
"upload_to_gpu: dataset ready",
);
Ok(super::gpu_dataset::GpuDataset {
snapshots_d,
file_offsets_d,
file_sizes_d,
frd_labels_d,
prng_d,
n_files,
total_snapshots: total_snaps,
max_horizon,
})
}
/// Iterator-style accessor for inference mode: walks every loaded
/// snapshot across every loaded file in chronological order.
/// Returns `Ok(None)` when the stream is exhausted. Unlike

1
crates/ml-alpha/src/data/mod.rs
View File

@@ -1,4 +1,5 @@
//! Phase A data path — predecoded MBP-10 -> Mbp10RawInput sequences + labels.
pub mod aggregation;
pub mod gpu_dataset;
pub mod loader;

708
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -986,6 +986,15 @@ pub struct IntegratedTrainer {
/// the same config + step sequence.
step_counter: u64,
/// When true, `step_synthetic` dispatches the encoder forward via
/// `forward_encoder_from_device()` (GPU-resident SoA buffers
/// pre-filled by the caller) instead of `forward_encoder(snapshots)`.
/// Set by `step_with_lobsim_gpu` before the K-loop, cleared after.
/// This avoids duplicating the 850-line reward + K-loop pipeline
/// between the host and GPU data paths. NOT a feature flag — it's
/// call-context state consumed within the same step.
gpu_encoder_active: bool,
/// Host-side counter for spectral norm amortization. Spectral norm
/// (power iteration on full weight matrices) costs 175us x 3 heads =
/// 843us/step. Running every step is wasteful because per-step weight
@@ -2464,6 +2473,7 @@ impl IntegratedTrainer {
last_step_stats: IntegratedStepStats::default(),
last_k_updates: 0,
step_counter: 0,
gpu_encoder_active: false,
spectral_norm_counter: 0,
frd_head,
frd_hidden_d,
@@ -3657,6 +3667,27 @@ impl IntegratedTrainer {
&mut self,
snapshots: &[Mbp10RawInput],
) -> Result<IntegratedStepStats> {
self.step_synthetic_read_deferred_stats()?;
if self.gpu_encoder_active {
// GPU data path: SoA buffers pre-filled by gather kernels.
let _ = self
.perception
.forward_encoder_from_device()
.context("forward_encoder_from_device")?;
} else {
let _ = self
.perception
.forward_encoder(snapshots)
.context("forward_encoder")?;
}
self.step_synthetic_body()
}
/// Read previous step's losses from mapped-pinned memory and
/// compose `self.last_step_stats`. Shared preamble for both
/// `step_synthetic` (host data) and `step_synthetic_gpu` (device
/// data). Must be called before the encoder forward each step.
fn step_synthetic_read_deferred_stats(&mut self) -> Result<()> {
let b_size = self.cfg.perception.n_batch;
if b_size == 0 {
anyhow::bail!("step_synthetic: empty batch (b_size = 0)");
@@ -3731,12 +3762,16 @@ impl IntegratedTrainer {
lambdas: prev_lambdas,
};
}
Ok(())
}
// ── Step 1: encoder forward ──────────────────────────────────
let _ = self
.perception
.forward_encoder(snapshots)
.context("forward_encoder")?;
/// Post-encoder body shared by `step_synthetic` and
/// `step_synthetic_gpu`. Runs the LR controller, all head
/// forwards/backwards, Adam steps, and loss composition. The
/// encoder must have been forwarded before this call (either via
/// `forward_encoder` or `forward_encoder_from_device`).
fn step_synthetic_body(&mut self) -> Result<IntegratedStepStats> {
let b_size = self.cfg.perception.n_batch;
// ── Step 2: per-head LR controller emit ──────────────────────
// The controller emits the bootstrap target on first call
@@ -5733,6 +5768,26 @@ impl IntegratedTrainer {
.step_fill_from_market_targets(last_snap.ts_ns)
.context("step_with_lobsim: lobsim.step_fill_from_market_targets")?;
self.step_with_lobsim_reward_and_train(lobsim, snapshots, b_size)
}
/// Post-fill reward pipeline, controllers, advantage computation,
/// PER push/sample, K-loop (step_synthetic + dqn_replay), and
/// target-net soft update. Shared by both `step_with_lobsim` (host
/// data) and `step_with_lobsim_gpu` (GPU data). The `snapshots`
/// parameter is passed through to `step_synthetic` in the K-loop
/// for the host path; when `self.gpu_encoder_active` is true,
/// `step_synthetic` ignores snapshots and uses
/// `forward_encoder_from_device`.
fn step_with_lobsim_reward_and_train(
&mut self,
lobsim: &mut dyn RlLobBackend,
snapshots: &[Mbp10RawInput],
b_size: usize,
) -> Result<IntegratedStepStats> {
let b_size_i = b_size as i32;
let pos_bytes_i = lobsim.pos_bytes() as i32;
// ── Graph B: post-fill reward/EMA/controller pipeline ─────────
// ~20 kernels from extract_realized_pnl_delta through
// var_over_abs_mean. Same three-state machine as Graph A/A2 —
@@ -6585,6 +6640,649 @@ impl IntegratedTrainer {
Ok(stats)
}
/// GPU-resident data path variant of `step_with_lobsim`. Runs one
/// full RL training step driven by `LobSimCuda` (via `RlLobBackend`)
/// using GPU-resident `GpuDataset` + `GpuDataLoader` for data
/// loading. Zero CPU data work per step — sampling, window gathering,
/// FRD label gathering, and AoS->SoA scatter all run as GPU kernels.
///
/// Key differences from `step_with_lobsim`:
/// * Encoder forwards read from SoA buffers pre-filled by the GPU
/// gather kernels (`forward_encoder_from_device` replaces
/// `forward_encoder(snapshots)`).
/// * `lobsim.apply_snapshot` is skipped — the lobsim's book state
/// carries forward from the previous step (one-tick stale). This
/// is acceptable for RL training: fill pricing is approximate by
/// design (the agent learns on noisy fills).
/// * FRD labels are gathered by `gpu_gather_frd_labels` directly
/// into `self.frd_labels_d` (no mapped-pinned staging).
/// * `ts_ns` for `step_fill_from_market_targets` is DtoD'd from
/// the SoA's last-snapshot position; the multires kernel reads
/// `self.ts_ns_d` which is filled by the same DtoD.
///
/// Call sequence per step:
/// 1. `sample_and_gather` — samples (file, anchor), gathers s_t
/// 2. `gather_next_window` — gathers s_{t+1} using saved anchor
/// 3. `forward_encoder_from_device` on s_{t+1} -> h_{t+1}
/// 4. DtoD h_t -> h_tp1
/// 5. `gather_current_window` — re-gathers s_t (saved anchor, no PRNG)
/// 6. `forward_encoder_from_device` on s_t -> h_t
/// 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
pub fn step_with_lobsim_gpu(
&mut self,
gpu_loader: &mut crate::data::gpu_dataset::GpuDataLoader,
gpu_dataset: &crate::data::gpu_dataset::GpuDataset,
lobsim: &mut dyn RlLobBackend,
) -> Result<IntegratedStepStats> {
let b_size = self.cfg.perception.n_batch;
let seq_len = self.cfg.perception.seq_len;
if b_size == 0 {
anyhow::bail!("step_with_lobsim_gpu: empty batch (n_batch = 0)");
}
// ── 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))?;
}
// ── Step 0: bump device-resident step counter (ISV[548]). ──
{
let mut args = RawArgs::new();
args.push_ptr(self.isv_dev_ptr);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_increment_step_fn.cu_function(),
(1, 1, 1), (1, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_increment_step: {:?}", e))?;
}
}
// ── Step 1a: GPU data loading + encoder forward on s_{t+1}. ──
//
// 1. sample_and_gather: samples (file, anchor), gathers s_t to
// SoA, saves anchor in device scratch.
// 2. gather_next_window: gathers s_{t+1} to SoA (overwrites s_t)
// using the saved anchor.
// 3. forward_encoder_from_device: runs encoder on SoA (= s_{t+1}).
// 4. DtoD h_t -> h_tp1: saves h_{t+1}.
let soa = self.perception.soa_buffer_ptrs();
gpu_loader
.sample_and_gather(gpu_dataset, &soa, seq_len, b_size)
.context("step_with_lobsim_gpu: sample_and_gather")?;
gpu_loader
.gather_next_window(gpu_dataset, &soa, seq_len, b_size)
.context("step_with_lobsim_gpu: gather_next_window")?;
let _ = self
.perception
.forward_encoder_from_device()
.context("step_with_lobsim_gpu: forward_encoder_from_device(s_tp1)")?;
{
let nbytes = b_size * HIDDEN_DIM * std::mem::size_of::<f32>();
unsafe {
raw_memcpy_dtod_async(
self.hot.h_tp1, self.hot.h_t_encoder, nbytes, self.raw_stream,
)
.map_err(|e| anyhow::anyhow!("step_with_lobsim_gpu: DtoD h_t -> h_tp1: {:?}", e))?;
}
}
// ── Step 1b: re-gather s_t + encoder forward on current snap. ──
//
// gather_current_window re-populates the SoA with s_t using the
// saved (file_offset, anchor) from step 1a — no PRNG re-sampling.
gpu_loader
.gather_current_window(gpu_dataset, &soa, seq_len, b_size)
.context("step_with_lobsim_gpu: gather_current_window")?;
// FRD labels gathered by GPU kernel into trainer's frd_labels_d.
gpu_loader
.gather_frd_labels(gpu_dataset, seq_len, b_size, self.frd_labels_d.raw_ptr())
.context("step_with_lobsim_gpu: gather_frd_labels")?;
{
self.perception.encoder_ctx_trade_ptr = self.hot.trade_context;
self.perception.encoder_ctx_multires_ptr = self.hot.multires_output;
}
let _ = self
.perception
.forward_encoder_from_device()
.context("step_with_lobsim_gpu: forward_encoder_from_device(s_t)")?;
self.step_counter = self.step_counter.wrapping_add(1);
// ── From here, the kernel pipeline is identical to
// step_with_lobsim. The only remaining difference is that
// apply_snapshot + ts_ns staging are handled differently:
//
// * apply_snapshot is SKIPPED (lobsim book carries forward from
// previous step — one-tick-stale pricing, acceptable for RL).
// * ts_ns for step_fill_from_market_targets is DtoD'd from the
// 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)
}
/// Shared body for `step_with_lobsim_gpu`. Runs the graph pipelines,
/// lobsim fill, reward/controller chain, PER, K-loop, and target-net
/// update. Called after both encoder forwards have completed and h_t /
/// h_tp1 are populated.
///
/// `soa` is passed for the ts_ns DtoD (avoids reconstructing the
/// pointer struct). `b_size` and `seq_len` are passed to avoid
/// re-reading from cfg.
fn step_with_lobsim_gpu_body(
&mut self,
lobsim: &mut dyn RlLobBackend,
soa: &crate::trainer::perception::SoaBufferPtrs,
b_size: usize,
seq_len: usize,
) -> Result<IntegratedStepStats> {
// ── Graph A: pre-snapshot kernel pipeline ──────────────────────
if self.prefill_graph.is_some() {
unsafe {
raw_graph_launch(
self.prefill_graph.as_ref().unwrap().cu_graph_exec,
self.raw_stream,
).map_err(|e| anyhow::anyhow!("prefill graph launch: {:?}", e))?;
}
} else {
let capturing_prefill = self.graph_warmup_done;
if !self.graph_warmup_done {
self.graph_warmup_done = true;
}
if capturing_prefill {
self.stream
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
.map_err(|e| anyhow::anyhow!("prefill begin_capture: {e}"))?;
}
// ── Step 2: Q + V forwards on h_t for action sampling. ────────
let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM);
debug_assert_eq!(self.h_tp1_d.len(), b_size * HIDDEN_DIM);
self.frd_head
.forward(
h_t_borrow,
&mut self.frd_hidden_d,
&mut self.frd_logits_d,
b_size,
)
.context("frd_head.forward in step_with_lobsim_gpu")?;
self.dqn_head
.forward(h_t_borrow, b_size, &mut self.q_logits_d)
.context("step_with_lobsim_gpu: dqn_head.forward(h_t)")?;
self.dqn_head
.forward(&self.h_tp1_d, b_size, &mut self.q_logits_tp1_d)
.context("step_with_lobsim_gpu: dqn_head.forward(h_tp1)")?;
self.value_head
.forward(h_t_borrow, &self.isv_dev_ptr, b_size, &mut self.v_pred_d)
.context("step_with_lobsim_gpu: value_head.forward(h_t)")?;
self.value_head
.forward(&self.h_tp1_d, &self.isv_dev_ptr, b_size, &mut self.v_pred_tp1_d)
.context("step_with_lobsim_gpu: value_head.forward(h_tp1)")?;
// IQN forward
{
let n_tau = self.iqn_head.n_tau();
self.iqn_head
.forward(
&mut self.iqn_prng_state_d,
h_t_borrow,
&mut self.iqn_tau_d,
b_size,
n_tau,
&mut self.iqn_q_values_d,
)
.context("step_with_lobsim_gpu: iqn_head.forward(h_t)")?;
self.iqn_head
.expected_q(
&self.iqn_q_values_d,
b_size,
n_tau,
&mut self.iqn_expected_q_d,
)
.context("step_with_lobsim_gpu: iqn_head.expected_q")?;
}
// Ensemble Q
{
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.q_logits_d.raw_ptr());
args.push_ptr(self.atom_supports_d.raw_ptr());
args.push_ptr(self.iqn_expected_q_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_ptr(self.ensemble_q_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_ensemble_action_value_fn.cu_function(),
(b_size as u32, 1, 1), (N_ACTIONS as u32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_ensemble_action_value: {:?}", e))?;
}
}
// NoisyNet exploration
self.noisy_exploration
.resample_noise()
.context("step_with_lobsim_gpu: noisy_exploration.resample_noise")?;
self.noisy_exploration
.forward(h_t_borrow, b_size, &mut self.noisy_output_d)
.context("step_with_lobsim_gpu: noisy_exploration.forward")?;
{
let n_elems = (b_size * N_ACTIONS) as i32;
let grid_x = ((n_elems as u32) + 255) / 256;
let mut args = RawArgs::new();
args.push_ptr(self.ensemble_q_d.raw_ptr());
args.push_ptr(self.noisy_output_d.raw_ptr());
args.push_i32(n_elems);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.aux_vec_add_fn.cu_function(),
(grid_x, 1, 1), (256, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("aux_vec_add_inplace (noisy -> ensemble_q): {:?}", e))?;
}
}
// pi logits
self.policy_head
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
.context("step_with_lobsim_gpu: policy_head.forward_logits")?;
// Q-vs-pi agree
{
let block_dim = (b_size.next_power_of_two() as u32).max(1);
let smem = (block_dim as usize * std::mem::size_of::<f32>()) as u32;
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.q_logits_d.raw_ptr());
args.push_ptr(self.pi_logits_d.raw_ptr());
args.push_ptr(self.atom_supports_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_q_pi_agree_b_fn.cu_function(),
(1, 1, 1), (block_dim, 1, 1), smem,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_q_pi_agree_b: {:?}", e))?;
}
}
// pi-driven action selection
{
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.pi_logits_d.raw_ptr());
args.push_ptr(self.prng_state_d.raw_ptr());
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_pi_action_kernel_fn.cu_function(),
(b_size as u32, 1, 1), (1, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_pi_action_kernel: {:?}", e))?;
}
}
// argmax expected Q for Bellman target
{
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.q_logits_tp1_d.raw_ptr());
args.push_ptr(self.atom_supports_d.raw_ptr());
args.push_ptr(self.next_actions_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.argmax_expected_q_fn.cu_function(),
(b_size as u32, 1, 1), (N_ACTIONS as u32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("argmax_expected_q (h_tp1): {:?}", e))?;
}
}
// log pi_old at sampled action
{
let grid_x = ((b_size as u32) + 31) / 32;
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.pi_logits_d.raw_ptr());
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.log_pi_old_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.log_pi_at_action_fn.cu_function(),
(grid_x.max(1), 1, 1), (32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("log_pi_at_action: {:?}", e))?;
}
}
// Confidence gate
{
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
let b_size_i = b_size as i32;
let pos_bytes_i = lobsim.pos_bytes() as i32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.q_logits_d.raw_ptr());
args.push_ptr(self.atom_supports_d.raw_ptr());
args.push_ptr(pos_d_ref.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
args.push_i32(pos_bytes_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_confidence_gate_fn.cu_function(),
(b_size as u32, 1, 1), (1, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_confidence_gate: {:?}", e))?;
}
}
// FRD gate
{
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
let b_size_i = b_size as i32;
let pos_bytes_i = lobsim.pos_bytes() as i32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.frd_logits_d.raw_ptr());
args.push_ptr(pos_d_ref.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
args.push_i32(pos_bytes_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_frd_gate_fn.cu_function(),
(b_size as u32, 1, 1), (1, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_frd_gate: {:?}", e))?;
}
}
if capturing_prefill {
let graph = self
.stream
.end_capture(
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
)
.context("prefill end_capture")?
.ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?;
self.prefill_graph = Some(graph);
}
} // end else (warmup / capture dispatch)
// ── Step 5: GPU-pure env step (GPU data path). ─────────────────
// apply_snapshot SKIPPED: the lobsim's book state carries forward
// from the previous step. For RL training the one-tick-stale book
// pricing is acceptable — the agent learns on approximate fills.
//
// ts_ns: DtoD from the SoA's last-snapshot ts_ns into self.ts_ns_d
// for the multires kernel. The lobsim fill kernel receives 0 as
// the timestamp (diagnostic-only, PnL tracking timestamps).
{
let last_ts_offset = ((seq_len - 1) * std::mem::size_of::<u64>()) as u64;
unsafe {
raw_memcpy_dtod_async(
self.ts_ns_d.raw_ptr(),
soa.ts_ns + last_ts_offset,
std::mem::size_of::<u64>(),
self.raw_stream,
).map_err(|e| anyhow::anyhow!("ts_ns DtoD from SoA: {:?}", e))?;
}
}
let pos_bytes_i = lobsim.pos_bytes() as i32;
let b_size_i = b_size as i32;
// ── Graph A2: post-snapshot / pre-fill kernel pipeline ─────────
if self.postfill_graph.is_some() {
unsafe {
raw_graph_launch(
self.postfill_graph.as_ref().unwrap().cu_graph_exec,
self.raw_stream,
).map_err(|e| anyhow::anyhow!("postfill graph launch: {:?}", e))?;
}
} else {
let capturing_postfill = self.prefill_graph.is_some();
if capturing_postfill {
self.stream
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
.map_err(|e| anyhow::anyhow!("postfill begin_capture: {e}"))?;
}
// Session risk check
{
let b_size_i = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.rewards_d.raw_ptr());
args.push_ptr(self.dones_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_session_risk_check_fn.cu_function(),
(1, 1, 1), (1, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_session_risk_check: {:?}", e))?;
}
}
// Min hold check
{
let pos_d_ref: &CudaSlice<u8> = lobsim.pos_d();
let grid_x = ((b_size as u32) + 31) / 32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.steps_since_done_d.raw_ptr());
args.push_ptr(pos_d_ref.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
args.push_i32(pos_bytes_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_min_hold_check_fn.cu_function(),
(grid_x, 1, 1), (32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_min_hold_check: {:?}", e))?;
}
}
// Asymmetric trail decay
{
let bid_px_d = lobsim.bid_px_d();
let ask_px_d = lobsim.ask_px_d();
let mut args = RawArgs::new();
args.push_ptr(self.unit_trail_distance_d.raw_ptr());
args.push_ptr(self.unit_active_d.raw_ptr());
args.push_ptr(self.unit_entry_price_d.raw_ptr());
args.push_ptr(self.unit_initial_r_d.raw_ptr());
args.push_ptr(self.unit_lots_d.raw_ptr());
args.push_ptr(bid_px_d.raw_ptr());
args.push_ptr(ask_px_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_asymmetric_trail_decay_fn.cu_function(),
(b_size as u32, 1, 1), (4, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_asymmetric_trail_decay: {:?}", e))?;
}
}
// Trail mutate
{
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_ptr(self.unit_active_d.raw_ptr());
args.push_ptr(self.unit_trail_distance_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_trail_mutate_fn.cu_function(),
(b_size as u32, 1, 1), (4, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_trail_mutate: {:?}", e))?;
}
}
// Trail stop check
{
let bid_px_d = lobsim.bid_px_d();
let ask_px_d = lobsim.ask_px_d();
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(bid_px_d.raw_ptr());
args.push_ptr(ask_px_d.raw_ptr());
args.push_ptr(self.unit_active_d.raw_ptr());
args.push_ptr(self.unit_entry_price_d.raw_ptr());
args.push_ptr(self.unit_lots_d.raw_ptr());
args.push_ptr(self.unit_trail_distance_d.raw_ptr());
args.push_ptr(self.pyramid_units_count_d.raw_ptr());
args.push_ptr(self.close_unit_index_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_trail_stop_check_fn.cu_function(),
(b_size as u32, 1, 1), (4, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_trail_stop_check: {:?}", e))?;
}
}
// Position heat cap
{
let pos_d_ref_heat: &CudaSlice<u8> = lobsim.pos_d();
let block = (b_size as u32).min(256);
let grid = ((b_size as u32) + block - 1) / block;
let smem = 256 * std::mem::size_of::<i32>() as u32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(pos_d_ref_heat.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_i32(b_size_i);
args.push_i32(pos_bytes_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.rl_position_heat_check_fn.cu_function(),
(grid.max(1), 1, 1), (block.max(1), 1, 1), smem,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("rl_position_heat_check: {:?}", e))?;
}
}
// actions_to_market_targets
{
let (pos_d_ref, bid_px_d, ask_px_d, market_targets_d) =
lobsim.pos_book_and_market_targets_mut();
let grid_x = ((b_size as u32) + 31) / 32;
let mut args = RawArgs::new();
args.push_ptr(self.actions_d.raw_ptr());
args.push_ptr(pos_d_ref.raw_ptr());
args.push_ptr(market_targets_d.raw_ptr());
args.push_ptr(self.isv_dev_ptr);
args.push_ptr(bid_px_d.raw_ptr());
args.push_ptr(ask_px_d.raw_ptr());
args.push_ptr(self.unit_entry_price_d.raw_ptr());
args.push_ptr(self.unit_active_d.raw_ptr());
args.push_ptr(self.unit_lots_d.raw_ptr());
args.push_ptr(self.pyramid_units_count_d.raw_ptr());
args.push_ptr(self.close_unit_index_d.raw_ptr());
args.push_ptr(self.outcome_ema_d.raw_ptr());
args.push_i32(b_size_i);
args.push_i32(pos_bytes_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.actions_to_market_targets_fn.cu_function(),
(grid_x, 1, 1), (32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("actions_to_market_targets: {:?}", e))?;
}
}
if capturing_postfill {
let graph = self
.stream
.end_capture(
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
)
.context("postfill end_capture")?
.ok_or_else(|| anyhow::anyhow!("postfill end_capture returned None"))?;
self.postfill_graph = Some(graph);
}
} // 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")?;
// Activate GPU encoder path for the K-loop's step_synthetic call.
// step_synthetic checks this flag and dispatches
// forward_encoder_from_device() instead of forward_encoder().
// The SoA buffers were last filled by gather_current_window in
// step 1b, so the encoder re-forward inside step_synthetic
// produces the same h_t. Cleared after the helper returns.
self.gpu_encoder_active = true;
// Empty slice: step_synthetic ignores snapshots when
// gpu_encoder_active is set.
let empty: &[Mbp10RawInput] = &[];
let result = self.step_with_lobsim_reward_and_train(lobsim, empty, b_size);
self.gpu_encoder_active = false;
result
}
/// Launch `rl_lr_controller` to emit per-head learning rates into
/// `ISV[412..417]` via ReduceLROnPlateau-style monotone decay.
/// Per-head observed loss (l_q / l_pi / l_v from step_synthetic's

161
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -236,6 +236,24 @@ pub fn auto_horizon_weights(_seq_len: usize, _horizons: &[usize; N_HORIZONS]) ->
[1.0f32; N_HORIZONS]
}
/// Raw device pointers to the 10 SoA input buffers that `snap_feature_assemble`
/// reads from. Used by the GPU-resident data loader (`GpuDataLoader`) to write
/// directly into the perception trainer's buffers from the gather kernel,
/// bypassing the host-side AoS staging path.
#[derive(Clone, Copy, Debug)]
pub struct SoaBufferPtrs {
pub bid_px: u64,
pub bid_sz: u64,
pub ask_px: u64,
pub ask_sz: u64,
pub regime: u64,
pub prev_mid: u64,
pub trade_signed_vol: u64,
pub trade_count: u64,
pub ts_ns: u64,
pub prev_ts_ns: u64,
}
pub struct PerceptionTrainer {
cfg: PerceptionTrainerConfig,
stream: Arc<CudaStream>,
@@ -671,6 +689,13 @@ pub struct PerceptionTrainer {
/// no Adam, no reduce-axis0). The first forward_only call runs
/// eager; the second captures; the third+ replays.
forward_warmed: bool,
/// CUDA Graph for the GPU-resident data path (skips AoS scatter).
/// Separate from `forward_graph` because the kernel chain differs
/// (no `snapshot_aos_to_soa` at the front — the SoA buffers are
/// populated by `gpu_sample_and_gather` before this graph runs).
forward_graph_no_scatter: Option<CudaGraph>,
/// Warmup flag for the no-scatter forward path.
forward_no_scatter_warmed: bool,
/// cuBLAS warmup done — required before stream capture can safely
/// include cuBLAS calls.
cublas_warmed: bool,
@@ -2254,6 +2279,8 @@ impl PerceptionTrainer {
train_graph: None,
forward_graph: None,
forward_warmed: false,
forward_graph_no_scatter: None,
forward_no_scatter_warmed: false,
cublas_warmed: false,
// AoS staging — single mapped-pinned buffer for B*K Mbp10RawInput.
@@ -3821,6 +3848,114 @@ impl PerceptionTrainer {
Ok(&self.h_t_d)
}
/// GPU-resident data path: run the encoder forward when SoA buffers
/// have already been populated by an external GPU kernel (e.g.
/// `gpu_sample_and_gather`). Skips the mapped-pinned AoS staging
/// and the `snapshot_aos_to_soa` scatter kernel — the SoA buffers
/// `bid_px_all_d`, etc. already contain valid data.
///
/// Same post-condition as `forward_encoder`: `h_t_d` holds the
/// encoder representation at slot K-1 (DtoD-copied from
/// `h_new_per_k_d`).
pub fn forward_encoder_from_device(&mut self) -> Result<&CudaSlice<f32>> {
let b_sz = self.cfg.n_batch;
let k_seq = self.cfg.seq_len;
let total_snaps = b_sz * k_seq;
// Three-state machine identical to forward_only_dispatch, but
// without the AoS staging + scatter. The dispatch_forward_kernels
// method starts with launch_aos_scatter which reads from
// stg_aos_snapshots — but when using the GPU data path, the
// SoA buffers are already filled. We need a variant that skips
// the scatter.
//
// Approach: call dispatch_forward_kernels_no_scatter (a new
// helper that is dispatch_forward_kernels minus the scatter).
// However, the CUDA graph captures the scatter kernel as part
// of the forward graph. With the GPU data path, the scatter
// is a wasted no-op (reads stale stg_aos_snapshots) and gets
// overwritten by the gather kernel's output. This is harmless
// — the SoA buffers are written by gather BEFORE dispatch, and
// the scatter inside the captured graph re-writes them from
// stale stg_aos_snapshots (which is NOT what we want).
//
// Correct approach: for the GPU data path, we skip the graph
// entirely and call dispatch_forward_kernels_no_scatter each
// step. The kernel chain (snap_feature_assemble → Mamba2 → CfC
// → heads) without graph capture is ~0.2ms overhead vs graph
// replay, but eliminates the stale-scatter corruption.
//
// Alternative: invalidate the forward graph when switching to
// GPU data path. The next forward_encoder_from_device call does
// 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() {
self.forward_graph_no_scatter
.as_ref()
.unwrap()
.launch()
.context("forward_no_scatter graph launch")?;
} else if !self.forward_no_scatter_warmed {
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 {
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
let begin = self.stream.begin_capture(
CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
);
if let Err(e) = begin {
return Err(anyhow::anyhow!("forward_no_scatter begin_capture: {e}"));
}
let dispatch_result =
self.dispatch_forward_kernels_no_scatter(b_sz, k_seq, total_snaps);
let graph_result = self.stream.end_capture(
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
dispatch_result.context("forward_no_scatter dispatch (during capture)")?;
let graph = graph_result
.context("forward_no_scatter end_capture")?
.ok_or_else(|| anyhow::anyhow!(
"forward_no_scatter end_capture returned None"
))?;
self.forward_graph_no_scatter = Some(graph);
}
// DtoD copy h_new_per_k_d at slot K-1 into h_t_d (same as
// forward_encoder).
let slot_floats = b_sz * HIDDEN_DIM;
let slot_bytes = slot_floats * std::mem::size_of::<f32>();
unsafe {
let src_at_last = self.h_new_per_k_d.raw_ptr() + ((k_seq - 1) * slot_bytes) as u64;
raw_memcpy_dtod_async(self.h_t_d.raw_ptr(), src_at_last, slot_bytes, self.raw_stream)
.map_err(|e| anyhow::anyhow!("forward_encoder_from_device: h_t DtoD: {:?}", e))?;
}
Ok(&self.h_t_d)
}
/// Return raw device pointers to the SoA input buffers so an
/// external GPU kernel (e.g. `gpu_sample_and_gather`) can write
/// directly into them, bypassing the mapped-pinned AoS staging.
pub fn soa_buffer_ptrs(&self) -> SoaBufferPtrs {
SoaBufferPtrs {
bid_px: self.bid_px_all_d.raw_ptr(),
bid_sz: self.bid_sz_all_d.raw_ptr(),
ask_px: self.ask_px_all_d.raw_ptr(),
ask_sz: self.ask_sz_all_d.raw_ptr(),
regime: self.regime_all_d.raw_ptr(),
prev_mid: self.prev_mid_all_d.raw_ptr(),
trade_signed_vol: self.trade_signed_vol_all_d.raw_ptr(),
trade_count: self.trade_count_all_d.raw_ptr(),
ts_ns: self.ts_ns_all_d.raw_ptr(),
prev_ts_ns: self.prev_ts_ns_all_d.raw_ptr(),
}
}
/// Phase E.2 RL hook: borrow the `h_t_d` buffer without re-running
/// the encoder. Returns the slice last written by `forward_encoder`
/// (slot K-1 of `h_new_per_k_d`). Stable until the next
@@ -6374,6 +6509,20 @@ impl PerceptionTrainer {
/// device addresses + `self.trunk.*` weights); output is written to
/// `self.probs_per_k_d`. No host branches, no scalar-arg-changes,
/// no host-mallocs inside per pearl_no_host_branches_in_captured_graph.
/// GPU-resident data path variant: forward kernel chain WITHOUT the
/// AoS→SoA scatter (the SoA buffers were already populated by
/// `gpu_sample_and_gather`). Delegates to
/// `dispatch_forward_kernels_core` which is the shared implementation
/// per `feedback_single_source_of_truth_no_duplicates`.
fn dispatch_forward_kernels_no_scatter(
&mut self,
b_sz: usize,
k_seq: usize,
total_snaps: usize,
) -> Result<()> {
self.dispatch_forward_kernels_core(b_sz, k_seq, total_snaps)
}
fn dispatch_forward_kernels(
&mut self,
b_sz: usize,
@@ -6383,7 +6532,19 @@ impl PerceptionTrainer {
// AoS→SoA scatter: GPU kernel reads contiguous Mbp10RawInput from
// mapped-pinned buffer, writes the 10 SoA device buffers.
self.launch_aos_scatter(total_snaps)?;
self.dispatch_forward_kernels_core(b_sz, k_seq, total_snaps)
}
/// Core of the forward kernel chain: snap_feature_assemble → Mamba2 ×2
/// → LN ×2 → attn-pool → CfC K-loop → heads. Called by both
/// `dispatch_forward_kernels` (after AoS scatter) and
/// `dispatch_forward_kernels_no_scatter` (GPU data path).
fn dispatch_forward_kernels_core(
&mut self,
b_sz: usize,
k_seq: usize,
total_snaps: usize,
) -> Result<()> {
// snap_feature_assemble_batched: raw MBP-10 SoA → window_tensor [B,K,FEATURE_DIM].
let tick_size = ES_TICK_SIZE;
let n_total_i32 = total_snaps as i32;