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

feat(rl): encoder input expansion 40 → 56 dims (trade-context + multires)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-24 22:35:26 +02:00
parent 233894a4bf
commit 914a6e8e72
6 changed files with 87 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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

39
crates/ml-alpha/cuda/rl_encoder_context_broadcast.cu Normal file
View File

@@ -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 <20><><EFBFBD> this kernel broadcasts them.
//
// One thread per (batch × K) row. No shared memory.
// Per `feedback_no_atomicadd`, `feedback_cpu_is_read_only`.
#include <stdint.h>
#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];
}

12
crates/ml-alpha/cuda/snap_feature_assemble.cu
View File

@@ -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;
}

5
crates/ml-alpha/src/cfc/snap_features.rs
View File

@@ -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

2
crates/ml-alpha/src/cfc/trunk.rs
View File

@@ -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,

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

@@ -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<f32> = (0..FEATURE_DIM * FEATURE_DIM)
let vsn_scale = (1.0_f32 / ENCODER_INPUT_DIM as f32).sqrt();
let vsn_w_init: Vec<f32> = (0..ENCODER_INPUT_DIM * ENCODER_INPUT_DIM)
.map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect();
let vsn_b_init: Vec<f32> = vec![0.0; FEATURE_DIM];
let vsn_b_init: Vec<f32> = 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::<f32>(
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::<f32>(
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::<f32>(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::<f32>(cfg.n_batch * FEATURE_DIM)
let vsn_step_gates_d = stream.alloc_zeros::<f32>(cfg.n_batch * ENCODER_INPUT_DIM)
.context("vsn_step_gates_d alloc")?;
let ln_a_step_out_d = stream.alloc_zeros::<f32>(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::<i32>(cfg.n_batch)?;
let step_ts_ns_d = stream.alloc_zeros::<i64>(cfg.n_batch)?;
let step_prev_ts_ns_d = stream.alloc_zeros::<i64>(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::<f32>(cfg.n_batch * k * FEATURE_DIM)?,
vsn_grad_x_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * FEATURE_DIM)?,
grad_vsn_w_d: stream.alloc_zeros::<f32>(FEATURE_DIM * FEATURE_DIM)?,
grad_vsn_b_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
vsn_gates_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * ENCODER_INPUT_DIM)?,
vsn_grad_x_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * ENCODER_INPUT_DIM)?,
grad_vsn_w_d: stream.alloc_zeros::<f32>(ENCODER_INPUT_DIM * ENCODER_INPUT_DIM)?,
grad_vsn_b_d: stream.alloc_zeros::<f32>(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);
{