feat(ml): ReplayBufferType::GpuPrioritized variant with full dispatch

New variant delegates to GpuReplayBuffer for all operations. Sample
returns BatchSample with gpu_batch populated (empty CPU vecs). Added
update_priorities_gpu() for tensor-based priority updates, as_gpu_buffer()
for direct access from trainer. Existing tests unchanged.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml/src/dqn/replay_buffer_type.rs

@@ -60,6 +60,9 @@ pub enum ReplayBufferType {
     Uniform(Arc<Mutex<ExperienceReplayBuffer>>),
     /// Prioritized sampling based on TD errors (Rainbow DQN)
     Prioritized(Arc<PrioritizedReplayBuffer>),
+    /// GPU-resident prioritized replay — all sampling on GPU (zero CPU round-trips)
+    #[cfg(feature = "cuda")]
+    GpuPrioritized(Arc<Mutex<crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>>),
 }
 
 impl std::fmt::Debug for ReplayBufferType {
@@ -67,6 +70,8 @@ impl std::fmt::Debug for ReplayBufferType {
         match self {
             Self::Uniform(_) => write!(f, "ReplayBufferType::Uniform"),
             Self::Prioritized(_) => write!(f, "ReplayBufferType::Prioritized"),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(_) => write!(f, "ReplayBufferType::GpuPrioritized"),
         }
     }
 }
@@ -102,6 +107,36 @@ impl ReplayBufferType {
         Ok(Self::Prioritized(Arc::new(buffer)))
     }
 
+    /// Create GPU-resident prioritized replay buffer.
+    ///
+    /// All experience data lives as contiguous GPU tensors. Sampling and
+    /// priority updates happen entirely on GPU — zero CPU round-trips.
+    #[cfg(feature = "cuda")]
+    pub fn new_gpu_prioritized(
+        capacity: usize,
+        state_dim: usize,
+        alpha: f64,
+        beta: f64,
+        beta_max: f64,
+        beta_annealing_steps: usize,
+        device: &candle_core::Device,
+    ) -> Result<Self, MLError> {
+        use crate::cuda_pipeline::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
+
+        let config = GpuReplayBufferConfig {
+            capacity,
+            state_dim,
+            alpha: alpha as f32,
+            beta_start: beta as f32,
+            beta_max: beta_max as f32,
+            beta_annealing_steps,
+            epsilon: 1e-6,
+        };
+
+        let buffer = GpuReplayBuffer::new(config, device)?;
+        Ok(Self::GpuPrioritized(Arc::new(Mutex::new(buffer))))
+    }
+
     /// Sample a batch from the buffer
     pub fn sample(&self, batch_size: usize) -> Result<BatchSample, MLError> {
         match self {
@@ -126,6 +161,17 @@ impl ReplayBufferType {
                     gpu_batch: None,
                 })
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => {
+                let buf = buffer.lock();
+                let gpu_batch = buf.sample_proportional(batch_size)?;
+                Ok(BatchSample {
+                    experiences: vec![], // Empty — data is on GPU
+                    weights: vec![],     // Empty — weights on GPU
+                    indices: vec![],     // Empty — indices on GPU
+                    gpu_batch: Some(gpu_batch),
+                })
+            }
         }
     }
 
@@ -140,6 +186,11 @@ impl ReplayBufferType {
             Self::Prioritized(buffer) => {
                 buffer.push(experience)
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(_) => {
+                // No-op: GPU buffer uses insert_batch() from GPU experience collector
+                Ok(())
+            }
         }
     }
 
@@ -164,6 +215,11 @@ impl ReplayBufferType {
                 }
                 Ok(())
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(_) => {
+                // No-op: GPU buffer uses insert_batch() with tensors directly
+                Ok(())
+            }
         }
     }
 
@@ -176,21 +232,49 @@ impl ReplayBufferType {
                 let priorities: Vec<f32> = td_errors.iter().map(|&td| td.abs()).collect();
                 buffer.update_priorities(indices, &priorities)
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(_) => {
+                // Use update_priorities_gpu() with tensor inputs instead
+                Ok(())
+            }
+        }
+    }
+
+    /// Update priorities from GPU-resident TD error tensors (GpuPrioritized only).
+    ///
+    /// Takes tensor indices and TD errors directly — no `to_vec1()` needed.
+    #[cfg(feature = "cuda")]
+    pub fn update_priorities_gpu(
+        &self,
+        indices: &candle_core::Tensor,
+        td_errors: &candle_core::Tensor,
+    ) -> Result<(), MLError> {
+        match self {
+            Self::GpuPrioritized(buffer) => {
+                let mut buf = buffer.lock();
+                buf.update_priorities_gpu(indices, td_errors)
+            }
+            _ => Ok(()), // No-op for non-GPU buffers
         }
     }
 
     /// Step training counter for beta annealing (PER only)
     pub fn step(&self) {
-        if let Self::Prioritized(buffer) = self {
-            buffer.step();
+        match self {
+            Self::Prioritized(buffer) => buffer.step(),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => buffer.lock().step(),
+            _ => {}
         }
     }
 
     /// Get current beta value (PER only)
     pub fn current_beta(&self) -> f32 {
         match self {
-            Self::Uniform(_) => 1.0, // No importance sampling for uniform
+            Self::Uniform(_) => 1.0,
             Self::Prioritized(buffer) => buffer.current_beta(),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => buffer.lock().current_beta(),
         }
     }
 
@@ -199,6 +283,8 @@ impl ReplayBufferType {
         match self {
             Self::Uniform(buffer) => buffer.lock().len(),
             Self::Prioritized(buffer) => buffer.len(),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => buffer.lock().len(),
         }
     }
 
@@ -217,6 +303,10 @@ impl ReplayBufferType {
             Self::Prioritized(buffer) => {
                 buffer.clear();
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => {
+                buffer.lock().clear();
+            }
         }
     }
 
@@ -225,6 +315,8 @@ impl ReplayBufferType {
         match self {
             Self::Uniform(buffer) => buffer.lock().len() >= batch_size,
             Self::Prioritized(buffer) => buffer.can_sample(batch_size),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => buffer.lock().can_sample(batch_size),
         }
     }
 
@@ -233,6 +325,12 @@ impl ReplayBufferType {
         matches!(self, Self::Prioritized(_))
     }
 
+    /// Check if using GPU-resident prioritized replay
+    #[cfg(feature = "cuda")]
+    pub fn is_gpu_prioritized(&self) -> bool {
+        matches!(self, Self::GpuPrioritized(_))
+    }
+
     /// Adaptive buffer sizing based on epsilon decay
     ///
     /// Resizes buffer capacity using threshold-based strategy to reduce memory usage
@@ -294,6 +392,11 @@ impl ReplayBufferType {
                 // PER resize not yet implemented (requires rebuilding sum tree)
                 return Ok(false);
             }
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(_) => {
+                // GPU buffer is fixed-capacity (pre-allocated tensors)
+                return Ok(false);
+            }
         }
 
         tracing::info!(
@@ -312,6 +415,17 @@ impl ReplayBufferType {
         match self {
             Self::Uniform(buffer) => buffer.lock().capacity(),
             Self::Prioritized(buffer) => buffer.capacity(),
+            #[cfg(feature = "cuda")]
+            Self::GpuPrioritized(buffer) => buffer.lock().capacity(),
+        }
+    }
+
+    /// Get a locked reference to the inner GPU replay buffer (GpuPrioritized only).
+    #[cfg(feature = "cuda")]
+    pub fn as_gpu_buffer(&self) -> Option<parking_lot::MutexGuard<'_, crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>> {
+        match self {
+            Self::GpuPrioritized(buffer) => Some(buffer.lock()),
+            _ => None,
         }
     }