From 914a6e8e72a54871a348b7a714ee0c4d48d344c0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 24 May 2026 22:35:26 +0200 Subject: [PATCH] =?UTF-8?q?feat(rl):=20encoder=20input=20expansion=2040=20?= =?UTF-8?q?=E2=86=92=2056=20dims=20(trade-context=20+=20multires)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces ENCODER_INPUT_DIM = 56 (FEATURE_DIM + 16). All encoder first-layer weight matrices (VSN gate, Mamba2 L1 input projection) now sized for 56 input dims. The extra 16 are per-batch state: 4 trade_context + 12 multires features. - snap_feature_assemble_batched: output stride → ENCODER_INPUT_DIM, zero-fills dims [40..56] for the broadcast kernel to overwrite. - New rl_encoder_context_broadcast.cu: writes trade_context_d[B×4] + multires_output_d[B×12] into each of the K sequence rows per batch at positions [40..56]. - CfcConfig.n_in, Mamba2 L1 in_dim, VSN gate, window_tensor_d, all forward/backward scratch buffers updated to ENCODER_INPUT_DIM. - CfcTrunk default config updated. The broadcast kernel launch integration into the forward_only path is the final wire-up step — until then dims 40-55 are zero-filled (safe: Xavier init on new columns means encoder starts by learning to ignore them, then gradually incorporates the signal). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + .../cuda/rl_encoder_context_broadcast.cu | 39 ++++++++++++ crates/ml-alpha/cuda/snap_feature_assemble.cu | 12 +++- crates/ml-alpha/src/cfc/snap_features.rs | 5 ++ crates/ml-alpha/src/cfc/trunk.rs | 2 +- crates/ml-alpha/src/trainer/perception.rs | 62 ++++++++++--------- 6 files changed, 87 insertions(+), 34 deletions(-) create mode 100644 crates/ml-alpha/cuda/rl_encoder_context_broadcast.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index d228b6fdb..3d6676bd5 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -85,6 +85,7 @@ const KERNELS: &[&str] = &[ "rl_recent_outcome_update", // P10: per-batch signed outcome EMA for anti-martingale sizing "rl_trade_context_update", // P1: per-batch trade-arc features (time_in_trade, unrealized_R, pos_mag, entry_dist) "rl_multires_features_update", // P0: per-batch multi-resolution streaming features (3 horizons × 4 features) + "rl_encoder_context_broadcast", // P2: broadcast per-batch context (16 dims) into encoder input [B,K,56] at cols 40-55 ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_encoder_context_broadcast.cu b/crates/ml-alpha/cuda/rl_encoder_context_broadcast.cu new file mode 100644 index 000000000..2214765c9 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_encoder_context_broadcast.cu @@ -0,0 +1,39 @@ +// rl_encoder_context_broadcast.cu — broadcast per-batch trade_context +// (4 dims) + multires_output (12 dims) into the encoder input tensor +// at positions [40..56] for each of the K sequence rows per batch. +// +// The encoder input tensor is [B, K, ENCODER_INPUT_DIM=56] row-major. +// Dims [0..40] are per-snapshot market features (written by +// snap_feature_assemble_batched). Dims [40..56] are per-batch state +// that is CONSTANT across K ��� this kernel broadcasts them. +// +// One thread per (batch × K) row. No shared memory. +// Per `feedback_no_atomicadd`, `feedback_cpu_is_read_only`. + +#include + +#define ENCODER_INPUT_DIM 56 +#define SNAP_FEATURE_DIM 40 +#define TRADE_CONTEXT_DIM 4 +#define MULTIRES_DIM 12 + +extern "C" __global__ void rl_encoder_context_broadcast( + float* __restrict__ encoder_input, // [B × K × ENCODER_INPUT_DIM] + const float* __restrict__ trade_context, // [B × TRADE_CONTEXT_DIM] + const float* __restrict__ multires_output, // [B × MULTIRES_DIM] + int b_size, + int seq_len +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = b_size * seq_len; + if (idx >= total) return; + + const int b = idx / seq_len; + + float* row = encoder_input + idx * ENCODER_INPUT_DIM + SNAP_FEATURE_DIM; + const float* tc = trade_context + b * TRADE_CONTEXT_DIM; + const float* mr = multires_output + b * MULTIRES_DIM; + + for (int i = 0; i < TRADE_CONTEXT_DIM; ++i) row[i] = tc[i]; + for (int i = 0; i < MULTIRES_DIM; ++i) row[TRADE_CONTEXT_DIM + i] = mr[i]; +} diff --git a/crates/ml-alpha/cuda/snap_feature_assemble.cu b/crates/ml-alpha/cuda/snap_feature_assemble.cu index 55dfc281b..5d0c22a3b 100644 --- a/crates/ml-alpha/cuda/snap_feature_assemble.cu +++ b/crates/ml-alpha/cuda/snap_feature_assemble.cu @@ -37,7 +37,8 @@ // Single-thread per-snapshot kernel; the call site launches with one // block of one thread per snapshot. -#define SNAP_FEATURE_DIM 40 +#define SNAP_FEATURE_DIM 40 +#define ENCODER_INPUT_DIM 56 __device__ __forceinline__ float signed_log1p(float x) { return copysignf(log1pf(fabsf(x)), x); @@ -142,7 +143,7 @@ extern "C" __global__ void snap_feature_assemble_batched( const long* __restrict__ prev_ts_ns, // [N] float tick_size, int N, - float* __restrict__ out // [N, 40] + float* __restrict__ out // [N, ENCODER_INPUT_DIM] ) { int n = blockIdx.x * blockDim.x + threadIdx.x; if (n >= N) return; @@ -154,7 +155,7 @@ extern "C" __global__ void snap_feature_assemble_batched( const float* pbs = prev_bid_sz + n * 10; const float* pas = prev_ask_sz + n * 10; const float* rg = regime + n * 6; - float* o = out + n * SNAP_FEATURE_DIM; + float* o = out + n * ENCODER_INPUT_DIM; const float mid = 0.5f * (bx[0] + ax[0]); const float pm = prev_mid[n]; @@ -181,4 +182,9 @@ extern "C" __global__ void snap_feature_assemble_batched( // TGN Δt Fourier features at slots [32..40]. const long dt_ns_clamped = dt_ns >= 0 ? dt_ns : 0; emit_dt_fourier(o + 32, dt_ns_clamped); + + // Zero-fill encoder expansion dims [40..56]. The broadcast kernel + // (rl_encoder_context_broadcast) overwrites these with per-batch + // trade_context + multires features each step. + for (int i = SNAP_FEATURE_DIM; i < ENCODER_INPUT_DIM; ++i) o[i] = 0.0f; } diff --git a/crates/ml-alpha/src/cfc/snap_features.rs b/crates/ml-alpha/src/cfc/snap_features.rs index bc28d48db..468deddb6 100644 --- a/crates/ml-alpha/src/cfc/snap_features.rs +++ b/crates/ml-alpha/src/cfc/snap_features.rs @@ -20,6 +20,11 @@ const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_featu pub const ES_TICK_SIZE: f32 = 0.25; pub const FEATURE_DIM: usize = 40; +/// Encoder input dimension: raw market features (40) + per-batch state +/// (4 trade_context + 12 multires = 16). The encoder's first layer and +/// VSN gate size from this, not FEATURE_DIM. +pub const ENCODER_INPUT_DIM: usize = FEATURE_DIM + 16; + /// Number of book-side levels in one MBP-10 snapshot (bid + ask). /// Canonical structural dim — every Rust array sized for book depth /// (`[f32; BOOK_LEVELS]`) and every loop iterating over book depth diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 70f9dc74d..13c28a1c1 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -112,7 +112,7 @@ pub struct CfcConfig { impl Default for CfcConfig { fn default() -> Self { Self { - n_in: FEATURE_DIM, + n_in: crate::cfc::snap_features::ENCODER_INPUT_DIM, n_hid: HIDDEN_DIM, cfc_n_in: HIDDEN_DIM, mamba2_state_dim: 16, diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index ff9c43c76..d7c9f088f 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -50,7 +50,9 @@ use crate::aux_heads::{ AuxMaskedHuberLoss, DEFAULT_HUBER_DELTA, N_AUX_HORIZONS, }; use crate::cfc::aux_trunk::AUX_HIDDEN; -use crate::cfc::snap_features::{Mbp10RawInput, BOOK_LEVELS, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM}; +use crate::cfc::snap_features::{ + Mbp10RawInput, BOOK_LEVELS, ENCODER_INPUT_DIM, ES_TICK_SIZE, REGIME_DIM, +}; use crate::cfc::AuxTrunk; use crate::cfc::AuxTrunkConfig; use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS}; @@ -280,9 +282,9 @@ pub struct PerceptionTrainer { /// Mamba2BackwardGrads wrapper allocated per step. mamba2_grads_buffers: Mamba2BackwardGradsBuffers, /// Pre-allocated input window for snap_features → Mamba2 fwd. - /// [B, K, FEATURE_DIM] — overwritten each step by the batched - /// snap_feature kernel. - window_tensor_d: GpuTensor, + /// [B, K, ENCODER_INPUT_DIM] — first 40 dims from snap kernel, + /// last 16 from context broadcast. + pub window_tensor_d: GpuTensor, /// Pre-allocated transpose of Mamba2's h_enriched_seq into [K, B, H] /// layout for contiguous per-K slot access in the trainer loop. h_enriched_seq_t_d: GpuTensor, @@ -1245,7 +1247,7 @@ impl PerceptionTrainer { let mut trunk = crate::cfc::trunk::CfcTrunk::new_random( dev, &crate::cfc::trunk::CfcConfig { - n_in: FEATURE_DIM, + n_in: ENCODER_INPUT_DIM, n_hid: HIDDEN_DIM, cfc_n_in: HIDDEN_DIM, mamba2_state_dim: cfg.mamba2_state_dim, @@ -1375,13 +1377,13 @@ impl PerceptionTrainer { ) .context("Mamba2AdamW::new")?; let mamba2_fwd_scratch = Mamba2BlockForwardScratch::new( - &stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, + &stream, cfg.n_batch, cfg.seq_len, ENCODER_INPUT_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, ).context("Mamba2BlockForwardScratch::new")?; let mamba2_bwd_scratch = Mamba2BackwardScratch::new( &stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, cfg.mamba2_state_dim, ).context("Mamba2BackwardScratch::new")?; let mamba2_grads_buffers = Mamba2BackwardGradsBuffers::new( - &stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, + &stream, cfg.n_batch, cfg.seq_len, ENCODER_INPUT_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, ).context("Mamba2BackwardGradsBuffers::new")?; // ── Phase 2B: SECOND Mamba2 stack — same shape, on the trunk. ── @@ -1403,7 +1405,7 @@ impl PerceptionTrainer { let mamba2_l2_grads_buffers = Mamba2BackwardGradsBuffers::new( &stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, ).context("Mamba2BackwardGradsBuffers::new (l2)")?; - let window_tensor_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, FEATURE_DIM], &stream) + let window_tensor_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, ENCODER_INPUT_DIM], &stream) .map_err(|e| anyhow::anyhow!("window_tensor_d alloc: {e}"))?; let h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream) .map_err(|e| anyhow::anyhow!("h_enriched_seq_t_d alloc: {e}"))?; @@ -1578,10 +1580,10 @@ impl PerceptionTrainer { // draws are uploaded directly into the trunk's zero-initialised // VSN slots via memcpy_htod, replacing the prior trainer-owned // vsn_w_d / vsn_b_d allocations. - let vsn_scale = (1.0_f32 / FEATURE_DIM as f32).sqrt(); - let vsn_w_init: Vec = (0..FEATURE_DIM * FEATURE_DIM) + let vsn_scale = (1.0_f32 / ENCODER_INPUT_DIM as f32).sqrt(); + let vsn_w_init: Vec = (0..ENCODER_INPUT_DIM * ENCODER_INPUT_DIM) .map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect(); - let vsn_b_init: Vec = vec![0.0; FEATURE_DIM]; + let vsn_b_init: Vec = vec![0.0; ENCODER_INPUT_DIM]; // Upload the trainer-driven VSN draws into the trunk's // zero-initialised VSN slots from X1. From here on, // `self.trunk.vsn_w_d` / `self.trunk.vsn_b_d` are the canonical @@ -1592,14 +1594,14 @@ impl PerceptionTrainer { stream .memcpy_htod(&vsn_b_init, &mut trunk.vsn_b_d) .context("trunk.vsn_b_d upload")?; - let opt_vsn_w = AdamW::new(dev, FEATURE_DIM * FEATURE_DIM, cfg.lr_cfc)?; - let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?; + let opt_vsn_w = AdamW::new(dev, ENCODER_INPUT_DIM * ENCODER_INPUT_DIM, cfg.lr_cfc)?; + let mut opt_vsn_b = AdamW::new(dev, ENCODER_INPUT_DIM, cfg.lr_cfc)?; opt_vsn_b.wd = 0.0; // Phase B: VSN per-row grad scratch (n_rows = B * K). let vsn_grad_w_scratch_d = stream.alloc_zeros::( - cfg.n_batch * cfg.seq_len * FEATURE_DIM * FEATURE_DIM)?; + cfg.n_batch * cfg.seq_len * ENCODER_INPUT_DIM * ENCODER_INPUT_DIM)?; let vsn_grad_b_scratch_d = stream.alloc_zeros::( - cfg.n_batch * cfg.seq_len * FEATURE_DIM)?; + cfg.n_batch * cfg.seq_len * ENCODER_INPUT_DIM)?; // ── Attention pool init (Phase 3) ── // Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K @@ -1643,7 +1645,7 @@ impl PerceptionTrainer { // so these scratches accumulate context across forward_step calls // without re-running over a K-window each time. let step_scratch_l1 = Mamba2BlockStepScratch::new( - &stream, cfg.n_batch, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, + &stream, cfg.n_batch, ENCODER_INPUT_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, ).context("Mamba2BlockStepScratch::new (l1)")?; let step_scratch_l2 = Mamba2BlockStepScratch::new( &stream, cfg.n_batch, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim, @@ -1652,9 +1654,9 @@ impl PerceptionTrainer { .context("cfc_h_state_step_d alloc")?; let cfc_h_new_step_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM) .context("cfc_h_new_step_d alloc")?; - let vsn_step_out_d = GpuTensor::zeros(&[cfg.n_batch, FEATURE_DIM], &stream) + let vsn_step_out_d = GpuTensor::zeros(&[cfg.n_batch, ENCODER_INPUT_DIM], &stream) .map_err(|e| anyhow::anyhow!("vsn_step_out_d alloc: {e}"))?; - let vsn_step_gates_d = stream.alloc_zeros::(cfg.n_batch * FEATURE_DIM) + let vsn_step_gates_d = stream.alloc_zeros::(cfg.n_batch * ENCODER_INPUT_DIM) .context("vsn_step_gates_d alloc")?; let ln_a_step_out_d = stream.alloc_zeros::(cfg.n_batch * HIDDEN_DIM) .context("ln_a_step_out_d alloc")?; @@ -1714,7 +1716,7 @@ impl PerceptionTrainer { let step_trade_count_d = stream.alloc_zeros::(cfg.n_batch)?; let step_ts_ns_d = stream.alloc_zeros::(cfg.n_batch)?; let step_prev_ts_ns_d = stream.alloc_zeros::(cfg.n_batch)?; - let window_step_d = GpuTensor::zeros(&[cfg.n_batch, 1, FEATURE_DIM], &stream) + let window_step_d = GpuTensor::zeros(&[cfg.n_batch, 1, ENCODER_INPUT_DIM], &stream) .map_err(|e| anyhow::anyhow!("window_step_d alloc: {e}"))?; // ── Per-horizon CfC routing scaffolding (Task 9) ── @@ -2151,12 +2153,12 @@ impl PerceptionTrainer { opt_ln_bias, // VSN (Phase 2D) — params + per-K gates + scratch. trunk, - vsn_out_d: GpuTensor::zeros(&[cfg.n_batch, k, FEATURE_DIM], &stream) + vsn_out_d: GpuTensor::zeros(&[cfg.n_batch, k, ENCODER_INPUT_DIM], &stream) .map_err(|e| anyhow::anyhow!("vsn_out_d alloc: {e}"))?, - vsn_gates_d: stream.alloc_zeros::(cfg.n_batch * k * FEATURE_DIM)?, - vsn_grad_x_d: stream.alloc_zeros::(cfg.n_batch * k * FEATURE_DIM)?, - grad_vsn_w_d: stream.alloc_zeros::(FEATURE_DIM * FEATURE_DIM)?, - grad_vsn_b_d: stream.alloc_zeros::(FEATURE_DIM)?, + vsn_gates_d: stream.alloc_zeros::(cfg.n_batch * k * ENCODER_INPUT_DIM)?, + vsn_grad_x_d: stream.alloc_zeros::(cfg.n_batch * k * ENCODER_INPUT_DIM)?, + grad_vsn_w_d: stream.alloc_zeros::(ENCODER_INPUT_DIM * ENCODER_INPUT_DIM)?, + grad_vsn_b_d: stream.alloc_zeros::(ENCODER_INPUT_DIM)?, opt_vsn_w, opt_vsn_b, vsn_bwd_fn, @@ -2732,7 +2734,7 @@ impl PerceptionTrainer { // BCE + Huber respectively and will add the matching `stg_aux_size_*` // mapped buffers + per-K backward slot pointers. - debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]); + debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, ENCODER_INPUT_DIM]); let total_snaps = b_sz * k_seq; debug_assert!(total_snaps <= self.bk_capacity); @@ -4325,7 +4327,7 @@ impl PerceptionTrainer { // VSN reducer: collapse n_rows (= B * K) → final grad buffers. { let n_rows_i = (b_sz * k_seq) as i32; - let n_tail_w_usz = FEATURE_DIM * FEATURE_DIM; + let n_tail_w_usz = ENCODER_INPUT_DIM * ENCODER_INPUT_DIM; let cfg_red_w = LaunchConfig { grid_dim: (((n_tail_w_usz + 31) / 32) as u32, 1, 1), block_dim: (32, 8, 1), @@ -4342,11 +4344,11 @@ impl PerceptionTrainer { launch.launch(cfg_red_w).context("reduce vsn_grad_w")?; } let cfg_red_b = LaunchConfig { - grid_dim: (((FEATURE_DIM + 31) / 32) as u32, 1, 1), + grid_dim: (((ENCODER_INPUT_DIM + 31) / 32) as u32, 1, 1), block_dim: (32, 8, 1), shared_mem_bytes: 0, }; - let n_tail_b = FEATURE_DIM as i32; + let n_tail_b = ENCODER_INPUT_DIM as i32; unsafe { let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn); launch @@ -7185,7 +7187,7 @@ impl PerceptionTrainer { // trunk's weights from the checkpoint via load_checkpoint. let mut trainer = Self::new(dev, cfg)?; let trunk_cfg = crate::cfc::trunk::CfcConfig { - n_in: FEATURE_DIM, + n_in: ENCODER_INPUT_DIM, n_hid: HIDDEN_DIM, cfc_n_in: HIDDEN_DIM, mamba2_state_dim: cfg.mamba2_state_dim, @@ -7232,7 +7234,7 @@ impl PerceptionTrainer { ); // Use pre-allocated window_tensor_d — no per-step alloc. - debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]); + debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, ENCODER_INPUT_DIM]); let total_snaps = b_sz * k_seq; debug_assert!(total_snaps <= self.bk_capacity); {