Files
foxhunt/docs/superpowers/plans/2026-04-06-per-tree-fix-and-zero-sync-readback.md
jgrusewski d4269252c6 docs: H100 maximum performance spec and implementation plan
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-06 11:15:24 +02:00

26 KiB
Raw Blame History

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

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Fix broken PER segment tree updates, remove all dead code from H100 hang fixes, and consolidate all GPU→CPU readbacks into a single pinned host buffer with zero synchronous copies in the training hot loop.

Architecture: Delete the orphaned per_update_kernel.cu and route PER priority updates through the existing seg_tree_update kernel (which correctly propagates the segment tree). Extend the existing 12-byte pinned readback buffer to 64 bytes to cover Q-stats and causal sensitivity. Replace the causal sensitivity sync readback with a GPU-side mean reduction kernel.

Tech Stack: Rust 1.85, CUDA 12.4 (cudarc 0.19), half (bf16), parking_lot::Mutex


Task 1: Delete per_update_kernel.cu and All References

Files:

  • Delete: crates/ml/src/cuda_pipeline/per_update_kernel.cu

  • Modify: crates/ml/build.rs:75

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:65,496,2230,2726,3677-3710,6030-6038

  • Step 1: Delete the kernel file

rm crates/ml/src/cuda_pipeline/per_update_kernel.cu
  • Step 2: Remove from build.rs

In crates/ml/build.rs, remove the line:

        "per_update_kernel.cu",

from the CUDA source files list (line 75).

  • Step 3: Remove PER_UPDATE_CUBIN include

In crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs, delete line 65:

static PER_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_update_kernel.cubin"));
  • Step 4: Remove per_update_kernel field from struct

In gpu_dqn_trainer.rs, delete the field at line 496:

    per_update_kernel: CudaFunction,
  • Step 5: Remove per_update_kernel compilation in constructor

In gpu_dqn_trainer.rs, delete line 2230:

        let per_update_kernel = compile_per_update_kernel(&stream)?;

And delete the assignment at line 2726:

            per_update_kernel,
  • Step 6: Delete compile_per_update_kernel function

In gpu_dqn_trainer.rs, delete the entire function at lines 6030-6038:

fn compile_per_update_kernel(stream: &Arc<CudaStream>) -> Result<CudaFunction, MLError> {
    let context = stream.context();
    let module = context.load_cubin(PER_UPDATE_CUBIN.to_vec()).map_err(|e| {
        MLError::ModelError(format!("per_update cubin load: {e}"))
    })?;
    module.load_function("per_update_priorities_kernel").map_err(|e| {
        MLError::ModelError(format!("per_update_priorities_kernel load: {e}"))
    })
}
  • Step 7: Delete update_priorities_cuda_raw method

In gpu_dqn_trainer.rs, delete the entire method at lines 3677-3710:

    pub fn update_priorities_cuda_raw(
        &mut self,
        indices_ptr: u64,
        priorities_ptr: u64,
        capacity: usize,
        alpha: f32,
        epsilon: f32,
    ) -> Result<(), MLError> {
        // ... entire method body ...
    }
  • Step 8: Delete batch_max_buf field and allocation

In gpu_dqn_trainer.rs:

  • Delete field declaration at line 630:

        batch_max_buf: CudaSlice<f32>,
    
  • Delete allocation at line 2330:

        let batch_max_buf: CudaSlice<f32> = stream.alloc_zeros(1)
            .map_err(|e| MLError::ModelError(format!("alloc batch_max: {e}")))?;
    
  • Delete assignment in struct init (around line 2805):

                batch_max_buf,
    
  • Step 9: Delete batch_max_readback_and_reset method

In gpu_dqn_trainer.rs, delete the entire method at lines 3712-3720:

    pub fn batch_max_readback_and_reset(&mut self) -> Result<f32, MLError> {
        let mut host = [0.0_f32; 1];
        self.stream.memcpy_dtoh(&self.batch_max_buf, &mut host)
            .map_err(|e| MLError::ModelError(format!("batch_max DtoH: {e}")))?;
        self.stream
            .memset_zeros(&mut self.batch_max_buf)
            .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
        Ok(host[0])
    }
  • Step 10: Build and verify

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 Expected: Compilation errors in fused_training.rs (calls to deleted methods) — these are fixed in Task 2.

  • Step 11: Commit
git add -u crates/ml/src/cuda_pipeline/per_update_kernel.cu crates/ml/build.rs crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "refactor: delete orphaned per_update_kernel and batch_max_buf

The per_update_priorities_kernel only wrote to the priorities[] buffer
without updating the PER segment tree — sampling used stale tree values.
All PER priority updates now go through seg_tree_update which correctly
propagates from leaf to root."

Task 2: Route PER Updates Through seg_tree_update

Files:

  • Modify: crates/ml/src/trainers/dqn/fused_training.rs:753-768

  • Modify: crates/ml/src/trainers/dqn/config.rs:638-644

  • Modify: crates/ml-dqn/src/replay_buffer_type.rs:358-373,388-401

  • Step 1: Add update_priorities_from_td to ReplayBufferType

In crates/ml-dqn/src/replay_buffer_type.rs, add a new method after update_priorities_gpu (after line 373):

    /// Update PER priorities using td_errors and the indices from the last sample.
    ///
    /// 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 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(_) => Ok(()),
        }
    }
  • Step 2: Add update_priorities_from_td passthrough on agent config

In crates/ml/src/trainers/dqn/config.rs, replace priorities_f32_ptr (lines 638-644) with:

    /// Update PER priorities from GPU-resident TD errors.
    ///
    /// 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)
    }
  • Step 3: Remove priorities_f32_ptr from ReplayBufferType

In crates/ml-dqn/src/replay_buffer_type.rs, delete priorities_f32_ptr method (lines 388-401):

    pub fn priorities_f32_ptr(&self) -> Option<(u64, usize)> {
        match self {
            Self::GpuPrioritized(buffer) => {
                let buf = buffer.lock();
                let slice = buf.gpu.priorities_slice();
                Some((slice.raw_ptr(), slice.len()))
            }
            Self::Uniform(_) | Self::Prioritized(_) => None,
        }
    }
  • Step 4: Update fused_training.rs PER update section

In crates/ml/src/trainers/dqn/fused_training.rs, replace lines 753-768:

        if step < 3 { eprintln!("H100_STEP: step {step} — PER priority update"); }
        // ── Step 6: PER priority update (outside graph) ──────────────────
        // Pass the f32 priorities slice directly — no GPU→CPU→GPU roundtrip.
        if let (Some((prio_ptr, prio_cap)), Some((alpha, epsilon))) =
            (agent.priorities_f32_ptr(), agent.per_alpha_epsilon())
        {
            self.trainer.update_priorities_cuda_raw(
                gpu_batch.indices_ptr,
                prio_ptr,
                prio_cap,
                alpha,
                epsilon,
            ).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?;
        }
        if step < 3 { eprintln!("H100_HANG4: step {step} — PER update done"); }

with:

        if step < 3 { eprintln!("H100_STEP: step {step} — PER priority update"); }
        // ── 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"); }
  • Step 5: Build and verify

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 Expected: Clean compilation (0 errors, 0 warnings).

  • Step 6: Run smoke tests

Run: SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result' Expected: test result: ok. 19 passed; 0 failed;

  • Step 7: Commit
git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/config.rs crates/ml-dqn/src/replay_buffer_type.rs
git commit -m "fix: route PER updates through seg_tree_update kernel

Fused training now calls agent.update_priorities_from_td() which goes
through the replay buffer's seg_tree_update kernel. This correctly
writes priorities, updates tree leaves, and propagates sums to root.

Previously per_update_priorities_kernel only wrote to priorities[]
without touching the segment tree — sampling was using stale values."

Task 3: Delete Dead Code (GpuTrainResult, cast helpers, stale annotations)

Files:

  • Modify: crates/ml-dqn/src/dqn.rs:824-873

  • Modify: crates/ml-dqn/src/gpu_replay_buffer.rs:119-148,175-178

  • Step 1: Delete GpuTrainResult struct and all methods

In crates/ml-dqn/src/dqn.rs, delete lines 824-873 (the entire pub struct GpuTrainResult and its impl block including loss_cuda_slice, grad_norm_cuda_slice, from_fused_scalars, gpu_tensor_to_cuda_slice, loss_scalar, grad_norm_scalar).

Also remove any use imports for types only needed by GpuTrainResult (check if GpuTensor, CudaStream, Arc are still used elsewhere in the file before removing).

  • Step 2: Delete get_cast_kernels_f32_to_u32

In crates/ml-dqn/src/gpu_replay_buffer.rs, delete the get_cast_kernels_f32_to_u32 function (around line 119-122):

pub fn get_cast_kernels_f32_to_u32(stream: &Arc<CudaStream>) -> Result<F32ToU32Caster, MLError> {
    let kernels = get_cast_kernels(stream)?;
    Ok(F32ToU32Caster { kernel: &kernels.f32_to_u32, bf16_to_f32_kernel: &kernels.bf16_to_f32 })
}

Also delete F32ToU32Caster struct if it exists and is only used here.

  • Step 3: Delete f32_slice_to_gpu_tensor_gpu

In crates/ml-dqn/src/gpu_replay_buffer.rs, delete the #[allow(dead_code)] annotated function f32_slice_to_gpu_tensor_gpu (around lines 125-148).

  • Step 4: Remove #[allow(dead_code)] from gather_f32

In crates/ml-dqn/src/gpu_replay_buffer.rs, check if gather_f32 field (line 176) is actually used. If unused, delete the field and its kernel load. If used, remove only the #[allow(dead_code)] annotation.

  • Step 5: Clean up unused imports

Run: SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | grep 'unused import' Remove any unused imports flagged by the compiler.

  • Step 6: Build and verify

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 Expected: Clean compilation.

  • Step 7: Run smoke tests

Run: SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result' Expected: test result: ok. 19 passed; 0 failed;

  • Step 8: Commit
git add -u crates/ml-dqn/
git commit -m "cleanup: remove dead GpuTrainResult, cast helpers, stale annotations

GpuTrainResult was replaced by FusedStepResult (CPU-only f32 scalars).
get_cast_kernels_f32_to_u32 and f32_slice_to_gpu_tensor_gpu were only
used by the deleted priorities_tensor() roundtrip path."

Task 4: Extend Pinned Readback Buffer (Q-Stats Migration)

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:870,2808-2815,920,4054-4125

  • Step 1: Extend pinned buffer from 3 to 16 floats

In gpu_dqn_trainer.rs, change the readback_pinned allocation (line 2808-2815) from 3 * size_of::<f32>() to 16 * size_of::<f32>():

readback_pinned: {
    let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
    unsafe {
        cudarc::driver::result::malloc_host(16 * std::mem::size_of::<f32>(), flags)
            .map_err(|e| MLError::ModelError(format!("pinned readback alloc: {e}")))?
            as *mut f32
    }
},

Add a comment above the field declaration documenting the layout:

/// Pinned host buffer [16 × f32] for async DtoH scalar readback.
/// Layout:
///   [0]=loss, [1]=mse_loss, [2]=grad_norm_sq          (per-step, from graph_adam)
///   [3]=q_mean, [4]=q_min, [5]=q_max, [6]=q_var, [7]=avg_max_q  (every 50 steps)
///   [8]=causal_mean_sens                               (every N steps)
///   [9..16]=nan_flags (as reinterpreted i32)           (per-step)
readback_pinned: *mut f32,
  • Step 2: Remove q_stats_pending field

In gpu_dqn_trainer.rs, delete the field at line 870:

    q_stats_pending: [f32; 5],

And its initialization in the constructor (search for q_stats_pending: [0.0; 5]).

  • Step 3: Update reduce_current_q_stats to use pinned buffer

In gpu_dqn_trainer.rs, replace the async DtoH at lines 4106-4113:

        // 3. Issue async DtoH for 5 f32s from q_stats_buf
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
                self.q_stats_pending.as_mut_ptr().cast(),
                self.q_stats_buf.raw_ptr(),
                5 * std::mem::size_of::<f32>(),
                self.stream.cu_stream(),
            );
        }

with:

        // 3. Issue async DtoH for 5 f32s into pinned buffer at offset 3
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
                self.readback_pinned.add(3).cast(),
                self.q_stats_buf.raw_ptr(),
                5 * std::mem::size_of::<f32>(),
                self.stream.cu_stream(),
            );
        }
  • Step 4: Update Q-stats result reading to use pinned buffer

In the reduce_current_q_stats method, replace all reads from self.q_stats_pending[N] with reads from the pinned buffer:

        let prev = if self.q_stats_ready {
            if let Some(ref event) = self.q_stats_event {
                if !event.is_complete() {
                    event.synchronize().map_err(|e|
                        MLError::ModelError(format!("q_stats event sync: {e}"))
                    )?;
                }
            }
            self.q_stats_ready = false;
            unsafe {
                QValueStatsResult {
                    avg_max_q: *self.readback_pinned.add(7) as f64,
                    q_min: *self.readback_pinned.add(4),
                    q_max: *self.readback_pinned.add(5),
                    q_mean: *self.readback_pinned.add(3),
                    q_variance: *self.readback_pinned.add(6),
                }
            }
        } else {
            QValueStatsResult {
                avg_max_q: 0.0, q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_variance: 0.0,
            }
        };
  • Step 5: Update flush_q_stats_readback similarly

In flush_q_stats_readback (line 4129+), replace reads from self.q_stats_pending with the same pinned buffer reads at offsets 3-7.

  • Step 6: Build and verify

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 Expected: Clean compilation.

  • Step 7: Run smoke tests

Run: SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result' Expected: test result: ok. 19 passed; 0 failed;

  • Step 8: Commit
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "perf: migrate Q-stats readback to consolidated pinned buffer

Extend readback_pinned from 12 to 64 bytes. Q-stats async DtoH now
targets the pinned buffer at offset 3, fixing a latent H100 hang where
cuMemcpyDtoHAsync into non-pinned q_stats_pending degrades to sync."

Task 5: Causal Sensitivity GPU-Native Mean Reduction

Files:

  • Modify: crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3335-3430

  • Step 1: Add causal_mean_reduce kernel

In crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu, add after the existing causal_q_delta_reduce kernel:

/**
 * Reduce causal_sensitivity_buf[0..n_features] to a single mean scalar.
 * Writes result to out[0]. Single block, 256 threads.
 */
extern "C" __global__ void causal_mean_reduce(
    const float* __restrict__ sensitivity,  /* [market_dim] */
    float* __restrict__ out,                /* [1] — mean sensitivity */
    int n_features)
{
    __shared__ float s_sum[256];
    int tid = threadIdx.x;
    float val = (tid < n_features) ? sensitivity[tid] : 0.0f;
    s_sum[tid] = val;
    __syncthreads();

    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) s_sum[tid] += s_sum[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        out[0] = s_sum[0] / (float)n_features;
    }
}
  • Step 2: Load the new kernel in gpu_dqn_trainer.rs

In the constructor, after loading other utility kernels, add:

        let causal_mean_reduce_kernel = utility_module
            .load_function("causal_mean_reduce")
            .map_err(|e| MLError::ModelError(format!("causal_mean_reduce load: {e}")))?;

Add the field to the struct:

    causal_mean_reduce_kernel: CudaFunction,

And assign it in the struct initializer.

  • Step 3: Replace sync readback with GPU reduction + async pinned readback

In gpu_dqn_trainer.rs, replace the sync section at lines 3420-3430:

        // Single DtoH readback at the end — read mean sensitivity
        let n_features = market_dim.min(14);
        let mut host_sens = vec![0.0_f32; market_dim];
        self.stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("causal final sync: {e}")))?;
        self.stream.memcpy_dtoh(&self.causal_sensitivity_buf, &mut host_sens[..market_dim])
            .map_err(|e| MLError::ModelError(format!("causal readback: {e}")))?;
        let total: f32 = host_sens[..n_features].iter().sum();
        Ok(total / n_features as f32)

with:

        // GPU-side mean reduction → single scalar async readback to pinned buffer
        let n_features = market_dim.min(14) as i32;
        unsafe {
            self.stream
                .launch_builder(&self.causal_mean_reduce_kernel)
                .arg(&self.causal_sensitivity_buf)
                .arg(&self.readback_pinned.add(8))  // offset 8 in pinned buffer
                .arg(&n_features)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("causal_mean_reduce: {e}")))?;
        }

        // The result will be available at readback_pinned[8] after the stream
        // catches up. The caller reads it on the next invocation (double-buffered).
        // Return previous result for now.
        let prev = unsafe { *self.readback_pinned.add(8) };
        Ok(prev)

Note: readback_pinned is host-mapped (CU_MEMHOSTALLOC_DEVICEMAP), so the kernel can write directly to it via its device-mapped address. However, the kernel arg needs the device pointer to the pinned memory. We need to get the device pointer via cuMemHostGetDevicePointer. Let me adjust — instead, use cuMemcpyDtoHAsync from a 1-element GPU scratch buffer:

Replace the above with:

        // GPU-side mean reduction into a scratch scalar
        let n_features_i32 = market_dim.min(14) as i32;
        // Reuse batch_max scratch (1 f32) — we deleted the old one, allocate a small one
        // Actually, use q_stats_buf[0] as scratch (it's [5] f32, we only need [0])
        self.stream.memset_zeros(&mut self.causal_mean_scratch)
            .map_err(|e| MLError::ModelError(format!("zero causal_mean: {e}")))?;
        unsafe {
            self.stream
                .launch_builder(&self.causal_mean_reduce_kernel)
                .arg(&self.causal_sensitivity_buf)
                .arg(&self.causal_mean_scratch)
                .arg(&n_features_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("causal_mean_reduce: {e}")))?;
        }

        // Async DtoH: 1 f32 from causal_mean_scratch → pinned buffer offset 8
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
                self.readback_pinned.add(8).cast(),
                self.causal_mean_scratch.raw_ptr(),
                std::mem::size_of::<f32>(),
                self.stream.cu_stream(),
            );
        }

        // Return previous result (double-buffered — current result arrives by next call)
        let prev = unsafe { *self.readback_pinned.add(8) };
        Ok(prev)
  • Step 4: Add causal_mean_scratch field

In the struct, add:

    causal_mean_scratch: CudaSlice<f32>,

In the constructor, allocate:

        let causal_mean_scratch = stream.alloc_zeros::<f32>(1)
            .map_err(|e| MLError::ModelError(format!("alloc causal_mean_scratch: {e}")))?;

And assign in the struct initializer.

  • Step 5: Build and verify

Run: SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5 Expected: Clean compilation.

  • Step 6: Run smoke tests

Run: SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result' Expected: test result: ok. 19 passed; 0 failed;

  • Step 7: Commit
git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "perf: GPU-native causal sensitivity mean, zero sync in hot loop

Replace stream.synchronize() + sync memcpy_dtoh of per-feature
sensitivities with a single-block GPU mean reduction kernel + async
DtoH to pinned buffer. Eliminates the last synchronous readback in
the training hot loop."

Task 6: Migrate NaN Flags to Pinned Buffer

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

  • Step 1: Find read_nan_flags and reset_nan_flags

Search for these methods in gpu_dqn_trainer.rs. They use sync memcpy_dtoh for 8 × i32 (32 bytes). Since NaN flags are checked per-step (in run_nan_checks_post_forward), migrate the readback to the pinned buffer at offset 9 (slots 9-16 reinterpreted as i32).

  • Step 2: Update read_nan_flags to read from pinned buffer

Replace the sync memcpy_dtoh with a read from readback_pinned offsets 9-16 (reinterpreted as *const i32). The async write to pinned should be done in run_nan_checks_post_forward after the check kernel launches.

  • Step 3: Add async DtoH of nan_flags to pinned buffer

After the NaN check kernel launch, add:

unsafe {
    cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
        self.readback_pinned.add(9).cast(),
        self.nan_flags_buf.raw_ptr(),
        8 * std::mem::size_of::<i32>(),
        self.stream.cu_stream(),
    );
}

Note: 8 × i32 = 32 bytes = 8 × f32 slots. The pinned buffer has slots 9-16 available (7 slots × 4 bytes = 28 bytes). Check if 8 × i32 fits — it's 32 bytes but we have 7 × f32 = 28 bytes. Adjust pinned buffer to 17 floats if needed, or reduce nan_flags to 7.

Actually, nan_flags is [i32; 8] = 32 bytes. Pinned buffer slots 9-16 = 7 × f32 = 28 bytes. Not enough. Either extend to 17 f32 (68 bytes), or keep nan_flags as sync (it's only read on NaN detection, which is rare).

Decision: Keep read_nan_flags as sync. It's only called when the training guard detects NaN — a rare error path, not the hot loop. The spec says "nan_flags per-step" but the actual read is conditional on gr.halt_nan.

  • Step 4: Skip this task — nan_flags readback is cold path

The NaN check kernel writes flags per-step, but the CPU only reads them when halt_nan fires (rare error). No hot-path sync to eliminate.

  • Step 5: Commit (if any changes)

No changes needed — nan_flags stays as-is.


Task 7: Final Verification and Deploy

Files: None (testing only)

  • Step 1: Run full smoke test suite

Run: SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result' Expected: test result: ok. 19 passed; 0 failed;

  • Step 2: Run full workspace check for warnings

Run: SQLX_OFFLINE=true cargo check --workspace 2>&1 | grep -E 'warning|error' | head -20 Expected: No errors. Fix any warnings in modified files.

  • Step 3: Push and deploy to H100
git push origin main
./scripts/argo-train.sh dqn --baseline --watch

Monitor for:

  • H100_HANG4: step 0 — PER update done (confirms seg_tree_update completes)

  • H100_LOOP: step 0 — fused done, running guard (confirms no hang)

  • Epoch 1 appearing (confirms training progresses past step 0)

  • Step 4: Verify Q-stats still reported

In the H100 logs, check for Q-value statistics in epoch summaries:

Epoch N/200: Q-value range=[...], mean=...
  • Step 5: Verify causal sensitivity runs

Check for causal intervention log entries (if causal_config.interval triggers during early epochs).