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

feat: OFI embed MLP forward (18→10 cuBLAS) — learned order flow compression

Browse Source 
Extracts [raw_ofi(8); delta_ofi(8); book_aggression(1); log_duration(1)]
= 18-dim from state vector, compresses to 10-dim via cuBLAS SGEMM +
bias+ReLU. Xavier init. Runs in training forward path before Mamba2
and attention, making the embedding available for temporal enrichment.

190 trainable params (18×10 + 10 bias). ~0.05ms per step.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 00:49:18 +02:00
parent 9b61e46a71
commit 6575eea8da
2 changed files with 106 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5790,3 +5790,31 @@ extern "C" __global__ void plan_noise_inject(
plan_params[i * 6 + p] *= (1.0f + noise);
}
}
/* ── OFI embedding MLP input construction ─────────────────────────────────
*
* ofi_embed_build_input — Extract 18-dim OFI input from states_buf.
* raw OFI at state[66..74), delta OFI at state[74..82),
* book_aggression at state[82], log_bar_duration at state[83].
* Output: row-major [B, 18] for cuBLAS sgemm_f32 (row-major trick).
* Grid: ceil(B/256), Block: 256.
*/
extern "C" __global__ void ofi_embed_build_input(
const float* __restrict__ states,
float* __restrict__ output,
int B, int state_dim
) {
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 > 66 + k) ? s[66 + k] : 0.0f;
/* Delta OFI [8..16) */
for (int k = 0; k < 8; k++)
output[b * 18 + 8 + k] = (state_dim > 74 + k) ? s[74 + k] : 0.0f;
/* Book aggression [16] */
output[b * 18 + 16] = (state_dim > 82) ? s[82] : 0.0f;
/* Log bar duration [17] */
output[b * 18 + 17] = (state_dim > 83) ? s[83] : 0.0f;
}

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

@@ -1591,6 +1591,13 @@ 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_build_input_kernel: CudaFunction,
// ── Speculative inference cache ──
/// Cached trunk output [B, SH2] from between-bar speculative forward.
speculative_h_s2: CudaSlice<f32>,
@@ -2534,6 +2541,49 @@ 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).
pub(crate) fn launch_ofi_embed_forward(&self, batch_size: usize) -> Result<(), MLError> {
let sd = self.config.state_dim as i32;
let b_i32 = batch_size as i32;
// Step 1: Build input [B, 18] 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);
unsafe {
self.stream.launch_builder(&self.ofi_embed_build_input_kernel)
.arg(&states_ptr)
.arg(&input_ptr)
.arg(&b_i32)
.arg(&sd)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.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]
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",
)?;
// Step 3: Add bias + ReLU in-place on output[B, 10]
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,
)?;
Ok(())
}
/// Recursive confidence backward: MSE loss gradient into trunk + conf weight gradients.
/// Accumulates into grad_buf (same buffer Adam reads) and bw_d_h_s2 trunk gradient.
pub(crate) fn launch_recursive_confidence_backward(&self, batch_size: usize) -> Result<(), MLError> {
@@ -5000,7 +5050,9 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("trade_plan_activate load: {e}")))?;
let plan_noise_kernel = exp_module_for_mag.load_function("plan_noise_inject")
.map_err(|e| MLError::ModelError(format!("plan_noise_inject load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward + fill_gamma_buf + trade_plan_activate + plan_noise kernels loaded");
let ofi_embed_build_input_kernel = exp_module_for_mag.load_function("ofi_embed_build_input")
.map_err(|e| MLError::ModelError(format!("ofi_embed_build_input load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate + adaptive_atom + atom_grad + q_anchor + regime_dropout + G5/G6/G10/G12 + risk_budget + isv_signal_update + isv_forward + fill_gamma_buf + trade_plan_activate + plan_noise + ofi_embed_build_input kernels loaded");
// ── G5: Epistemic-gated magnitude — pinned var_ema threshold ─
let (var_ema_pinned, var_ema_dev_ptr) = {
@@ -6296,6 +6348,24 @@ 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 mut ofi_embed_w = alloc_f32(&stream, 10 * 18, "ofi_embed_w")?;
// Xavier init: scale = sqrt(2.0 / (fan_in + fan_out)) = sqrt(2.0 / 28)
{
let scale = (2.0_f32 / 28.0).sqrt();
let init_data: Vec<f32> = (0..10 * 18).map(|j| {
let hash = (j as u32).wrapping_mul(2654435761).wrapping_add(0xCAFEBEEF);
let u = (hash as f32) / (u32::MAX as f32) * 2.0 - 1.0;
u * scale
}).collect();
stream.memcpy_htod(&init_data, &mut ofi_embed_w)
.map_err(|e| MLError::ModelError(format!("ofi_embed_w xavier init: {e}")))?;
}
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);
// ── Speculative inference cache ──
let speculative_h_s2 = stream.alloc_zeros::<f32>(b * sh2)
.map_err(|e| MLError::ModelError(format!("speculative_h_s2 alloc: {e}")))?;
@@ -6790,6 +6860,11 @@ impl GpuDqnTrainer {
trade_plan_pre_out_buf,
trade_plan_activate_kernel,
plan_noise_kernel,
ofi_embed_input_buf,
ofi_embed_output_buf,
ofi_embed_w,
ofi_embed_b,
ofi_embed_build_input_kernel,
speculative_h_s2,
speculative_features,
speculative_valid: false,
@@ -9246,6 +9321,8 @@ impl GpuDqnTrainer {
self.launch_isv_feature_gate(batch_size)?;
self.launch_recursive_confidence_forward(batch_size)?;
self.launch_trade_plan_forward(batch_size)?;
// OFI embedding: extract [raw_ofi; delta_ofi; book_aggression; log_duration] → MLP → 10-dim
self.launch_ofi_embed_forward(batch_size)?;
}
// ── 1c. Temporal pipeline (enriches h_s2 before loss) ──