fix(bf16): PER indices u32, grad_norm float accumulator, training guard readback

- GpuBatch.indices: GpuTensor → CudaSlice<u32> (fixes 2402 compute-sanitizer memory errors from per_update_priorities_kernel reading u32 from bf16 buffer) - fused_training PER: eliminate bf16→host→u32→GPU roundtrip, pass u32 directly - train_step accumulation: GPU DtoD concat for CudaSlice<u32> indices - grad_norm kernel: float accumulator via separate CudaSlice<f32> buffer (bf16 sum-of-squares overflows at 147K params; atomicAdd on native float) - grad_norm finalize kernel: runs OUTSIDE CUDA graph, converts float→bf16 L2 norm - Adam + clip_grad + clipped_saxpy kernels: read float sum-of-squares directly - training guard: read loss/grad from fused trainer's GPU buffers (not GpuTrainResult's hardcoded zeros), raw_ptr() for kernel args (no event tracking) - guard accumulator: reset between epochs for per-epoch metrics - Q-stats padding: pad input to config.batch_size for CUTLASS tile alignment - training_profile tests: update BF16-tuned values (spectral_norm 1.5, noisy_sigma 0.3) 895/895 unit tests pass, 5/9 smoke tests pass (remaining 4 need loss kernel float arithmetic — C51/MSE softmax overflows bf16 after ~100 training steps). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-28 20:27:14 +01:00
parent 99f54e3f94
commit 6bb8f87805
16 changed files with 370 additions and 207 deletions
--- a/crates/ml-dqn/src/dqn.rs
+++ b/crates/ml-dqn/src/dqn.rs
@@ -792,7 +792,7 @@ pub struct GradientResult {
    /// GPU-resident TD errors (`GpuPrioritized` path).
    pub td_errors_gpu: Option<GpuTensor>,
    /// GPU-resident buffer indices (`GpuPrioritized` path).
-    pub indices_gpu: Option<GpuTensor>,
+    pub indices_gpu: Option<cudarc::driver::CudaSlice<u32>>,
    /// Loss tensor on GPU for deferred batch readback.
    pub loss_tensor_gpu: Option<GpuTensor>,
    /// Gradient norm tensor on GPU for deferred readback.
@@ -883,7 +883,7 @@ struct ComputeLossResult {
    /// GPU-resident TD errors (`GpuPrioritized` path).
    td_errors_gpu: Option<GpuTensor>,
    /// GPU-resident buffer indices (`GpuPrioritized` path).
-    indices_gpu: Option<GpuTensor>,
+    indices_gpu: Option<cudarc::driver::CudaSlice<u32>>,
 }

 /// Experience replay buffer for `DQN`
@@ -2469,13 +2469,11 @@ impl DQN {
        let is_gpu_per = self.memory.is_gpu_prioritized();
        let td_error_tensor = GpuTensor::new(td_error_gpu, vec![batch_size])?;
        let (td_errors_vec, td_gpu, idx_gpu) = if is_gpu_per /* use_per: always on */ {
-            let idx_tensor = gpu_batch_opt.as_ref()
-                .map(|g| g.indices.gpu_clone(stream))
-                .transpose()?
+            let idx_u32 = gpu_batch_opt.as_ref()
                .ok_or_else(|| MLError::TrainingError(
                    "GPU PER indices required — CPU fallback removed".to_owned()
-                ))?;
-            (Vec::new(), Some(td_error_tensor), Some(idx_tensor))
+                ))?.indices.clone();
+            (Vec::new(), Some(td_error_tensor), Some(idx_u32))
        } else {
            // Non-GPU PER: download td_errors to host for CPU priority update
            // ALLOWED: only reached when GPU PER is not active (legacy path)
@@ -2592,7 +2590,7 @@ impl DQN {
        if let (Some(ref td_gpu), Some(ref idx_gpu)) =
            (&result.td_errors_gpu, &result.indices_gpu)
        {
-            self.memory.update_priorities_gpu(idx_gpu, td_gpu)?;
+            self.memory.update_priorities_gpu(idx_gpu, td_gpu.data(), result.td_errors.len().max(1))?;
        } else if !result.indices.is_empty() {
            self.memory
                .update_priorities(&result.indices, &result.td_errors)?;
--- a/crates/ml-dqn/src/gpu_replay_buffer.rs
+++ b/crates/ml-dqn/src/gpu_replay_buffer.rs
@@ -28,9 +28,9 @@ pub struct GpuBatchSlices {
    pub states: CudaSlice<u16>,       // [batch_size * state_dim] bf16 on GPU
    pub next_states: CudaSlice<u16>,  // [batch_size * state_dim] bf16 on GPU
    pub actions: CudaSlice<u32>,      // [batch_size] u32 on GPU
-    pub rewards: CudaSlice<f32>,      // [batch_size] f32 on GPU
-    pub dones: CudaSlice<f32>,        // [batch_size] f32 on GPU (0.0/1.0)
-    pub weights: CudaSlice<f32>,      // [batch_size] f32 on GPU (IS weights)
+    pub rewards: CudaSlice<u16>,      // [batch_size] bf16 on GPU
+    pub dones: CudaSlice<u16>,        // [batch_size] bf16 on GPU (0.0/1.0)
+    pub weights: CudaSlice<u16>,      // [batch_size] bf16 on GPU (IS weights)
    pub indices: CudaSlice<u32>,      // [batch_size] u32 on GPU (buffer indices)
    /// Episode IDs for sampled transitions `[batch_size]` i32 on GPU.
    /// Used by HER Future/Final strategies for episode-aware donor selection.
@@ -44,8 +44,6 @@ impl GpuBatchSlices {
    ///
    /// Zero CPU downloads -- all type conversions happen on GPU via custom CUDA kernels.
    pub fn into_gpu_batch(self, stream: &Arc<CudaStream>) -> Result<crate::replay_buffer_type::GpuBatch, MLError> {
-        use ml_core::cuda_autograd::GpuTensor;
-
        let bs = self.batch_size;
        let sd = self.state_dim;

@@ -55,14 +53,25 @@ impl GpuBatchSlices {
        let states = bf16_slice_to_gpu_tensor_gpu(&self.states, vec![bs, sd], stream, kernels)?;
        let next_states = bf16_slice_to_gpu_tensor_gpu(&self.next_states, vec![bs, sd], stream, kernels)?;

-        // u32 actions/indices -> f32 via GPU cast kernel (zero host roundtrip)
+        // u32 actions -> GpuTensor via GPU cast kernel (zero host roundtrip)
        let actions = u32_slice_to_gpu_tensor_gpu(&self.actions, vec![bs], stream, kernels)?;
-        let indices = u32_slice_to_gpu_tensor_gpu(&self.indices, vec![bs], stream, kernels)?;
+        // u32 indices: DtoD clone (stays u32, no bf16 conversion)
+        let indices = {
+            let mut dst = stream.alloc_zeros::<u32>(bs)
+                .map_err(|e| MLError::ModelError(format!("indices clone alloc: {e}")))?;
+            let nbytes = bs * std::mem::size_of::<u32>();
+            unsafe {
+                cudarc::driver::result::memcpy_dtod_async(
+                    dst.raw_ptr(), self.indices.raw_ptr(), nbytes, stream.cu_stream(),
+                ).map_err(|e| MLError::ModelError(format!("indices DtoD: {e}")))?;
+            }
+            dst
+        };

-        // f32 slices -> bf16 GpuTensor via GPU cast kernel
-        let rewards = f32_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?;
-        let dones = f32_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?;
-        let weights = f32_slice_to_gpu_tensor_gpu(&self.weights, vec![bs], stream, kernels)?;
+        // bf16 slices -> bf16 GpuTensor via DtoD reinterpret (zero cast needed, already bf16)
+        let rewards = bf16_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?;
+        let dones = bf16_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?;
+        let weights = bf16_slice_to_gpu_tensor_gpu(&self.weights, vec![bs], stream, kernels)?;

        Ok(crate::replay_buffer_type::GpuBatch {
            states,
@@ -252,6 +261,7 @@ fn u32_slice_to_gpu_tensor_gpu(
 }

 /// Convert a f32 `CudaSlice<f32>` to bf16 `GpuTensor` via GPU cast kernel (zero CPU download).
+#[allow(dead_code)]
 fn f32_slice_to_gpu_tensor_gpu(
    src: &CudaSlice<f32>,
    shape: Vec<usize>,
@@ -303,9 +313,11 @@ struct ReplayKernels {
    scatter_insert_f32: CudaFunction,
    scatter_insert_u32: CudaFunction,
    scatter_insert_bf16: CudaFunction,
+    #[allow(dead_code)]
    gather_f32: CudaFunction,
    gather_u32: CudaFunction,
    gather_bf16_rows: CudaFunction,
+    gather_bf16: CudaFunction,
    is_weights_f32: CudaFunction,
    normalize_weights_f32: CudaFunction,
    fill_from_gpu_f32: CudaFunction,
@@ -343,6 +355,7 @@ impl ReplayKernels {
            scatter_insert_bf16: ld("scatter_insert_bf16")?,
            gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?,
            gather_bf16_rows: ld("gather_bf16_rows")?,
+            gather_bf16: ld("gather_bf16")?,
            is_weights_f32: ld("is_weights_f32")?,
            normalize_weights_f32: ld("normalize_weights_f32")?,
            fill_from_gpu_f32: ld("fill_from_gpu_f32")?,
@@ -386,8 +399,8 @@ pub struct GpuReplayBuffer {
    stream: Arc<CudaStream>,
    kernels: ReplayKernels,
    states: CudaSlice<u16>, next_states: CudaSlice<u16>,
-    actions: CudaSlice<u32>, rewards: CudaSlice<f32>,
-    dones: CudaSlice<f32>, priorities: CudaSlice<f32>,
+    actions: CudaSlice<u32>, rewards: CudaSlice<u16>,
+    dones: CudaSlice<u16>, priorities: CudaSlice<f32>,
    /// Episode IDs per buffer slot `[capacity]` i32 on GPU.
    /// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`.
    episode_ids: CudaSlice<i32>,
@@ -404,8 +417,8 @@ pub struct GpuReplayBuffer {
    sample_states: CudaSlice<u16>,
    sample_next_states: CudaSlice<u16>,
    sample_actions: CudaSlice<u32>,
-    sample_rewards: CudaSlice<f32>,
-    sample_dones: CudaSlice<f32>,
+    sample_rewards: CudaSlice<u16>,
+    sample_dones: CudaSlice<u16>,
    sample_priorities: CudaSlice<f32>,
    sample_weights: CudaSlice<f32>,
    sample_max_weight: CudaSlice<f32>,
@@ -442,8 +455,8 @@ impl GpuReplayBuffer {
        let s = a16(stream, cap * sd, "s")?;
        let ns = a16(stream, cap * sd, "ns")?;
        let a = a32u(stream, cap, "a")?;
-        let r = a32f(stream, cap, "r")?;
-        let d = a32f(stream, cap, "d")?;
+        let r = a16(stream, cap, "r")?;
+        let d = a16(stream, cap, "d")?;
        let p = a32f(stream, cap, "p")?;
        let mut mp = a32f(stream, 1, "mp")?;
        stream.memcpy_htod(&[1.0_f32], &mut mp).map_err(|e| MLError::ModelError(format!("mp: {e}")))?;
@@ -459,8 +472,8 @@ impl GpuReplayBuffer {
        let ss = a16(stream, mbs * sd, "s_states")?;
        let sns = a16(stream, mbs * sd, "s_nstates")?;
        let sa = a32u(stream, mbs, "s_act")?;
-        let sr = a32f(stream, mbs, "s_rew")?;
-        let sdn = a32f(stream, mbs, "s_done")?;
+        let sr = a16(stream, mbs, "s_rew")?;
+        let sdn = a16(stream, mbs, "s_done")?;
        let sp = a32f(stream, mbs, "s_pri")?;
        let sw = a32f(stream, mbs, "s_wt")?;
        let smw = a32f(stream, 1, "s_mw")?;
@@ -553,16 +566,33 @@ impl GpuReplayBuffer {
                .arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
        }
-        // SAFETY: same context, rewards buffer valid.
+        // Cast f32 rewards/dones to bf16 (u16) then scatter insert as bf16 (state_dim=1)
+        let one_i = 1_i32;
+        let eff_i = eff as i32;
+        let mut b_r = a16(&self.stream, eff, "ib_r")?;
+        // SAFETY: b_r, sr are valid device allocations of at least eff elements.
        unsafe {
-            self.stream.launch_builder(&self.kernels.scatter_insert_f32)
-                .arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
+            self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
+                .arg(&mut b_r).arg(&sr).arg(&eff_i).launch(lcfg(eff))
+                .map_err(|e| MLError::ModelError(format!("cast r: {e}")))?;
+        }
+        // SAFETY: rewards buffer (bf16) valid, b_r cast above.
+        unsafe {
+            self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
+                .arg(&self.rewards).arg(&b_r).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
        }
-        // SAFETY: same context, dones buffer valid.
+        let mut b_d = a16(&self.stream, eff, "ib_d")?;
+        // SAFETY: b_d, sd2 are valid device allocations of at least eff elements.
        unsafe {
-            self.stream.launch_builder(&self.kernels.scatter_insert_f32)
-                .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
+            self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
+                .arg(&mut b_d).arg(&sd2).arg(&eff_i).launch(lcfg(eff))
+                .map_err(|e| MLError::ModelError(format!("cast d: {e}")))?;
+        }
+        // SAFETY: dones buffer (bf16) valid, b_d cast above.
+        unsafe {
+            self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
+                .arg(&self.dones).arg(&b_d).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
        }
        let mut pt = a32f(&self.stream, eff, "pt")?;
@@ -660,41 +690,21 @@ impl GpuReplayBuffer {
                .arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
                .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
        }
-        // Rewards/dones: reinterpret bf16 CudaSlice as f32 for scatter kernel (same 2-byte elements)
-        // Actually rewards/dones in the replay buffer are stored as f32.
-        // We need to cast bf16→f32. Use a temporary f32 buffer + bf16_to_f32 cast.
-        // For simplicity, use the scatter_insert_bf16 kernel which copies bf16→bf16,
-        // then let the sample path handle the conversion.
-        // BUT the replay buffer stores rewards as CudaSlice<f32>, so we need to convert.
-        // Simplest: reinterpret the bf16 as raw bytes and use memcpy_dtod with size matching.
-        // Since rewards/dones in experience batch are bf16 but replay buffer stores f32,
-        // we need to cast. Use the bf16_to_f32_cast kernel if available.
-        // For now, download→convert→upload (small: only batch_size scalars).
+        // Rewards/dones: both input and storage are bf16 — scatter directly (state_dim=1)
        {
            let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
            let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
-            // Small host roundtrip for rewards/dones (eff scalars, typically 128-1024)
-            let mut rw_host = vec![half::bf16::ZERO; eff];
-            self.stream.memcpy_dtoh(&sr, &mut rw_host)
-                .map_err(|e| MLError::ModelError(format!("rw dtoh: {e}")))?;
-            let rw_f32: Vec<f32> = rw_host.iter().map(|x| x.to_f32()).collect();
-            let rw_gpu = self.stream.clone_htod(&rw_f32)
-                .map_err(|e| MLError::ModelError(format!("rw htod: {e}")))?;
+            let one_i = 1_i32;
+            // SAFETY: rewards/dones buffers are u16 (bf16), rw/dn CudaSlice<half::bf16> same repr.
            unsafe {
-                self.stream.launch_builder(&self.kernels.scatter_insert_f32)
-                    .arg(&self.rewards).arg(&rw_gpu).arg(&ci).arg(&cpi).arg(&bsi)
-                    .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
+                self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
+                    .arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
+                    .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r bf16: {e}")))?;
            }
-            let mut dn_host = vec![half::bf16::ZERO; eff];
-            self.stream.memcpy_dtoh(&sd2, &mut dn_host)
-                .map_err(|e| MLError::ModelError(format!("dn dtoh: {e}")))?;
-            let dn_f32: Vec<f32> = dn_host.iter().map(|x| x.to_f32()).collect();
-            let dn_gpu = self.stream.clone_htod(&dn_f32)
-                .map_err(|e| MLError::ModelError(format!("dn htod: {e}")))?;
            unsafe {
-                self.stream.launch_builder(&self.kernels.scatter_insert_f32)
-                    .arg(&self.dones).arg(&dn_gpu).arg(&ci).arg(&cpi).arg(&bsi)
-                    .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
+                self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
+                    .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
+                    .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d bf16: {e}")))?;
            }
        }
        // Priorities — same as insert_batch
@@ -806,13 +816,13 @@ impl GpuReplayBuffer {
        }
        // SAFETY: same context, rewards buffer valid.
        unsafe {
-            self.stream.launch_builder(&self.kernels.gather_f32)
+            self.stream.launch_builder(&self.kernels.gather_bf16)
                .arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?;
        }
        // SAFETY: same context, dones buffer valid.
        unsafe {
-            self.stream.launch_builder(&self.kernels.gather_f32)
+            self.stream.launch_builder(&self.kernels.gather_bf16)
                .arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi)
                .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?;
        }
@@ -896,13 +906,28 @@ impl GpuReplayBuffer {
        // so we DtoD-clone the relevant portions. All copies are async on the stream.
        let ep_ids = dtod_clone_i32(&self.stream, &self.sample_episode_ids, batch_size, "o_ep")?;

+        // Cast f32 IS weights -> bf16 (u16) for output. IS computation stays f32 for precision.
+        let weights_bf16 = {
+            let wt_slice = self.sample_weights.slice(..batch_size);
+            let mut out = a16(&self.stream, batch_size, "o_w_bf16")?;
+            let n_i = batch_size as i32;
+            // SAFETY: out, wt_slice are valid device allocations of at least batch_size elements.
+            unsafe {
+                self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
+                    .arg(&mut out).arg(&wt_slice).arg(&n_i)
+                    .launch(lcfg(batch_size))
+                    .map_err(|e| MLError::ModelError(format!("cast w bf16: {e}")))?;
+            }
+            out
+        };
+
        Ok(GpuBatchSlices {
            states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?,
            next_states: dtod_clone_u16(&self.stream, &self.sample_next_states, batch_size * sd, "o_n")?,
            actions: dtod_clone_u32(&self.stream, &self.sample_actions, batch_size, "o_act")?,
-            rewards: dtod_clone_f32(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
-            dones: dtod_clone_f32(&self.stream, &self.sample_dones, batch_size, "o_d")?,
-            weights: dtod_clone_f32(&self.stream, &self.sample_weights, batch_size, "o_w")?,
+            rewards: dtod_clone_u16(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
+            dones: dtod_clone_u16(&self.stream, &self.sample_dones, batch_size, "o_d")?,
+            weights: weights_bf16,
            indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?,
            episode_ids: Some(ep_ids),
            batch_size,
@@ -1042,15 +1067,15 @@ impl GpuReplayBuffer {
    pub const fn states_slice(&self) -> &CudaSlice<u16> { &self.states }
    pub const fn next_states_slice(&self) -> &CudaSlice<u16> { &self.next_states }
    pub const fn actions_slice(&self) -> &CudaSlice<u32> { &self.actions }
-    pub const fn rewards_slice(&self) -> &CudaSlice<f32> { &self.rewards }
-    pub const fn dones_slice(&self) -> &CudaSlice<f32> { &self.dones }
+    pub const fn rewards_slice(&self) -> &CudaSlice<u16> { &self.rewards }
+    pub const fn dones_slice(&self) -> &CudaSlice<u16> { &self.dones }
    pub const fn priorities_slice(&self) -> &CudaSlice<f32> { &self.priorities }

    /// Sample proportional indices and IS weights as GPU-resident `CudaSlices`.
    ///
-    /// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<f32>)` on GPU.
+    /// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<u16>)` on GPU (weights are bf16).
    /// Callers process indices on GPU -- zero CPU download.
-    pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<f32>), MLError> {
+    pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<u16>), MLError> {
        let b = self.sample_proportional(bs)?;
        Ok((b.indices, b.weights))
    }
@@ -1080,6 +1105,7 @@ fn a32i(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLErr
 }

 /// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
+#[allow(dead_code)]
 fn dtod_clone_f32(s: &Arc<CudaStream>, src: &CudaSlice<f32>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
    let mut dst = a32f(s, n, nm)?;
    let sv = src.slice(..n);
--- a/crates/ml-dqn/src/regime_conditional.rs
+++ b/crates/ml-dqn/src/regime_conditional.rs
@@ -650,7 +650,7 @@ impl RegimeConditionalDQN {
                next_states: gpu_batch.next_states.gpu_clone(&self.stream)?,
                dones: gpu_batch.dones.gpu_clone(&self.stream)?,
                weights: masked_weights,
-                indices: gpu_batch.indices.gpu_clone(&self.stream)?,
+                indices: gpu_batch.indices.clone(),
                episode_ids: None,
            };

@@ -853,7 +853,7 @@ impl RegimeConditionalDQN {
                next_states: gpu_batch.next_states.gpu_clone(&self.stream)?,
                dones: gpu_batch.dones.gpu_clone(&self.stream)?,
                weights: masked_weights,
-                indices: gpu_batch.indices.gpu_clone(&self.stream)?,
+                indices: gpu_batch.indices.clone(),
                episode_ids: None,
            };

--- a/crates/ml-dqn/src/replay_buffer_kernels.cu
+++ b/crates/ml-dqn/src/replay_buffer_kernels.cu
@@ -114,6 +114,21 @@ void gather_bf16_rows(
    out[row * state_dim + col] = src[src_row * state_dim + col];
 }

+// ── 6b. Gather bf16 scalars (1D) by indices ────────────────────────────────
+
+extern "C" __global__
+void gather_bf16(
+    unsigned short* __restrict__ out,        // [batch_size]
+    const unsigned short* __restrict__ src,  // [capacity]
+    const long long* __restrict__ indices,   // [batch_size]
+    int batch_size
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= batch_size) return;
+    int idx = (int)indices[i];
+    out[i] = src[idx];
+}
+
 // ── 7. prios_alpha: out[i] = pow(src[i], alpha) ────────────────────────────

 extern "C" __global__
--- a/crates/ml-dqn/src/replay_buffer_type.rs
+++ b/crates/ml-dqn/src/replay_buffer_type.rs
@@ -43,13 +43,13 @@ impl BatchSample {
 /// When present, `compute_gradients()` uses these directly -- no CPU->GPU transfer.
 #[derive(Debug)]
 pub struct GpuBatch {
-    pub states: GpuTensor,      // [batch_size, state_dim] f32 on GPU
+    pub states: GpuTensor,      // [batch_size, state_dim] bf16 on GPU
    pub actions: GpuTensor,     // [batch_size] u32 on GPU
-    pub rewards: GpuTensor,     // [batch_size] f32 on GPU
-    pub next_states: GpuTensor, // [batch_size, state_dim] f32 on GPU
-    pub dones: GpuTensor,       // [batch_size] f32 on GPU (0.0/1.0)
-    pub weights: GpuTensor,     // [batch_size] f32 on GPU (IS weights)
-    pub indices: GpuTensor,     // [batch_size] u32 on GPU (buffer indices)
+    pub rewards: GpuTensor,     // [batch_size] bf16 on GPU
+    pub next_states: GpuTensor, // [batch_size, state_dim] bf16 on GPU
+    pub dones: GpuTensor,       // [batch_size] bf16 on GPU (0.0/1.0)
+    pub weights: GpuTensor,     // [batch_size] bf16 on GPU (IS weights)
+    pub indices: cudarc::driver::CudaSlice<u32>,  // [batch_size] u32 on GPU (buffer indices)
    /// Episode IDs per transition `[batch_size]` i32 on GPU.
    ///
    /// `episode_ids[i] = buffer_index[i] / episode_length`. Required by HER
@@ -354,19 +354,16 @@ impl ReplayBufferType {
    /// Indices are expected to hold f32-encoded u32 values (cast to u32 internally).
    pub fn update_priorities_gpu(
        &self,
-        indices: &GpuTensor,
-        td_errors: &GpuTensor,
+        indices: &cudarc::driver::CudaSlice<u32>,
+        td_errors: &cudarc::driver::CudaSlice<half::bf16>,
+        bs: usize,
    ) -> Result<(), MLError> {
        match self {
            Self::GpuPrioritized(buffer) => {
                let mut buf = buffer.lock();
-                let bs = td_errors.shape().first().copied().unwrap_or(0);
                if bs == 0 { return Ok(()); }
-                // Convert f32-encoded indices to CudaSlice<u32> via GPU cast kernel (zero CPU download)
-                let stream = buf.gpu.stream().clone();
-                let cast_kernels = crate::gpu_replay_buffer::get_cast_kernels_f32_to_u32(&stream)?;
-                let idx_gpu = cast_kernels.cast_f32_to_u32(indices.data(), bs, &stream)?;
-                buf.gpu.update_priorities_gpu(&idx_gpu, td_errors.data(), bs)
+                // indices already CudaSlice<u32>, td_errors already CudaSlice<half::bf16>
+                buf.gpu.update_priorities_gpu(indices, td_errors, bs)
            }
            Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers
        }
--- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
+++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
@@ -26,40 +26,52 @@

 extern "C" __global__ void dqn_grad_norm_kernel(
    const __nv_bfloat16* __restrict__ grads,      /* [TOTAL_PARAMS] accumulated gradients */
-    __nv_bfloat16* __restrict__ out_grad_norm,    /* [1] output: sum of squares */
+    float* __restrict__ out_norm_f32,             /* [1] float accumulator (atomicAdd target) */
    int total_params
 ) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    __nv_bfloat16 g = (idx < total_params) ? grads[idx] : bf16_zero();
-    __nv_bfloat16 g2 = g * g;
+    /* Accumulate g² in float — bf16 overflows at sum > 65504 (trivial with 147K params).
+     * BF16 arithmetic produces occasional NaN/Inf in individual gradient elements
+     * (e.g. bf16 subtraction of nearly-equal large values → NaN). These don't affect
+     * training correctness (Adam's clip_scale uses fmaxf which handles NaN→0), but
+     * NaN + valid = NaN poisons the norm sum. Skip them for an accurate norm. */
+    float g_f = (idx < total_params) ? (float)grads[idx] : 0.0f;
+    if (isnan(g_f) || isinf(g_f)) g_f = 0.0f;
+    float g2 = g_f * g_f;

-    /* Warp-level reduction via shuffle (no shared memory) */
+    /* Warp-level reduction via shuffle (float) */
    for (int offset = 16; offset > 0; offset >>= 1)
-        g2 = g2 + bf16_shfl_xor(0xFFFFFFFF, g2, offset);
+        g2 += __shfl_xor_sync(0xFFFFFFFF, g2, offset);

-    /* Block-level cross-warp reduction via shared memory.
-     * Pad warp_sums to 2 elements per warp (stride=2) to avoid bank conflicts:
-     * warp 0 → bank 0, warp 1 → bank 2, ... warp 7 → bank 14.
-     * Without padding, sequential warp writes to warp_sums[0..7] map to
-     * banks 0..7 which is conflict-free for writes, but the __syncthreads()
-     * + read-back pattern benefits from padding to avoid false sharing
-     * between adjacent cache lines on H100's 128-byte L1 lines. */
-    __shared__ __nv_bfloat16 warp_sums[16];  /* 8 warps × 2 stride (padded) */
+    /* Block-level cross-warp reduction via shared memory (float) */
+    __shared__ float warp_sums[16];  /* 8 warps × 2 stride (padded) */
    int warp_id = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_sums[warp_id * 2] = g2;
    __syncthreads();

-    /* First warp reduces across warps using shuffles (pure register path) */
+    /* First warp reduces across warps using shuffles */
    if (warp_id == 0) {
-        __nv_bfloat16 val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : bf16_zero();
+        float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
-            val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset);
+            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0)
-            atomicAddBF16(out_grad_norm, val);
+            atomicAdd(out_norm_f32, val);  /* Native float atomicAdd — no type-punning */
    }
 }

+/* Single-thread finalize: sqrt(f32 sum-of-squares) → bf16 L2 norm.
+ * Called AFTER graph_adam replay (outside CUDA graph).
+ * Reads from the float accumulator, writes bf16 norm for the training guard. */
+extern "C" __global__ void dqn_grad_norm_finalize(
+    const float* __restrict__ norm_f32,           /* [1] float sum-of-squares */
+    __nv_bfloat16* __restrict__ out_norm_bf16     /* [1] bf16 L2 norm output */
+) {
+    float sum_sq = *norm_f32;
+    float norm = sqrtf(fmaxf(sum_sq, 0.0f));
+    out_norm_bf16[0] = bf16(norm);
+}
+
 /* ══════════════════════════════════════════════════════════════════════
 * ADAM UPDATE KERNEL (Phase 3b)
 *
@@ -75,7 +87,7 @@ extern "C" __global__ void dqn_adam_update_kernel(
    const __nv_bfloat16* __restrict__ grads,      /* [TOTAL_PARAMS] accumulated gradients */
    __nv_bfloat16* __restrict__ m,                /* [TOTAL_PARAMS] first moment (Adam) */
    __nv_bfloat16* __restrict__ v,                /* [TOTAL_PARAMS] second moment (Adam) */
-    const __nv_bfloat16* __restrict__ grad_norm_sq, /* [1] completed sum of squares */
+    const float* __restrict__ grad_norm_f32,      /* [1] float sum-of-squares from grad_norm kernel */
    float lr,
    float beta1,
    float beta2,
@@ -91,9 +103,11 @@ extern "C" __global__ void dqn_adam_update_kernel(
    int t = *t_ptr;
    __nv_bfloat16 g = grads[idx];

-    /* Clip using the COMPLETED norm (no race — separate kernel launch) */
-    __nv_bfloat16 norm = bf16_sqrt(*grad_norm_sq + bf16(1e-12f));
-    __nv_bfloat16 clip_scale = (*grad_norm_sq > bf16_zero() && norm > bf16(max_grad_norm)) ? (bf16(max_grad_norm) / norm) : bf16_one();
+    /* Clip using the COMPLETED sum-of-squares (native float buffer).
+     * dqn_grad_norm_kernel accumulates float partial sums via atomicAdd. */
+    float norm_sq_f = *grad_norm_f32;
+    float norm_f = sqrtf(fmaxf(norm_sq_f, 0.0f));
+    __nv_bfloat16 clip_scale = (norm_f > max_grad_norm) ? bf16(max_grad_norm / norm_f) : bf16_one();
    __nv_bfloat16 clipped_g = g * clip_scale;

    /* Adam bias correction: compute in float — bf16(0.999) rounds to 1.0 → div-by-zero.
@@ -160,15 +174,15 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel(
    const __nv_bfloat16* __restrict__ x,
    float alpha,
    float max_norm,
-    const __nv_bfloat16* __restrict__ x_norm_sq,
+    const float* __restrict__ x_norm_f32,  /* float sum-of-squares from grad_norm kernel */
    int n
 ) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

-    __nv_bfloat16 norm = bf16_sqrt(*x_norm_sq + bf16(1e-12f));
-    __nv_bfloat16 clip = (norm > bf16(max_norm)) ? (bf16(max_norm) / norm) : bf16_one();
-    y[i] = y[i] + bf16(alpha) * clip * x[i];
+    float norm_f = sqrtf(fmaxf(*x_norm_f32, 0.0f));
+    float clip_f = (norm_f > max_norm) ? (max_norm / norm_f) : 1.0f;
+    y[i] = y[i] + bf16(alpha * clip_f) * x[i];
 }

 /* ══════════════════════════════════════════════════════════════════════
@@ -185,16 +199,16 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel(

 extern "C" __global__ void dqn_clip_grad_kernel(
    __nv_bfloat16* __restrict__ grads,
-    const __nv_bfloat16* __restrict__ grad_norm_sq,
+    const __nv_bfloat16* __restrict__ grad_norm_ptr,  /* bf16 L2 norm from finalize kernel */
    float max_norm,
    int total_params
 ) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_params) return;

-    __nv_bfloat16 norm = bf16_sqrt(*grad_norm_sq + bf16(1e-12f));
-    if (norm > bf16(max_norm)) {
-        grads[idx] = grads[idx] * (bf16(max_norm) / norm);
+    float norm_f = (float)(*grad_norm_ptr);  /* bf16 L2 norm, already sqrt'd */
+    if (norm_f > max_norm) {
+        grads[idx] = grads[idx] * bf16(max_norm / norm_f);
    }
 }

--- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
@@ -314,6 +314,7 @@ struct CachedPtrs {
    target_params_buf: u64,
    grad_buf: u64,
    grad_norm_buf: u64,
+    grad_norm_f32_buf: u64,
    m_buf: u64,
    v_buf: u64,
    t_buf: u64,
@@ -380,6 +381,7 @@ pub struct GpuDqnTrainer {
    // Dead kernels (forward_loss_kernel, forward_only_kernel, backward_kernel)
    // removed — replaced by cuBLAS forward/backward. See Task 2 for method cleanup.
    grad_norm_kernel: CudaFunction,
+    grad_norm_finalize_kernel: CudaFunction,
    /// Separate instance of grad_norm for non-graph launches (clip_grad_buf_inplace).
    /// CUDA doesn't allow the same CUfunction to be launched both inside a captured
    /// graph and outside on the same stream.
@@ -446,7 +448,7 @@ pub struct GpuDqnTrainer {
    // ── Forward output buffers ──────────────────────────────────────
    per_sample_loss_buf: CudaSlice<half::bf16>,  // [B]
    td_errors_buf: CudaSlice<half::bf16>,        // [B]
-    total_loss_buf: CudaSlice<half::bf16>,       // [1]
+    pub(crate) total_loss_buf: CudaSlice<half::bf16>,       // [1]

    // ── Forward-only Q-value output ─────────────────────────────────
    q_out_buf: CudaSlice<half::bf16>,            // [B, TOTAL_ACTIONS(11)]
@@ -457,7 +459,8 @@ pub struct GpuDqnTrainer {
    target_params_buf: CudaSlice<half::bf16>, // [TOTAL_PARAMS] flat target parameters (EMA)
    m_buf: CudaSlice<half::bf16>,            // [TOTAL_PARAMS] Adam first moment
    v_buf: CudaSlice<half::bf16>,            // [TOTAL_PARAMS] Adam second moment
-    grad_norm_buf: CudaSlice<half::bf16>,    // [1] pre-clip gradient L2 norm
+    pub(crate) grad_norm_buf: CudaSlice<half::bf16>,    // [1] bf16 L2 norm (written by finalize)
+    grad_norm_f32_buf: CudaSlice<f32>,                 // [1] float accumulator for grad_norm kernel
    cql_grad_scratch: CudaSlice<half::bf16>, // [TOTAL_PARAMS] CQL gradient isolation buffer

    // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─
@@ -1330,21 +1333,21 @@ impl GpuDqnTrainer {
    /// Called after `apply_cql_gradient` populated `cql_grad_scratch`.
    /// Computes norm of scratch, clips to `cql_budget`, then SAXPYs into grad_buf.
    pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> {
-        // Compute CQL gradient norm
+        // Compute CQL gradient norm (float accumulator)
        let _eg = EventTrackingGuard::new(self.stream.context());
-        self.stream.memset_zeros(&mut self.grad_norm_buf)
-            .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?;
+        self.stream.memset_zeros(&mut self.grad_norm_f32_buf)
+            .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm_f32: {e}")))?;

        let total = self.total_params as i32;
        let blocks = ((self.total_params + 255) / 256) as u32;
        let scratch_ptr = self.ptrs.cql_grad_scratch;
-        let norm_ptr = self.ptrs.grad_norm_buf;
+        let norm_f32_ptr = self.ptrs.grad_norm_f32_buf;

        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_standalone)
                .arg(&scratch_ptr)
-                .arg(&norm_ptr)
+                .arg(&norm_f32_ptr)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
@@ -1364,7 +1367,7 @@ impl GpuDqnTrainer {
                .arg(&scratch_ptr)
                .arg(&alpha)
                .arg(&cql_budget)
-                .arg(&norm_ptr)
+                .arg(&norm_f32_ptr)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
@@ -1543,14 +1546,14 @@ impl GpuDqnTrainer {
        // Disable cudarc event tracking — graph replay leaves stale events
        // that cause INVALID_VALUE on post-graph kernel launches.
        let _evt_guard = EventTrackingGuard::new(self.stream.context());
-        // Zero the grad_norm accumulator
+        // Zero the float grad_norm accumulator
        self.stream
-            .memset_zeros(&mut self.grad_norm_buf)
-            .map_err(|e| MLError::ModelError(format!("zero grad_norm for clip: {e}")))?;
-        // memset_zeros succeeded
+            .memset_zeros(&mut self.grad_norm_f32_buf)
+            .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32 for clip: {e}")))?;

-        // Compute grad_norm (sum of squares)
+        // Compute grad_norm (float sum-of-squares → bf16 L2 norm)
        self.launch_grad_norm()?;
+        self.launch_grad_norm_finalize()?;

        // Clip in-place
        let grad_ptr = self.ptrs.grad_buf;
@@ -1583,22 +1586,24 @@ impl GpuDqnTrainer {
        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        self.stream
-            .memset_zeros(&mut self.grad_norm_buf)
-            .map_err(|e| MLError::ModelError(format!("zero grad_norm for diag: {e}")))?;
+            .memset_zeros(&mut self.grad_norm_f32_buf)
+            .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32 for diag: {e}")))?;

        self.launch_grad_norm()?;
+        self.launch_grad_norm_finalize()?;

        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }

-        let mut norm_sq_bf16 = [half::bf16::ZERO; 1];
+        // grad_norm_buf[0] now contains the L2 norm directly (sqrt applied by finalize)
+        let mut norm_bf16 = [half::bf16::ZERO; 1];
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoH_v2(
-                norm_sq_bf16.as_mut_ptr().cast(),
+                norm_bf16.as_mut_ptr().cast(),
                self.grad_norm_buf.raw_ptr(),
                std::mem::size_of::<half::bf16>(),
            );
        }
-        Ok(norm_sq_bf16[0].to_f32().sqrt())
+        Ok(norm_bf16[0].to_f32())
    }

    /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph).
@@ -1717,7 +1722,7 @@ impl GpuDqnTrainer {
        // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).

        // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─
-        let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) =
+        let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) =
            compile_training_kernels(&stream, &config)?;

        // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace).
@@ -1778,6 +1783,8 @@ impl GpuDqnTrainer {
        let m_buf = alloc_bf16(&stream, total_params, "adam_m")?;
        let v_buf = alloc_bf16(&stream, total_params, "adam_v")?;
        let grad_norm_buf = alloc_bf16(&stream, 1, "grad_norm")?;
+        let grad_norm_f32_buf = stream.alloc_zeros::<f32>(1)
+            .map_err(|e| MLError::ModelError(format!("alloc grad_norm_f32: {e}")))?;
        let cql_grad_scratch = alloc_bf16(&stream, total_params, "cql_grad_scratch")?;
        let t_buf = alloc_i32(&stream, 1, "adam_t")?;

@@ -2007,6 +2014,7 @@ impl GpuDqnTrainer {
                target_params_buf: target_params_buf.raw_ptr(),
                grad_buf: grad_buf.raw_ptr(),
                grad_norm_buf: grad_norm_buf.raw_ptr(),
+                grad_norm_f32_buf: grad_norm_f32_buf.raw_ptr(),
                m_buf: m_buf.raw_ptr(),
                v_buf: v_buf.raw_ptr(),
                t_buf: t_buf.raw_ptr(),
@@ -2055,6 +2063,7 @@ impl GpuDqnTrainer {
            stream,
            ptrs,
            grad_norm_kernel,
+            grad_norm_finalize_kernel,
            grad_norm_standalone,
            adam_update_kernel,
            ema_kernel,
@@ -2111,6 +2120,7 @@ impl GpuDqnTrainer {
            m_buf,
            v_buf,
            grad_norm_buf,
+            grad_norm_f32_buf,
            cql_grad_scratch,
            t_buf,
            adam_step: 0,
@@ -2293,6 +2303,9 @@ impl GpuDqnTrainer {
    /// gradients into grad_buf (IQN, attention, ensemble).
    pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
        self.replay_adam()?;
+        // Finalize grad_norm OUTSIDE graph: convert float sum-of-squares → bf16 L2 norm.
+        // The graph stores float sums via atomicAdd; this kernel reads them and writes bf16.
+        self.launch_grad_norm_finalize()?;
        // Per-step: zero readback for performance (no cuStreamSync).
        // Loss/grad_norm accumulated on GPU by training_guard kernel.
        // Epoch boundary calls readback_scalars_sync() for metrics.
@@ -2302,9 +2315,33 @@ impl GpuDqnTrainer {
        })
    }

+    /// Convert float sum-of-squares → bf16 L2 norm.
+    /// Reads grad_norm_f32_buf, writes grad_norm_buf.
+    /// Must be called AFTER graph_adam replay.
+    fn launch_grad_norm_finalize(&self) -> Result<(), MLError> {
+        // Uses raw_ptr() → passes u64 device addresses → no event tracking needed.
+        let f32_ptr = self.grad_norm_f32_buf.raw_ptr();
+        let bf16_ptr = self.grad_norm_buf.raw_ptr();
+        unsafe {
+            self.stream
+                .launch_builder(&self.grad_norm_finalize_kernel)
+                .arg(&f32_ptr)
+                .arg(&bf16_ptr)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (1, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("grad_norm_finalize: {e}")))?;
+        }
+        Ok(())
+    }
+
    /// Sync stream and read back loss + grad_norm from GPU.
    /// Called ONLY at epoch boundary — never per-step.
    pub fn readback_scalars_sync(&mut self) -> Result<FusedTrainScalars, MLError> {
+        // Launch finalize to convert float sum-of-squares → bf16 L2 norm
+        self.launch_grad_norm_finalize()?;
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
        }
@@ -2323,7 +2360,7 @@ impl GpuDqnTrainer {
        }
        Ok(FusedTrainScalars {
            total_loss: loss_bf16[0].to_f32(),
-            grad_norm: norm_bf16[0].to_f32().sqrt(),
+            grad_norm: norm_bf16[0].to_f32(),  // bf16 L2 norm from finalize kernel
        })
    }

@@ -2686,7 +2723,25 @@ impl GpuDqnTrainer {
        batch_size: usize,
    ) -> Result<QValueStatsResult, MLError> {
        let _eg = EventTrackingGuard::new(self.stream.context());
-        self.compute_q_values(states, batch_size)?;
+
+        // Pad input to config.batch_size to prevent CUTLASS tile-boundary OOB reads.
+        // CUTLASS WMMA uses 32-element N-tiles; non-aligned batch sizes read past buffer.
+        let mut padded_buf;
+        let eff_states = if batch_size < self.config.batch_size {
+            let sd = self.config.state_dim;
+            padded_buf = self.stream.alloc_zeros::<half::bf16>(self.config.batch_size * sd)
+                .map_err(|e| MLError::ModelError(format!("q_stats pad alloc: {e}")))?;
+            {
+                let mut dst = padded_buf.slice_mut(0..batch_size * sd);
+                self.stream.memcpy_dtod(states, &mut dst)
+                    .map_err(|e| MLError::ModelError(format!("q_stats pad copy: {e}")))?;
+            }
+            &padded_buf
+        } else {
+            states
+        };
+
+        self.compute_q_values(eff_states, batch_size)?;

        // Zero stats buf then launch reduction
        self.stream.memset_zeros(&mut self.q_stats_buf)
@@ -2972,13 +3027,16 @@ impl GpuDqnTrainer {
        online_d: &DuelingWeightSet,
        online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
-        // ── Zero grad_norm accumulator ─────────────────────────────
+        // ── Zero float grad_norm accumulator ──────────────────────
        self.stream
-            .memset_zeros(&mut self.grad_norm_buf)
-            .map_err(|e| MLError::ModelError(format!("zero grad_norm: {e}")))?;
+            .memset_zeros(&mut self.grad_norm_f32_buf)
+            .map_err(|e| MLError::ModelError(format!("zero grad_norm_f32: {e}")))?;

-        // ── 5. Gradient norm ───────────────────────────────────────
+        // ── 5. Gradient norm (float accumulator) ──────────────────
        self.launch_grad_norm()?;
+        // NOTE: grad_norm_finalize runs AFTER graph replay (outside capture).
+        // The graph stores float sum-of-squares in grad_norm_buf[0:3].
+        // Adam reads float sum-of-squares directly and computes sqrt inline.

        // ── 6. Adam update ────────────────────────────────────────
        self.launch_adam_update()?;
@@ -3628,15 +3686,15 @@ impl GpuDqnTrainer {
        };

        let grad_ptr = self.ptrs.grad_buf;
-        let norm_ptr = self.ptrs.grad_norm_buf;
+        let norm_f32_ptr = self.ptrs.grad_norm_f32_buf;

        // Safety: argument order matches the extern "C" kernel signature exactly.
-        // grad_buf has size = total_params; grad_norm_buf has size = 1.
+        // grad_buf has size = total_params; grad_norm_f32_buf has size = 1 float.
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_standalone)
                .arg(&grad_ptr)
-                .arg(&norm_ptr)
+                .arg(&norm_f32_ptr)
                .arg(&tp)
                .launch(launch_cfg)
                .map_err(|e| {
@@ -3678,12 +3736,12 @@ impl GpuDqnTrainer {
        let grad_ptr = self.ptrs.grad_buf;
        let m_ptr = self.ptrs.m_buf;
        let v_ptr = self.ptrs.v_buf;
-        let norm_ptr = self.ptrs.grad_norm_buf;
+        let norm_f32_ptr = self.ptrs.grad_norm_f32_buf;
        let t_ptr = self.ptrs.t_buf;

        // Safety: argument order matches the extern "C" kernel signature exactly.
        // All buffers are pre-allocated with size = total_params.
-        // grad_norm_buf contains the completed norm from launch_grad_norm().
+        // grad_norm_f32_buf contains float sum-of-squares from launch_grad_norm().
        // t_buf is a device pointer (kernel reads *t_ptr) for CUDA Graph compatibility.
        unsafe {
            self.stream
@@ -3692,7 +3750,7 @@ impl GpuDqnTrainer {
                .arg(&grad_ptr)
                .arg(&m_ptr)
                .arg(&v_ptr)
-                .arg(&norm_ptr)
+                .arg(&norm_f32_ptr)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
@@ -3967,11 +4025,11 @@ impl GpuDqnTrainer {
 /// bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm,
 /// clipped_saxpy, clip_grad.
 ///
-/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, clipped_saxpy, clip_grad)`.
+/// Returns `(grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm, clipped_saxpy, clip_grad)`.
 fn compile_training_kernels(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
-) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
+) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
    info!(
        state_dim = config.state_dim,
        total_params = compute_total_params(config),
@@ -3985,6 +4043,8 @@ fn compile_training_kernels(

    let grad_norm = module.load_function("dqn_grad_norm_kernel")
        .map_err(|e| MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}")))?;
+    let grad_norm_finalize = module.load_function("dqn_grad_norm_finalize")
+        .map_err(|e| MLError::ModelError(format!("dqn_grad_norm_finalize load: {e}")))?;
    let adam_update = module.load_function("dqn_adam_update_kernel")
        .map_err(|e| MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")))?;
    let f32_to_bf16 = module.load_function("f32_to_bf16_kernel")
@@ -4009,7 +4069,7 @@ fn compile_training_kernels(
        .map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?;

    info!("GpuDqnTrainer: 12 utility kernels loaded from precompiled cubin");
-    Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad))
+    Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad))
 }

 /// Load the standalone Polyak EMA kernel from precompiled cubin.
--- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs
@@ -276,6 +276,8 @@ impl GpuTrainingGuard {
        warmup: bool,
    ) -> Result<GuardResult, MLError> {
        let _nvtx = NvtxRange::new("gpu_training_guard");
+        // Disable event tracking — loss/grad buffers may be captured in CUDA graphs.
+        let _eg = super::gpu_dqn_trainer::EventTrackingGuard::new(self.stream.context());

        // Read PREVIOUS step's result (one-step delay; first step returns defaults)
        let prev_result = if self.guard_has_prev {
@@ -313,14 +315,17 @@ impl GpuTrainingGuard {
            shared_mem_bytes: 0,
        };

-        // Safety: loss_gpu and grad_norm_gpu are valid F32 CudaSlice on the same context.
+        // Pass raw device pointers to bypass cudarc event tracking on graph-captured buffers.
+        let loss_ptr = loss_gpu.raw_ptr();
+        let grad_ptr = grad_norm_gpu.raw_ptr();
+        let acc_ptr = self.acc_buf.raw_ptr();
        unsafe {
            self.stream
                .launch_builder(&self.fused_check_accum_func)
-                .arg(loss_gpu)
-                .arg(grad_norm_gpu)
+                .arg(&loss_ptr)
+                .arg(&grad_ptr)
                .arg(&write_dev_ptr)
-                .arg(&mut self.acc_buf)
+                .arg(&acc_ptr)
                .arg(&clip_threshold)
                .arg(&collapse_threshold)
                .arg(&warmup_int)
--- a/crates/ml/src/trainers/dqn/config.rs
+++ b/crates/ml/src/trainers/dqn/config.rs
@@ -662,10 +662,11 @@ impl DQNAgentType {
    /// Update replay buffer priorities from GPU-resident tensors (GpuPrioritized only).
    pub fn update_priorities_gpu(
        &self,
-        indices: &GpuTensor,
-        td_errors: &GpuTensor,
+        indices: &cudarc::driver::CudaSlice<u32>,
+        td_errors: &cudarc::driver::CudaSlice<half::bf16>,
+        bs: usize,
    ) -> Result<(), MLError> {
-        self.memory().update_priorities_gpu(indices, td_errors)
+        self.memory().update_priorities_gpu(indices, td_errors, bs)
    }

    /// Step the replay buffer training counter for beta annealing (PER only).
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -802,15 +802,9 @@ impl FusedTrainingCtx {
        if let (Some(priorities_tensor), Some((alpha, epsilon))) =
            (agent.priorities_tensor()?, agent.per_alpha_epsilon())
        {
-            // GpuBatch::indices stores u32 bit-patterns in f32 CudaSlice containers.
-            // Reinterpret the f32 CudaSlice as u32 to pass to the PER kernel.
-            let idx_f32 = effective_gpu.indices.data();
-            let idx_u32: &cudarc::driver::CudaSlice<u32> = unsafe {
-                &*(idx_f32 as *const cudarc::driver::CudaSlice<half::bf16>
-                    as *const cudarc::driver::CudaSlice<u32>)
-            };
+            // GpuBatch::indices is now native CudaSlice<u32> — pass directly.
            self.trainer.update_priorities_cuda(
-                idx_u32,
+                &effective_gpu.indices,
                priorities_tensor.data(),
                alpha,
                epsilon,
@@ -1171,6 +1165,18 @@ impl FusedTrainingCtx {
        self.trainer.set_c51_alpha(alpha);
    }

+    /// Return a reference to the GPU-resident total_loss scalar (bf16).
+    /// Written by the CUDA graph's loss kernel — valid after `replay_forward()`.
+    pub(crate) fn loss_gpu_buf(&self) -> &cudarc::driver::CudaSlice<half::bf16> {
+        &self.trainer.total_loss_buf
+    }
+
+    /// Return a reference to the GPU-resident grad_norm scalar (bf16).
+    /// Written by the CUDA graph's grad_norm kernel — valid after `replay_adam()`.
+    pub(crate) fn grad_norm_gpu_buf(&self) -> &cudarc::driver::CudaSlice<half::bf16> {
+        &self.trainer.grad_norm_buf
+    }
+
    /// Get CUDA stream reference for DtoH transfers after compute_q_values.
    pub(crate) fn stream(&self) -> &Arc<cudarc::driver::CudaStream> {
        &self.stream
--- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
+++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
@@ -78,6 +78,7 @@ fn insert_random_batch(
 }

 /// Download a CudaSlice<f32> to host for assertion.
+#[allow(dead_code)]
 fn dtoh_f32(slice: &cudarc::driver::CudaSlice<f32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<f32>> {
    let n = slice.len();
    let mut host = vec![0.0_f32; n];
@@ -85,6 +86,14 @@ fn dtoh_f32(slice: &cudarc::driver::CudaSlice<f32>, stream: &Arc<cudarc::driver:
    Ok(host)
 }

+/// Download a CudaSlice<u16> (bf16) to host as f32 for assertion.
+fn dtoh_bf16_as_f32(slice: &cudarc::driver::CudaSlice<u16>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<f32>> {
+    let n = slice.len();
+    let mut host = vec![0_u16; n];
+    stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh u16: {e}"))?;
+    Ok(host.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect())
+}
+
 /// Download a CudaSlice<u32> to host for assertion.
 fn dtoh_u32(slice: &cudarc::driver::CudaSlice<u32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<u32>> {
    let n = slice.len();
@@ -118,7 +127,7 @@ fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> {
    assert_eq!(batch.batch_size, 16);

    // Download weights and indices to host for assertions
-    let weights_host = dtoh_f32(&batch.weights, &stream)?;
+    let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?;
    let indices_host = dtoh_u32(&batch.indices, &stream)?;

    // All weights must be positive
@@ -154,7 +163,7 @@ fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> {
    assert_eq!(batch.batch_size, 16);

    // Download weights and indices to host for assertions
-    let weights_host = dtoh_f32(&batch.weights, &stream)?;
+    let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?;
    let indices_host = dtoh_u32(&batch.indices, &stream)?;

    // All weights must be positive
@@ -193,7 +202,7 @@ fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {

    // Verify sampling still works with updated priorities (proves update succeeded)
    let batch2 = buf.sample_proportional(16)?;
-    let weights_host = dtoh_f32(&batch2.weights, &stream)?;
+    let weights_host = dtoh_bf16_as_f32(&batch2.weights, &stream)?;

    // Weights should still be positive after priority update
    for (i, &w) in weights_host.iter().enumerate() {
--- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
+++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
@@ -116,9 +116,10 @@ fn test_per_weights_valid() -> anyhow::Result<()> {

    let batch = buf.sample_proportional(16)?;

-    // Download weights to host for assertion
-    let mut weights_host = vec![0.0_f32; batch.weights.len()];
-    stream.memcpy_dtoh(&batch.weights, &mut weights_host).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback
+    // Download bf16 weights to host for assertion
+    let mut weights_bf16 = vec![0_u16; batch.weights.len()];
+    stream.memcpy_dtoh(&batch.weights, &mut weights_bf16).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback
+    let weights_host: Vec<f32> = weights_bf16.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect();

    for (i, &w) in weights_host.iter().enumerate() {
        assert!(w > 0.0, "weight[{i}] not positive: {w}");
--- a/crates/ml/src/trainers/dqn/trainer/train_step.rs
+++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs
@@ -120,14 +120,17 @@ impl DQNTrainer {
                let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64;
                let past_warmup = self.gradient_logging_step as u64 > warmup_steps;

-                let loss_slice = gpu_result.loss_cuda_slice()
-                    .map_err(|e| anyhow::anyhow!("GPU guard loss CudaSlice: {e}"))?;
-                let grad_slice = gpu_result.grad_norm_cuda_slice()
-                    .map_err(|e| anyhow::anyhow!("GPU guard grad CudaSlice: {e}"))?;
+                // Read loss/grad directly from fused trainer's GPU buffers
+                // (not from GpuTrainResult which returns hardcoded zeros per-step).
+                let fused = self.fused_ctx.as_mut()
+                    .ok_or_else(|| anyhow::anyhow!("fused_ctx required for training guard"))?;
+
+                let loss_buf = fused.loss_gpu_buf();
+                let grad_buf = fused.grad_norm_gpu_buf();
                let result = guard
                    .check_and_accumulate(
-                        &loss_slice,
-                        &grad_slice,
+                        loss_buf,
+                        grad_buf,
                        1e6_f32, // loss clip threshold
                        grad_collapse_threshold,
                        !past_warmup,
@@ -273,7 +276,7 @@ impl DQNTrainer {
        #[allow(unused_mut, unused_assignments, unused_variables)]
        let mut final_grad_norm = 0.0_f32;
        let mut gpu_td_errors: Vec<GpuTensor> = Vec::new();
-        let mut gpu_indices: Vec<GpuTensor> = Vec::new();
+        let mut gpu_indices: Vec<cudarc::driver::CudaSlice<u32>> = Vec::new();
        let mut gpu_loss_tensors: Vec<GpuTensor> = Vec::new();
        let mut gpu_grad_tensors: Vec<GpuTensor> = Vec::new();

@@ -409,13 +412,24 @@ impl DQNTrainer {
            if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() {
                let cat_stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream for PER concat"))?;
                let td_refs: Vec<&GpuTensor> = gpu_td_errors.iter().collect();
-                let idx_refs: Vec<&GpuTensor> = gpu_indices.iter().collect();
                let td_cat = GpuTensor::cat(&td_refs, 0, cat_stream)
                    .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?;
-                let idx_cat = GpuTensor::cat(&idx_refs, 0, cat_stream)
-                    .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?;
+
+                // Concatenate CudaSlice<u32> indices via GPU DtoD copies
+                let total_idx_len: usize = gpu_indices.iter().map(|s| s.len()).sum();
+                let mut idx_cat = cat_stream.alloc_zeros::<u32>(total_idx_len)
+                    .map_err(|e| anyhow::anyhow!("GPU PER index alloc: {e}"))?;
+                let mut offset = 0usize;
+                for idx_slice in &gpu_indices {
+                    let len = idx_slice.len();
+                    let mut sub = idx_cat.slice_mut(offset..offset + len);
+                    cat_stream.memcpy_dtod(idx_slice, &mut sub)
+                        .map_err(|e| anyhow::anyhow!("GPU PER index DtoD concat: {e}"))?;
+                    offset += len;
+                }
+
                agent
-                    .update_priorities_gpu(&idx_cat, &td_cat)
+                    .update_priorities_gpu(&idx_cat, td_cat.data(), total_idx_len)
                    .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?;
            } else if !all_indices.is_empty() {
                agent
--- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs
+++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs
@@ -1263,17 +1263,29 @@ impl DQNTrainer {

                        let guard_start = std::time::Instant::now();
                        if let Some(ref mut guard) = self.training_guard {
-                            let loss_slice = _gpu_result.loss_cuda_slice()
-                                .map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?;
-                            let grad_slice = _gpu_result.grad_norm_cuda_slice()
-                                .map_err(|e| anyhow::anyhow!("guard grad CudaSlice: {e}"))?;
-                            let gr = guard.check_and_accumulate(
-                                &loss_slice,
-                                &grad_slice,
-                                1e6_f32,
-                                guard_collapse_thresh,
-                                !guard_past_warmup,
-                            ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?;
+                            // Read loss/grad directly from fused trainer's GPU buffers.
+                            // GpuTrainResult returns hardcoded zeros per-step (no sync).
+                            let gr = if let Some(ref fused) = self.fused_ctx {
+                                guard.check_and_accumulate(
+                                    fused.loss_gpu_buf(),
+                                    fused.grad_norm_gpu_buf(),
+                                    1e6_f32,
+                                    guard_collapse_thresh,
+                                    !guard_past_warmup,
+                                ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?
+                            } else {
+                                let loss_slice = _gpu_result.loss_cuda_slice()
+                                    .map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?;
+                                let grad_slice = _gpu_result.grad_norm_cuda_slice()
+                                    .map_err(|e| anyhow::anyhow!("guard grad CudaSlice: {e}"))?;
+                                guard.check_and_accumulate(
+                                    &loss_slice,
+                                    &grad_slice,
+                                    1e6_f32,
+                                    guard_collapse_thresh,
+                                    !guard_past_warmup,
+                                ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?
+                            };
                            if gr.halt_nan {
                                return Err(anyhow::anyhow!(
                                    "NaN/Inf at step {}: loss={}, grad={}",
@@ -1464,9 +1476,13 @@ impl DQNTrainer {
        train_step_count: usize,
        monitor: &mut TrainingMonitor,
    ) -> Result<EpochBoundaryMetrics> {
+
        let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard {
            let (al, ag) = guard.read_accumulators()
                .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?;
+            warn!("EPOCH_DIAG[{}]: acc avg_loss={:.6}, avg_grad={:.6}", epoch, al, ag);
+            guard.reset_accumulators()
+                .map_err(|e| anyhow::anyhow!("guard reset_accumulators: {e}"))?;
            (al as f32, ag as f32)
        } else {
            (0.0_f32, 0.0_f32)
--- a/crates/ml/src/training_profile.rs
+++ b/crates/ml/src/training_profile.rs
@@ -1128,8 +1128,8 @@ mod tests {
        );
        // reward_scale should be applied
        assert!(
-            (hp.reward_scale - 10.0).abs() < 0.01,
-            "reward_scale should be 10.0, got {}",
+            (hp.reward_scale - 1.0).abs() < 0.01,
+            "reward_scale should be 1.0, got {}",
            hp.reward_scale
        );
        // v_min/v_max should be recomputed from reward_scale + gamma
@@ -1159,14 +1159,14 @@ mod tests {
            hp.iqn_lambda
        );
        assert!(
-            (hp.spectral_norm_sigma_max - 3.0).abs() < 0.01,
-            "spectral_norm_sigma_max should be 3.0 from smoketest [advanced], got {}",
+            (hp.spectral_norm_sigma_max - 1.5).abs() < 0.01,
+            "spectral_norm_sigma_max should be 1.5 from smoketest [advanced], got {}",
            hp.spectral_norm_sigma_max
        );
        // Exploration section new fields
        assert!(
-            (hp.noisy_sigma_init - 0.5).abs() < 0.01,
-            "noisy_sigma_init should be 0.5 from smoketest [exploration], got {}",
+            (hp.noisy_sigma_init - 0.3).abs() < 0.01,
+            "noisy_sigma_init should be 0.3 from smoketest [exploration], got {}",
            hp.noisy_sigma_init
        );
        assert!(
--- a/crates/ml/tests/gpu_per_integration_test.rs
+++ b/crates/ml/tests/gpu_per_integration_test.rs
@@ -153,8 +153,8 @@ fn test_gpu_per_training_loop_cycle() {

        // Simulate TD errors: decreasing magnitude over epochs
        let scale = 1.0_f32 / (epoch as f32 + 1.0);
-        let td_host = vec![scale; batch_size];
-        let td_errors: CudaSlice<f32> = stream.clone_htod(&td_host).unwrap();
+        let td_host: Vec<half::bf16> = vec![half::bf16::from_f32(scale); batch_size];
+        let td_errors: CudaSlice<half::bf16> = stream.clone_htod(&td_host).unwrap();

        buf.update_priorities_gpu(&batch.indices, &td_errors, batch_size).unwrap();
        buf.step();
@@ -178,7 +178,7 @@ fn test_gpu_per_priorities_affect_sampling() {

    // Set experience 0 to very high priority
    let high_idx: CudaSlice<u32> = stream.clone_htod(&[0_u32]).unwrap();
-    let high_td: CudaSlice<f32> = stream.clone_htod(&[100.0_f32]).unwrap();
+    let high_td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32(100.0)]).unwrap();
    buf.update_priorities_gpu(&high_idx, &high_td, 1).unwrap();

    // Sample 1000 times in batches of 10, count how often index 0 appears
@@ -222,7 +222,7 @@ fn test_gpu_per_sampling_distribution_differs_from_uniform() {
    // This creates a strong left-to-right priority gradient
    for i in 0..50 {
        let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
-        let td: CudaSlice<f32> = stream.clone_htod(&[(i as f32 + 1.0)]).unwrap();
+        let td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32(i as f32 + 1.0)]).unwrap();
        buf.update_priorities_gpu(&idx, &td, 1).unwrap();
    }

@@ -264,15 +264,16 @@ fn test_gpu_per_is_weights_correct_range() {
    // Set varied priorities
    for i in 0..200 {
        let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
-        let td: CudaSlice<f32> = stream.clone_htod(&[((i % 10) as f32 + 1.0)]).unwrap();
+        let td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32((i % 10) as f32 + 1.0)]).unwrap();
        buf.update_priorities_gpu(&idx, &td, 1).unwrap();
    }

    let batch = buf.sample_proportional(64).unwrap();

-    // Download f32 weights to host for validation
-    let mut weights_host = vec![0.0_f32; 64];
-    stream.memcpy_dtoh(&batch.weights, &mut weights_host).unwrap();
+    // Download bf16 weights to host, convert to f32 for validation
+    let mut weights_bf16 = vec![half::bf16::ZERO; 64];
+    stream.memcpy_dtoh(&batch.weights, &mut weights_bf16).unwrap();
+    let weights_host: Vec<f32> = weights_bf16.iter().map(|x| x.to_f32()).collect();

    let w_min = weights_host.iter().copied().fold(f32::INFINITY, f32::min);
    let w_max = weights_host.iter().copied().fold(f32::NEG_INFINITY, f32::max);