Files
foxhunt/docs/plans/2026-02-28-gpu-pipeline-phase2b-implementation.md
jgrusewski 5c23c79a1c docs: add Phase 2b implementation plan — 10 tasks for zero-roundtrip kernel
Detailed task-by-task plan for implementing the full GPU experience
collection pipeline:

Task 1-3: CUDA kernel (Q-network, curiosity, barrier, diversity, reward)
Task 4: Weight extraction module (VarMap → CudaSlice)
Task 5: GPU experience collector (Rust wrapper + buffer management)
Task 6: Unit tests for weights and kernel compilation
Task 7: DQN trainer integration (GPU round-based collection)
Task 8: Periodic curiosity training with weight sync
Task 9: Workspace verification (2392+ tests, clippy)
Task 10: Final commit and branch completion

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-28 11:58:30 +01:00

1119 lines
40 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# GPU Pipeline Phase 2b Implementation Plan
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
**Goal:** Replace the DQN trainer's CPU inner loop with a single zero-roundtrip CUDA kernel that runs 128 parallel episodes of experience collection entirely on GPU — including Dueling Q-network inference, curiosity inference, barrier tracking, diversity entropy, and reward combination.
**Architecture:** One monolithic CUDA kernel (`dqn_full_experience_kernel`) runs N=128 parallel threads, each processing L=500 timesteps independently. Market features and Q-network weights are pre-uploaded. The kernel outputs N×L experiences downloaded in a single PCIe transfer. Weight synchronization happens after Candle gradient updates.
**Tech Stack:** CUDA (NVRTC runtime compilation), cudarc 0.17.3 (via `candle_core::cuda_backend::cudarc`), Candle (gradient updates, curiosity training)
**Design doc:** `docs/plans/2026-02-28-gpu-pipeline-phase2b-design.md`
---
### Task 1: CUDA Device Functions — Dueling Q-Network Forward Pass
**Files:**
- Create: `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu`
This is the new kernel file. Phase 2's `experience_kernels.cu` is kept for backward compatibility. The new kernel is a complete replacement.
**Step 1: Write the Q-network device functions**
Create `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu` with the Dueling Q-network forward pass and helper functions:
```cuda
/**
* DQN Full Experience Collection Kernel
*
* Zero-roundtrip CUDA kernel for DQN training experience collection.
* Each thread runs an independent episode: Q-network inference,
* epsilon-greedy action selection, portfolio simulation, barrier tracking,
* diversity entropy, curiosity inference, and reward combination.
*
* Architecture: N threads × L timesteps, one kernel launch per training round.
*/
// ============================================================
// Constants
// ============================================================
#define STATE_DIM 54
#define MARKET_DIM 51
#define PORTFOLIO_DIM 3
#define NUM_ACTIONS 45
#define SHARED_H1 256
#define SHARED_H2 256
#define VALUE_H 128
#define ADV_H 128
#define CUR_INPUT 35
#define CUR_HIDDEN 64
#define CUR_OUTPUT 32
#define DIVERSITY_WINDOW 100
#define PORTFOLIO_STATE_SIZE 8
#define BARRIER_STATE_SIZE 5
// ============================================================
// RNG — Linear Congruential Generator per thread
// ============================================================
__device__ __forceinline__ float gpu_random(unsigned int* state) {
*state = *state * 1664525u + 1013904223u;
return (float)(*state & 0x7FFFFF) / (float)0x7FFFFF;
}
// ============================================================
// Activation: LeakyReLU(x, alpha=0.01)
// ============================================================
__device__ __forceinline__ float leaky_relu(float x) {
return (x > 0.0f) ? x : 0.01f * x;
}
// ============================================================
// Matrix-vector multiply: out[out_dim] = W[out_dim, in_dim] * in[in_dim] + b[out_dim]
// Then apply LeakyReLU if activate=true.
// Weight layout: Candle convention [out_features, in_features] (row-major).
// ============================================================
__device__ void matvec_leaky_relu(
const float* W, // [out_dim, in_dim]
const float* b, // [out_dim]
const float* input, // [in_dim]
float* output, // [out_dim]
int in_dim,
int out_dim,
bool activate
) {
for (int j = 0; j < out_dim; j++) {
float sum = b[j];
for (int i = 0; i < in_dim; i++) {
sum += input[i] * W[j * in_dim + i];
}
output[j] = activate ? leaky_relu(sum) : sum;
}
}
// ============================================================
// Dueling Q-Network forward pass
//
// Architecture:
// Shared: [54] → [256] LReLU → [256] LReLU
// Value: [256] → [128] LReLU → [1]
// Advantage: [256] → [128] LReLU → [45]
// Q(s,a) = V(s) + A(s,a) - mean(A)
//
// Weight pointers point into contiguous GPU weight buffer.
// ============================================================
__device__ void q_forward_dueling(
const float* state, // [STATE_DIM]
// Shared pathway weights
const float* sw1, const float* sb1, // [SHARED_H1, STATE_DIM], [SHARED_H1]
const float* sw2, const float* sb2, // [SHARED_H2, SHARED_H1], [SHARED_H2]
// Value stream weights
const float* vw, const float* vb, // [VALUE_H, SHARED_H2], [VALUE_H]
const float* vow, const float* vob, // [1, VALUE_H], [1]
// Advantage stream weights
const float* aw, const float* ab, // [ADV_H, SHARED_H2], [ADV_H]
const float* aow, const float* aob, // [NUM_ACTIONS, ADV_H], [NUM_ACTIONS]
// Scratch space (caller-allocated)
float* h1, // [SHARED_H1]
float* h2, // [SHARED_H2]
float* vh, // [VALUE_H]
float* ah, // [ADV_H]
// Output
float* q_values // [NUM_ACTIONS]
) {
// Shared pathway
matvec_leaky_relu(sw1, sb1, state, h1, STATE_DIM, SHARED_H1, true);
matvec_leaky_relu(sw2, sb2, h1, h2, SHARED_H1, SHARED_H2, true);
// Value stream → scalar
matvec_leaky_relu(vw, vb, h2, vh, SHARED_H2, VALUE_H, true);
float value;
matvec_leaky_relu(vow, vob, vh, &value, VALUE_H, 1, false);
// Advantage stream → [NUM_ACTIONS]
matvec_leaky_relu(aw, ab, h2, ah, SHARED_H2, ADV_H, true);
matvec_leaky_relu(aow, aob, ah, q_values, ADV_H, NUM_ACTIONS, false);
// Compute mean advantage
float adv_mean = 0.0f;
for (int i = 0; i < NUM_ACTIONS; i++) adv_mean += q_values[i];
adv_mean /= (float)NUM_ACTIONS;
// Q(s,a) = V(s) + A(s,a) - mean(A)
for (int i = 0; i < NUM_ACTIONS; i++) {
q_values[i] = value + q_values[i] - adv_mean;
}
}
// ============================================================
// Argmax over NUM_ACTIONS Q-values
// ============================================================
__device__ int argmax_q(const float* q_values) {
int best = 0;
float best_val = q_values[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
if (q_values[i] > best_val) {
best_val = q_values[i];
best = i;
}
}
return best;
}
```
**Step 2: Verify kernel compiles (syntax check)**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | head -5`
This won't compile the .cu file yet (it's not included). Just ensure the Rust crate still compiles after adding the file.
**Step 3: Commit**
```bash
git add crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
git commit -m "feat(cuda): add Dueling Q-network device functions for experience kernel"
```
---
### Task 2: CUDA Device Functions — Portfolio, Barrier, Diversity, Curiosity, Reward
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu`
Append the remaining device functions to the kernel file.
**Step 1: Add portfolio simulation helpers**
These already exist in `experience_kernels.cu`. Copy and adapt `action_to_exposure` and `action_to_tx_cost` into the new file. Then add the barrier, diversity, curiosity, and reward functions:
```cuda
// ============================================================
// Action helpers (from Phase 2 kernel — copied for self-containment)
// ============================================================
__device__ __forceinline__ float action_to_exposure(int action_idx) {
int exposure_level = action_idx / 9;
switch (exposure_level) {
case 0: return -1.0f;
case 1: return -0.5f;
case 2: return 0.0f;
case 3: return 0.5f;
case 4: return 1.0f;
default: return 0.0f;
}
}
__device__ __forceinline__ float action_to_tx_cost(int action_idx) {
int order_type = (action_idx / 3) % 3;
switch (order_type) {
case 0: return 0.0001f;
case 1: return 0.00005f;
case 2: return 0.00015f;
default: return 0.0001f;
}
}
// ============================================================
// Barrier tracking
//
// barrier_state[5]: entry_price, upper_barrier, lower_barrier, expiry_step, result
// barrier_config[4]: profit_target_bps, stop_loss_bps, max_holding_bars, reward_scale
// result: 0=none, 1=profit, -1=stop_loss, 2=time_expiry
// ============================================================
__device__ void barrier_init(
float* barrier_state,
const float* barrier_config,
float entry_price,
int current_step
) {
float profit_bps = barrier_config[0]; // e.g., 100 = 1%
float stop_bps = barrier_config[1]; // e.g., 50 = 0.5%
float max_bars = barrier_config[2]; // e.g., 500
barrier_state[0] = entry_price;
barrier_state[1] = entry_price * (1.0f + profit_bps / 10000.0f); // upper
barrier_state[2] = entry_price * (1.0f - stop_bps / 10000.0f); // lower
barrier_state[3] = (float)(current_step + (int)max_bars); // expiry
barrier_state[4] = 0.0f; // no result yet
}
__device__ float barrier_check(
float* barrier_state,
float price,
int current_step,
float position
) {
// Only check if we have a position and no result yet
if (fabsf(position) < 0.001f || barrier_state[4] != 0.0f) return 0.0f;
if (price >= barrier_state[1]) {
barrier_state[4] = 1.0f; // profit target hit
return 1.0f;
}
if (price <= barrier_state[2]) {
barrier_state[4] = -1.0f; // stop loss hit
return -1.0f;
}
if ((float)current_step >= barrier_state[3]) {
barrier_state[4] = 2.0f; // time expiry
return 0.0f; // neutral for time
}
return 0.0f; // no barrier hit
}
__device__ void barrier_reset(float* barrier_state) {
barrier_state[0] = 0.0f;
barrier_state[1] = 0.0f;
barrier_state[2] = 0.0f;
barrier_state[3] = 0.0f;
barrier_state[4] = 0.0f;
}
// ============================================================
// Action diversity — Shannon entropy over sliding window
//
// diversity_window[DIVERSITY_WINDOW]: circular buffer of action categories (0-2)
// diversity_meta[2]: write_cursor, count
// Maps 45-action to 3 categories: 0=Short, 1=Flat, 2=Long
// ============================================================
__device__ float diversity_entropy(
int* diversity_window,
int* diversity_meta,
int action_idx
) {
// Map to 3 categories
int exposure = action_idx / 9;
int category = (exposure <= 1) ? 0 : (exposure == 2) ? 1 : 2;
// Insert into circular buffer
int cursor = diversity_meta[0];
int count = diversity_meta[1];
diversity_window[cursor] = category;
diversity_meta[0] = (cursor + 1) % DIVERSITY_WINDOW;
if (count < DIVERSITY_WINDOW) diversity_meta[1] = count + 1;
// Compute histogram
int n = diversity_meta[1];
if (n < 3) return 0.0f; // Not enough data
int hist[3] = {0, 0, 0};
for (int i = 0; i < n; i++) {
int c = diversity_window[i];
if (c >= 0 && c < 3) hist[c]++;
}
// Shannon entropy: H = -sum(p * log2(p))
float entropy = 0.0f;
float fn = (float)n;
for (int c = 0; c < 3; c++) {
if (hist[c] > 0) {
float p = (float)hist[c] / fn;
entropy -= p * log2f(p);
}
}
// Penalty if entropy < 1.0
return (entropy < 1.0f) ? -0.1f : 0.0f;
}
// ============================================================
// Curiosity forward model — inference only
//
// Architecture: [35] → [64] LeakyReLU → [32]
// Input: [32 state features (first 32 of STATE_DIM)] + [3 action one-hot]
// Output: predicted next state embedding [32]
// Return: clamped MSE between predicted and actual next state
// ============================================================
__device__ float curiosity_inference(
const float* state, // [STATE_DIM] — use first 32 dims
const float* next_state, // [STATE_DIM] — use first 32 dims
int action_idx,
const float* cw1, const float* cb1, // [CUR_HIDDEN, CUR_INPUT], [CUR_HIDDEN]
const float* cw2, const float* cb2, // [CUR_OUTPUT, CUR_HIDDEN], [CUR_OUTPUT]
float* cur_h, // [CUR_HIDDEN] scratch
float max_reward
) {
// Build input: [32 state features, 3 one-hot action]
float input[CUR_INPUT];
for (int i = 0; i < 32; i++) input[i] = state[i];
// One-hot action encoding
int exposure = action_idx / 9;
int cat = (exposure <= 1) ? 0 : (exposure == 2) ? 1 : 2;
input[32] = (cat == 0) ? 1.0f : 0.0f; // Short
input[33] = (cat == 1) ? 1.0f : 0.0f; // Flat
input[34] = (cat == 2) ? 1.0f : 0.0f; // Long
// Forward pass: input → hidden
matvec_leaky_relu(cw1, cb1, input, cur_h, CUR_INPUT, CUR_HIDDEN, true);
// Hidden → prediction [32]
float predicted[CUR_OUTPUT];
matvec_leaky_relu(cw2, cb2, cur_h, predicted, CUR_HIDDEN, CUR_OUTPUT, false);
// MSE between predicted and actual next state (first 32 dims)
float mse = 0.0f;
for (int i = 0; i < CUR_OUTPUT; i++) {
float diff = predicted[i] - next_state[i];
mse += diff * diff;
}
mse /= (float)CUR_OUTPUT;
return fminf(mse, max_reward);
}
```
**Step 2: Commit**
```bash
git add crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
git commit -m "feat(cuda): add barrier, diversity, curiosity, portfolio device functions"
```
---
### Task 3: CUDA Main Kernel — Parallel Episode Loop
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu`
Append the main kernel function that ties all device functions together.
**Step 1: Write the main kernel**
Append to `dqn_experience_kernel.cu`:
```cuda
// ============================================================
// Main kernel: DQN full experience collection
//
// Launch config: grid=(ceil(N/32), 1, 1), block=(32, 1, 1)
// Each thread processes one independent episode of L timesteps.
// ============================================================
extern "C" __global__ void dqn_full_experience_kernel(
// Market data (pre-uploaded, read-only)
const float* __restrict__ market_features, // [total_bars, MARKET_DIM]
const float* __restrict__ targets, // [total_bars, 4]
const int* __restrict__ episode_starts, // [N]
// Online Q-network weights (read-only, all Candle [out,in] layout)
const float* sw1, const float* sb1, // shared layer 1
const float* sw2, const float* sb2, // shared layer 2
const float* vw, const float* vb, // value fc
const float* vow, const float* vob, // value out
const float* aw, const float* ab, // advantage fc
const float* aow, const float* aob, // advantage out
// Target Q-network weights (read-only, same layout)
const float* tsw1, const float* tsb1,
const float* tsw2, const float* tsb2,
const float* tvw, const float* tvb,
const float* tvow, const float* tvob,
const float* taw, const float* tab,
const float* taow, const float* taob,
// Curiosity weights (read-only)
const float* cw1, const float* cb1,
const float* cw2, const float* cb2,
// Per-episode state (IN/OUT) — indexed by tid
float* portfolio_states, // [N, PORTFOLIO_STATE_SIZE]
float* barrier_states, // [N, BARRIER_STATE_SIZE]
int* diversity_windows, // [N, DIVERSITY_WINDOW]
int* diversity_meta, // [N, 2]
// Config (read-only)
const float* barrier_config, // [4]
float epsilon,
float max_position,
int episode_length,
int total_bars,
int L, // timesteps to simulate
float gamma, // discount factor
float curiosity_max_reward,
int N, // number of parallel episodes
// RNG (IN/OUT)
unsigned int* rng_states, // [N]
// Outputs — [N, L] layout (thread-major)
float* states_out, // [N * L * STATE_DIM]
int* actions_out, // [N * L]
float* rewards_out, // [N * L]
int* done_out, // [N * L]
float* target_q_out, // [N * L]
float* td_error_out // [N * L]
) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= N) return;
// ---- Per-thread scratch (local memory) ----
float h1[SHARED_H1], h2[SHARED_H2];
float vh[VALUE_H], ah[ADV_H];
float q_values[NUM_ACTIONS];
float tgt_q[NUM_ACTIONS];
float cur_h[CUR_HIDDEN];
float state[STATE_DIM];
float next_state[STATE_DIM];
// ---- Load per-episode state from global memory ----
int ps_off = tid * PORTFOLIO_STATE_SIZE;
float cash = portfolio_states[ps_off + 0];
float position = portfolio_states[ps_off + 1];
float entry_price = portfolio_states[ps_off + 2];
float initial_cap = portfolio_states[ps_off + 3];
float spread = portfolio_states[ps_off + 4];
float last_price = portfolio_states[ps_off + 5];
float reserve_pct = portfolio_states[ps_off + 6];
float cum_costs = portfolio_states[ps_off + 7];
int bs_off = tid * BARRIER_STATE_SIZE;
float local_barrier[BARRIER_STATE_SIZE];
for (int i = 0; i < BARRIER_STATE_SIZE; i++)
local_barrier[i] = barrier_states[bs_off + i];
int dw_off = tid * DIVERSITY_WINDOW;
int local_div_window[DIVERSITY_WINDOW];
for (int i = 0; i < DIVERSITY_WINDOW; i++)
local_div_window[i] = diversity_windows[dw_off + i];
int dm_off = tid * 2;
int local_div_meta[2];
local_div_meta[0] = diversity_meta[dm_off + 0];
local_div_meta[1] = diversity_meta[dm_off + 1];
unsigned int rng = rng_states[tid];
int start_bar = episode_starts[tid];
int step_in_episode = 0;
// ============================================================
// Main loop: L timesteps per episode
// ============================================================
for (int t = 0; t < L; t++) {
int global_bar = start_bar + t;
if (global_bar >= total_bars - 1) {
// Out of data — mark remaining as done with zero reward
int out_idx = tid * L + t;
for (int d = 0; d < STATE_DIM; d++) states_out[out_idx * STATE_DIM + d] = 0.0f;
actions_out[out_idx] = 0;
rewards_out[out_idx] = 0.0f;
done_out[out_idx] = 1;
target_q_out[out_idx] = 0.0f;
td_error_out[out_idx] = 0.0f;
continue;
}
// ---- 1. Read market features ----
int mf_off = global_bar * MARKET_DIM;
for (int i = 0; i < MARKET_DIM; i++)
state[i] = market_features[mf_off + i];
// ---- 2. Compute portfolio features ----
int t_off = global_bar * 4;
float current_close_raw = targets[t_off + 2];
float current_close = targets[t_off + 0];
float price = (current_close_raw != 0.0f) ? current_close_raw : current_close;
if (price <= 0.0f) price = 1.0f;
float current_value = cash + position * price;
float current_norm = current_value / initial_cap;
float max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f;
float pos_norm = position / max_pos_norm;
state[MARKET_DIM + 0] = current_norm; // normalized portfolio value
state[MARKET_DIM + 1] = pos_norm; // normalized position
state[MARKET_DIM + 2] = spread; // spread
// ---- 3. Online Q-network forward ----
q_forward_dueling(state,
sw1, sb1, sw2, sb2, vw, vb, vow, vob, aw, ab, aow, aob,
h1, h2, vh, ah, q_values);
// ---- 4. Epsilon-greedy action selection ----
int action;
if (gpu_random(&rng) < epsilon) {
action = (int)(gpu_random(&rng) * (float)NUM_ACTIONS) % NUM_ACTIONS;
} else {
action = argmax_q(q_values);
}
float online_q_selected = q_values[action];
// ---- 5. Portfolio simulation (from Phase 2 kernel) ----
float target_exposure = action_to_exposure(action);
float target_position = target_exposure * max_position;
float tx_rate = action_to_tx_cost(action);
bool is_reversal = (position > 0.0f && target_position < 0.0f) ||
(position < 0.0f && target_position > 0.0f);
if (is_reversal) {
float close_cash = position * price;
float close_cost = fabsf(position) * price * tx_rate;
cash += close_cash - close_cost;
cum_costs += close_cost;
float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f;
float affordable = fmaxf(cash - reserve, 0.0f);
float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f;
max_contracts = floorf(max_contracts);
float actual = fminf(max_contracts, fabsf(target_position));
if (actual > 0.0f) {
float new_pos = (target_position > 0.0f) ? actual : -actual;
float open_cost = actual * price * tx_rate;
cash -= new_pos * price + open_cost;
cum_costs += open_cost;
position = new_pos;
entry_price = price;
// Init barrier for new position
barrier_init(local_barrier, barrier_config, price, global_bar);
} else {
position = 0.0f;
entry_price = 0.0f;
barrier_reset(local_barrier);
}
} else {
float delta = target_position - position;
if (fabsf(delta) > 0.0f) {
float trade_cost = fabsf(delta) * price * tx_rate;
cum_costs += trade_cost;
cash -= trade_cost;
if (delta > 0.0f && reserve_pct > 0.0f) {
float pv = cash + position * price;
float reserve = pv * (reserve_pct / 100.0f);
float buy_cost = delta * price;
if (cash - buy_cost < reserve) {
float afford = fmaxf(cash - reserve, 0.0f);
delta = fminf(delta, floorf(afford / price));
}
}
if (delta > 0.0f) {
entry_price = price;
// Init barrier if opening new position
if (fabsf(position) < 0.001f) {
barrier_init(local_barrier, barrier_config, price, global_bar);
}
} else if (target_position == 0.0f) {
entry_price = 0.0f;
barrier_reset(local_barrier);
}
cash -= delta * price;
position += delta;
}
}
last_price = price;
// ---- 6. Barrier tracking ----
float barrier_label = barrier_check(local_barrier, price, global_bar, position);
float barrier_scale = 0.0f;
float b_reward_scale = barrier_config[3]; // reward_scale factor
if (barrier_label == 1.0f) barrier_scale = 0.5f * b_reward_scale;
else if (barrier_label == -1.0f) barrier_scale = -0.5f * b_reward_scale;
// ---- 7. Diversity entropy ----
float diversity_bonus = diversity_entropy(local_div_window, local_div_meta, action);
// ---- 8. Mark-to-market at next price ----
float next_close_raw = targets[t_off + 3];
float next_close = targets[t_off + 1];
float next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close;
if (next_price <= 0.0f) next_price = price;
float next_value = cash + position * next_price;
float next_norm = next_value / initial_cap;
// Build next_state for curiosity
int next_mf_off = (global_bar + 1) * MARKET_DIM;
for (int i = 0; i < MARKET_DIM; i++)
next_state[i] = market_features[next_mf_off + i];
float next_max_pos = (next_price > 0.0f) ? initial_cap / next_price : 1.0f;
next_state[MARKET_DIM + 0] = next_norm;
next_state[MARKET_DIM + 1] = position / next_max_pos;
next_state[MARKET_DIM + 2] = spread;
// ---- 9. Curiosity inference ----
float curiosity_bonus = curiosity_inference(
state, next_state, action,
cw1, cb1, cw2, cb2, cur_h,
curiosity_max_reward
);
// ---- 10. Reward combination ----
float raw_pnl = 0.0f;
if (current_norm > 0.0f) {
raw_pnl = (next_norm - current_norm) / current_norm;
}
float reward = raw_pnl * (1.0f + barrier_scale);
// Risk penalty
float abs_pos = fabsf(pos_norm);
if (abs_pos > 0.8f) {
reward -= (abs_pos - 0.8f) * 5.0f * 0.1f;
}
reward += diversity_bonus;
reward += curiosity_bonus;
// ---- 11. Target Q-network → TD error ----
q_forward_dueling(next_state,
tsw1, tsb1, tsw2, tsb2, tvw, tvb, tvow, tvob, taw, tab, taow, taob,
h1, h2, vh, ah, tgt_q);
float max_tgt_q = tgt_q[0];
for (int i = 1; i < NUM_ACTIONS; i++)
if (tgt_q[i] > max_tgt_q) max_tgt_q = tgt_q[i];
// ---- 12. Episode boundary ----
step_in_episode++;
int time_done = (step_in_episode >= episode_length) ? 1 : 0;
int barrier_done = (barrier_label != 0.0f) ? 1 : 0;
int data_done = (global_bar + 1 >= total_bars - 1) ? 1 : 0;
int done = (time_done || barrier_done || data_done) ? 1 : 0;
// TD target and error
float td_target = reward + gamma * max_tgt_q * (float)(1 - done);
float td_error = fabsf(td_target - online_q_selected);
// ---- 13. Write experience ----
int out_idx = tid * L + t;
for (int d = 0; d < STATE_DIM; d++)
states_out[out_idx * STATE_DIM + d] = state[d];
actions_out[out_idx] = action;
rewards_out[out_idx] = reward;
done_out[out_idx] = done;
target_q_out[out_idx] = max_tgt_q;
td_error_out[out_idx] = td_error;
// ---- 14. Episode reset ----
if (done) {
cash = initial_cap;
position = 0.0f;
entry_price = 0.0f;
cum_costs = 0.0f;
barrier_reset(local_barrier);
// Reset diversity window
for (int i = 0; i < DIVERSITY_WINDOW; i++) local_div_window[i] = 0;
local_div_meta[0] = 0;
local_div_meta[1] = 0;
step_in_episode = 0;
}
}
// ---- Write back per-episode state ----
portfolio_states[ps_off + 0] = cash;
portfolio_states[ps_off + 1] = position;
portfolio_states[ps_off + 2] = entry_price;
portfolio_states[ps_off + 3] = initial_cap;
portfolio_states[ps_off + 4] = spread;
portfolio_states[ps_off + 5] = last_price;
portfolio_states[ps_off + 6] = reserve_pct;
portfolio_states[ps_off + 7] = cum_costs;
for (int i = 0; i < BARRIER_STATE_SIZE; i++)
barrier_states[bs_off + i] = local_barrier[i];
for (int i = 0; i < DIVERSITY_WINDOW; i++)
diversity_windows[dw_off + i] = local_div_window[i];
diversity_meta[dm_off + 0] = local_div_meta[0];
diversity_meta[dm_off + 1] = local_div_meta[1];
rng_states[tid] = rng;
}
```
**Step 2: Commit**
```bash
git add crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
git commit -m "feat(cuda): add main dqn_full_experience_kernel with parallel episodes"
```
---
### Task 4: Rust Weight Extraction Module
**Files:**
- Create: `crates/ml/src/cuda_pipeline/gpu_weights.rs`
- Modify: `crates/ml/src/cuda_pipeline/mod.rs` (add module declaration)
This module handles extracting weights from Candle VarMap and uploading them as flat GPU buffers.
**Step 1: Write the weight extraction module**
Create `crates/ml/src/cuda_pipeline/gpu_weights.rs` with:
1. `DuelingWeightSet` struct — holds 12 `CudaSlice<f32>` buffers (6 weight + 6 bias for shared/value/advantage)
2. `CuriosityWeightSet` struct — holds 4 `CudaSlice<f32>` buffers
3. `extract_dueling_weights(vars: &VarMap, stream: &Arc<CudaStream>) -> Result<DuelingWeightSet>` — extracts from VarMap key paths: `shared_0.weight`, `shared_0.bias`, `shared_1.weight`, `shared_1.bias`, `value_fc.weight`, `value_fc.bias`, `value_out.weight`, `value_out.bias`, `advantage_fc.weight`, `advantage_fc.bias`, `advantage_out.weight`, `advantage_out.bias`
4. `extract_curiosity_weights(vars: &VarMap, stream: &Arc<CudaStream>) -> Result<CuriosityWeightSet>` — extracts `fc1.weight`, `fc1.bias`, `fc2.weight`, `fc2.bias`
The extraction pattern for each tensor:
```rust
let vars_data = vars.data().lock().map_err(|e| MLError::ModelError(format!("Lock: {e}")))?;
let tensor = vars_data.get("shared_0.weight")
.ok_or_else(|| MLError::ModelError("Missing shared_0.weight".into()))?
.as_tensor();
let flat: Vec<f32> = tensor.flatten_all()?.to_vec1::<f32>()?;
let mut buf = stream.alloc_zeros::<f32>(flat.len()).map_err(...)?;
stream.memcpy_htod(&flat, &mut buf).map_err(...)?;
```
**Step 2: Add module declaration to mod.rs**
In `crates/ml/src/cuda_pipeline/mod.rs`, add after the `gpu_portfolio` module declaration:
```rust
#[cfg(feature = "cuda")]
pub mod gpu_weights;
```
**Step 3: Run compilation check**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5`
Expected: Clean compile
**Step 4: Commit**
```bash
git add crates/ml/src/cuda_pipeline/gpu_weights.rs crates/ml/src/cuda_pipeline/mod.rs
git commit -m "feat(cuda): add weight extraction module for Q-network and curiosity GPU upload"
```
---
### Task 5: Rust GPU Experience Collector
**Files:**
- Create: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
- Modify: `crates/ml/src/cuda_pipeline/mod.rs` (add module declaration)
This is the main Rust wrapper that compiles the new kernel, manages all GPU buffers, and provides a `collect_experiences()` method.
**Step 1: Write the experience collector struct and constructor**
Create `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` with:
1. Constants: `MAX_EPISODES = 256`, `MAX_TIMESTEPS = 1000`, `STATE_DIM = 54`
2. `GpuExperienceCollector` struct — holds:
- `stream: Arc<CudaStream>`, `kernel_func: CudaFunction`
- `online_weights: DuelingWeightSet`, `target_weights: DuelingWeightSet`
- `curiosity_weights: CuriosityWeightSet`
- Per-episode state buffers: `portfolio_states: CudaSlice<f32>`, `barrier_states: CudaSlice<f32>`, `diversity_windows: CudaSlice<i32>`, `diversity_meta: CudaSlice<i32>`
- Output buffers: `states_out`, `actions_out`, `rewards_out`, `done_out`, `target_q_out`, `td_error_out`
- Config: `episode_starts: CudaSlice<i32>`, `barrier_config: CudaSlice<f32>`, `rng_states: CudaSlice<u32>`
3. `new()` — compiles kernel via `include_str!("dqn_experience_kernel.cu")`, allocates all buffers
4. `ExperienceCollectorConfig` struct — holds `max_position`, `episode_length`, `total_bars`, `epsilon`, `gamma`, `curiosity_max_reward`, `n_episodes`, `timesteps_per_episode`
**Step 2: Write the `collect_experiences` method**
```rust
pub fn collect_experiences(
&mut self,
market_features_buf: &CudaSlice<f32>,
targets_buf: &CudaSlice<f32>,
episode_starts: &[i32],
config: &ExperienceCollectorConfig,
) -> Result<ExperienceBatch, MLError>
```
This method:
1. Uploads `episode_starts` to GPU
2. Computes launch config: `grid_dim = (ceil(N/32), 1, 1)`, `block_dim = (32, 1, 1)`
3. Launches kernel with all 40+ arguments using `stream.launch_builder(&func).arg(...)` chain
4. Downloads all output buffers
5. Returns `ExperienceBatch` struct with `states`, `actions`, `rewards`, `done_flags`, `target_q_values`, `td_errors`, `n_episodes`, `timesteps`
**Step 3: Write weight sync methods**
```rust
pub fn sync_online_weights(&mut self, vars: &VarMap) -> Result<(), MLError>
pub fn sync_target_weights(&mut self, vars: &VarMap) -> Result<(), MLError>
pub fn sync_curiosity_weights(&mut self, vars: &VarMap) -> Result<(), MLError>
```
Each extracts weights from VarMap and uploads to the corresponding GPU buffers.
**Step 4: Add module declaration**
In `crates/ml/src/cuda_pipeline/mod.rs`:
```rust
#[cfg(feature = "cuda")]
pub mod gpu_experience_collector;
```
**Step 5: Run compilation check**
Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5`
Expected: Clean compile
**Step 6: Commit**
```bash
git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/cuda_pipeline/mod.rs
git commit -m "feat(cuda): add GpuExperienceCollector with zero-roundtrip kernel launch"
```
---
### Task 6: Unit Tests for Weight Extraction and Kernel Compilation
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/mod.rs` (add tests)
- Modify: `crates/ml/src/cuda_pipeline/gpu_weights.rs` (add tests)
**Step 1: Write weight extraction test**
In `gpu_weights.rs`, add a `#[cfg(test)]` module with:
```rust
#[test]
fn test_dueling_weight_key_paths() {
// Create a DuelingQNetwork with known config
// Verify all 12 VarMap keys exist
// Verify weight tensor shapes match expected dimensions
}
```
This test creates a DuelingQNetwork on CPU device, then checks that all VarMap keys (`shared_0.weight`, `shared_0.bias`, etc.) exist with correct shapes.
**Step 2: Write kernel compilation test**
In `gpu_experience_collector.rs` or `mod.rs`, add:
```rust
#[test]
fn test_dqn_experience_kernel_compiles() {
// Verify NVRTC can compile the new kernel
let src = include_str!("dqn_experience_kernel.cu");
assert!(src.contains("dqn_full_experience_kernel"));
// On CPU device: just verify source is loadable
// On CUDA device: verify NVRTC compilation succeeds
}
```
**Step 3: Write experience collector config test**
```rust
#[test]
fn test_experience_batch_sizing() {
let config = ExperienceCollectorConfig {
n_episodes: 128,
timesteps_per_episode: 500,
..Default::default()
};
assert_eq!(config.total_experiences(), 64_000);
}
```
**Step 4: Run tests**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib cuda_pipeline -j8 2>&1 | tail -10`
Expected: All tests pass (10 existing + new tests)
**Step 5: Commit**
```bash
git add crates/ml/src/cuda_pipeline/
git commit -m "test(cuda): add weight extraction and kernel compilation tests"
```
---
### Task 7: DQN Trainer Integration — GPU Experience Collection Round
**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer.rs`
This is the critical integration: replacing the CPU inner loop with GPU round-based experience collection.
**Step 1: Add GPU experience collector field to DQNTrainer**
Add to the struct definition (behind `#[cfg(feature = "cuda")]`):
```rust
#[cfg(feature = "cuda")]
gpu_experience_collector: Option<crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector>,
```
Initialize to `None` in the constructor.
**Step 2: Initialize GPU experience collector after data upload**
In `train_with_data_full_loop`, after the existing GPU data upload section (Phase 1b), add initialization of the `GpuExperienceCollector`:
1. Get the `CudaStream` from candle's device
2. Extract online/target Q-network weights from VarMap
3. Extract curiosity weights
4. Create `GpuExperienceCollector::new(stream, online_weights, target_weights, curiosity_weights, config)`
5. Store in `self.gpu_experience_collector`
**Step 3: Add GPU experience collection round method**
Add a new method `collect_gpu_experiences`:
```rust
#[cfg(feature = "cuda")]
fn collect_gpu_experiences(
&mut self,
total_bars: usize,
epoch: usize,
) -> Result<usize, MLError>
```
This method:
1. Generates N random episode start indices using `fastrand`
2. Calls `self.gpu_experience_collector.collect_experiences()`
3. Iterates over the returned `ExperienceBatch`
4. For each experience `(state, action, reward, done, target_q, td_error)`:
- Creates a `TradingState` from the state vector
- Stores in replay buffer via `self.store_experience()`
5. Returns the number of experiences collected
**Step 4: Wire GPU round into training loop**
In `train_with_data_full_loop`, replace the Phase 1 batch loop with:
```rust
#[cfg(feature = "cuda")]
if self.gpu_experience_collector.is_some() {
// GPU path: N rounds of parallel episode collection
let num_rounds = (total_bars + n_experiences_per_round - 1) / n_experiences_per_round;
for round in 0..num_rounds {
let n_collected = self.collect_gpu_experiences(total_bars, epoch)?;
total_experiences += n_collected;
}
} else {
// CPU fallback: original per-bar inner loop
// ... existing code ...
}
```
**Step 5: Add weight sync after gradient updates**
After the Phase 2 gradient update loop, add:
```rust
#[cfg(feature = "cuda")]
if let Some(ref mut collector) = self.gpu_experience_collector {
collector.sync_online_weights(&self.q_network_vars())?;
if should_sync_target {
collector.sync_target_weights(&self.target_network_vars())?;
}
}
```
**Step 6: Run DQN tests**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib dqn -j8 -- --test-threads=1 2>&1 | tail -10`
Expected: 427+ tests pass (existing tests use CPU fallback path since test device is CPU)
**Step 7: Commit**
```bash
git add crates/ml/src/trainers/dqn/trainer.rs
git commit -m "feat(dqn): integrate GPU experience collection with parallel episodes"
```
---
### Task 8: Periodic Curiosity Training
**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer.rs`
**Step 1: Add curiosity training round**
After the gradient update phase, add periodic curiosity training:
```rust
// Every CURIOSITY_TRAIN_INTERVAL rounds, train curiosity model on Candle
if round_idx % CURIOSITY_TRAIN_INTERVAL == 0 {
if let Some(ref curiosity) = self.curiosity_module {
// Sample batch from replay buffer
let batch = self.replay_buffer.sample(1024)?;
// Train curiosity model via Candle backward + optimizer step
curiosity.train_batch(&batch)?;
// Sync updated weights to GPU
#[cfg(feature = "cuda")]
if let Some(ref mut collector) = self.gpu_experience_collector {
collector.sync_curiosity_weights(&curiosity.vars())?;
}
}
}
```
**Step 2: Run DQN tests**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib dqn -j8 -- --test-threads=1 2>&1 | tail -10`
Expected: 427+ tests pass
**Step 3: Commit**
```bash
git add crates/ml/src/trainers/dqn/trainer.rs
git commit -m "feat(dqn): add periodic curiosity training with GPU weight sync"
```
---
### Task 9: Workspace Verification
**Files:** None (verification only)
**Step 1: Full workspace compilation**
Run: `SQLX_OFFLINE=true cargo check --workspace -j8 2>&1 | tail -5`
Expected: `Finished` with no errors
**Step 2: Clippy on ml crate**
Run: `SQLX_OFFLINE=true cargo clippy -p ml --lib -- -D warnings 2>&1 | tail -5`
Expected: Zero warnings
**Step 3: cuda_pipeline tests**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib cuda_pipeline -j8 2>&1 | tail -10`
Expected: All tests pass (10 existing + new tests from Task 6)
**Step 4: Full DQN test suite**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib dqn -j8 -- --test-threads=1 2>&1 | tail -10`
Expected: 427+ tests pass
**Step 5: Full ML test suite**
Run: `SQLX_OFFLINE=true cargo test -p ml --lib -j8 2>&1 | tail -5`
Expected: 2392+ tests pass, 0 failures
**Step 6: Commit any fixups**
If any issues found, fix and commit.
---
### Task 10: Final Commit and Summary
**Step 1: Review all changes**
Run: `git log --oneline feature/cuda-pipeline-phase2..HEAD`
Verify all commits are clean and descriptive.
**Step 2: Run final verification**
Run: `SQLX_OFFLINE=true cargo check --workspace -j8 && SQLX_OFFLINE=true cargo clippy -p ml --lib -- -D warnings`
Expected: Clean
**Step 3: Summary of files created/modified**
New files:
- `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu` — Full CUDA kernel (~500 lines)
- `crates/ml/src/cuda_pipeline/gpu_weights.rs` — Weight extraction and GPU upload
- `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` — Rust wrapper for kernel launch
Modified files:
- `crates/ml/src/cuda_pipeline/mod.rs` — Module declarations
- `crates/ml/src/trainers/dqn/trainer.rs` — GPU experience collection integration
**Step 4: Use superpowers:finishing-a-development-branch to complete**