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

```rust
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

```rust
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:
```rust
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`:
```cuda
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