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

Merge: TLOB Phase B — OFI_DIM 20→32, +12 real microstructure features, kernel-read gap closed

Browse Source
This commit is contained in:
jgrusewski
2026-04-23 11:37:40 +02:00
parent 8d4c2c3b03 79578bbaf6
commit 07b70ccff5
18 changed files with 312 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
crates/ml-core/src/state_layout.rs
View File

@@ -4,30 +4,52 @@
/// All systems use these constants instead of configurable state_dim fields.
///
/// Layout (grouped by update frequency, fast-changing first):
/// [0..42) Market features (OHLCV + technicals)
/// [42..62) OFI (20-dim order flow from MBP-10)
/// [62..78) MTF (4 lookbacks x 4 features)
/// [78..86) Portfolio base (8 features)
/// [86..92) Plan/ISV (6 features)
/// [92..96) Zero padding (tensor core alignment)
/// [0..42) Market features (OHLCV + technicals)
/// [42..74) OFI (32-dim order flow from MBP-10 — see OFI slot map below)
/// [74..90) MTF (4 lookbacks x 4 features)
/// [90..98) Portfolio base (8 features)
/// [98..104) Plan/ISV (6 features)
/// [104..104) Zero padding (empty — layout is natively 8-aligned)
///
/// OFI slot map (OFI_DIM = 32, indices relative to OFI_START):
/// [0..8) raw OFI — ofi_level1, ofi_level5, depth_imbalance, vpin,
/// kyle_lambda, bid_slope, ask_slope, trade_imbalance
/// [8..16) lag-1 deltas of slots [0..8)
/// [16] book_aggression (MBP-10 center-of-mass asymmetry)
/// [17] log_bar_duration (ln(bar_secs) / 10)
/// [18] ofi_acceleration (EMA of ΔOFI_L1 — MicrostructureState)
/// [19] toxicity_gradient (EMA of ΔVPIN — MicrostructureState)
/// [20] ofi_trajectory (online-linreg slope of OFI_L1 — MicrostructureState)
/// [21] realized_variance (Σ log-returns² — MicrostructureState)
/// [22] hawkes_intensity (Hawkes trade process — MicrostructureState)
/// [23] book_pressure (weighted exp(-0.3·l) 10-level — MicrostructureState)
/// [24] spread_dynamics ((max-min)/mean — MicrostructureState)
/// [25] aggression_ratio (trades_at_ask / total — MicrostructureState)
/// [26] queue_depletion_asymmetry (EMA bid/ask depletion — MicrostructureState)
/// [27] order_count_flux (ct_inc / (ct_inc + ct_dec) — MicrostructureState)
/// [28] intra_bar_momentum (sign-weighted half-bar product — MicrostructureState)
/// [29] regime_score (sigmoid of aggr·spread·thin_book — MicrostructureState)
/// [30] order_count_imbalance ((Σbid_ct − Σask_ct) / Σ(bid_ct + ask_ct) — Mbp10Snapshot)
/// [31] microprice_residual ((weighted_mid − mid) / mid — Mbp10Snapshot)
pub const STATE_DIM: usize = 96;
pub const STATE_DIM: usize = 104;
pub const STATE_DIM_PADDED: usize = 128; // pad128(STATE_DIM) for cuBLAS K-tile alignment
pub const MARKET_DIM: usize = 42;
pub const OFI_DIM: usize = 20;
pub const OFI_DIM: usize = 32;
pub const MTF_DIM: usize = 16;
pub const PORTFOLIO_BASE_DIM: usize = 8;
pub const PORTFOLIO_PLAN_DIM: usize = 6;
pub const PADDING_DIM: usize = 4;
pub const PADDING_DIM: usize = 0;
pub const MARKET_START: usize = 0;
pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42
pub const MTF_START: usize = OFI_START + OFI_DIM; // 62
pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 78
pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 86
pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 92
pub const MTF_START: usize = OFI_START + OFI_DIM; // 74
pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 90
pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 98
pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 104
const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM");
const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores");
const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS");
const _: () = assert!(STATE_DIM <= STATE_DIM_PADDED, "STATE_DIM must fit within padded width");

75
crates/ml/examples/precompute_features.rs
View File

@@ -348,8 +348,9 @@ async fn main() -> Result<()> {
.collect();
// ── Step 4: Compute OFI from MBP-10 + trades ────────────────────────────
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
let t2 = Instant::now();
let ofi: Vec<[f64; 20]> = if let Some(ref mbp10_path) = mbp10_dir {
let ofi: Vec<[f64; OFI_DIM]> = if let Some(ref mbp10_path) = mbp10_dir {
use ml::features::mbp10_loader::load_ofi_features_parallel;
use ml::features::trades_loader::load_trades_sync;
use ml::features::ofi_calculator::OFICalculator;
@@ -360,7 +361,7 @@ async fn main() -> Result<()> {
if mbp10_files.is_empty() {
info!("No MBP-10 files found, OFI will be zeros");
vec![[0.0; 20]; n]
vec![[0.0; OFI_DIM]; n]
} else {
// Load MBP-10 snapshots (parallel)
use data::providers::databento::dbn_parser::DbnParser;
@@ -463,23 +464,51 @@ async fn main() -> Result<()> {
let arr8 = f.to_array();
let micro_12 = micro_state.snapshot();
// v4 OFI layout: [0..8) raw OFI, [8..16) deltas (filled below),
// [16] book_aggression (filled below), [17] log_bar_duration (filled below),
// [18..20) ofi_acceleration + toxicity_gradient from microstructure
let mut arr20 = [0.0_f64; 20];
arr20[..8].copy_from_slice(&arr8);
// Slots [8..16) reserved for OFI deltas (computed after all bars)
// Slot [16] reserved for book_aggression (computed below)
// Slot [17] reserved for log_bar_duration (computed below)
// Keep micro_12[10..12) (ofi_acceleration, toxicity_gradient) at [18..20)
arr20[18] = micro_12[10]; // ofi_acceleration
arr20[19] = micro_12[11]; // toxicity_gradient
ofi_per_bar.push(arr20);
// v5 OFI layout (see state_layout.rs OFI slot map):
// [0..8) raw OFI (8 features)
// [8..16) lag-1 deltas (filled in post-loop)
// [16] book_aggression (filled below)
// [17] log_bar_duration (filled in post-loop)
// [18] ofi_acceleration (micro_12[10])
// [19] toxicity_gradient (micro_12[11])
// [20..30) MicrostructureState::snapshot()[0..10]
// [30] order_count_imbalance ((Σbid_ct − Σask_ct) / Σ(bid_ct+ask_ct))
// [31] microprice_residual ((weighted_mid − mid) / mid)
let mut ofi_row = [0.0_f64; OFI_DIM];
ofi_row[..8].copy_from_slice(&arr8);
ofi_row[18] = micro_12[10]; // ofi_acceleration
ofi_row[19] = micro_12[11]; // toxicity_gradient
for k in 0..10 {
ofi_row[20 + k] = micro_12[k];
}
// TLOB-novel slots: derive directly from Mbp10Snapshot counts/prices.
// order_count_imbalance: count-based analog to depth_imbalance.
let (bid_ct_sum, ask_ct_sum) = snap.levels.iter().fold((0u64, 0u64), |(b, a), l| {
(b + l.bid_ct as u64, a + l.ask_ct as u64)
});
let total_ct = bid_ct_sum + ask_ct_sum;
let order_count_imbalance = if total_ct > 0 {
(bid_ct_sum as f64 - ask_ct_sum as f64) / total_ct as f64
} else {
0.0
};
ofi_row[30] = order_count_imbalance.clamp(-1.0, 1.0);
// microprice_residual: (weighted_mid − mid) / mid; dimensionless.
// Mbp10Snapshot::weighted_mid_price() uses top-of-book volume-weighted mid.
let mid = snap.mid_price();
let wmid = snap.weighted_mid_price();
let microprice_residual = if mid > 0.0 && mid.is_finite() && wmid.is_finite() {
((wmid - mid) / mid).clamp(-0.01, 0.01)
} else {
0.0
};
ofi_row[31] = microprice_residual;
ofi_per_bar.push(ofi_row);
}
_ => ofi_per_bar.push([0.0; 20]),
_ => ofi_per_bar.push([0.0; OFI_DIM]),
}
} else {
ofi_per_bar.push([0.0; 20]);
ofi_per_bar.push([0.0; OFI_DIM]);
}
// ── Book aggression: center-of-mass asymmetry from MBP-10 depth ──
@@ -500,7 +529,9 @@ async fn main() -> Result<()> {
let buy_com = if buy_total > 0.0 { buy_weighted_sum / buy_total } else { 5.5 };
let sell_com = if sell_total > 0.0 { sell_weighted_sum / sell_total } else { 5.5 };
let book_aggression = (sell_com - buy_com) / 10.0; // normalize to [-1, 1]
ofi_per_bar.last_mut().unwrap()[16] = book_aggression;
if let Some(last) = ofi_per_bar.last_mut() {
last[16] = book_aggression;
}
}
}
@@ -528,8 +559,10 @@ async fn main() -> Result<()> {
let delta_nonzero = ofi_per_bar.iter().filter(|f| f[8..16].iter().any(|&v| v != 0.0)).count();
let book_nonzero = ofi_per_bar.iter().filter(|f| f[16] != 0.0).count();
info!("OFI v4: deltas_nonzero={}/{}, book_aggression_nonzero={}/{}",
delta_nonzero, n, book_nonzero, n);
let micro_nonzero = ofi_per_bar.iter().filter(|f| f[20..30].iter().any(|&v| v != 0.0)).count();
let tlob_nonzero = ofi_per_bar.iter().filter(|f| f[30] != 0.0 || f[31] != 0.0).count();
info!("OFI v5: deltas_nonzero={}/{}, book_aggression_nonzero={}/{}, microstructure_nonzero={}/{}, tlob_novel_nonzero={}/{}",
delta_nonzero, n, book_nonzero, n, micro_nonzero, n, tlob_nonzero, n);
// Fill targets[i][5] with MBP-10 midpoint at bar open (mid_price_open)
let mut mid_fill_count = 0_usize;
@@ -558,7 +591,7 @@ async fn main() -> Result<()> {
}
} else {
info!("No MBP-10 directory, OFI will be zeros (log_bar_duration still computed)");
let mut ofi_per_bar = vec![[0.0_f64; 20]; n];
let mut ofi_per_bar = vec![[0.0_f64; OFI_DIM]; n];
// Even without MBP-10 data, compute log_bar_duration from bar timestamps
for i in 0..n {
let duration_secs = if i > 0 {
@@ -638,7 +671,7 @@ async fn main() -> Result<()> {
println!("Bars: {}", total_len);
println!("Features: 42-dim");
println!("Targets: 6-dim");
println!("OFI: 20-dim ({})", if has_ofi { "from MBP-10" } else { "zero-padded" });
println!("OFI: {}-dim ({})", OFI_DIM, if has_ofi { "from MBP-10" } else { "zero-padded" });
println!("Format: f32 (v{}), OFI_DIM={}", ml::fxcache::FXCACHE_VERSION, ml::fxcache::OFI_DIM);
println!("Cache key: {}", hex_key);
println!("Output: {}", output_path.display());

2
crates/ml/examples/train_baseline_rl.rs
View File

@@ -599,7 +599,7 @@ fn run_training(args: &Args) -> Result<Vec<RlTrainingResult>> {
let targets: Vec<[f64; 6]> = aligned_bars.iter()
.map(|b| [b.close, b.close, b.close, b.close, b.open, b.open])
.collect();
let ofi = vec![[0.0_f64; 20]; n];
let ofi = vec![[0.0_f64; ml_core::state_layout::OFI_DIM]; n];
ml::fxcache::FxCacheData {
timestamps,

29
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -6235,11 +6235,15 @@ extern "C" __global__ void plan_noise_inject(
/* ── OFI embedding MLP input construction ─────────────────────────────────
*
* ofi_embed_build_input — Extract 18-dim OFI input from states_buf.
* Canonical layout via state_layout.cuh:
* raw OFI at state[SL_OFI_START..SL_OFI_START+8), delta at [+8..+16),
* book_aggression at [+16], log_bar_duration at [+17].
* Output: row-major [B, 18] for cuBLAS sgemm_f32 (row-major trick).
* ofi_embed_build_input — Extract the full SL_OFI_DIM OFI input from states_buf.
* Canonical layout via state_layout.cuh — see the OFI slot map in
* crates/ml-core/src/state_layout.rs for the current assignment.
* Previously this kernel read only the first 18 slots (raw 8 + delta 8 +
* book_aggression + log_duration) and silently discarded ofi_acceleration
* (slot 18) and toxicity_gradient (slot 19) even though they were persisted.
* It now reads the full SL_OFI_DIM range, feeding every persisted OFI slot
* into the embedding MLP.
* Output: row-major [B, SL_OFI_DIM] for cuBLAS sgemm_f32 (row-major trick).
* Grid: ceil(B/256), Block: 256.
*/
extern "C" __global__ void ofi_embed_build_input(
@@ -6250,16 +6254,11 @@ extern "C" __global__ void ofi_embed_build_input(
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
const float* s = states + (long long)b * state_dim;
/* Raw OFI [0..8) */
for (int k = 0; k < 8; k++)
output[b * 18 + k] = (state_dim > SL_OFI_START + k) ? s[SL_OFI_START + k] : 0.0f;
/* Delta OFI [8..16) */
for (int k = 0; k < 8; k++)
output[b * 18 + 8 + k] = (state_dim > SL_OFI_START + 8 + k) ? s[SL_OFI_START + 8 + k] : 0.0f;
/* Book aggression [16] */
output[b * 18 + 16] = (state_dim > SL_OFI_START + 16) ? s[SL_OFI_START + 16] : 0.0f;
/* Log bar duration [17] */
output[b * 18 + 17] = (state_dim > SL_OFI_START + 17) ? s[SL_OFI_START + 17] : 0.0f;
/* Full OFI [0..SL_OFI_DIM) */
for (int k = 0; k < SL_OFI_DIM; k++) {
output[b * SL_OFI_DIM + k] =
(state_dim > SL_OFI_START + k) ? s[SL_OFI_START + k] : 0.0f;
}
}
/* ================================================================== */

129
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -82,6 +82,13 @@ static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR")
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension
const OFI_EMBED_DIM: usize = 10; // OFI embedding width appended to history rows
/// OFI embed MLP input width — mirrors `ml_core::state_layout::OFI_DIM` so the
/// MLP consumes every OFI slot persisted to fxcache (no silent [OFI_DIM..) gap).
const OFI_EMBED_IN: usize = ml_core::state_layout::OFI_DIM;
/// OFI embed MLP W tensor element count: [OFI_EMBED_DIM, OFI_EMBED_IN] row-major.
const OFI_EMBED_W_COUNT: usize = OFI_EMBED_DIM * OFI_EMBED_IN;
/// OFI embed MLP total parameter count: W[OFI_EMBED_W_COUNT] + b[OFI_EMBED_DIM].
const OFI_EMBED_TOTAL_PARAMS: usize = OFI_EMBED_W_COUNT + OFI_EMBED_DIM;
/// Introspective State Vector (ISV) configuration.
const ISV_K: usize = 4; // Temporal ISV history length
@@ -1859,26 +1866,26 @@ pub struct GpuDqnTrainer {
trade_plan_activate_kernel: CudaFunction, // bias + sigmoid/scaling activations
plan_noise_kernel: CudaFunction, // plan_noise_inject: ±5% temporal diversity
// ── OFI embed MLP (18→10) ──
ofi_embed_input_buf: CudaSlice<f32>, // [B, 18] scratch for input extraction
ofi_embed_output_buf: CudaSlice<f32>, // [B, 10] embedding output
ofi_embed_w: CudaSlice<f32>, // [10, 18] weight matrix (row-major)
ofi_embed_b: CudaSlice<f32>, // [10] bias
// ── OFI embed MLP (OFI_EMBED_IN→OFI_EMBED_DIM) ──
ofi_embed_input_buf: CudaSlice<f32>, // [B, OFI_EMBED_IN] scratch for input extraction
ofi_embed_output_buf: CudaSlice<f32>, // [B, OFI_EMBED_DIM] embedding output
ofi_embed_w: CudaSlice<f32>, // [OFI_EMBED_DIM, OFI_EMBED_IN] weight matrix (row-major)
ofi_embed_b: CudaSlice<f32>, // [OFI_EMBED_DIM] bias
ofi_embed_build_input_kernel: CudaFunction,
// ── OFI embed MLP backward + Adam ──
/// Contiguous param buffer [190] = W[180] + b[10] for Adam update.
/// Contiguous param buffer [OFI_EMBED_TOTAL_PARAMS] = W[OFI_EMBED_W_COUNT] + b[OFI_EMBED_DIM] for Adam update.
ofi_embed_params: CudaSlice<f32>,
/// Weight gradient [180] (10×18).
/// Weight gradient [OFI_EMBED_W_COUNT] (OFI_EMBED_DIM × OFI_EMBED_IN).
ofi_embed_grad_w: CudaSlice<f32>,
/// Bias gradient [10].
/// Bias gradient [OFI_EMBED_DIM].
ofi_embed_grad_b: CudaSlice<f32>,
/// Contiguous gradient buffer [190] for Adam (assembled from grad_w + grad_b).
/// Contiguous gradient buffer [OFI_EMBED_TOTAL_PARAMS] for Adam (assembled from grad_w + grad_b).
ofi_embed_grad: CudaSlice<f32>,
/// Combined gradient from Mamba2 + attention [OFI_EMBED_DIM * B].
d_ofi_embed_combined: CudaSlice<f32>,
/// Adam first moment [190].
/// Adam first moment [OFI_EMBED_TOTAL_PARAMS].
ofi_embed_adam_m: CudaSlice<f32>,
/// Adam second moment [190].
/// Adam second moment [OFI_EMBED_TOTAL_PARAMS].
ofi_embed_adam_v: CudaSlice<f32>,
/// Adam step counter.
ofi_embed_adam_step: i32,
@@ -1891,7 +1898,7 @@ pub struct GpuDqnTrainer {
/// Pinned device-mapped step counter for Adam.
ofi_embed_t_pinned: *mut i32,
ofi_embed_t_dev_ptr: u64,
/// Bias grad reduce partial sums: [bias_num_blocks, 10].
/// Bias grad reduce partial sums: [bias_num_blocks, OFI_EMBED_DIM].
ofi_embed_bias_grad_partials: CudaSlice<f32>,
// ── Speculative inference cache ──
@@ -3604,14 +3611,17 @@ impl GpuDqnTrainer {
Ok(())
}
/// OFI embedding MLP forward: extract 18-dim OFI features from states, project to 10-dim.
/// Input: states_buf [B, state_dim] → extract [raw_ofi(8); delta_ofi(8); book_aggression; log_duration]
/// Output: ofi_embed_output_buf [B, 10] (bias + ReLU activated).
/// OFI embedding MLP forward: extract OFI_EMBED_IN OFI features from
/// states, project to OFI_EMBED_DIM.
///
/// Input: states_buf [B, state_dim] → extracts the full OFI slice at
/// [SL_OFI_START..SL_OFI_START+OFI_EMBED_IN).
/// Output: ofi_embed_output_buf [B, OFI_EMBED_DIM] (bias + ReLU activated).
pub(crate) fn launch_ofi_embed_forward(&self, batch_size: usize) -> Result<(), MLError> {
let sd = ml_core::state_layout::STATE_DIM as i32;
let b_i32 = batch_size as i32;
// Step 1: Build input [B, 18] from states_buf
// Step 1: Build input [B, OFI_EMBED_IN] from states_buf
let states_ptr = self.states_buf.raw_ptr();
let input_ptr = self.ofi_embed_input_buf.raw_ptr();
let blocks = ((batch_size as u32 + 255) / 256).max(1);
@@ -3629,19 +3639,20 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("ofi_embed_build_input: {e}")))?;
}
// Step 2: cuBLAS SGEMM: output[B, 10] = input[B, 18] @ W^T[18, 10]
// sgemm_f32 uses row-major trick: W[10, 18] row-major, input [B, 18], output [B, 10]
// Step 2: cuBLAS SGEMM: output[B, OFI_EMBED_DIM] = input[B, OFI_EMBED_IN] @ W^T[OFI_EMBED_IN, OFI_EMBED_DIM]
// sgemm_f32 uses row-major trick: W[OFI_EMBED_DIM, OFI_EMBED_IN] row-major,
// input [B, OFI_EMBED_IN], output [B, OFI_EMBED_DIM]
let w_ptr = self.ofi_embed_w.raw_ptr();
let output_ptr = self.ofi_embed_output_buf.raw_ptr();
self.cublas_forward.sgemm_f32(
&self.stream, w_ptr, input_ptr, output_ptr,
10, batch_size, 18, "ofi_embed_fwd",
OFI_EMBED_DIM, batch_size, OFI_EMBED_IN, "ofi_embed_fwd",
)?;
// Step 3: Add bias + ReLU in-place on output[B, 10]
// Step 3: Add bias + ReLU in-place on output[B, OFI_EMBED_DIM]
let bias_ptr = self.ofi_embed_b.raw_ptr();
self.cublas_forward.launch_add_bias_relu_f32_raw(
&self.stream, output_ptr, bias_ptr, 10, batch_size,
&self.stream, output_ptr, bias_ptr, OFI_EMBED_DIM, batch_size,
)?;
Ok(())
@@ -3650,15 +3661,15 @@ impl GpuDqnTrainer {
/// OFI embed MLP backward pass.
///
/// Accumulates gradients from two sources:
/// - `d_ofi_embed_mamba2 [10, B]` — from Mamba2 backward (self.d_ofi_embed_mamba2)
/// - `d_ofi_embed_attn [10, B]` — from attention backward (passed as raw pointer)
/// - `d_ofi_embed_mamba2 [OFI_EMBED_DIM, B]` — from Mamba2 backward (self.d_ofi_embed_mamba2)
/// - `d_ofi_embed_attn [OFI_EMBED_DIM, B]` — from attention backward (passed as raw pointer)
///
/// Pipeline:
/// 1. d_combined = d_mamba2 + d_attn (graph-safe copy + SAXPY via add kernel)
/// 2. ReLU mask: d_combined *= (ofi_embed_output > 0)
/// 3. dW[10, 18] = d_combined[10, B] @ ofi_input[18, B]^T (cuBLAS backward)
/// 4. db[10] = sum_B(d_combined) (2-phase bias grad reduce)
/// 5. Assemble contiguous grad[190] = [grad_w[180]; grad_b[10]]
/// 3. dW[OFI_EMBED_DIM, OFI_EMBED_IN] = d_combined[OFI_EMBED_DIM, B] @ ofi_input[OFI_EMBED_IN, B]^T (cuBLAS backward)
/// 4. db[OFI_EMBED_DIM] = sum_B(d_combined) (2-phase bias grad reduce)
/// 5. Assemble contiguous grad[OFI_EMBED_TOTAL_PARAMS] = [grad_w[OFI_EMBED_W_COUNT]; grad_b[OFI_EMBED_DIM]]
/// 6. Copy params W+b back to ofi_embed_w/ofi_embed_b (keep in sync)
pub(crate) fn launch_ofi_embed_backward(&mut self, batch_size: usize, d_ofi_attn_ptr: Option<u64>) -> Result<(), MLError> {
let b = batch_size;
@@ -3706,13 +3717,13 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("ofi_embed relu_mask: {e}")))?;
}
// Step 3: dW[10, 18] = d_combined[10, B] @ ofi_input[18, B]^T
// Step 3: dW[OFI_EMBED_DIM, OFI_EMBED_IN] = d_combined[OFI_EMBED_DIM, B] @ ofi_input[OFI_EMBED_IN, B]^T
//
// Forward used row-major W[10,18], input[B,18], output[B,10].
// Forward used row-major W[OFI_EMBED_DIM, OFI_EMBED_IN], input[B, OFI_EMBED_IN], output[B, OFI_EMBED_DIM].
// In col-major cuBLAS terms:
// d_combined is col-major [10, B], ofi_input is col-major [18, B]
// dW_col[10, 18] = d_combined[10, B] @ ofi_input[18, B]^T
// → TRANSA=N, TRANSB=T, m=10, n=18, k=B
// d_combined is col-major [OFI_EMBED_DIM, B], ofi_input is col-major [OFI_EMBED_IN, B]
// dW_col[OFI_EMBED_DIM, OFI_EMBED_IN] = d_combined @ ofi_input^T
// → TRANSA=N, TRANSB=T, m=OFI_EMBED_DIM, n=OFI_EMBED_IN, k=B
let grad_w_ptr = self.ofi_embed_grad_w.raw_ptr();
let input_ptr = self.ofi_embed_input_buf.raw_ptr();
{
@@ -3727,8 +3738,8 @@ impl GpuDqnTrainer {
let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;
let op_n: i32 = 0; // CUBLAS_OP_N
let op_t: i32 = 1; // CUBLAS_OP_T
let m = OFI_EMBED_DIM; // 10
let n = 18_usize;
let m = OFI_EMBED_DIM;
let n = OFI_EMBED_IN;
let k = b;
unsafe {
@@ -3747,15 +3758,15 @@ impl GpuDqnTrainer {
std::mem::size_of::<i32>(),
).map_err(|e| MLError::ModelError(format!("ofi_embed dW TRANSB: {e:?}")))?;
// A = d_combined [10, B] col-major, lda=10
// A = d_combined [OFI_EMBED_DIM, B] col-major, lda=OFI_EMBED_DIM
let a_layout = cublaslt_result::create_matrix_layout(
f32_type, m as u64, k as u64, m as i64,
).map_err(|e| MLError::ModelError(format!("ofi_embed dW A layout: {e:?}")))?;
// B = ofi_input [18, B] col-major (transposed), physical [18, B], ldb=18
// B = ofi_input [OFI_EMBED_IN, B] col-major (transposed), physical [OFI_EMBED_IN, B], ldb=OFI_EMBED_IN
let b_layout = cublaslt_result::create_matrix_layout(
f32_type, n as u64, k as u64, n as i64,
).map_err(|e| MLError::ModelError(format!("ofi_embed dW B layout: {e:?}")))?;
// C = dW [10, 18] col-major, ldc=10
// C = dW [OFI_EMBED_DIM, OFI_EMBED_IN] col-major, ldc=OFI_EMBED_DIM
let c_layout = cublaslt_result::create_matrix_layout(
f32_type, m as u64, n as u64, m as i64,
).map_err(|e| MLError::ModelError(format!("ofi_embed dW C layout: {e:?}")))?;
@@ -3787,7 +3798,7 @@ impl GpuDqnTrainer {
}
}
// Step 4: db[10] = sum_B(d_combined) — 2-phase bias grad reduce
// Step 4: db[OFI_EMBED_DIM] = sum_B(d_combined) — 2-phase bias grad reduce
let bias_num_blocks = ((b + 255) / 256) as i32;
let partials_ptr = self.ofi_embed_bias_grad_partials.raw_ptr();
let grad_b_ptr = self.ofi_embed_grad_b.raw_ptr();
@@ -3819,19 +3830,21 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p2: {e}")))?;
}
// Step 5: Assemble contiguous grad[190] = [grad_w[180]; grad_b[10]]
// Step 5: Assemble contiguous grad[OFI_EMBED_TOTAL_PARAMS]
// = [grad_w[OFI_EMBED_W_COUNT]; grad_b[OFI_EMBED_DIM]]
let grad_ptr = self.ofi_embed_grad.raw_ptr();
self.graph_safe_copy_f32(grad_ptr, grad_w_ptr, 180 * 4, "ofi_embed_grad_w→grad")?;
self.graph_safe_copy_f32(grad_ptr + (180 * 4) as u64, grad_b_ptr, 10 * 4, "ofi_embed_grad_b→grad")?;
let w_bytes = OFI_EMBED_W_COUNT * std::mem::size_of::<f32>();
let b_bytes = OFI_EMBED_DIM * std::mem::size_of::<f32>();
self.graph_safe_copy_f32(grad_ptr, grad_w_ptr, w_bytes, "ofi_embed_grad_w→grad")?;
self.graph_safe_copy_f32(grad_ptr + w_bytes as u64, grad_b_ptr, b_bytes, "ofi_embed_grad_b→grad")?;
Ok(())
}
/// Adam optimizer step for ofi_embed_params [190]. Same pattern as step_denoise_adam.
/// Adam optimizer step for ofi_embed_params [OFI_EMBED_TOTAL_PARAMS]. Same pattern as step_denoise_adam.
pub(crate) fn step_ofi_embed_adam(&mut self) -> Result<(), MLError> {
self.ofi_embed_adam_step += 1;
let step_val = self.ofi_embed_adam_step;
const OFI_EMBED_TOTAL_PARAMS: usize = 190; // 10*18 + 10
let n = OFI_EMBED_TOTAL_PARAMS as i32;
// Write step counter to pinned device-mapped memory
@@ -3917,8 +3930,10 @@ impl GpuDqnTrainer {
// so forward pass uses the latest weights.
let w_dst = self.ofi_embed_w.raw_ptr();
let b_dst = self.ofi_embed_b.raw_ptr();
self.graph_safe_copy_f32(w_dst, params_ptr, 180 * 4, "ofi_embed_params→w")?;
self.graph_safe_copy_f32(b_dst, params_ptr + (180 * 4) as u64, 10 * 4, "ofi_embed_params→b")?;
let w_bytes = OFI_EMBED_W_COUNT * std::mem::size_of::<f32>();
let b_bytes = OFI_EMBED_DIM * std::mem::size_of::<f32>();
self.graph_safe_copy_f32(w_dst, params_ptr, w_bytes, "ofi_embed_params→w")?;
self.graph_safe_copy_f32(b_dst, params_ptr + w_bytes as u64, b_bytes, "ofi_embed_params→b")?;
Ok(())
}
@@ -8339,33 +8354,35 @@ impl GpuDqnTrainer {
let trade_plan_hidden_buf = alloc_f32(&stream, b * config.adv_h, "trade_plan_hidden")?;
let trade_plan_pre_out_buf = alloc_f32(&stream, b * 6, "trade_plan_pre_out")?;
// ── OFI embed MLP buffers (18→10) ──
let ofi_embed_input_buf = alloc_f32(&stream, b * 18, "ofi_embed_input")?;
let ofi_embed_output_buf = alloc_f32(&stream, b * 10, "ofi_embed_output")?;
let ofi_embed_w = alloc_f32(&stream, 10 * 18, "ofi_embed_w")?;
// ── OFI embed MLP buffers (OFI_EMBED_IN → OFI_EMBED_DIM) ──
let ofi_embed_input_buf = alloc_f32(&stream, b * OFI_EMBED_IN, "ofi_embed_input")?;
let ofi_embed_output_buf = alloc_f32(&stream, b * OFI_EMBED_DIM, "ofi_embed_output")?;
let ofi_embed_w = alloc_f32(&stream, OFI_EMBED_W_COUNT, "ofi_embed_w")?;
// Zero-init OFI embed weights so Mamba2/attention start with clean h_s2.
// The embed produces zeros at init → h_history = [h_s2; zeros] → no corruption.
// Weights learn from zero as gradients flow back from Mamba2/attention.
// (Previous Xavier init with scale ~0.27 injected noise that corrupted trunk features.)
let ofi_embed_b = alloc_f32(&stream, 10, "ofi_embed_b")?; // zero-init bias
info!("GpuDqnTrainer: OFI embed MLP buffers allocated (18→10, {} params)", 10 * 18 + 10);
let ofi_embed_b = alloc_f32(&stream, OFI_EMBED_DIM, "ofi_embed_b")?; // zero-init bias
info!(
"GpuDqnTrainer: OFI embed MLP buffers allocated ({}→{}, {} params)",
OFI_EMBED_IN, OFI_EMBED_DIM, OFI_EMBED_TOTAL_PARAMS,
);
// ── OFI embed MLP backward + Adam buffers ──
const OFI_EMBED_TOTAL_PARAMS: usize = 190; // 10*18 + 10
// Contiguous params buffer: copy W then b for Adam
let ofi_embed_params = alloc_f32(&stream, OFI_EMBED_TOTAL_PARAMS, "ofi_embed_params")?;
{
// Copy W[180] from ofi_embed_w into params[0..180]
// Copy W[OFI_EMBED_W_COUNT] from ofi_embed_w into params[0..OFI_EMBED_W_COUNT]
let w_src = ofi_embed_w.raw_ptr();
let p_dst = ofi_embed_params.raw_ptr();
let w_bytes = 180 * std::mem::size_of::<f32>();
let w_bytes = OFI_EMBED_W_COUNT * std::mem::size_of::<f32>();
unsafe {
cudarc::driver::sys::cuMemcpyDtoDAsync_v2(p_dst, w_src, w_bytes, stream.cu_stream());
}
// b[10] is zero-init, params[180..190] already zeroed by alloc_zeros
// b[OFI_EMBED_DIM] is zero-init, tail of params already zeroed by alloc_zeros
}
let ofi_embed_grad_w = alloc_f32(&stream, 180, "ofi_embed_grad_w")?;
let ofi_embed_grad_b = alloc_f32(&stream, 10, "ofi_embed_grad_b")?;
let ofi_embed_grad_w = alloc_f32(&stream, OFI_EMBED_W_COUNT, "ofi_embed_grad_w")?;
let ofi_embed_grad_b = alloc_f32(&stream, OFI_EMBED_DIM, "ofi_embed_grad_b")?;
let ofi_embed_grad = alloc_f32(&stream, OFI_EMBED_TOTAL_PARAMS, "ofi_embed_grad")?;
let d_ofi_embed_combined = alloc_f32(&stream, OFI_EMBED_DIM * b, "d_ofi_embed_combined")?;
let ofi_embed_adam_m = alloc_f32(&stream, OFI_EMBED_TOTAL_PARAMS, "ofi_embed_adam_m")?;

2
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -837,7 +837,7 @@ impl GpuExperienceCollector {
let (market_dim, num_atoms_max) = kernel_dims;
let state_dim = ml_core::state_layout::STATE_DIM;
let _portfolio_dim: usize = 12; // 8 base + 4 plan progress features
let ofi_dim: usize = 20;
let ofi_dim: usize = ml_core::state_layout::OFI_DIM;
// Branch sizes for 4-branch hierarchical DQN (always enabled).
// [direction(4), magnitude(3), order(3), urgency(3)]

13
crates/ml/src/cuda_pipeline/gpu_walk_forward.rs
View File

@@ -547,11 +547,10 @@ impl GpuWalkForwardData {
/// On H100 (80 GB), a typical dataset of 1M bars uses:
/// features: 1M × 42 × 4 = 168 MB
/// targets: 1M × 4 × 4 = 16 MB
/// ofi: 1M × 20 × 4 = 80 MB
/// total: 216 MB (~0.27% of VRAM)
/// ofi: 1M × OFI_DIM × 4 (~128 MB at OFI_DIM=32)
pub fn upload(
training_data: &[(FeatureVector, Vec<f64>)],
ofi_data: Option<&[[f64; 20]]>,
ofi_data: Option<&[[f64; crate::fxcache::OFI_DIM]]>,
wf_config: &GpuWalkForwardConfig,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
@@ -591,19 +590,21 @@ impl GpuWalkForwardData {
// Upload OFI (optional)
let ofi_gpu = if let Some(ofi) = ofi_data {
let ofi_dim: usize = crate::fxcache::OFI_DIM;
let n = ofi.len().min(total_bars);
let mut flat_ofi = Vec::with_capacity(total_bars * 20);
let mut flat_ofi = Vec::with_capacity(total_bars * ofi_dim);
let zero_row = vec![0.0_f32; ofi_dim];
for i in 0..total_bars {
if i < n {
for &v in ofi[i].iter() {
flat_ofi.push(v as f32);
}
} else {
flat_ofi.extend_from_slice(&[0.0_f32; 20]);
flat_ofi.extend_from_slice(&zero_row);
}
}
let buf = super::clone_htod_f32(stream, &flat_ofi)?;
vram += total_bars * 20 * 4;
vram += total_bars * ofi_dim * 4;
Some(buf)
} else {
None

33
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -249,7 +249,7 @@ impl DqnGpuData {
pub fn upload_slices(
features: &[[f64; 42]],
targets: &[[f64; 6]],
ofi: &[[f64; 20]],
ofi: &[[f64; crate::fxcache::OFI_DIM]],
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
let num_bars = features.len();
@@ -291,7 +291,7 @@ impl DqnGpuData {
ofi.len(), num_bars,
)));
}
let mut flat_ofi = Vec::with_capacity(num_bars * 20);
let mut flat_ofi = Vec::with_capacity(num_bars * crate::fxcache::OFI_DIM);
for row in ofi {
for &v in row.iter() {
flat_ofi.push(v as f32);
@@ -309,13 +309,14 @@ impl DqnGpuData {
})
}
/// Upload OFI features (20 dims per bar) from MBP-10 order book data.
/// Upload OFI features (`OFI_DIM` dims per bar) from MBP-10 order book data.
///
/// Must be called after `upload()` with a slice matching `num_bars` length.
/// Features are: OFI, VPIN, Kyle's Lambda, trade imbalance, etc.
/// Features are: OFI, VPIN, Kyle's Lambda, trade imbalance, microstructure
/// state, etc. — see `ml_core::state_layout` for the full slot map.
pub fn upload_ofi(
&mut self,
ofi_data: &[[f64; 20]],
ofi_data: &[[f64; crate::fxcache::OFI_DIM]],
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
if ofi_data.len() != self.num_bars {
@@ -325,7 +326,7 @@ impl DqnGpuData {
)));
}
let mut flat = Vec::with_capacity(self.num_bars * 20);
let mut flat = Vec::with_capacity(self.num_bars * crate::fxcache::OFI_DIM);
for row in ofi_data {
for &v in row.iter() {
flat.push(v as f32);
@@ -348,7 +349,7 @@ impl DqnGpuData {
/// Estimated VRAM usage in bytes.
pub fn vram_bytes(&self) -> usize {
let ofi_dim: usize = 20;
let ofi_dim: usize = crate::fxcache::OFI_DIM;
estimate_vram_bytes(self.num_bars * (self.feature_dim + 6 + ofi_dim))
}
@@ -423,7 +424,7 @@ impl DqnGpuData {
portfolio_features: &[f32; 3],
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
let ofi_dim: usize = 20;
let ofi_dim: usize = crate::fxcache::OFI_DIM;
let raw_dim = self.feature_dim + 3 + ofi_dim;
let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);
@@ -438,12 +439,12 @@ impl DqnGpuData {
// HtoD copy: portfolio features -> dst[feature_dim..feature_dim+3] (3 scalars, tiny)
Self::htod_copy_into(portfolio_features, &mut dst, self.feature_dim, stream)?;
// DtoD copy: OFI features -> dst[feature_dim+3..feature_dim+3+20]
// DtoD copy: OFI features -> dst[feature_dim+3..feature_dim+3+ofi_dim]
let ofi = self.ofi_features.as_ref().ok_or_else(|| {
MLError::ModelError("OFI features missing on GPU — model requires order flow data".into())
})?;
let ofi_gpu = Self::d2d_subrange(ofi, bar_idx * 20, 20, stream)?;
Self::dtod_copy_into(&ofi_gpu, &mut dst, self.feature_dim + 3, 20, stream)?;
let ofi_gpu = Self::d2d_subrange(ofi, bar_idx * ofi_dim, ofi_dim, stream)?;
Self::dtod_copy_into(&ofi_gpu, &mut dst, self.feature_dim + 3, ofi_dim, stream)?;
Ok(dst)
}
@@ -463,7 +464,7 @@ impl DqnGpuData {
if count == 0 {
return Err(MLError::ModelError("Empty batch for state construction".to_owned()));
}
let ofi_dim: usize = 20;
let ofi_dim: usize = crate::fxcache::OFI_DIM;
let raw_dim = self.feature_dim + 3 + ofi_dim;
let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);
@@ -477,11 +478,11 @@ impl DqnGpuData {
// Get contiguous market features [count * feature_dim]
let market_gpu = self.batch_features(start, count, stream)?;
// Get contiguous OFI features [count * 20]
// Get contiguous OFI features [count * ofi_dim]
let ofi_src = self.ofi_features.as_ref().ok_or_else(|| {
MLError::ModelError("OFI features missing on GPU — model requires order flow data".into())
})?;
let ofi_gpu = Self::d2d_subrange(ofi_src, start * 20, count * 20, stream)?;
let ofi_gpu = Self::d2d_subrange(ofi_src, start * ofi_dim, count * ofi_dim, stream)?;
// For each row, DtoD copy: market, portfolio, OFI into the interleaved output
for i in 0..count {
@@ -499,9 +500,9 @@ impl DqnGpuData {
// OFI features: contiguous in ofi_gpu
Self::dtod_copy_into_at_offset(
&ofi_gpu, i * 20,
&ofi_gpu, i * ofi_dim,
&mut dst, row_offset + self.feature_dim + 3,
20, stream,
ofi_dim, stream,
)?;
}

26
crates/ml/src/cuda_pipeline/state_layout.cuh
View File

@@ -5,13 +5,13 @@
#pragma once
// ── Dimensions ──
#define SL_STATE_DIM 96
#define SL_MARKET_DIM 42
#define SL_OFI_DIM 20
#define SL_MTF_DIM 16
#define SL_PORTFOLIO_BASE_DIM 8
#define SL_PORTFOLIO_PLAN_DIM 6
#define SL_PADDING_DIM 4
#define SL_STATE_DIM 104
#define SL_MARKET_DIM 42
#define SL_OFI_DIM 32
#define SL_MTF_DIM 16
#define SL_PORTFOLIO_BASE_DIM 8
#define SL_PORTFOLIO_PLAN_DIM 6
#define SL_PADDING_DIM 0
// ── Offsets ──
#define SL_MARKET_START 0
@@ -29,6 +29,8 @@ static_assert(SL_PADDING_START + SL_PADDING_DIM == SL_STATE_DIM,
"State layout dimensions must sum to SL_STATE_DIM");
static_assert(SL_STATE_DIM % 8 == 0,
"SL_STATE_DIM must be 8-aligned for tensor core cuBLAS");
static_assert(SL_STATE_DIM <= SL_STATE_DIM_PADDED,
"SL_STATE_DIM must fit within SL_STATE_DIM_PADDED");
// ── Shared state assembly function ──
// Called by both experience_state_gather (training) and backtest_state_gather (validation).
@@ -49,23 +51,23 @@ __device__ __forceinline__ void assemble_state(
for (int k = 0; k < SL_MARKET_DIM; k++)
out[SL_MARKET_START + k] = market[k];
// OFI [42..62)
// OFI [42..74)
for (int k = 0; k < SL_OFI_DIM; k++)
out[SL_OFI_START + k] = ofi[k];
// MTF [62..78)
// MTF [74..90)
for (int k = 0; k < SL_MTF_DIM; k++)
out[SL_MTF_START + k] = mtf[k];
// Portfolio base [78..86)
// Portfolio base [90..98)
for (int k = 0; k < SL_PORTFOLIO_BASE_DIM; k++)
out[SL_PORTFOLIO_START + k] = portfolio[k];
// Plan/ISV [86..92)
// Plan/ISV [98..104)
for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++)
out[SL_PLAN_ISV_START + k] = plan_isv[k];
// Padding [92..96) — zero
// Padding [104..104) — empty (layout is natively 8-aligned)
for (int k = 0; k < SL_PADDING_DIM; k++)
out[SL_PADDING_START + k] = 0.0f;
}

51
crates/ml/src/fxcache.rs
View File

@@ -10,17 +10,17 @@
//! ┌───────────────────────────────────────────────────┐
//! │ FxCacheHeader (64 bytes) │
//! │ magic [u8; 8] = b"FXCACHE\0" │
//! │ version u16 = 4 (f32) │
//! │ version u16 = 5 (f32) │
//! │ feat_dim u16 = 42 │
//! │ target_dim u16 = 6 │
//! │ ofi_dim u16 = 20 │
//! │ ofi_dim u16 = 32 │
//! │ bar_count u64 │
//! │ cache_key [u8; 32] (SHA256 raw bytes) │
//! │ reserved [u8; 8] │
//! └───────────────────────────────────────────────────┘
//! │ Body (bar_count records) │
//! │ Each record starts with an i64 timestamp (ns). │
//! │ Version 4: [i64 ts][68 × f32] = 280 bytes/bar │
//! │ Version 5: [i64 ts][80 × f32] = 328 bytes/bar │
//! └───────────────────────────────────────────────────┘
//! ```
@@ -40,7 +40,12 @@ const HEADER_SIZE: usize = 64;
/// Cache format version. Bump on ANY format change (OFI_DIM, FEAT_DIM, etc.).
/// Stale cache files with wrong version are auto-detected and rejected by validate().
/// The ensure-fxcache Argo step catches the error and regenerates.
pub const FXCACHE_VERSION: u16 = 4; // v4: precomputed OFI deltas + book_aggression + log_duration in ofi[8..18)
/// v5: OFI_DIM 20→32 — adds 10 MicrostructureState slots (ofi_trajectory,
/// realized_variance, hawkes_intensity, book_pressure, spread_dynamics,
/// aggression_ratio, queue_depletion_asymmetry, order_count_flux,
/// intra_bar_momentum, regime_score) + 2 TLOB-novel slots
/// (order_count_imbalance, microprice_residual).
pub const FXCACHE_VERSION: u16 = 5;
/// Feature vector dimensionality.
const FEAT_DIM: usize = 42;
@@ -48,8 +53,9 @@ const FEAT_DIM: usize = 42;
/// Target vector dimensionality (close, next_close, raw_close, raw_next, raw_open, mid_price_open).
const TARGET_DIM: usize = 6;
/// OFI vector dimensionality (8 base + 12 extended microstructure features).
pub const OFI_DIM: usize = 20;
/// OFI vector dimensionality — mirrors `ml_core::state_layout::OFI_DIM`
/// (20 legacy slots + 12 new microstructure slots).
pub const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
/// Total f64 values per record: features + targets + OFI = 42 + 6 + 20 = 68.
const RECORD_F64_COUNT: usize = FEAT_DIM + TARGET_DIM + OFI_DIM;
@@ -64,13 +70,13 @@ const RECORD_F32_COUNT: usize = RECORD_F64_COUNT;
pub struct FxCacheHeader {
/// Magic bytes: `b"FXCACHE\0"`.
pub magic: [u8; 8],
/// Format version: 1 = f32.
/// Format version — see `FXCACHE_VERSION`.
pub version: u16,
/// Feature dimension (42).
pub feat_dim: u16,
/// Target dimension (6; legacy v2 files have 4).
/// Target dimension (6).
pub target_dim: u16,
/// OFI dimension (20).
/// OFI dimension (see `OFI_DIM`).
pub ofi_dim: u16,
/// Number of bars (records) in the file.
pub bar_count: u64,
@@ -99,8 +105,9 @@ impl FxCacheHeader {
/// Validate header integrity.
///
/// Accepts both v2 (target_dim=4) and v3 (target_dim=6) files.
/// v2 files are read with legacy padding (tgt[4]=raw_close, tgt[5]=raw_close).
/// Strictly enforces the current `FXCACHE_VERSION`. Any older version
/// (including v4 with OFI_DIM=20) is rejected — regenerate via
/// `precompute_features`.
pub fn validate(&self) -> Result<()> {
if self.magic != FXCACHE_MAGIC {
bail!(
@@ -200,7 +207,7 @@ pub struct FxCacheData {
pub features: Vec<[f64; FEAT_DIM]>,
/// Target vectors, one per bar (6 elements each).
pub targets: Vec<[f64; TARGET_DIM]>,
/// OFI vectors, one per bar (20 elements each).
/// OFI vectors, one per bar (`OFI_DIM` elements each).
pub ofi: Vec<[f64; OFI_DIM]>,
/// SHA256 cache key (raw 32 bytes).
pub cache_key: [u8; 32],
@@ -220,7 +227,7 @@ pub struct FxCacheData {
/// * `path` — Output file path (parent directories are created automatically)
/// * `features` — Slice of 42-element feature vectors
/// * `targets` — Slice of 6-element target vectors
/// * `ofi` — Slice of 20-element OFI vectors
/// * `ofi` — Slice of `OFI_DIM`-element OFI vectors
/// * `timestamps` — Per-bar timestamps (nanoseconds since Unix epoch)
/// * `cache_key` — SHA256 key (raw 32 bytes)
/// * `has_ofi` — If true, OFI was computed from real MBP-10 data; false means zero-filled
@@ -277,8 +284,10 @@ pub fn write_fxcache(
let total_bytes = HEADER_SIZE as u64 + body_bytes;
info!(
"FxCache written: {} bars, v1 (f32), {:.2} MB -> {:?}",
"FxCache written: {} bars, v{} (f32, OFI_DIM={}), {:.2} MB -> {:?}",
bar_count,
FXCACHE_VERSION,
OFI_DIM,
total_bytes as f64 / 1_048_576.0,
path
);
@@ -286,7 +295,7 @@ pub fn write_fxcache(
Ok(total_bytes)
}
/// Write body in f32 format: [i64 ts] + (42+6+20) × f32 = 280 bytes per bar.
/// Write body in f32 format: [i64 ts] + (FEAT_DIM+TARGET_DIM+OFI_DIM) × f32 per bar.
fn write_body_f32(
writer: &mut BufWriter<std::fs::File>,
features: &[[f64; FEAT_DIM]],
@@ -346,7 +355,7 @@ pub fn load_fxcache(path: &Path) -> Result<FxCacheData> {
let bar_count = header.bar_count as usize;
// Sanity-check file size (v4 only — legacy versions rejected by validate())
// Sanity-check file size — current version only; legacy versions rejected by validate()
let on_disk_record_f32 = FEAT_DIM + TARGET_DIM + OFI_DIM;
let expected_body = bar_count as u64 * (8 + on_disk_record_f32 as u64 * 4);
let expected_total = HEADER_SIZE as u64 + expected_body;
@@ -378,7 +387,7 @@ pub fn load_fxcache(path: &Path) -> Result<FxCacheData> {
})
}
/// Read body in f32 format (v3/v4 / current), converting f32 back to f64 on load.
/// Read body in f32 format (current `FXCACHE_VERSION`), converting f32 back to f64 on load.
fn read_body_f32(
reader: &mut BufReader<std::fs::File>,
bar_count: usize,
@@ -419,14 +428,6 @@ fn read_body_f32(
Ok((timestamps, features, targets, ofi))
}
/// Read body from legacy v2 files (target_dim=4), padding to 6-slot targets.
///
/// v2 layout: [preproc_close, preproc_next, raw_close, raw_next]
/// Padded to: [preproc_close, preproc_next, raw_close, raw_next, raw_close, raw_close]
/// tgt[4] = raw_close (fallback for raw_open — unavailable in v2)
/// tgt[5] = raw_close (fallback for mid_price_open — unavailable in v2)
// ── Finder ───────────────────────────────────────────────────────────────────
#[cfg(test)]

32
crates/ml/src/hyperopt/adapters/dqn.rs
View File

@@ -47,6 +47,9 @@
use ml_core::device::MlDevice;
use anyhow::Context;
/// OFI dimension — mirrors `ml_core::state_layout::OFI_DIM`.
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
use serde::{Deserialize, Serialize};
use std::fs::OpenOptions;
use std::io::Write as IoWrite;
@@ -454,9 +457,12 @@ pub struct DQNMetrics {
/// ## Fixed Architecture
///
/// The following parameters are fixed for consistency:
/// - `state_dim`: 66 without OFI (42 market + 8 portfolio + 16 MTF), or 86 with MBP-10 data (+20 OFI)
/// - `state_dim`: canonical layout at `STATE_DIM` (see `ml_core::state_layout`) —
/// 42 market + `OFI_DIM` OFI + 16 MTF + 8 portfolio + 6 plan/ISV.
/// - Market features (42): OHLCV, technical indicators, price patterns, volume, time, statistical, ADX, CUSUM direction
/// - OFI features (20, optional): 8 base OFI from MBP-10 order book + 12 extended microstructure
/// - OFI features (`OFI_DIM`): 8 raw OFI + 8 deltas + book_aggression + log_duration
/// + ofi_acceleration + toxicity_gradient + 10 MicrostructureState snapshot slots
/// + order_count_imbalance + microprice_residual
/// - Portfolio features (3): Position, unrealized PnL, drawdown
/// - `num_actions`: 9 (9 exposure levels: S100/S75/S50/S25/Flat/L25/L50/L75/L100)
/// - `hidden_dims`: [256, 128, 64]
@@ -535,7 +541,7 @@ pub struct DQNTrainer {
trades_data_dir: Option<String>,
/// Preloaded OFI features from MBP-10 data (8 features per bar).
/// Loaded during `preload_data()` and injected into each trial's DQNTrainer.
preloaded_ofi_features: Option<Arc<[[f64; 20]]>>,
preloaded_ofi_features: Option<Arc<[[f64; OFI_DIM]]>>,
/// Trading instrument symbol (e.g. "ES.FUT") — used in cache key to prevent
/// cross-instrument collisions when multiple symbols share the same data directory.
symbol: String,
@@ -917,10 +923,10 @@ impl DQNTrainer {
targets.push(t);
}
let has_ofi = ofi_features.is_some();
let ofi: Vec<[f64; 20]> = if let Some(ref ofi_arc) = ofi_features {
let ofi: Vec<[f64; OFI_DIM]> = if let Some(ref ofi_arc) = ofi_features {
ofi_arc.iter().cloned().collect()
} else {
vec![[0.0_f64; 20]; total]
vec![[0.0_f64; OFI_DIM]; total]
};
let timestamps = vec![0_i64; total];
@@ -959,7 +965,7 @@ impl DQNTrainer {
if self.preloaded_ofi_features.is_none() {
self.preloaded_ofi_features = Some(Arc::from(fxcache.ofi.as_slice()));
let ofi_nonzero = fxcache.ofi.iter().filter(|o| o.iter().any(|&v| v != 0.0)).count();
info!("OFI features from fxcache: {} bars x 20 dims ({} non-zero)", fxcache.ofi.len(), ofi_nonzero);
info!("OFI features from fxcache: {} bars x {} dims ({} non-zero)", fxcache.ofi.len(), OFI_DIM, ofi_nonzero);
}
let elapsed = preload_start.elapsed();
@@ -1203,7 +1209,7 @@ impl DQNTrainer {
/// Attempt to load OFI features from MBP10 data.
/// Returns None if MBP10 data is not available or on error.
fn load_ofi_features(&self) -> Option<Arc<[[f64; 20]]>> {
fn load_ofi_features(&self) -> Option<Arc<[[f64; OFI_DIM]]>> {
use crate::features::mbp10_loader::load_ofi_features_parallel;
let mbp10_dir = if let Some(ref dir) = self.mbp10_data_dir {
@@ -1223,15 +1229,19 @@ impl DQNTrainer {
.map(|p| p.join("trades"))
};
// This fallback loader only fills the 8 raw OFI slots (ofi8). The
// remaining [8..OFI_DIM) slots stay zero — the canonical signal path
// is precompute_features, which fills every slot from full MBP-10
// history. This path is used when precompute has not been run.
load_ofi_features_parallel(&mbp10_dir, trades_dir.as_deref(), |dir| {
collect_dbn_files(dir)
}).map(|ofi8| {
let ofi20: Vec<[f64; 20]> = ofi8.into_iter().map(|row| {
let mut wide = [0.0_f64; 20];
let ofi_wide: Vec<[f64; OFI_DIM]> = ofi8.into_iter().map(|row| {
let mut wide = [0.0_f64; OFI_DIM];
wide[..8].copy_from_slice(&row);
wide
}).collect();
Arc::from(ofi20)
Arc::from(ofi_wide)
})
}
@@ -1432,7 +1442,7 @@ impl DQNTrainer {
max_leverage: internal_trainer.hyperparams().max_leverage as f32,
// OFI reorder in gather kernel: produces [market, portfolio, OFI, pad]
// directly, eliminating the Candle narrow+cat closure.
ofi_dim: 20,
ofi_dim: OFI_DIM,
holding_cost_rate: internal_trainer.hyperparams().holding_cost_rate as f32,
churn_threshold: internal_trainer.hyperparams().churn_threshold_bars as f32,
churn_penalty_scale: internal_trainer.hyperparams().churn_penalty_scale as f32,

11
crates/ml/src/hyperopt/adapters/ppo.rs
View File

@@ -32,6 +32,9 @@
use ml_core::device::MlDevice;
use serde::{Deserialize, Serialize};
/// OFI dimension — mirrors `ml_core::state_layout::OFI_DIM`.
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
use std::fs::OpenOptions;
use std::io::Write as IoWrite;
use std::sync::atomic::{AtomicUsize, Ordering};
@@ -1495,7 +1498,7 @@ impl PPOTrainer {
/// Attempt to load OFI features from MBP10 data.
/// Returns None if MBP10 data is not available or on error.
fn load_ofi_features(&self) -> Option<Vec<[f64; 20]>> {
fn load_ofi_features(&self) -> Option<Vec<[f64; OFI_DIM]>> {
use crate::features::mbp10_loader::load_ofi_features_parallel;
let mbp10_dir = self
@@ -1506,11 +1509,15 @@ impl PPOTrainer {
// Derive trades directory as sibling of the OHLCV data dir
let trades_dir = self.dbn_data_dir.parent().map(|p| p.join("trades"));
// This fallback loader only fills the 8 raw OFI slots (ofi8). The
// remaining [8..OFI_DIM) slots stay zero — the canonical signal path
// is precompute_features, which fills every slot from full MBP-10
// history.
load_ofi_features_parallel(&mbp10_dir, trades_dir.as_deref(), |dir| {
collect_dbn_files_recursive(dir)
}).map(|ofi8| {
ofi8.into_iter().map(|row| {
let mut wide = [0.0_f64; 20];
let mut wide = [0.0_f64; OFI_DIM];
wide[..8].copy_from_slice(&row);
wide
}).collect()

22
crates/ml/src/trainers/dqn/data_loading.rs
View File

@@ -15,6 +15,9 @@ use crate::training_pipeline::FinancialFeatures;
use crate::features::extraction::FeatureVector;
use super::trainer::DQNTrainer;
/// OFI dimension — mirrors `ml_core::state_layout::OFI_DIM`.
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
/// Resolve a data directory path — handles relative paths by checking against
/// the workspace root (via CARGO_MANIFEST_DIR) when the raw path doesn't exist.
/// Used for all data paths: dbn_data_dir, mbp10_data_dir, trades_data_dir.
@@ -411,8 +414,8 @@ impl DQNTrainer {
Ok(features) => {
if features.is_valid() {
let arr8 = features.to_array();
let mut arr20 = [0.0_f64; 20];
arr20[..8].copy_from_slice(&arr8);
let mut ofi_row = [0.0_f64; OFI_DIM];
ofi_row[..8].copy_from_slice(&arr8);
// Book aggression from MBP-10 depth
let mut buy_ws = 0.0_f64;
let mut buy_t = 0.0_f64;
@@ -428,16 +431,19 @@ impl DQNTrainer {
}
let buy_com = if buy_t > 0.0 { buy_ws / buy_t } else { 5.5 };
let sell_com = if sell_t > 0.0 { sell_ws / sell_t } else { 5.5 };
arr20[16] = (sell_com - buy_com) / 10.0;
ofi_per_bar.push(arr20);
ofi_row[16] = (sell_com - buy_com) / 10.0;
// Slots [17..OFI_DIM) stay zero in this DBN-fallback path;
// the canonical signal source is precompute_features
// which fills every slot from full MBP-10 history.
ofi_per_bar.push(ofi_row);
} else {
ofi_per_bar.push([0.0; 20]);
ofi_per_bar.push([0.0; OFI_DIM]);
}
}
Err(_) => ofi_per_bar.push([0.0; 20]),
Err(_) => ofi_per_bar.push([0.0; OFI_DIM]),
}
} else {
ofi_per_bar.push([0.0; 20]);
ofi_per_bar.push([0.0; OFI_DIM]);
}
}
@@ -511,7 +517,7 @@ impl DQNTrainer {
tracing::info!("max_bars={}: truncating {} → {} bars (incl. OFI)", self.hyperparams.max_bars, training_data.len(), self.hyperparams.max_bars);
training_data.truncate(self.hyperparams.max_bars);
if let Some(ofi_arc) = self.ofi_features.take() {
let mut ofi_vec: Vec<[f64; 20]> = ofi_arc.iter().copied().collect();
let mut ofi_vec: Vec<[f64; OFI_DIM]> = ofi_arc.iter().copied().collect();
ofi_vec.truncate(self.hyperparams.max_bars);
self.ofi_features = Some(Arc::from(ofi_vec));
}

9
crates/ml/src/trainers/dqn/trainer/metrics.rs
View File

@@ -505,8 +505,9 @@ impl DQNTrainer {
// ── Lazy-init the GPU evaluator (once per fold, reused across epochs) ──
if self.gpu_evaluator.is_none() {
let market_dim: usize = 42;
let feature_dim: usize = 62; // 42 market + 20 OFI (portfolio+MTF in state_dim, not feature_dim)
let market_dim: usize = ml_core::state_layout::MARKET_DIM;
let ofi_dim: usize = ml_core::state_layout::OFI_DIM;
let feature_dim: usize = market_dim + ofi_dim; // 42 market + OFI (portfolio+MTF in state_dim, not feature_dim)
// Build a single window from all val_data
let mut prices: Vec<[f32; 4]> = Vec::with_capacity(self.val_data.len());
@@ -520,7 +521,7 @@ impl DQNTrainer {
let close = if target.len() >= 2 { target[0] as f32 } else { fv[3] as f32 };
prices.push([close, close, close, close]);
// Features: 42 market features + 20 OFI
// Features: 42 market features + OFI_DIM OFI
let fv_slice = &fv[..market_dim.min(fv.len())];
let mut fv_f32: Vec<f32> = fv_slice.iter().map(|&v| v as f32).collect();
let ofi_idx = self.ofi_val_offset + i;
@@ -546,7 +547,7 @@ impl DQNTrainer {
contract_multiplier: hp.contract_multiplier as f32,
margin_pct: hp.margin_pct as f32,
max_leverage: hp.max_leverage as f32,
ofi_dim: 20,
ofi_dim,
holding_cost_rate: hp.holding_cost_rate as f32,
churn_threshold: hp.churn_threshold_bars as f32,
churn_penalty_scale: hp.churn_penalty_scale as f32,

13
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -20,6 +20,9 @@ use tokio::sync::RwLock;
use tracing::info;
use uuid::Uuid;
/// OFI dimension — mirrors `ml_core::state_layout::OFI_DIM`.
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
use crate::dqn::action_space::FactoredAction;
use crate::dqn::circuit_breaker::CircuitBreaker;
use crate::dqn::curiosity::CuriosityModule;
@@ -512,7 +515,7 @@ pub struct DQNTrainer {
/// Populated during data loading when MBP-10 order book data is available.
/// Passed as `regime_features` in `TradingState::from_normalized()`.
/// Arc-shared to avoid copy per hyperopt trial.
pub ofi_features: Option<Arc<[[f64; 20]]>>,
pub ofi_features: Option<Arc<[[f64; OFI_DIM]]>>,
/// Number of training bars (OFI offset for validation data).
/// val_data[i] corresponds to ofi_features[ofi_val_offset + i].
pub(crate) ofi_val_offset: usize,
@@ -876,7 +879,7 @@ impl DQNTrainer {
// Upload FULL dataset to GPU FIRST so we can use GPU prefix sums for fold generation
let ofi = self.ofi_features.as_ref()
.ok_or_else(|| anyhow::anyhow!("OFI features required — model is worthless without order flow data. Check fxcache has_ofi flag and --mbp10-data-dir."))?;
let ofi_slice: Vec<[f64; 20]> = ofi.to_vec();
let ofi_slice: Vec<[f64; OFI_DIM]> = ofi.to_vec();
self.init_from_fxcache(&features, &targets, &ofi_slice).await?;
// Generate regime-stratified folds using GPU prefix sums (O(1) range queries)
@@ -1350,7 +1353,7 @@ impl DQNTrainer {
&mut self,
features: &[[f64; 42]],
targets: &[[f64; 6]],
ofi: &[[f64; 20]],
ofi: &[[f64; OFI_DIM]],
) -> Result<()> {
let stream = self.cuda_stream.as_ref()
.ok_or_else(|| anyhow::anyhow!("CUDA stream required for init_from_fxcache"))?;
@@ -1360,14 +1363,14 @@ impl DQNTrainer {
features, targets, ofi, stream,
).map_err(|e| anyhow::anyhow!("DqnGpuData::upload_slices failed: {e}"))?;
let raw_dim: usize = 65; // 42 market + 3 portfolio + 20 OFI
let raw_dim: usize = 42 + 3 + OFI_DIM; // market + portfolio (legacy 3-slot) + OFI
let aligned_dim = (raw_dim + 7) & !7;
gpu_data.set_aligned_state_dim(aligned_dim);
tracing::info!(
"init_from_fxcache: {} bars uploaded to GPU ({:.1} MB, OFI=true)",
features.len(),
(features.len() * (42 + 6 + 20) * 4) as f64 / 1_048_576.0,
(features.len() * (42 + 6 + OFI_DIM) * 4) as f64 / 1_048_576.0,
);
self.gpu_data = Some(gpu_data);

20
crates/ml/tests/fxcache_roundtrip_test.rs
View File

@@ -3,7 +3,7 @@
//! Covers f32 serialization, cache-key lookup, header validation,
//! and the empty-bars error path.
use ml::fxcache::{find_fxcache, load_fxcache, write_fxcache, FxCacheData};
use ml::fxcache::{find_fxcache, load_fxcache, write_fxcache, FxCacheData, FXCACHE_VERSION, OFI_DIM};
use tempfile::TempDir;
// ── Helpers ─────────────────────────────────────────────────────────────────
@@ -35,12 +35,12 @@ fn make_targets(n: usize) -> Vec<[f64; 6]> {
}
/// Deterministic OFI vector for bar `i`.
fn make_ofi(n: usize) -> Vec<[f64; 20]> {
fn make_ofi(n: usize) -> Vec<[f64; OFI_DIM]> {
(0..n)
.map(|i| {
let mut row = [0.0_f64; 20];
for j in 0..20 {
row[j] = (i * 20 + j) as f64 * 0.001;
let mut row = [0.0_f64; OFI_DIM];
for j in 0..OFI_DIM {
row[j] = (i * OFI_DIM + j) as f64 * 0.001;
}
row
})
@@ -119,7 +119,7 @@ fn test_fxcache_f32_roundtrip() {
data.targets[i][j]
);
}
for j in 0..20 {
for j in 0..OFI_DIM {
let diff = (data.ofi[i][j] - ofi[i][j]).abs();
assert!(
diff < 1e-4,
@@ -187,7 +187,7 @@ fn test_fxcache_empty() {
// Writer rejects 0 bars.
let features: Vec<[f64; 42]> = vec![];
let targets: Vec<[f64; 6]> = vec![];
let ofi: Vec<[f64; 20]> = vec![];
let ofi: Vec<[f64; OFI_DIM]> = vec![];
let timestamps: Vec<i64> = vec![];
let key = test_cache_key();
@@ -199,13 +199,13 @@ fn test_fxcache_empty() {
);
// Reader also rejects a header with bar_count=0.
// Hand-craft a valid-magic header with bar_count=0.
// Hand-craft a valid-magic header with current version + bar_count=0.
let mut header = [0u8; 64];
header[0..8].copy_from_slice(b"FXCACHE\0");
header[8..10].copy_from_slice(&3u16.to_le_bytes()); // version=3
header[8..10].copy_from_slice(&FXCACHE_VERSION.to_le_bytes()); // current version
header[10..12].copy_from_slice(&42u16.to_le_bytes()); // feat_dim
header[12..14].copy_from_slice(&6u16.to_le_bytes()); // target_dim
header[14..16].copy_from_slice(&20u16.to_le_bytes()); // ofi_dim
header[14..16].copy_from_slice(&(OFI_DIM as u16).to_le_bytes()); // ofi_dim
header[16..24].copy_from_slice(&0u64.to_le_bytes()); // bar_count=0
// cache_key + reserved stay zeroed.

5
crates/ml/tests/smoke_test_real_data.rs
View File

@@ -441,7 +441,8 @@ mod gpu_smoke {
assert!(n > 200, "Need at least 200 bars for GPU test, got {n}");
// 3. Optionally load OFI features from MBP-10
let ofi_features: Option<Vec<[f64; 20]>> = mbp10_dir.and_then(|dir| {
const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
let ofi_features: Option<Vec<[f64; OFI_DIM]>> = mbp10_dir.and_then(|dir| {
let files: Vec<_> = std::fs::read_dir(dir)
.ok()?
.filter_map(|e| e.ok())
@@ -455,7 +456,7 @@ mod gpu_smoke {
info!(file = %file.display(), "Loading OFI");
ml::features::mbp10_loader::compute_ofi_from_file(file).ok().map(|ofi8| {
ofi8.into_iter().map(|row| {
let mut wide = [0.0_f64; 20];
let mut wide = [0.0_f64; OFI_DIM];
wide[..8].copy_from_slice(&row);
wide
}).collect::<Vec<_>>()

BIN
test_data/feature-cache/9412eb8eddebfb1935378b13c5c87e05ead68044198c67912a3a06cdc20514fb.fxcache
View File

Binary file not shown.