plan: OFI momentum + dense micro-reward — 9 tasks, full gradient flow
9-task implementation plan covering: - fxcache target_dim 4→6 (raw_open + mid_price_open from MBP-10) - OFI data pipeline fix (ofi_gpu wiring, indices 42→66, PORTFOLIO_STRIDE 30→38) - Dense micro-reward (OFI momentum × price confirm × adaptive cost tolerance) - OFI embed MLP (16→8 cuBLAS) with full backward gradient flow - Mamba2 history enrichment (SH2→SH2+8) + d_h_history backward - Attention enrichment (D→D+8) + d_input_scratch exposure Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
635
docs/superpowers/plans/2026-04-19-ofi-momentum-dense-reward.md
Normal file
635
docs/superpowers/plans/2026-04-19-ofi-momentum-dense-reward.md
Normal file
@@ -0,0 +1,635 @@
|
||||
# OFI Momentum + Dense Micro-Reward Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Add dense per-bar micro-reward using OFI momentum × price confirmation with adaptive cost tolerance, feed OFI embedding into Mamba2 history + attention input with full backward gradient flow.
|
||||
|
||||
**Architecture:** Three layers: (1) Fix broken OFI data pipeline + expand fxcache targets to include open/mid prices. (2) Implement dense micro-reward in env_step_batch using OFI momentum, price confirmation (open→close), and performance-adaptive spread cost. (3) Add learned OFI embedding MLP (16→8) feeding into Mamba2 + attention with full backward pass gradient flow via 2 new cuBLAS GEMMs.
|
||||
|
||||
**Tech Stack:** CUDA 12.4, cuBLAS SGEMM (TF32), cudarc vendored, existing 2-phase reduce pattern, fxcache binary format v3.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Expand fxcache targets (target_dim 4→6) + FXCACHE_VERSION bump
|
||||
|
||||
Add `raw_open` and `mid_price_open` (from MBP-10 best bid/ask midpoint) to the targets buffer. Bump FXCACHE_VERSION from 2 to 3 so the Argo `ensure-fxcache` step automatically rebuilds.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/fxcache.rs` — FXCACHE_VERSION 2→3, target_dim references
|
||||
- Modify: `crates/ml/src/trainers/dqn/data_loading.rs` — populate tgt[4]=raw_open, tgt[5]=mid_price_open
|
||||
- Modify: `crates/ml/examples/precompute_features.rs` — target construction
|
||||
- Modify: `crates/ml/examples/train_baseline_rl.rs` — any TARGET_DIM=4 constants
|
||||
- Modify: `crates/ml/examples/evaluate_baseline.rs` — any TARGET_DIM=4 constants
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` — update target slot comments
|
||||
- Modify: `crates/ml/tests/fxcache_roundtrip_test.rs` — update test from 4→6
|
||||
|
||||
- [ ] **Step 1: Bump FXCACHE_VERSION in fxcache.rs**
|
||||
|
||||
In `crates/ml/src/fxcache.rs` line 43:
|
||||
```rust
|
||||
pub const FXCACHE_VERSION: u16 = 3; // v3: target_dim 4→6 (added raw_open + mid_price_open)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update target construction in data_loading.rs**
|
||||
|
||||
Find the `targets` construction (around line 677 in `data_loading.rs`) where `vec![close, next_close, close, next_close]` is built. Add `raw_open` and `mid_price_open`:
|
||||
|
||||
```rust
|
||||
// targets = [preproc_close, preproc_next, raw_close, raw_next, raw_open, mid_price_open]
|
||||
let targets = vec![
|
||||
preproc_close, preproc_next_close,
|
||||
raw_close, raw_next_close,
|
||||
raw_open, // bar open price from OHLCV
|
||||
mid_price_at_open, // MBP-10 best bid/ask midpoint at bar formation
|
||||
];
|
||||
```
|
||||
|
||||
The `raw_open` is already parsed (line 650). For `mid_price_at_open`, extract from the MBP-10 data if available, otherwise use `raw_open` as fallback (for OHLCV-only mode).
|
||||
|
||||
- [ ] **Step 3: Update all target_dim=4 references**
|
||||
|
||||
Search for hardcoded `4` used as target dimension in:
|
||||
- `fxcache.rs`: FxCacheHeader target_dim field, reader validation
|
||||
- `precompute_features.rs`: target construction loop
|
||||
- `train_baseline_rl.rs` / `evaluate_baseline.rs`: any `target_dim` or `tgt[3]` bounds
|
||||
- `experience_kernels.cu`: update comment at kernel parameter docs (tgt[0..3] → tgt[0..5])
|
||||
- `fxcache_roundtrip_test.rs` line 207: change `4u16` to `6u16`
|
||||
|
||||
- [ ] **Step 4: Build and run fxcache roundtrip test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- fxcache --nocapture
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/fxcache.rs crates/ml/src/trainers/dqn/data_loading.rs \
|
||||
crates/ml/examples/precompute_features.rs crates/ml/examples/train_baseline_rl.rs \
|
||||
crates/ml/examples/evaluate_baseline.rs crates/ml/src/cuda_pipeline/experience_kernels.cu \
|
||||
crates/ml/tests/fxcache_roundtrip_test.rs
|
||||
git commit -m "feat: target_dim 4→6 — add raw_open + mid_price_open from MBP-10"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Fix OFI data pipeline + PORTFOLIO_STRIDE expansion
|
||||
|
||||
Fix 3 pre-existing bugs: (1) `ofi_gpu` uploaded but never wired to kernels, (2) `concat_ofi_features` reads wrong indices (42 vs 66), (3) PORTFOLIO_STRIDE too small for prev-OFI storage.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` — PORTFOLIO_STRIDE 30→38, add `batch_states` param to env_step_batch, OFI delta precompute in state[74..82)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/trade_stats_kernel.cu` — PORTFOLIO_STRIDE 30→38
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` — PORTFOLIO_STRIDE const, wire ofi_gpu to state_gather, pass batch_states to env_step, update portfolio_init
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — fix concat_ofi ofi_base 42→66, 45→69
|
||||
|
||||
- [ ] **Step 1: Expand PORTFOLIO_STRIDE in all .cu files**
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/experience_kernels.cu` around line 65:
|
||||
```c
|
||||
#define PORTFOLIO_STRIDE 38 /* was 30. Slots 30-37: prev-bar OFI for delta computation */
|
||||
```
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/trade_stats_kernel.cu` line 12:
|
||||
```c
|
||||
#define PORTFOLIO_STRIDE 38
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update PORTFOLIO_STRIDE in Rust + all allocation sites**
|
||||
|
||||
In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` line 47:
|
||||
```rust
|
||||
const PORTFOLIO_STRIDE: usize = 38; // was 30. Slots 30-37: prev-bar OFI for delta
|
||||
```
|
||||
|
||||
Search for all `PORTFOLIO_STRIDE` uses in the file — the allocation at line 901, init loops at lines 905-907, 2525-2528 — all reference the constant so they auto-update.
|
||||
|
||||
- [ ] **Step 3: Add `batch_states` parameter to env_step_batch kernel signature**
|
||||
|
||||
In `experience_kernels.cu`, add to the `experience_env_step` kernel signature (after existing params):
|
||||
```c
|
||||
const float* __restrict__ batch_states, /* [N, state_dim] full state vector with OFI at [66..82) */
|
||||
int state_dim, /* state vector width (96) */
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Wire ofi_gpu into state_gather and pass batch_states to env_step**
|
||||
|
||||
In `gpu_experience_collector.rs`:
|
||||
- Find `state_gather_kernel` launch (~line 1988). Add `ofi_gpu` buffer as a source for indices 66-73 in the gathered states.
|
||||
- Find `env_step_kernel` launch (~line 2356). Add the `batch_states_buf` pointer and `state_dim` as new kernel args.
|
||||
|
||||
- [ ] **Step 5: Fix concat_ofi_features indices**
|
||||
|
||||
In `gpu_dqn_trainer.rs`, find `launch_concat_ofi_features` calls (~line 2082-2083). Change:
|
||||
```rust
|
||||
// Order branch: was ofi_base=42, now 66
|
||||
let ofi_base_order = 66_i32;
|
||||
// Urgency branch: was ofi_base=45, now 69
|
||||
let ofi_base_urgency = 69_i32;
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Add OFI delta precomputation in env_step_batch**
|
||||
|
||||
In `experience_kernels.cu`, after the micro-reward computation (or in a dedicated section), add:
|
||||
```c
|
||||
/* Store current OFI in ps[30..37] for next bar's delta computation */
|
||||
if (batch_states != NULL && state_dim >= 82) {
|
||||
const float* ofi_cur = batch_states + (long long)i * state_dim + 66;
|
||||
for (int k = 0; k < 8; k++) ps[30 + k] = ofi_cur[k];
|
||||
}
|
||||
|
||||
/* Write OFI deltas to state vector at indices 74-81 for training replay */
|
||||
if (batch_states != NULL && state_dim >= 82) {
|
||||
float* state_out = /* output state buffer */ ;
|
||||
const float* ofi_cur = batch_states + (long long)i * state_dim + 66;
|
||||
for (int k = 0; k < 8; k++) {
|
||||
state_out[74 + k] = ofi_cur[k] - ps[30 + k]; /* delta = current - previous */
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 7: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "fix: wire OFI data pipeline + PORTFOLIO_STRIDE 30→38 + fix concat_ofi indices 42→66"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Dense micro-reward + rank normalization fix + config
|
||||
|
||||
Implement the OFI momentum × price confirmation micro-reward with adaptive cost tolerance. Fix rank normalization threshold. Enable via config.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` — micro-reward body in env_step_batch
|
||||
- Modify: `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu` — threshold 0.001→1e-5
|
||||
- Modify: `config/training/dqn-production.toml` — micro_reward_scale = 0.1
|
||||
|
||||
- [ ] **Step 1: Implement micro-reward formula in env_step_batch**
|
||||
|
||||
Replace the holding cost section (lines 1715-1721) in `experience_kernels.cu`:
|
||||
|
||||
```c
|
||||
if (!segment_complete && fabsf(position) > 0.001f && micro_reward_scale > 0.0f
|
||||
&& batch_states != NULL && state_dim >= 74) {
|
||||
const float* ofi_cur = batch_states + (long long)i * state_dim + 66;
|
||||
float sign_pos = (position > 0.0f) ? 1.0f : -1.0f;
|
||||
|
||||
/* OFI momentum: delta from previous bar (stored in ps[30..37]) */
|
||||
float delta_depth = ofi_cur[1] - ps[31];
|
||||
float delta_vpin = ofi_cur[6] - ps[36];
|
||||
float delta_queue = ofi_cur[2] - ps[32];
|
||||
float delta_tarr = ofi_cur[5] - ps[35];
|
||||
float flow_momentum = sign_pos * 0.25f * (delta_depth + delta_vpin + delta_queue + delta_tarr);
|
||||
|
||||
/* Price confirmation: intra-bar move from MBP-10 open/mid prices */
|
||||
float raw_open = tgt[4];
|
||||
float mid_open = tgt[5];
|
||||
float bar_move = (raw_close - raw_open) / fmaxf(atr_abs, 1e-6f);
|
||||
float price_confirm = sign_pos * bar_move;
|
||||
|
||||
/* Adaptive cost tolerance: profitable models trade more freely */
|
||||
float spread_at_entry = fabsf(raw_open - mid_open);
|
||||
float capital_ratio = equity / fmaxf(initial_capital, 1.0f);
|
||||
float sharpe_ema = ps[DSR_A_SLOT];
|
||||
float cost_tolerance = fmaxf(capital_ratio * fmaxf(sharpe_ema, 0.0f), 0.1f);
|
||||
float spread_penalty = spread_at_entry / (fmaxf(atr_abs, 1e-6f) * cost_tolerance);
|
||||
|
||||
/* Informed flow intensity */
|
||||
float vol_informed = ofi_cur[6] * ofi_cur[7];
|
||||
|
||||
/* Composite quality — performance-adaptive */
|
||||
float quality = flow_momentum * (1.0f + vol_informed)
|
||||
+ 0.5f * price_confirm
|
||||
- spread_penalty;
|
||||
reward = micro_reward_scale * tanhf(quality / 3.0f);
|
||||
} else if (!segment_complete && fabsf(position) > 0.001f) {
|
||||
/* Fallback: original holding cost when micro_reward_scale=0 or no OFI */
|
||||
reward = -0.0001f * fabsf(position);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Lower rank normalization threshold**
|
||||
|
||||
In `reward_shaping_kernel.cu`, in the `reward_scatter_rank` kernel:
|
||||
```c
|
||||
int is_trade = (fabsf(sharpe) > 1e-5f) ? 1 : 0; /* was 0.001f — lower for dense micro-rewards */
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Enable micro_reward_scale in config**
|
||||
|
||||
In `config/training/dqn-production.toml`, in the `[reward]` section:
|
||||
```toml
|
||||
micro_reward_scale = 0.1
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: dense micro-reward — OFI momentum × price confirm × adaptive cost tolerance"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: OFI embed MLP forward (16→8 cuBLAS)
|
||||
|
||||
Add the learned OFI embedding MLP that compresses [raw_ofi(8) ; delta_ofi(8)] into 8-dim embedding. Runs during training forward pass via cuBLAS SGEMM.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — W_ofi params, scratch buffers, cuBLAS SGEMM call, `ofi_embed_build_input` kernel load
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` — add `ofi_embed_build_input` kernel
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — call OFI embed in forward path
|
||||
|
||||
- [ ] **Step 1: Write `ofi_embed_build_input` CUDA kernel**
|
||||
|
||||
Add to `experience_kernels.cu` (or a dedicated OFI kernel file):
|
||||
```c
|
||||
/* ofi_embed_build_input — Extract [raw_ofi(8); delta_ofi(8)] from states_buf.
|
||||
* raw OFI at state[66..74), delta OFI at state[74..82) (precomputed in env_step).
|
||||
* Grid: ceil(B/256), Block: 256.
|
||||
*/
|
||||
extern "C" __global__ void ofi_embed_build_input(
|
||||
const float* __restrict__ states, /* [B, state_dim] */
|
||||
float* __restrict__ output, /* [16, B] col-major for cuBLAS */
|
||||
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;
|
||||
for (int k = 0; k < 8; k++) {
|
||||
output[k + b * 16] = s[66 + k]; /* raw OFI */
|
||||
output[8 + k + b * 16] = s[74 + k]; /* delta OFI */
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add OFI embed buffers + params to GpuDqnTrainer**
|
||||
|
||||
In the struct:
|
||||
```rust
|
||||
ofi_embed_input_buf: CudaSlice<f32>, // [16, B] scratch
|
||||
ofi_embed_output_buf: CudaSlice<f32>, // [8, B] ofi embedding
|
||||
ofi_embed_w: CudaSlice<f32>, // [8, 16] weight matrix
|
||||
ofi_embed_b: CudaSlice<f32>, // [8] bias
|
||||
ofi_embed_grad_w: CudaSlice<f32>, // [8, 16] weight gradient
|
||||
ofi_embed_grad_b: CudaSlice<f32>, // [8] bias gradient
|
||||
ofi_embed_adam_m: CudaSlice<f32>, // [136] Adam first moment
|
||||
ofi_embed_adam_v: CudaSlice<f32>, // [136] Adam second moment
|
||||
ofi_embed_build_input_kernel: CudaFunction,
|
||||
```
|
||||
|
||||
In the constructor: allocate all buffers, Xavier-init W_ofi, zero-init bias.
|
||||
|
||||
- [ ] **Step 3: Implement `launch_ofi_embed_forward`**
|
||||
|
||||
```rust
|
||||
pub(crate) fn launch_ofi_embed_forward(&self, batch_size: usize) -> Result<(), MLError> {
|
||||
// Step 1: build input from states_buf
|
||||
let states_ptr = self.ptrs.states_buf;
|
||||
let input_ptr = self.ofi_embed_input_buf.raw_ptr();
|
||||
let output_ptr = self.ofi_embed_output_buf.raw_ptr();
|
||||
// ... launch ofi_embed_build_input kernel ...
|
||||
|
||||
// Step 2: cuBLAS SGEMM: output[8,B] = W[8,16] @ input[16,B]
|
||||
// ... cublas_forward.sgemm_f32(...) ...
|
||||
|
||||
// Step 3: add_bias_relu_f32 kernel on output
|
||||
// ... launch add_bias_relu_f32_raw(...) ...
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
Follow the exact pattern of `launch_trade_plan_forward()`.
|
||||
|
||||
- [ ] **Step 4: Wire into training forward path**
|
||||
|
||||
In `fused_training.rs`, add OFI embed call inside `submit_forward_ops_main()` delegation, after trunk forward and before Mamba2/attention:
|
||||
```rust
|
||||
self.trainer.launch_ofi_embed_forward(self.batch_size)?;
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: OFI embed MLP forward (16→8 cuBLAS) — learned order flow compression"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Mamba2 history enrichment (SH2→SH2+8)
|
||||
|
||||
Widen Mamba2 history to carry OFI embedding alongside trunk features. The temporal scan then learns sequential patterns in order flow momentum.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu` — wider update_history kernel
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — wider h_history buffer, wider W_A/W_B params, recreate GEMM descs
|
||||
|
||||
- [ ] **Step 1: Modify `mamba2_update_history` to accept ofi_embed**
|
||||
|
||||
In `mamba2_temporal_kernel.cu`, update the kernel to copy `[h_s2(SH2); ofi_embed(8)]`:
|
||||
```c
|
||||
extern "C" __global__ void mamba2_update_history(
|
||||
float* __restrict__ h_history,
|
||||
const float* __restrict__ h_s2,
|
||||
const float* __restrict__ ofi_embed, /* [B, 8] — NEW param */
|
||||
int N, int K, int sh2, int embed_dim /* embed_dim=8 — NEW param */
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total_width = sh2 + embed_dim;
|
||||
int total = N * total_width;
|
||||
if (idx >= total) return;
|
||||
int i = idx / total_width;
|
||||
int j = idx % total_width;
|
||||
/* Shift history: copy slot t to t-1 for t=1..K-1 */
|
||||
for (int t = 0; t < K - 1; t++) {
|
||||
h_history[((long long)i * K + t) * total_width + j] =
|
||||
h_history[((long long)i * K + t + 1) * total_width + j];
|
||||
}
|
||||
/* Insert current at position K-1 */
|
||||
if (j < sh2) {
|
||||
h_history[((long long)i * K + K - 1) * total_width + j] = h_s2[(long long)i * sh2 + j];
|
||||
} else {
|
||||
h_history[((long long)i * K + K - 1) * total_width + j] = ofi_embed[i * embed_dim + (j - sh2)];
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Resize h_history buffer and W_A/W_B params**
|
||||
|
||||
In `gpu_dqn_trainer.rs` constructor:
|
||||
- `mamba2_h_history`: allocate `b * K * (sh2 + 8)` instead of `b * K * sh2`
|
||||
- `mamba2_params` (W_A + W_B + W_C): W_A and W_B grow from `[sh2, STATE_D]` to `[sh2+8, STATE_D]`. W_C stays `[sh2, STATE_D]`. Total params: `2 * (sh2+8) * STATE_D + sh2 * STATE_D`
|
||||
- Recreate `mamba2_gemm_proj` GEMM descriptor with k=sh2+8 (was k=sh2)
|
||||
- Recreate `mamba2_gemm_dw` GEMM descriptor with m=sh2+8 (was m=sh2) for dW_A and dW_B
|
||||
|
||||
- [ ] **Step 3: Update mamba2_forward() to use wider h_history**
|
||||
|
||||
The projection GEMMs `W_A @ h_history_flat` and `W_B @ h_history_flat` now have k=sh2+8. The GEMM descriptors handle this after recreation. The scan kernel (`mamba2_scan_projected_fwd`) operates on pre-projected values — unchanged.
|
||||
|
||||
- [ ] **Step 4: Update mamba2_update_history() Rust launch**
|
||||
|
||||
Pass `ofi_embed_output_buf` pointer and `embed_dim=8` as new args. Update grid to `total = N * (sh2 + 8)`.
|
||||
|
||||
- [ ] **Step 5: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: Mamba2 history enrichment SH2→SH2+8 — temporal scan learns OFI momentum patterns"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Mamba2 backward — add d_h_history gradient
|
||||
|
||||
Add gradient backpropagation to Mamba2's input (h_history), enabling the OFI embed MLP to learn from the temporal pipeline.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — 2 new cuBLAS GEMMs for d_h_history, elementwise add, extract d_ofi_embed
|
||||
|
||||
- [ ] **Step 1: Add d_h_history scratch buffers**
|
||||
|
||||
```rust
|
||||
d_h_history_buf: CudaSlice<f32>, // [(SH2+8), B*K] scratch for input gradient
|
||||
d_ofi_embed_mamba2: CudaSlice<f32>, // [8, B] accumulated OFI embed gradient from Mamba2
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add 2 cuBLAS GEMMs in mamba2_backward()**
|
||||
|
||||
After the existing dW_A, dW_B, dW_C GEMMs, add:
|
||||
```
|
||||
// d_h_from_A[(SH2+8), B*K] = W_A[(SH2+8), STATE_D] @ d_gate[STATE_D, B*K]
|
||||
// d_h_from_B[(SH2+8), B*K] = W_B[(SH2+8), STATE_D] @ d_x[STATE_D, B*K]
|
||||
// d_h_history = d_h_from_A + d_h_from_B (elementwise add, or use beta=1 on second GEMM)
|
||||
```
|
||||
|
||||
Use beta=0 on first GEMM (W_A @ d_gate → d_h_history), then beta=1 on second (W_B @ d_x, accumulate into d_h_history).
|
||||
|
||||
Create a new GEMM descriptor: TRANSA=N, TRANSB=N, m=SH2+8, n=B*K, k=STATE_D.
|
||||
|
||||
- [ ] **Step 3: Extract d_ofi_embed from d_h_history**
|
||||
|
||||
Write a small kernel `extract_ofi_embed_grad` that reads the last 8 rows of d_h_history and accumulates across K timesteps:
|
||||
```c
|
||||
/* d_ofi_embed_mamba2[d, b] = sum_{k=0}^{K-1} d_h_history[(SH2+d), b*K+k] */
|
||||
```
|
||||
Grid: (ceil(B/256), 8), Block: 256. Sum-reduce over K=8 steps per sample.
|
||||
|
||||
- [ ] **Step 4: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: Mamba2 d_h_history backward — gradient flows to OFI embed MLP"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Attention enrichment + backward gradient exposure
|
||||
|
||||
Widen attention input from D=256 to D+8=264, carrying OFI embedding into the cross-feature attention mechanism. Expose d_input_scratch for gradient backprop.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs` — wider W_Q/K/V, wider saved_input, wider d_input_scratch, public accessor
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` — concatenate ofi_embed to attention input
|
||||
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` — wider backward dimensions
|
||||
|
||||
- [ ] **Step 1: Widen W_Q, W_K, W_V input dimension**
|
||||
|
||||
In `GpuAttention::new()`:
|
||||
- `total_params` changes: W_Q/K/V become `[D, D+8]` instead of `[D, D]`. Total: `3*D*(D+8) + D*D + 6*D`
|
||||
- Forward GEMM descriptors: k changes from D to D+8 for Q/K/V projections
|
||||
- `saved_input` buffer: allocate `[D+8, B]` instead of `[D, B]`
|
||||
- `d_input_scratch` buffer: allocate `[D+8, B]` instead of `[D, B]`
|
||||
|
||||
- [ ] **Step 2: Concatenate ofi_embed to attention input**
|
||||
|
||||
Before calling `attn.forward()`, build `attn_input[D+8, B] = [h_s2[D, B]; ofi_embed[8, B]]`. Write a simple concat kernel or use two cuMemcpy operations. The `saved_input` in attention stores this concatenated input for backward.
|
||||
|
||||
- [ ] **Step 3: Update backward GEMM dimensions**
|
||||
|
||||
dW_Q/K/V GEMMs: k dimension of right-hand side (saved_input) changes from D to D+8.
|
||||
d_input_scratch: output width changes from D to D+8.
|
||||
The backward GEMMs `d_input = W_Q @ d_Q + W_K @ d_K + W_V @ d_V` now output `[D+8, B]`.
|
||||
|
||||
- [ ] **Step 4: Expose d_input_scratch via public accessor**
|
||||
|
||||
Add to `GpuAttention`:
|
||||
```rust
|
||||
pub fn d_input_ofi_embed_ptr(&self) -> u64 {
|
||||
// Return pointer to the OFI embed portion of d_input_scratch: offset D*B floats
|
||||
self.d_input_scratch.raw_ptr() + (self.config.state_dim * self.config.batch_size * 4) as u64
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: attention enrichment D→D+8 — OFI embed in cross-feature attention + backward gradient exposed"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: OFI embed MLP backward + gradient accumulation
|
||||
|
||||
Wire gradient flow from Mamba2 and attention back through the OFI embed MLP. This completes the full training loop for the OFI embedding.
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — gradient accumulation, MLP backward GEMM, bias reduce, Adam update
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — orchestrate backward calls
|
||||
|
||||
- [ ] **Step 1: Accumulate gradients from both sources**
|
||||
|
||||
After Mamba2 backward and attention backward, add elementwise:
|
||||
```rust
|
||||
// d_ofi_embed[8, B] = d_ofi_embed_mamba2[8, B] + d_ofi_embed_attn[8, B]
|
||||
```
|
||||
Use `add_f32_kernel` or write to `d_ofi_embed_mamba2` with beta=1 from attention source.
|
||||
|
||||
- [ ] **Step 2: MLP backward — dW and db**
|
||||
|
||||
```rust
|
||||
// dW[8, 16] = d_ofi_embed[8, B] @ ofi_input[16, B]^T
|
||||
// cuBLAS SGEMM: TRANSA=N, TRANSB=T, m=8, n=16, k=B
|
||||
self.cublas_forward.sgemm_f32(
|
||||
&self.stream, CUBLAS_OP_N, CUBLAS_OP_T,
|
||||
8, 16, batch_size,
|
||||
1.0, d_ofi_embed_ptr, 8, ofi_input_ptr, 16,
|
||||
0.0, grad_w_ptr, 8, "ofi_embed_dW"
|
||||
)?;
|
||||
|
||||
// db[8] = sum_B(d_ofi_embed) — 2-phase shared-memory reduce
|
||||
self.launch_ofi_embed_bias_grad(batch_size)?;
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Adam update on OFI embed params**
|
||||
|
||||
Use the existing `dqn_adam_update_kernel` pattern with 136 params (128 weights + 8 bias).
|
||||
|
||||
- [ ] **Step 4: Wire into fused_training.rs**
|
||||
|
||||
In `submit_post_aux_ops()` or a new phase after attention backward:
|
||||
```rust
|
||||
// Accumulate OFI embed gradients from Mamba2 + attention
|
||||
self.trainer.accumulate_ofi_embed_gradients()?;
|
||||
// MLP backward
|
||||
self.trainer.launch_ofi_embed_backward(self.batch_size)?;
|
||||
// Adam update
|
||||
self.trainer.step_ofi_embed_adam()?;
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Build and smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "feat: OFI embed MLP backward — full gradient flow from Mamba2 + attention"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: Deploy + validate
|
||||
|
||||
Push all changes, rebuild fxcache (automatic via Argo), deploy training run, validate metrics.
|
||||
|
||||
**Files:**
|
||||
- No code changes — deploy and monitor
|
||||
|
||||
- [ ] **Step 1: Push and wait for CI**
|
||||
|
||||
```bash
|
||||
git push origin main
|
||||
# Wait for ci-pipeline to succeed
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Deploy training run**
|
||||
|
||||
```bash
|
||||
./scripts/argo-train.sh dqn --baseline --epochs 200
|
||||
```
|
||||
|
||||
The `ensure-fxcache` step will automatically detect FXCACHE_VERSION mismatch (2→3) and rebuild the cache. First run takes ~5-10 min extra for cache regeneration.
|
||||
|
||||
- [ ] **Step 3: Monitor first 10 epochs**
|
||||
|
||||
Check:
|
||||
1. `atoms=%ent/%util` — should be 40-60% (working)
|
||||
2. Trade count — should be dropping below 1M (dense signal teaches hold)
|
||||
3. Val_Sharpe — should improve from -900 baseline
|
||||
4. OFI embed grad_norm — should be non-zero (gradient flowing)
|
||||
5. Step time — should stay <50ms (no regression)
|
||||
|
||||
- [ ] **Step 4: Monitor epoch 50+ for convergence**
|
||||
|
||||
Check:
|
||||
1. Trade count < 500K (model learned selectivity)
|
||||
2. Val_Sharpe > -200 (improving generalization)
|
||||
3. Atom utilization 40-60% (stable distributional learning)
|
||||
|
||||
---
|
||||
|
||||
## Dependency Graph
|
||||
|
||||
```
|
||||
Task 1 (fxcache targets 4→6)
|
||||
↓
|
||||
Task 2 (OFI pipeline fix + PORTFOLIO_STRIDE)
|
||||
↓
|
||||
Task 3 (Dense micro-reward + rank norm fix)
|
||||
↓
|
||||
Task 4 (OFI embed MLP forward)
|
||||
↓
|
||||
Task 5 (Mamba2 history enrichment) ──→ Task 6 (Mamba2 d_h_history backward)
|
||||
↓ ↓
|
||||
Task 7 (Attention enrichment) Task 8 (OFI embed backward)
|
||||
↓ ↑
|
||||
└──────────────────────────────────────┘
|
||||
↓
|
||||
Task 9 (Deploy)
|
||||
```
|
||||
|
||||
Tasks 1-3 can be deployed and validated independently — they provide the dense micro-reward without the temporal enrichment. Tasks 4-8 add the learned embedding pipeline. Task 9 validates everything.
|
||||
Reference in New Issue
Block a user