foxhunt/docs/superpowers/plans/2026-03-24-remaining-gpu-features.md

Remaining GPU Features — Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Wire and implement the 4 remaining unwired GPU features: multi-head attention integration, Decision Transformer CUDA kernels, ensemble training with diversity loss, and HER Future strategy with GPU episode boundaries.

Architecture: Each feature integrates into the existing fused CUDA training pipeline (GpuDqnTrainer + CUDA Graph). All kernels are GPU-native (zero CPU in hot path). Post-graph operations use the EventTrackingGuard RAII pattern for safe cudarc buffer access. Kernels are compiled via nvcc → cubin (cached at /tmp/.cubin_cache/).

CUDA Graph Integration Strategy: The CUDA Graph captures the full forward+loss+backward+Adam sequence. Features that modify intermediate activations (attention, ensemble heads) operate with a 1-step lag: they modify save_h_s2 AFTER graph replay, so the modification takes effect on the NEXT graph replay's unflatten step. This is the same pattern as the IQN trunk gradient and spectral normalization — standard in async gradient methods. The alternative (splitting the graph into trunk/head phases) is deferred until profiling shows the lag impacts convergence.

VRAM Budget:

Feature	RTX 3050 (4GB)	H100 (80GB)
Attention (params + activations)	~0.5MB	~2MB
Decision Transformer (params + context)	~4MB	~60MB
Ensemble (K=3 head copies)	~1.5MB	~6MB
HER episode tracking	~0.5MB	~4MB
Total	~6.5MB	~72MB

Tech Stack: Rust (cudarc 0.19.3), CUDA C (nvcc/cubin), cuBLAS SGEMM

---

File Structure

New Files

File	Responsibility
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu	Attention backward pass: LayerNorm backward + output projection backward + attention weight gradients
crates/ml/src/cuda_pipeline/dt_kernels.cu	Decision Transformer forward: causal attention, softmax, positional encoding, token embedding, FFN, cross-entropy loss. Backward: softmax_backward, attention_backward (dQ/dK/dV), ffn_backward, layernorm_backward, embed_backward — 10 kernels total
crates/ml/src/cuda_pipeline/ensemble_kernels.cu	Ensemble: multi-head Q-aggregation, KL-divergence diversity loss, Thompson sampling action selection
crates/ml/src/cuda_pipeline/her_episode_kernel.cu	HER: episode boundary tracking, Future strategy donor sampling, episode-end detection

Modified Files

File	Changes
crates/ml/src/cuda_pipeline/gpu_attention.rs	Add backward(), params_buf(), Adam state; wire into GpuDqnTrainer
crates/ml/src/cuda_pipeline/batched_forward.rs	Add attention call between h_s2 and value/advantage heads
crates/ml/src/cuda_pipeline/batched_backward.rs	Add attention backward in the gradient chain
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs	Conditional attention in forward/backward; ensemble head management
crates/ml/src/cuda_pipeline/decision_transformer.rs	Implement pretrain_step() with real CUDA kernels
crates/ml/src/cuda_pipeline/gpu_her.rs	Add Final/Future strategies, episode boundary buffer
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs	Add episode_ids to GpuExperienceBatch; fill via CUDA kernel (not CPU)
crates/ml/src/trainers/dqn/fused_training.rs	Wire attention, ensemble, HER strategies
crates/ml/src/trainers/dqn/config.rs	Add use_attention: bool to DQNHyperparameters (field exists in DQNParams hyperopt struct but not in the training config); wire from DQNParams during hyperopt conversion
crates/ml/src/trainers/dqn/trainer/training_loop.rs	Ensemble head rotation, DT pre-training phase, HER episode tracking
crates/ml/src/hyperopt/adapters/dqn.rs	Add DT pre-training epochs, ensemble params to search space

---

Task 1: Attention Forward Pass Integration

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

• Step 1: Add GpuAttention as optional field in FusedTrainingCtx

In fused_training.rs, add to FusedTrainingCtx:

pub(crate) gpu_attention: Option<GpuAttention>,

Initialize it in new() when hyperparams.use_attention is true (similar to gpu_iqn initialization pattern at line 233).

• Step 2: Wire attention forward BETWEEN trunk h_s2 and value/advantage heads

In gpu_dqn_trainer.rs, add a method:

pub fn apply_attention_forward(
    &self,
    attention: &mut GpuAttention,
    h_s2: &CudaSlice<f32>,
    batch_size: usize,
) -> Result<(), MLError>

This copies h_s2 into attention input, runs attention.forward(), then copies the attended output BACK into save_h_s2 (in-place replacement). The CUDA Graph captures the trunk forward FIRST, then attention runs OUTSIDE the graph (same pattern as IQN/spectral norm), then value/advantage heads run.

• Step 3: Call attention in run_full_step() after CUDA Graph replay

After Step 2 (graph replay) and before Step 5 (IQN):

if let Some(ref mut attn) = self.gpu_attention {
    let _evt_guard = EventTrackingGuard::new(self.stream.context());
    self.trainer.apply_attention_forward(attn, ...)?;
}

• Step 4: Compile check

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5

• Step 5: Commit

git add crates/ml/src/cuda_pipeline/gpu_attention.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: wire attention forward pass into DQN training (post-graph, frozen weights)"

---

Task 2: Attention Backward Pass Kernel

Files:

• Create: crates/ml/src/cuda_pipeline/attention_backward_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs

• Step 1: Write the attention backward CUDA kernel

The kernel computes dL/d(h_s2_input) given dL/d(h_s2_attended):

extern "C" __global__ void attention_backward_kernel(
    const float* d_output,      // [B, D] gradient from downstream
    const float* states_input,  // [B, D] saved input to attention
    const float* params,        // attention weights
    float* d_input,             // [B, D] gradient to upstream (trunk)
    float* d_params,            // [total_params] weight gradients (atomicAdd)
    int B
);

The backward flows through:

1. LayerNorm backward → d_prenorm
2. Residual: d_prenorm splits into d_projection + d_residual
3. Output projection backward: d_projection → d_concat, dW_O, db_O
4. Multi-head attention backward → dQ, dK, dV per head
5. Q/K/V projection backward → d_input, dW_Q, dW_K, dW_V

Since the kernel is complex (matching the 237-line forward kernel), implement in phases:

• Phase A: Residual passthrough only (d_input = d_output) — gradient flows, weights frozen
• Phase B: Full backward with weight gradients (unfreezes attention weights)

Start with Phase A (trivial kernel, verifiable).

• Step 2: Add backward() method to GpuAttention

pub fn backward(
    &self,
    d_output: &CudaSlice<f32>,
    d_input: &mut CudaSlice<f32>,
    batch_size: usize,
) -> Result<(), MLError>

For Phase A: simply memcpy_dtod(d_input, d_output) — residual gradient passthrough.

• Step 3: Wire backward into apply_iqn_trunk_gradient

After attention forward rewrites save_h_s2, the IQN trunk gradient flows through the attention backward before reaching the trunk layers:

// Before trunk backward: d_h_s2 → attention_backward → d_h_s2_input
if let Some(ref attn) = self.gpu_attention {
    attn.backward(&bw_d_h_s2, &mut bw_d_h_s2, batch_size)?;
}

• Step 4: Compile check + local test

SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5

• Step 5: Commit

git add crates/ml/src/cuda_pipeline/attention_backward_kernel.cu crates/ml/src/cuda_pipeline/gpu_attention.rs
git commit -m "feat: attention backward (Phase A — residual passthrough, gradient flow verified)"

---

Task 3: Decision Transformer CUDA Kernels

Files:

• Create: crates/ml/src/cuda_pipeline/dt_kernels.cu

• Modify: crates/ml/src/cuda_pipeline/decision_transformer.rs

• Step 1: Write the token embedding + positional encoding kernel

extern "C" __global__ void dt_embed_kernel(
    const float* trajectories,  // [B, T, state_dim+2] (return-to-go, state, action)
    const float* W_embed,       // [(state_dim+2), embed_dim]
    const float* b_embed,       // [embed_dim]
    const float* pos_embed,     // [T, embed_dim] (learned positional encoding)
    float* output,              // [B, T, embed_dim]
    int B, int T, int input_dim, int embed_dim
);

Each thread handles one (batch, timestep, embed_dim_idx) tuple:

• Linear projection: out[b][t][d] = sum_j(traj[b][t][j] * W[j][d]) + b[d]

• Add positional: out[b][t][d] += pos[t][d]

• Step 2: Write the causal self-attention kernel

extern "C" __global__ void dt_causal_attention_kernel(
    const float* Q,         // [B, T, embed_dim]
    const float* K,         // [B, T, embed_dim]
    const float* V,         // [B, T, embed_dim]
    float* output,          // [B, T, embed_dim]
    int B, int T, int D, int num_heads
);

Per-head computation with causal mask:

• attn[i][j] = Q[i] · K[j] / sqrt(D_h) for j ≤ i (causal), -inf otherwise
• Softmax over j dimension
• output[i] = sum_j(attn[i][j] * V[j])

Block-per-sample, threads handle different timesteps. Shared memory for K/V caching.

• Step 3: Write the FFN kernel (embed_dim → 4×embed_dim → embed_dim)

extern "C" __global__ void dt_ffn_kernel(
    const float* input,     // [B, T, embed_dim]
    const float* W1,        // [embed_dim, 4*embed_dim]
    const float* b1,        // [4*embed_dim]
    const float* W2,        // [4*embed_dim, embed_dim]
    const float* b2,        // [embed_dim]
    float* output,          // [B, T, embed_dim]
    int B, int T, int D
);

GELU activation between layers. Residual connection: output += input.

• Step 4: Write the cross-entropy loss kernel

extern "C" __global__ void dt_cross_entropy_kernel(
    const float* logits,        // [B, T, num_actions]
    const int* target_actions,  // [B, T]
    float* per_sample_loss,     // [B*T]
    float* total_loss,          // [1] atomicAdd
    int B, int T, int num_actions
);

Per-token cross-entropy: loss = -log(softmax(logits)[target_action]).

• Step 5: Implement pretrain_step() in decision_transformer.rs

Wire the 5 forward kernels into a training loop. Then:

• Step 5b: Write 5 backward CUDA kernels

// In dt_kernels.cu — add after forward kernels:
// dt_cross_entropy_backward_kernel: dL/d_logits = softmax - one_hot
// dt_ffn_backward_kernel: GELU backward + 2 linear backward
// dt_attention_backward_kernel: dQ,dK,dV from d_output (softmax Jacobian + causal mask)
// dt_layernorm_backward_kernel: chain through mean/variance
// dt_embed_backward_kernel: dW_embed from d_embedded

Each mirrors its forward kernel's thread/block layout. Total: 10 kernels.

• Step 5c: Implement pretrain_step() in decision_transformer.rs

1. Forward: Embed → {QKV → CausalAttn → LN → FFN → LN} × N → ActionHead → CE loss
2. Backward: CE backward → reverse chain through all N layers → embed backward
3. Optimizer: Reuse dqn_grad_norm_kernel + dqn_adam_update_kernel on DT params

DT action space = 9 exposure actions (branch_0 only, not full 81 factored).

• Step 6: Compile check

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5

• Step 7: Commit

git add crates/ml/src/cuda_pipeline/dt_kernels.cu crates/ml/src/cuda_pipeline/decision_transformer.rs
git commit -m "feat: Decision Transformer CUDA kernels (embed, causal attention, FFN, CE loss)"

---

Task 4: Decision Transformer Pre-Training Loop

Files:

• Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs

• Modify: crates/ml/src/trainers/dqn/config.rs

• Modify: crates/ml/src/hyperopt/adapters/dqn.rs

• Step 1: Add DT pre-training config fields

In DQNHyperparameters:

pub dt_pretrain_epochs: usize,   // 0 = disabled, 10-50 typical
pub dt_context_len: usize,       // 20 timesteps
pub dt_embed_dim: usize,         // 128
pub dt_num_layers: usize,        // 3
pub dt_target_return: f64,       // 2.0 (target Sharpe-like return)

• Step 2: Add DT pre-training phase at training start

In training_loop.rs, before the main training loop:

if self.hyperparams.dt_pretrain_epochs > 0 {
    let dt = DecisionTransformer::new(stream, dt_config)?;
    for epoch in 0..self.hyperparams.dt_pretrain_epochs {
        // Build trajectory batches from training_data
        // Run dt.pretrain_step(trajectories, target_actions, batch_size)
        // Log DT loss
    }
    // Transfer DT's learned representations to DQN trunk (optional)
}

• Step 3: Add DT params to hyperopt search space

In dqn.rs, add bounds for dt_pretrain_epochs, dt_context_len, dt_embed_dim.

• Step 4: Compile check + commit

SQLX_OFFLINE=true cargo check --workspace
git add crates/ml/src/cuda_pipeline/ crates/ml/src/trainers/dqn/ crates/ml/src/hyperopt/ && git commit -m "feat: Decision Transformer pre-training phase in training loop"

---

Task 5: Ensemble Multi-Head Q-Network

Files:

• Create: crates/ml/src/cuda_pipeline/ensemble_kernels.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

• Modify: crates/ml/src/trainers/dqn/config.rs

• Step 1: Write the ensemble aggregation kernel

extern "C" __global__ void ensemble_aggregate_kernel(
    const float* head_q_values,  // [K, B, num_actions] Q-values from K heads
    float* mean_q,               // [B, num_actions] mean across heads
    float* var_q,                // [B, num_actions] variance across heads
    int K, int B, int num_actions
);

Computes mean and variance of Q-values across K ensemble heads for uncertainty estimation.

• Step 2: Write the KL-divergence diversity loss kernel

extern "C" __global__ void ensemble_diversity_kernel(
    const float* head_logits,    // [K, B, num_atoms] per-head C51 logits
    float* diversity_loss,       // [1] total KL divergence (atomicAdd)
    int K, int B, int num_atoms
);

For each pair of heads (i, j): D_KL(p_i || p_j) averaged over the batch. Uses hierarchical reduction (warp reduce → block reduce → atomicAdd per block, same pattern as dqn_grad_norm_kernel) to avoid single-address atomicAdd serialization. The diversity loss ENCOURAGES disagreement (subtracted from total loss).

• Step 3: Implement multi-head architecture in GpuDqnTrainer

The ensemble shares the TRUNK (W_s1, W_s2) but has K SEPARATE value + advantage heads:

pub struct EnsembleHeads {
    heads: Vec<(DuelingWeightSet, BranchingWeightSet)>,  // K head weight sets
}

During training (heads run OUTSIDE the CUDA Graph, same pattern as IQN/attention):

1. CUDA Graph: forward through shared trunk → h_s2 (using head 0's weights in the graph)
2. Post-graph: for each head k=1..K-1, run cuBLAS forward through head_k's value/advantage layers on save_h_s2 (outside graph, with EventTrackingGuard)
3. Aggregate Q-values: mean ± std across K heads
4. C51 loss per head (head 0 inside graph, heads 1..K-1 outside)
5. Diversity regularization across heads
6. Backward: each head contributes trunk gradients via the existing IQN-style SGD/Adam correction

Note: The CUDA Graph remains captured for head 0 only. Heads 1..K-1 run as post-graph operations with their own cuBLAS forwards. This avoids re-capturing the graph for each head.

• Step 4: Add ensemble config and wiring

In DQNHyperparameters:

pub ensemble_count: usize,           // K heads (default 1 = no ensemble)
pub ensemble_diversity_weight: f64,  // 0.01 = mild diversity encouragement

In FusedTrainingCtx::run_full_step():

• If ensemble_count > 1: run K forward passes through different heads, average for action selection, sum losses with diversity term.

• Step 5: Compile check + commit

SQLX_OFFLINE=true cargo check --workspace
git add crates/ml/src/cuda_pipeline/ crates/ml/src/trainers/dqn/ crates/ml/src/hyperopt/ && git commit -m "feat: ensemble multi-head Q-network with KL diversity loss"

---

Task 6: HER Episode Boundary Tracking

Files:

• Create: crates/ml/src/cuda_pipeline/her_episode_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_her.rs

• Modify: crates/ml/src/cuda_pipeline/gpu_experience_collector.rs

• Step 1: Add episode index buffer to experience output

In GpuExperienceBatch, add:

pub episode_ids: CudaSlice<i32>,  // [N * L] episode index per transition

Fill via a simple CUDA kernel fill_episode_ids_kernel(int* ids, int L, int N) where thread i writes ids[i] = i / L. Launched as grid=(ceil(N*L/256)), block=256. Zero CPU fill — pure GPU.

• Step 2: Track episode boundaries in GPU replay buffer

When inserting transitions, also store the episode ID:

pub episode_id_buf: CudaSlice<i32>,  // [capacity] episode index per slot

This allows the HER kernel to find other transitions from the same episode.

• Step 3: Write the episode-aware donor sampling kernel

extern "C" __global__ void her_sample_future_donors(
    const int* episode_ids,        // [capacity] episode index per slot
    const int* source_indices,     // [her_batch_size] which to relabel
    int* donor_indices,            // [her_batch_size] output: future donors
    unsigned int* rng_states,      // [her_batch_size] per-sample RNG
    int capacity,
    int her_batch_size
);

Uses a GPU-resident episode_boundary_offsets auxiliary buffer (updated at insertion time) for O(log(N)) binary search instead of O(capacity) linear scan. Each thread binary-searches the boundary buffer to find its episode's [start, end) range, then uniformly samples from [source+1, end).

• Step 4: Write the episode-end detection kernel (Final strategy)

extern "C" __global__ void her_find_episode_end(
    const int* episode_ids,        // [capacity]
    const int* source_indices,     // [her_batch_size]
    int* end_indices,              // [her_batch_size] output: last index in same episode
    int capacity,
    int her_batch_size
);

For each source: binary search or linear scan to find the last index with the same episode_id.

• Step 5: Wire Future/Final strategies in GpuHer

pub fn relabel_batch_gpu(
    &mut self,
    strategy: HerGpuStrategy,
    episode_ids: &CudaSlice<i32>,
    ...
) -> Result<&HerBatch, MLError> {
    let donors = match strategy {
        HerGpuStrategy::Random => self.random_donors(her_batch_size),
        HerGpuStrategy::Future => self.sample_future_donors(episode_ids, source_indices)?,
        HerGpuStrategy::Final  => self.find_episode_ends(episode_ids, source_indices)?,
    };
    self.relabel_kernel(source_indices, donors, ...)?;
    Ok(&self.output)
}

• Step 6: Compile check + commit

SQLX_OFFLINE=true cargo check --workspace
git add crates/ml/src/cuda_pipeline/ crates/ml/src/trainers/dqn/ crates/ml/src/hyperopt/ && git commit -m "feat: HER Future/Final strategies with GPU episode boundary tracking"

---

Task 7: Integration Test — Full Pipeline

Files:

• Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs

• Step 1: Wire all 4 features into the training loop conditionally

// Attention: after trunk forward, before heads
if hyperparams.use_attention { ... }

// DT pre-training: before main training loop
if hyperparams.dt_pretrain_epochs > 0 { ... }

// Ensemble: K heads with diversity loss
if hyperparams.ensemble_count > 1 { ... }

// HER: Future strategy with episode boundaries
if hyperparams.her_ratio > 0.0 && hyperparams.her_strategy == "future" { ... }

• Step 2: Local test — all features enabled

rm -rf /tmp/.cubin_cache/
SQLX_OFFLINE=true cargo build --release -p ml --example hyperopt_baseline_rl
SQLX_OFFLINE=true timeout 120 ./target/release/examples/hyperopt_baseline_rl \
  --model dqn --phase fast --trials 1 --n-initial 1 --epochs 3 \
  --data-dir test_data/futures-baseline --symbol ES.FUT \
  --output /tmp/full_pipeline_test.json --seed 42

Expected: No NaN, no CUDA errors, train_loss < 50.

• Step 3: Commit

git add crates/ml/src/cuda_pipeline/ crates/ml/src/trainers/dqn/ crates/ml/src/hyperopt/ && git commit -m "feat: full pipeline integration — attention + DT + ensemble + HER Future"

---

Task 8: H100 Submission

• Step 1: Push to main

git push origin main

• Step 2: Submit H100 hyperopt run

50 epochs, 20 trials, with all features enabled via config:

use_attention = true
dt_pretrain_epochs = 10
ensemble_count = 3
her_ratio = 0.3
her_strategy = "future"

• Step 3: Monitor first trial

Check: no NaN, no CUDA OOM, loss decreasing, Q-gap > 0.

---

Dependency Graph

Task 1 (Attention Forward) → Task 2 (Attention Backward) → Task 7 (Integration)
Task 3 (DT Kernels) → Task 4 (DT Loop) → Task 7 (Integration)
Task 5 (Ensemble) → Task 7 (Integration)
Task 6 (HER Episodes) → Task 7 (Integration)
Task 7 (Integration) → Task 8 (H100)

Tasks 1-2, 3-4, 5, and 6 are INDEPENDENT and can be implemented in parallel.