foxhunt/docs/plans/2026-02-28-ppo-cuda-pipeline-design.md

PPO CUDA Pipeline Design — Zero-Roundtrip Experience Collection (Phase 2c)

Goal

Build a single monolithic CUDA kernel that runs the entire PPO experience collection pipeline — actor forward pass, softmax categorical sampling, critic forward pass, portfolio simulation, barrier tracking, diversity entropy, curiosity inference, reward combination, and GAE advantage computation — with zero CPU-GPU roundtrips per timestep. 128 parallel episodes × 500 timesteps = 64,000 training-ready experience tuples per kernel launch.

Architecture

Single entry point ppo_full_experience_kernel in ppo_experience_kernel.cu. Shared device functions (RNG, matmul, portfolio, barriers, diversity, curiosity) extracted from the DQN kernel into common_device_functions.cuh and reused by both kernels via source concatenation at NVRTC compile time.

Thread mapping: 1 thread = 1 episode, N=128 threads default. Each thread processes L=500 timesteps sequentially in two phases: forward rollout (collect experiences) then backward GAE scan (compute advantages and returns).

Kernel Phases

Phase A — Forward Rollout (L timesteps per thread)

For each timestep t:

1. Read 51 market features from global memory
2. Compute 3 portfolio features from simulation state (normalized value, normalized position, spread)
3. Actor forward: state[54] → h1[128] (ReLU) → h2[64] (ReLU) → logits[45]
4. Softmax + categorical sample: logits → exp → normalize → CDF scan → LCG random draw → action index + log(p[action])
5. Critic forward: state[54] → h1[512] (ReLU) → h2[384] (ReLU) → h3[256] (ReLU) → h4[128] (ReLU) → h5[64] (ReLU) → value[1]
6. Portfolio simulation (reused from DQN): action → exposure mapping, transaction costs, cash/position tracking, reversals
7. Barrier tracking (reused): profit-take, stop-loss, time-based exit detection
8. Diversity entropy (reused): action window penalty when entropy < 1.0
9. Curiosity inference (reused): forward dynamics model prediction error
10. Reward combination: PnL × (1 + barrier_scale) + diversity_penalty + curiosity_bonus - risk_penalty
11. Store per-timestep: state[54], action, log_prob, value, reward, done
12. Episode reset on done (reinitialize portfolio, barriers, diversity window)

Phase B — GAE Backward Scan (after all L timesteps)

Each thread scans backward over its own episode's stored rewards, values, and dones:

advantages[L-1] = delta[L-1]  // where delta = reward + gamma * V(next) * (1-done) - V(current)
for t = L-2 down to 0:
    delta[t] = rewards[t] + gamma * values[t+1] * (1 - dones[t]) - values[t]
    advantages[t] = delta[t] + gamma * lambda * (1 - dones[t]) * advantages[t+1]
returns[t] = advantages[t] + values[t]

Write advantages and returns to output buffers. Values buffer is NOT downloaded — only needed for GAE.

Shared Device Functions Header

Extract from dqn_experience_kernel.cu into common_device_functions.cuh:

Function	Purpose
gpu_random	LCG RNG, returns float in [0, 1)
leaky_relu	Activation with configurable alpha
matvec_leaky_relu	Matrix-vector multiply + bias + activation (use_leaky=0 for ReLU)
action_to_exposure	45-action index → position exposure float
action_to_tx_cost	45-action index → transaction cost rate
barrier_init/check/reset	Triple barrier state management
diversity_entropy	Action diversity penalty computation
curiosity_inference	Forward dynamics model inference

New PPO-specific functions in ppo_experience_kernel.cu:

Function	Purpose
ppo_actor_forward	MLP: 54→128(ReLU)→64(ReLU)→45 raw logits
softmax_sample	Stable softmax → CDF scan → categorical sample → (action, log_prob)
ppo_critic_forward	MLP: 54→512(ReLU)→384(ReLU)→256(ReLU)→128(ReLU)→64(ReLU)→1
compute_gae_backward	Backward scan: rewards + values + dones → advantages + returns

NVRTC compilation via source concatenation in Rust:

let common = include_str!("common_device_functions.cuh");
let kernel = include_str!("ppo_experience_kernel.cu");
let full_source = format!("{}\n{}", common, kernel);

DQN kernel refactored to use the same pattern.

Weight Layout

Actor (6 tensors, 18,093 parameters)

VarMap Key	Shape	CUDA name
policy_layer_0.weight	[128, 54]	pw1
policy_layer_0.bias	[128]	pb1
policy_layer_1.weight	[64, 128]	pw2
policy_layer_1.bias	[64]	pb2
policy_output.weight	[45, 64]	pw3
policy_output.bias	[45]	pb3

Critic (12 tensors, 346,369 parameters)

VarMap Key	Shape	CUDA name
value_layer_0.weight	[512, 54]	vw1
value_layer_0.bias	[512]	vb1
value_layer_1.weight	[384, 512]	vw2
value_layer_1.bias	[384]	vb2
value_layer_2.weight	[256, 384]	vw3
value_layer_2.bias	[256]	vb3
value_layer_3.weight	[128, 256]	vw4
value_layer_3.bias	[128]	vb4
value_layer_4.weight	[64, 128]	vw5
value_layer_4.bias	[64]	vb5
value_output.weight	[1, 64]	vw6
value_output.bias	[1]	vb6

Curiosity (4 tensors, reuses existing CuriosityWeightSet)

VarMap Key	Shape	CUDA name
fc1.weight	[64, 35]	cw1
fc1.bias	[64]	cb1
fc2.weight	[32, 64]	cw2
fc2.bias	[32]	cb2

Per-Thread Memory Budget

Array	Size	Bytes
state[54]	54 floats	216
logits[45] + probs[45]	90 floats	360
Actor scratch: h1[128], h2[64]	192 floats	768
Critic scratch: scratch_a[512], scratch_b[512] (ping-pong)	1,024 floats	4,096
next_state[54]	54 floats	216
Curiosity scratch: cur_h[64]	64 floats	256
Portfolio state: 8 floats	8 floats	32
Barrier state: 5 floats	5 floats	20
Diversity window: 100 ints	100 ints	400
GAE buffers: values_buf[L], rewards_buf[L], dones_buf[L]	3 × 500	6,000
Total per thread		~12.4 KB

At 128 threads: ~1.6 MB total local memory. Well within H100/L40S capacity.

Critic uses ping-pong pattern: two 512-wide scratch buffers alternate between layers, avoiding 5 separate arrays.

Output Buffers

Buffer	Shape	Type	Downloaded
states	[N × L × 54]	f32	Yes
actions	[N × L]	i32	Yes
old_log_probs	[N × L]	f32	Yes
advantages	[N × L]	f32	Yes
returns	[N × L]	f32	Yes
dones	[N × L]	i32	Yes
values	[N × L]	f32	No (only used for in-kernel GAE)

Total download per launch: N × L × (54 + 1 + 1 + 1 + 1 + 1) × 4 bytes = 128 × 500 × 59 × 4 = ~14.4 MB.

Trainer Integration

1. Add gpu_ppo_collector: Option<GpuPpoExperienceCollector> to PpoTrainer behind #[cfg(feature = "cuda")]
2. Initialize when CUDA device present and curiosity module available
3. GPU path in collect_rollouts(): kernel launch → download → build TrajectoryBatch
4. CPU fallback: existing code unchanged
5. Weight sync after update_mlp(): actor + critic + curiosity

Training flow:

GPU kernel launch (N=128 episodes × L=500 timesteps)
        ↓
Download 6 buffers (~14.4 MB)
        ↓
Build TrajectoryBatch from flat arrays
        ↓
update_mlp() [existing Candle code, unchanged]:
  For each PPO epoch (10):
    For each mini-batch (512):
      Recompute log_probs via Candle actor forward
      ratio = exp(new_log_probs - old_log_probs)
      PPO clipped surrogate loss + critic MSE loss
      Backward + optimizer step
        ↓
Sync actor + critic + curiosity weights → GPU

File Layout

crates/ml/src/cuda_pipeline/
├── common_device_functions.cuh      # NEW: extracted shared device functions
├── dqn_experience_kernel.cu         # MODIFIED: uses common header via concatenation
├── ppo_experience_kernel.cu         # NEW: PPO kernel entry point + actor/critic/GAE
├── gpu_weights.rs                   # MODIFIED: add PpoActorWeightSet, PpoCriticWeightSet
├── gpu_experience_collector.rs      # EXISTING: DQN collector (unchanged)
├── gpu_ppo_collector.rs             # NEW: PPO collector wrapper
├── gpu_portfolio.rs                 # EXISTING (unchanged)
└── mod.rs                           # MODIFIED: add new module declarations

crates/ml/src/trainers/ppo.rs        # MODIFIED: GPU collector integration
crates/ml/src/dqn/curiosity.rs       # ALREADY DONE: forward_model_vars() exists from Phase 2b

Testing Strategy

• CPU-only unit tests: actor/critic VarMap key paths, param counts, kernel source verification, config defaults
• DQN regression: extract shared header must not break existing DQN kernel compilation or 14 cuda_pipeline tests
• PPO trainer regression: 125 existing PPO tests pass with CPU fallback
• Full ML suite: 2396+ tests pass