foxhunt/docs/superpowers/specs/2026-03-10-zero-cpu-hotpath-design.md

Zero-CPU DQN Training Hot Path — Design Spec

Goal: Eliminate all GPU→CPU synchronization barriers (to_scalar(), to_vec1(), to_vec2(), memcpy_dtoh) from the DQN training hot path so the GPU pipeline never stalls between backward() and optimizer.step().

Architecture: Four new GPU components replace 10 CPU readback sites with on-device predicates, pinned-memory polling, and GPU-resident accumulators. All monitoring/safety checks (NaN detection, loss clipping, gradient collapse, Q-value divergence) execute as CUDA kernels writing boolean halt flags to pinned host memory — zero cudaStreamSynchronize barriers per training step.

Tech Stack: CUDA (NVRTC compiled), cudarc 0.17, Candle v0.9.1 (Rust ML framework), pinned host memory via cudarc::driver::CudaHostBuf

---

Readback Inventory (10 sites to eliminate)

#	File:Line	Pattern	Frequency	Replacement
1	trainer.rs:4531	stacked.to_vec1::<f32>()	Per train step	GpuTrainingGuard::check()
2	gradient_utils.rs:108	total_sum.to_scalar::<f32>()	Per accumulation boundary	GpuTrainingGuard NaN predicate
3	trainer.rs:4739	mean_all().to_scalar()	Per accumulation boundary	GpuTrainingGuard::read_accumulators()
4	trainer.rs:4751	mean_all().to_scalar()	Per accumulation boundary	GpuTrainingGuard::read_accumulators()
5	trainer.rs:4860	mean_all().to_scalar::<f32>()	Every 50 steps	GPU Q-value accumulator
6	dqn.rs:3618	to_vec2::<f32>()	Every 50 steps	GPU Q-value stats reduction
7	trainer.rs:4113	to_vec1::<u32>()	Per experience batch	GPU-resident action tensor
8	trainer.rs:4302-4303	to_vec1::<u32>()	Per experience batch	GPU-resident action tensor
9	trainer.rs:4334-4335	to_vec1::<u32>()	CPU fallback (non-cuda)	N/A (CPU-only path)
10	gpu_experience_collector.rs:769-772	memcpy_dtoh	Per GPU collection	Async deferred readback

Note: Site #9 is the CPU-only fallback path (no GPU selector). This only runs when #[cfg(not(feature = "cuda"))] or when GPU selector init fails. We leave this path as-is since it's not a GPU hot path.

---

Component A: GPU Training Guard

New files:

• crates/ml/src/cuda_pipeline/training_guard_kernel.cu
• crates/ml/src/cuda_pipeline/gpu_training_guard.rs

CUDA Kernel: training_guard_check

Single-thread kernel (1 block × 1 thread) that reads GPU-resident loss and grad_norm scalars and writes results to pinned host memory.

Inputs (device pointers):

• loss_val: *const f32 — scalar loss from train_step (GPU-resident)
• grad_norm_val: *const f32 — scalar grad_norm from clip_grad_norm (GPU-resident)
• loss_clip_threshold: f32 — 1e6 (constant)
• grad_collapse_threshold: f32 — LR × multiplier (passed from Rust)
• warmup_active: i32 — 1 if within warmup period (skip grad collapse check)

Outputs (pinned host memory, 7 × f32 = 28 bytes):

• [0]: halt_nan (1.0 if loss or grad_norm is NaN/Inf, else 0.0)
• [1]: halt_loss_clip (1.0 if loss > threshold, else 0.0)
• [2]: halt_grad_collapse (1.0 if grad_norm < threshold && !warmup, else 0.0)
• [3]: clipped_loss (min(loss, threshold))
• [4]: raw_loss (unclipped, for logging)
• [5]: raw_grad_norm (for logging)
• [6]: reserved (padding to 32 bytes)

On-device accumulators (persistent CudaSlice, not pinned):

• loss_sum: CudaSlice<f32> — running sum of losses across steps
• grad_norm_sum: CudaSlice<f32> — running sum of grad norms
• step_count: CudaSlice<i32> — number of steps accumulated

CUDA Kernel: training_guard_accumulate

Adds current loss and grad_norm to running sums, increments counter.

Rust Wrapper: GpuTrainingGuard

pub struct GpuTrainingGuard {
    check_func: CudaFunction,
    accumulate_func: CudaFunction,
    pinned_output: CudaHostBuf<f32>,  // 7 × f32 pinned
    loss_sum: CudaSlice<f32>,
    grad_norm_sum: CudaSlice<f32>,
    step_count: CudaSlice<i32>,
    device: Device,
}

Methods:

• new(device) -> Result<Self> — compile PTX, allocate pinned + device buffers
• check(loss_gpu: &Tensor, grad_norm_gpu: &Tensor, threshold: f32, warmup: bool) -> Result<GuardResult> — launches kernel, reads pinned memory (no cudaStreamSync)
• accumulate(loss_gpu: &Tensor, grad_norm_gpu: &Tensor) -> Result<()> — adds to running sums
• read_accumulators() -> Result<(f64, f64)> — single DtoH of mean loss + mean grad_norm (epoch boundary only)
• reset_accumulators() -> Result<()> — zeros device buffers

pub struct GuardResult {
    pub halt_nan: bool,
    pub halt_loss_clip: bool,
    pub halt_grad_collapse: bool,
    pub clipped_loss: f32,
    pub raw_grad_norm: f32,
}

Wiring (trainer.rs)

train_step_optimized (replaces sites #1):

// BEFORE: stacked.to_vec1::<f32>() — GPU sync barrier
// AFTER:
let guard_result = self.training_guard.check(
    &gpu_result.loss_gpu,
    &gpu_result.grad_norm_gpu,
    grad_collapse_threshold,
    !past_warmup,
)?;
if guard_result.halt_nan { return Err(...) }
if guard_result.halt_loss_clip { log warning }
if guard_result.halt_grad_collapse { increment counter, check patience }
self.training_guard.accumulate(&gpu_result.loss_gpu, &gpu_result.grad_norm_gpu)?;

train_step_with_accumulation (replaces sites #3, #4):

// BEFORE: gpu_loss_tensors concat → mean_all().to_scalar()
// AFTER: accumulate() called per sub-step, read_accumulators() at boundary
let (avg_loss, avg_grad_norm) = self.training_guard.read_accumulators()?;
self.training_guard.reset_accumulators()?;

check_gradients_finite (replaces site #2): Add #[cfg(feature = "cuda")] bypass: when CUDA is active, the training guard kernel already checks NaN/Inf on the loss and grad_norm scalars. The per-parameter NaN check in check_gradients_finite is redundant with the post-clip norm check — if any gradient was NaN, the norm would be NaN, which the guard catches.

---

Component B: GPU Q-Value Monitor

Extends training_guard_kernel.cu with a Q-value accumulation kernel.

CUDA Kernel: qvalue_accumulate

Inputs:

• q_values: *const f32 — [batch × actions] Q-value matrix
• batch_size: i32, num_actions: i32
• qvalue_sum: *mut f32 — device accumulator (running sum of max-Q per sample)
• qvalue_count: *mut i32 — device counter

Operation: Each thread handles one sample. Finds max Q across actions, atomicAdd to qvalue_sum.

CUDA Kernel: qvalue_divergence_check

Inputs:

• q_values: *const f32 — [1, actions] single-sample Q-values
• num_actions: i32
• divergence_threshold: f32 — 10000.0

Output (pinned, 4 × f32):

• [0]: q_min
• [1]: q_max
• [2]: q_mean
• [3]: divergence_detected (1.0 if |q_min| or |q_max| > threshold)

Wiring

estimate_avg_q_value_with_early_stopping (replaces site #5):

• Forward pass stays as-is (Candle tensor)
• Replace max(1) → mean_all() → to_scalar() with qvalue_accumulate kernel call
• Read accumulator at epoch boundary, not per 50-step check

log_q_values (replaces site #6):

• Replace to_vec2::<f32>() with qvalue_divergence_check kernel
• Pinned readback: 4 floats instead of entire Q-matrix
• Prometheus gauges set from pinned values

---

Component C: GPU-Resident Action Selection

No new kernels needed — select_actions(), select_actions_routed(), and select_actions_branching() already return GPU-resident Tensor. The readback happens in the caller (trainer.rs) which does to_vec1::<u32>() to iterate actions on CPU.

Solution: Keep actions on GPU, defer routing to experience storage

select_actions_batch_gpu (replaces sites #7, #8):

Currently at trainer.rs:4302:

let action_indices = action_indices_tensor.to_vec1::<u32>()?;  // GPU sync!
for &idx in &action_indices { ... route_action(exposure, spread) ... }

Replace with:

1. Keep action_indices_tensor on GPU
2. When use_routed is true: actions are already factored indices (0-44), store tensor directly
3. When use_routed is false AND use_branching is false: actions are exposure indices (0-4). Port route_action() to a batch CUDA kernel that converts exposure → factored index on-device.
4. When use_branching is true: actions are already factored (0-44) from branching kernel

New CUDA kernel: batch_route_exposure_to_factored

Add to epsilon_greedy_kernel.cu:

extern "C" __global__ void batch_route_exposure_to_factored(
    const unsigned int* exposure_actions,  // [N] exposure indices 0-4
    unsigned int* factored_actions,        // [N] output factored indices 0-44
    float spread, float median_spread,
    float volatility, float median_volatility,
    int N
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    // Port of OrderRouter::route()
    int order_type;
    if (median_spread <= 0.0f) order_type = 0;       // Market
    else if (spread < median_spread) order_type = 1;  // LimitMaker
    else if (spread > 2.0f * median_spread) order_type = 0; // Market
    else order_type = 2;                               // IoC

    int urgency;
    if (median_volatility <= 0.0f) urgency = 1;       // Normal
    else if (volatility > 1.5f * median_volatility) urgency = 2; // Aggressive
    else if (volatility < 0.5f * median_volatility) urgency = 0; // Patient
    else urgency = 1;                                  // Normal

    factored_actions[i] = exposure_actions[i] * 9 + order_type * 3 + urgency;
}

Downstream consumer change: The experience loop at trainer.rs:2184+ currently uses Vec<FactoredAction> from select_actions_batch_gpu(). With GPU-resident actions, we need to either:

• (a) Change the experience loop to accept GPU tensor indices + batch DtoH at the end, or
• (b) Keep actions as a Vec<FactoredAction> but derive it from a single batch readback at the loop boundary instead of per-batch

Chosen approach: (b) — Single batch readback. The experience loop needs CPU-side FactoredAction values for reward calculation (calculate_reward_for_action) and portfolio tracking. These are unavoidable CPU operations. The readback happens ONCE per experience collection (not per training step), so it's outside the critical train_step hot path. The key win is eliminating the per-step readbacks in the training loop (#1-#6).

---

Component D: Experience Collector Async Readback

gpu_experience_collector.rs:769-772 — memcpy_dtoh for rewards_cpu/actions_cpu.

These are ~0.5MB downloads needed for pnl_history (Sharpe ratio) and monitor.track_reward/track_action. They happen once per experience collection (every epoch), not per training step.

Solution: Async stream + deferred poll

1. Launch memcpy_dtoh_async on a separate CUDA stream
2. Return GpuExperienceBatch immediately (with empty rewards_cpu/actions_cpu)
3. Add poll_monitoring_data(&mut self) -> Option<(Vec<f32>, Vec<i32>)> that checks stream completion
4. Trainer calls poll_monitoring_data() at epoch boundary, updates pnl_history lazily

Alternative (simpler): Keep synchronous but move to AFTER insert_batch_tensors(). The DtoD copies for replay buffer insertion don't depend on the monitoring downloads. This reorders operations so the GPU pipeline isn't stalled during the insert.

Chosen approach: Reorder — simpler, same effect. The stream.synchronize() at line 777 already waits for DtoD copies. Move the monitoring downloads AFTER the sync, or use a separate stream for monitoring downloads that doesn't block the main pipeline.

---

Success Criteria

1. grep -rn 'to_scalar\|to_vec0\|to_vec1\|to_vec2' trainer.rs → 0 matches in #[cfg(feature = "cuda")] blocks of train_step_optimized and train_step_with_accumulation
2. check_gradients_finite bypassed under #[cfg(feature = "cuda")] (guard catches NaN)
3. log_q_values uses GPU reduction instead of to_vec2
4. All existing tests pass: SQLX_OFFLINE=true cargo test -p ml -p ml-dqn -p ml-core --lib
5. 0 clippy warnings in modified files
6. Training semantics unchanged (loss values, action distributions identical)

Non-Goals

• CPU-only fallback paths (#[cfg(not(feature = "cuda"))]) are unchanged
• CPU argmax fallback in select_actions_batch_gpu (line 4334) is for when GPU selector fails to init — not a hot path
• select_actions_batch (the older, non-GPU-tensor path at line 4090+) is the CPU experience fallback — readback there is unavoidable since it feeds CPU route_action loop
• Experience collector monitoring downloads are low-frequency (once per epoch) and can tolerate ~1ms latency