fix: route PER updates through seg_tree_update kernel

Replace the orphaned per_update_kernel call site in fused_training.rs
with agent.update_priorities_from_td() which routes through
GpuReplayBuffer::update_priorities_gpu_raw -> seg_tree_update.
This kernel correctly writes priorities AND propagates the segment tree
from leaf to root, fixing stale PER sampling.

Add update_priorities_from_td() on ReplayBufferType and DqnAgentConfig.
Remove priorities_f32_ptr() which exposed raw pointers for the deleted
kernel. Remove batch_max epoch-boundary flush since seg_tree_update
propagates max_priority internally.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/ml-dqn/src/replay_buffer_type.rs

@@ -13,7 +13,6 @@ use parking_lot::Mutex;
 use crate::dqn::ExperienceReplayBuffer;
 use crate::prioritized_replay::PrioritizedReplayBuffer;
 use crate::experience::Experience;
-use ml_core::cuda_autograd::GpuTensor;
 use ml_core::MLError;
 
 /// Batch sample with importance sampling weights
@@ -386,27 +385,23 @@ impl ReplayBufferType {
         }
     }
 
-    /// Direct reference to the GPU-resident priorities tensor (`GpuPrioritized` only).
+    /// Update PER priorities using td_errors and the indices from the last sample.
     ///
-    /// Returns `None` for non-GPU buffers. Used by fused CUDA training to pass
-    /// the priorities tensor to `GpuDqnTrainer::update_priorities_cuda()`.
-    pub fn priorities_tensor(&self) -> Option<GpuTensor> {
+    /// Uses the segment tree update kernel which writes priorities, updates
+    /// tree leaves, and propagates sums to root. The sample indices are stored
+    /// internally from the last `sample()` call.
+    pub fn update_priorities_from_td(
+        &self,
+        td_errors: &cudarc::driver::CudaSlice<half::bf16>,
+        batch_size: usize,
+    ) -> Result<(), MLError> {
         match self {
             Self::GpuPrioritized(buffer) => {
-                let buf = buffer.lock();
-                let slice = buf.gpu.priorities_slice();
-                let n = slice.len();
-                let stream = buf.gpu.stream().clone();
-                // Convert f32 priorities to bf16 GpuTensor via GPU cast kernel
-                let kernels = crate::gpu_replay_buffer::get_cast_kernels_f32_to_u32(&stream).ok()?;
-                // Use the f32 priorities slice directly -- download f32 to host and
-                // re-upload as bf16 (priorities are not hot path, called once per epoch).
-                let mut host_f32 = vec![0.0_f32; n];
-                stream.memcpy_dtoh(slice, &mut host_f32).ok()?;
-                let _ = kernels; // not needed for host path
-                GpuTensor::from_host(&host_f32, vec![n], &stream).ok()
+                let mut buf = buffer.lock();
+                if batch_size == 0 { return Ok(()); }
+                buf.gpu.update_priorities_gpu_raw(0, td_errors, batch_size)
             }
-            Self::Uniform(_) | Self::Prioritized(_) => None,
+            Self::Uniform(_) | Self::Prioritized(_) => Ok(()),
         }
     }
 

crates/ml/src/trainers/dqn/config.rs

@@ -635,12 +635,16 @@ impl DQNAgentType {
         self.memory().per_alpha_epsilon()
     }
 
-    /// Direct reference to the GPU-resident priorities tensor.
+    /// Update PER priorities from GPU-resident TD errors.
     ///
-    /// Returns `Ok(None)` for non-GPU buffers. Used by fused training to pass
-    /// the tensor to `GpuDqnTrainer::update_priorities_cuda()`.
-    pub fn priorities_tensor(&self) -> Result<Option<GpuTensor>, crate::MLError> {
-        Ok(self.memory().priorities_tensor())
+    /// Routes through the replay buffer's seg_tree_update kernel which
+    /// correctly propagates the segment tree from leaf to root.
+    pub fn update_priorities_from_td(
+        &self,
+        td_errors: &cudarc::driver::CudaSlice<half::bf16>,
+        batch_size: usize,
+    ) -> Result<(), crate::MLError> {
+        self.memory().update_priorities_from_td(td_errors, batch_size)
     }
 
     /// Set max_priority from a CPU scalar (epoch-boundary flush from CUDA kernel).

crates/ml/src/trainers/dqn/fused_training.rs

@@ -752,17 +752,13 @@ impl FusedTrainingCtx {
             .map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
 
         if step < 3 { eprintln!("H100_STEP: step {step} — PER priority update"); }
-        // ── Step 6: PER priority update (outside graph) ──────────────────
-        if let (Some(priorities_tensor), Some((alpha, epsilon))) =
-            (agent.priorities_tensor()?, agent.per_alpha_epsilon())
-        {
-            self.trainer.update_priorities_cuda(
-                gpu_batch.indices_ptr,
-                priorities_tensor.data(),
-                alpha,
-                epsilon,
-            ).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?;
-        }
+        // ── Step 6: PER priority update via seg_tree_update ─────────────
+        // Uses the replay buffer's segment tree kernel: writes priorities,
+        // updates tree leaves, propagates sums to root.
+        agent.update_priorities_from_td(
+            self.trainer.td_errors_buf(),
+            self.batch_size,
+        ).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?;
         if step < 3 { eprintln!("H100_HANG4: step {step} — PER update done"); }
 
         // Bookkeeping.
@@ -1280,14 +1276,7 @@ impl FusedTrainingCtx {
             target_vars, &self.target_branching, &self.stream,
         ).map_err(|e| anyhow::anyhow!("Reverse sync target branching: {e}"))?;
 
-        // Flush batch_max_buf: single scalar DtoH (4 bytes), once per epoch.
-        // Sets GpuReplayBuffer.max_priority so new experiences get sampled first.
-        let batch_max = self.trainer.batch_max_readback_and_reset()
-            .map_err(|e| anyhow::anyhow!("batch_max readback: {e}"))?;
-        if batch_max > 0.0 {
-            agent.apply_max_priority_scalar(batch_max)
-                .map_err(|e| anyhow::anyhow!("apply batch_max: {e}"))?;
-        }
+        // seg_tree_update propagates max_priority internally — no batch_max flush needed.
 
         self.steps_since_varmap_sync = 0;
         Ok(())