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>
40 KiB
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:
/**
* 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
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:
// ============================================================
// 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
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:
// ============================================================
// 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
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:
DuelingWeightSetstruct — holds 12CudaSlice<f32>buffers (6 weight + 6 bias for shared/value/advantage)CuriosityWeightSetstruct — holds 4CudaSlice<f32>buffersextract_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.biasextract_curiosity_weights(vars: &VarMap, stream: &Arc<CudaStream>) -> Result<CuriosityWeightSet>— extractsfc1.weight,fc1.bias,fc2.weight,fc2.bias
The extraction pattern for each tensor:
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:
#[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
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:
- Constants:
MAX_EPISODES = 256,MAX_TIMESTEPS = 1000,STATE_DIM = 54 GpuExperienceCollectorstruct — holds:stream: Arc<CudaStream>,kernel_func: CudaFunctiononline_weights: DuelingWeightSet,target_weights: DuelingWeightSetcuriosity_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>
new()— compiles kernel viainclude_str!("dqn_experience_kernel.cu"), allocates all buffersExperienceCollectorConfigstruct — holdsmax_position,episode_length,total_bars,epsilon,gamma,curiosity_max_reward,n_episodes,timesteps_per_episode
Step 2: Write the collect_experiences method
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:
- Uploads
episode_startsto GPU - Computes launch config:
grid_dim = (ceil(N/32), 1, 1),block_dim = (32, 1, 1) - Launches kernel with all 40+ arguments using
stream.launch_builder(&func).arg(...)chain - Downloads all output buffers
- Returns
ExperienceBatchstruct withstates,actions,rewards,done_flags,target_q_values,td_errors,n_episodes,timesteps
Step 3: Write weight sync methods
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:
#[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
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:
#[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:
#[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
#[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
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")]):
#[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:
- Get the
CudaStreamfrom candle's device - Extract online/target Q-network weights from VarMap
- Extract curiosity weights
- Create
GpuExperienceCollector::new(stream, online_weights, target_weights, curiosity_weights, config) - Store in
self.gpu_experience_collector
Step 3: Add GPU experience collection round method
Add a new method collect_gpu_experiences:
#[cfg(feature = "cuda")]
fn collect_gpu_experiences(
&mut self,
total_bars: usize,
epoch: usize,
) -> Result<usize, MLError>
This method:
- Generates N random episode start indices using
fastrand - Calls
self.gpu_experience_collector.collect_experiences() - Iterates over the returned
ExperienceBatch - For each experience
(state, action, reward, done, target_q, td_error):- Creates a
TradingStatefrom the state vector - Stores in replay buffer via
self.store_experience()
- Creates a
- 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:
#[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:
#[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
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:
// 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
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 uploadcrates/ml/src/cuda_pipeline/gpu_experience_collector.rs— Rust wrapper for kernel launch
Modified files:
crates/ml/src/cuda_pipeline/mod.rs— Module declarationscrates/ml/src/trainers/dqn/trainer.rs— GPU experience collection integration
Step 4: Use superpowers:finishing-a-development-branch to complete