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

40 KiB
Raw Blame History

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:

  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:

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:

  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

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

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:

  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:

#[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:

#[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 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