foxhunt/docs/superpowers/specs/2026-04-06-per-tree-fix-and-zero-sync-readback.md

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:

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:

self.trainer.update_priorities_cuda_raw(indices_ptr, prio_ptr, cap, alpha, epsilon)

with:

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