# PER Tree Fix, Dead Code Cleanup & Zero-Sync Readback

## Problem

Five issues found during H100 hang investigation:

1. **PER segment tree never updated** — `per_update_priorities_kernel` scatter-writes to the `priorities[]` buffer but never touches the segment tree (`tree[2*capacity]`). Sampling via `seg_tree_sample` walks the stale tree, producing incorrect priority distributions. The correct kernel (`seg_tree_update` in `gpu_replay_buffer.rs`) writes both priorities AND propagates the tree from leaf to root.

2. **Dead code from H100 hang fixes** — `GpuTrainResult` struct/methods, `per_update_kernel.cu`, `get_cast_kernels_f32_to_u32`, `f32_slice_to_gpu_tensor_gpu`, associated `#[allow(dead_code)]` annotations.

3. **Latent non-pinned async DtoH** — `q_stats_pending` is a regular `[f32; 5]` struct field used as the destination for `cuMemcpyDtoHAsync`. On H100 CUDA 13 driver, non-pinned host memory degrades async copies to synchronous, same class of bug as hang #3.

4. **Causal sensitivity sync** — `run_causal_intervention` calls `stream.synchronize()` + sync `memcpy_dtoh` to read back per-feature sensitivities. The mean reduction should happen on GPU; the scalar result should use async pinned readback.

5. **Scattered readback allocations** — readback scalars are spread across different buffers and mechanisms. A single consolidated pinned buffer simplifies management and eliminates all sync points.

## Design

### 1. PER Tree: Use `seg_tree_update` Directly

Delete the orphaned `per_update_kernel.cu` and all associated Rust code:
- `per_update_kernel.cu` — entire file
- `per_update_kernel` field in `GpuDqnTrainer` struct
- `per_update_kernel` PTX compilation in constructor
- `update_priorities_cuda_raw()` method
- `batch_max_buf` field, allocation, and `batch_max_readback_and_reset()` method (the replay buffer tracks batch_max internally via `seg_tree_update`'s atomicMax)

Add a new method to `ReplayBufferType`:
```rust
pub fn update_priorities_gpu(
    &self,
    indices: &CudaSlice<u32>,
    td_errors: &CudaSlice<half::bf16>,
    batch_size: usize,
) -> Result<(), MLError>
```
This acquires the inner Mutex, calls `GpuReplayBuffer::update_priorities_gpu()`, and releases. Matches the existing `sample()` / `step()` pattern.

In `fused_training.rs`, replace:
```rust
self.trainer.update_priorities_cuda_raw(indices_ptr, prio_ptr, cap, alpha, epsilon)
```
with:
```rust
agent.update_priorities_gpu(&indices_slice, &td_errors_buf, batch_size)
```

The `td_errors_buf` and sample indices are already on GPU in `GpuDqnTrainer` — expose them as `CudaSlice` references (not raw pointers) to pass through the agent API.

### 2. Dead Code Removal

| Item | File | Action |
|------|------|--------|
| `per_update_kernel.cu` | `cuda_pipeline/` | Delete file |
| `per_update_kernel` field + PTX | `gpu_dqn_trainer.rs` | Remove field, compile, load |
| `update_priorities_cuda_raw()` | `gpu_dqn_trainer.rs` | Delete method |
| `batch_max_buf` field + alloc + readback | `gpu_dqn_trainer.rs` | Delete field, alloc, `batch_max_readback_and_reset()` |
| `GpuTrainResult` struct + all methods | `ml-dqn/src/dqn.rs` | Delete struct, `from_fused_scalars`, `loss_cuda_slice`, `grad_norm_cuda_slice`, `loss_scalar`, `grad_norm_scalar`, `gpu_tensor_to_cuda_slice` |
| `get_cast_kernels_f32_to_u32()` | `gpu_replay_buffer.rs` | Delete function |
| `f32_slice_to_gpu_tensor_gpu()` | `gpu_replay_buffer.rs` | Delete function, remove `#[allow(dead_code)]` |
| `#[allow(dead_code)]` on `gather_f32` | `gpu_replay_buffer.rs` | Investigate — remove if truly unused |
| Stale doc comments referencing deleted items | Various | Update or remove |

### 3. Consolidated Pinned Readback Buffer

Extend `readback_pinned` from 12 bytes (3 x f32) to 64 bytes (16 x f32):

```
Offset  Size   Field              Written by              Frequency
0       4      total_loss         graph_adam readback      per step
4       4      mse_loss           graph_adam readback      per step
8       4      grad_norm_sq       graph_adam readback      per step
12      4      q_mean             q_stats reduction       every 50 steps
16      4      q_min              q_stats reduction       every 50 steps
20      4      q_max              q_stats reduction       every 50 steps
24      4      q_variance         q_stats reduction       every 50 steps
28      4      avg_max_q          q_stats reduction       every 50 steps
32      4      causal_mean_sens   causal reduction        every N steps
36      28     nan_flags[7]       nan check kernels       per step
```

Total: 64 bytes pinned via `cuMemHostAlloc(CU_MEMHOSTALLOC_DEVICEMAP)`.

All async DtoH writes target offsets within this buffer. Each readback records an event. CPU reads poll the event (usually already complete by next use).

### 4. Causal Sensitivity GPU-Native Mean

Replace `stream.synchronize()` + sync `memcpy_dtoh` of `[market_dim] f32` with:
- GPU reduction kernel: sum `causal_sensitivity_buf[0..n_features]` and divide by `n_features`, write scalar to `readback_pinned + 32` (offset for `causal_mean_sens`)
- Async `cuMemcpyDtoHAsync` of 1 x f32 (4 bytes) from the reduction output to the pinned buffer
- Record event
- Return previous step's result (double-buffered, same pattern as `reduce_current_q_stats`)

The per-feature sensitivity array stays on GPU — only the mean scalar exits.

### 5. Lock Safety

No structural change. The lock hierarchy is correct:
- Outer: Tokio `RwLock` on agent (`write().await` for training, `read().await` for experience collection)
- Inner: `parking_lot::Mutex` on replay buffer (microsecond holds, no GPU sync inside lock)

The new `ReplayBufferType::update_priorities_gpu()` method acquires the inner Mutex, launches the kernel (non-blocking), and releases. Same pattern as all other replay buffer methods.

`CudaSlice` allocations never reallocate, so raw device pointers remain valid across lock boundaries.

### 6. Q-Stats Pinned Migration

Change `q_stats_pending: [f32; 5]` from a regular struct field to an offset within `readback_pinned`:
- Write: `cuMemcpyDtoHAsync` targets `readback_pinned + 12` (q_stats offset)
- Read: unsafe pointer dereference at the same offset
- Event: reuse existing `q_stats_event` pattern

This fixes the latent H100 hang where async DtoH degrades to sync on non-pinned memory.

## Files Modified

| File | Changes |
|------|---------|
| `per_update_kernel.cu` | DELETE |
| `gpu_dqn_trainer.rs` | Remove per_update_kernel, batch_max_buf, update_priorities_cuda_raw; extend readback_pinned; migrate q_stats to pinned; expose td_errors/indices as CudaSlice refs |
| `fused_training.rs` | Call agent.update_priorities_gpu() instead of trainer.update_priorities_cuda_raw(); causal mean to GPU reduction + pinned readback |
| `dqn.rs` (ml-dqn) | Delete GpuTrainResult struct and all methods |
| `replay_buffer_type.rs` | Add update_priorities_gpu(&self) wrapper; delete get_cast_kernels_f32_to_u32, f32_slice_to_gpu_tensor_gpu |
| `gpu_replay_buffer.rs` | Remove #[allow(dead_code)] on items being deleted |
| `config.rs` | Add update_priorities_gpu() passthrough on agent; remove priorities_f32_ptr() |
| `dqn_utility_kernels.cu` | Add causal_mean_reduction kernel (sum + divide, single block) |

## Testing

- All 19 smoke tests must pass (validates PER sampling works correctly with seg_tree_update)
- `test_no_nan_after_graph_capture` — validates readback pipeline
- Manual H100 deployment — verify no hangs, training progresses past step 0
- Verify Q-stats still reported in epoch logs (pinned migration)
- Verify causal sensitivity still reported (GPU-native mean)

## Success Criteria

- Zero `stream.synchronize()` calls in the per-step training hot loop
- Zero sync `memcpy_dtoh` in the per-step training hot loop
- PER sampling uses correct tree priorities (seg_tree_update propagates)
- All readbacks via single consolidated pinned buffer + events
- No dead code, no `#[allow(dead_code)]` hiding unused items
- 19/19 smoke tests pass
