Files
foxhunt/docs/superpowers/plans/2026-04-01-h100-gpu-optimization.md
jgrusewski 171fb0832c docs: H100 GPU optimization implementation plan — 10 tasks
Phase 1: CUDA events, dynamic batch size, async q-stats, phase overlap
Phase 2: Multi-stream branches, target/online parallelism
Phase 3: CUDA Graph experience collection, lock removal
Validation: H100 benchmark with epoch time + SM utilization targets

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-01 21:41:21 +02:00

21 KiB
Raw Blame History

H100 GPU Pipeline Optimization 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: Reduce DQN epoch time from 2 min to <10s on H100 by eliminating GPU idle time from blocking synchronize calls, maximizing batch size, and enabling multi-stream parallelism.

Architecture: Replace stream.synchronize() with CUDA events for async scalar readback. Let AutoBatchSizer drive batch size (removes static cap). Capture experience collection as a CUDA Graph. Fork child streams for independent branch/target forward passes.

Tech Stack: Rust, cudarc 0.17 (CudaEvent, CudaStream::fork), CUDA 12.4, H100 PCIe 80GB.


File Structure

File Action Change
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs Modify CUDA events in replay_adam_and_readback, async q-stats, multi-stream forward
crates/ml/src/cuda_pipeline/batched_forward.rs Modify Multi-stream branch dispatch
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs Modify CUDA Graph capture for timestep loop
crates/ml/src/trainers/dqn/trainer/training_loop.rs Modify Async training loop, event-based phase transitions
config/gpu/h100.toml Modify Remove static batch_size cap
crates/ml/src/trainers/dqn/trainer/constructor.rs Modify Let AutoBatchSizer drive batch_size
infra/k8s/argo/precompute-features-template.yaml Done Already committed (data-dir/SYMBOL alignment)

Phase 1: Eliminate GPU Idle Time (biggest impact)

Task 1: CUDA Events replace per-step stream.synchronize()

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:2865-2893

The replay_adam_and_readback() method calls stream.synchronize() on every training step (341×/epoch). This is the primary bottleneck — H100 finishes both CUDA graphs in microseconds then sits idle.

  • Step 1: Add event and async readback state to GpuDqnTrainer struct

In gpu_dqn_trainer.rs, add fields to the GpuDqnTrainer struct (near line 479):

/// CUDA event for async scalar readback (replaces stream.synchronize per step)
readback_event: Option<Arc<CudaEvent>>,
/// Host-side buffers for deferred readback (loss, grad_norm from previous step)
pending_loss: f32,
pending_grad_norm: f32,
/// Whether there's an in-flight readback to collect
readback_pending: bool,

Initialize in new():

readback_event: None,
pending_loss: 0.0,
pending_grad_norm: 0.0,
readback_pending: false,
  • Step 2: Rewrite replay_adam_and_readback to use events

Replace the synchronous replay_adam_and_readback method:

/// Collect scalars from the PREVIOUS step (non-blocking poll),
/// then launch grad_norm + adam for the CURRENT step and record an event.
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    // 1. Collect previous step's scalars (if any)
    let prev_scalars = if self.readback_pending {
        // Poll: is the previous step's event done?
        if let Some(ref event) = self.readback_event {
            if !event.is_complete() {
                // Still in-flight — wait on just this event (not full stream)
                event.synchronize().map_err(|e|
                    MLError::ModelError(format!("readback event sync: {e}"))
                )?;
            }
        }
        // Now safe to read the scalars from the previous step
        self.readback_pending = false;
        FusedTrainScalars {
            total_loss: self.pending_loss,
            grad_norm: self.pending_grad_norm,
        }
    } else {
        FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }
    };

    // 2. Launch current step: grad_norm + adam
    self.compute_grad_norm_outside_graph()?;
    self.replay_adam()?;

    // 3. Issue async DtoH for current step's scalars
    unsafe {
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            (&mut self.pending_loss as *mut f32).cast(),
            self.total_loss_buf.raw_ptr(),
            std::mem::size_of::<f32>(),
            *self.stream.cu_stream(),
        );
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            (&mut self.pending_grad_norm as *mut f32).cast(),
            self.grad_norm_f32_buf.raw_ptr(),
            std::mem::size_of::<f32>(),
            *self.stream.cu_stream(),
        );
    }

    // 4. Record event marking when the async DtoH completes
    self.readback_event = Some(self.stream.record_event(None).map_err(|e|
        MLError::ModelError(format!("readback event record: {e}"))
    )?);
    self.readback_pending = true;

    // Return PREVIOUS step's scalars (current step's will be collected next call)
    self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm];
    Ok(prev_scalars)
}

/// Flush: block until the last in-flight readback completes.
/// Called once at epoch end.
pub fn flush_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    if !self.readback_pending {
        return Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 });
    }
    if let Some(ref event) = self.readback_event {
        if !event.is_complete() {
            event.synchronize().map_err(|e|
                MLError::ModelError(format!("flush readback sync: {e}"))
            )?;
        }
    }
    self.readback_pending = false;
    Ok(FusedTrainScalars {
        total_loss: self.pending_loss,
        grad_norm: self.pending_grad_norm,
    })
}
  • Step 3: Update training loop to flush at epoch end

In training_loop.rs, after the training step loop (after the for chunk_start in ... loop ends), add:

// Flush the last in-flight readback
if let Some(ref mut fused) = self.fused_ctx {
    let final_scalars = fused.flush_readback()?;
    // Use final_scalars for epoch-end logging if needed
}
  • Step 4: Verify compile

Run: SQLX_OFFLINE=true cargo check -p ml --lib Expected: compiles clean

  • Step 5: Run local smoke test to validate loss trajectory

Run: FOXHUNT_TEST_DATA=test_data/futures-baseline FOXHUNT_FEATURE_CACHE_DIR=test_data/feature-cache SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::training_stability::test_50_epoch_convergence --ignored --nocapture

Expected: test passes, loss values are finite, Sharpe is reasonable. Loss trajectory may differ slightly from synchronous path (1-step lag on PER priorities) but should converge similarly.

  • Step 6: Commit
git commit -m "perf: CUDA events replace per-step stream.synchronize in training loop"

Task 2: Async Q-stats readback

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3640-3660

The reduce_current_q_stats() calls memcpy_dtoh every 50 steps — implicit sync. Use the same event pattern.

  • Step 1: Add q-stats event and pending buffer

Add to GpuDqnTrainer struct:

/// Pending q-stats from previous reduction (deferred readback)
q_stats_event: Option<Arc<CudaEvent>>,
q_stats_pending: [f32; 5],
q_stats_ready: bool,

Initialize:

q_stats_event: None,
q_stats_pending: [0.0; 5],
q_stats_ready: false,
  • Step 2: Rewrite reduce_current_q_stats for async readback

Replace the synchronous memcpy_dtoh in reduce_current_q_stats():

pub fn reduce_current_q_stats(&mut self) -> Result<QValueStatsResult, MLError> {
    // Collect previous reduction if ready
    let result = 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;
        let h = self.q_stats_pending;
        QValueStatsResult {
            avg_max_q: h[0] as f64, q_min: h[1], q_max: h[2], q_mean: h[3], q_var: h[4],
        }
    } else {
        QValueStatsResult { avg_max_q: 0.0, q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_var: 0.0 }
    };

    // Launch new reduction kernel
    // ... (existing kernel launch code stays the same) ...

    // Async readback
    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(),
        );
    }
    self.q_stats_event = Some(self.stream.record_event(None).map_err(|e|
        MLError::ModelError(format!("q_stats event record: {e}"))
    )?);
    self.q_stats_ready = true;

    Ok(result)
}
  • Step 3: Verify compile + smoke test

Run: SQLX_OFFLINE=true cargo check -p ml --lib Run: smoke test (same as Task 1 Step 5)

  • Step 4: Commit
git commit -m "perf: async q-stats readback via CUDA events"

Task 3: Dynamic batch sizing — remove static cap

Files:

  • Modify: config/gpu/h100.toml:3

  • Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs (AutoBatchSizer logic)

  • Step 1: Remove static batch_size from h100.toml

Change config/gpu/h100.toml line 3:

# batch_size is auto-computed by AutoBatchSizer from available VRAM.
# Do NOT set a static value here — let the sizer maximize GPU utilization.
# batch_size = 1024  # REMOVED: was capping H100 far below optimal

Actually: the TOML profile must still have the field for the parser. Set it to 0 as a sentinel meaning "auto":

batch_size = 0  # 0 = auto-compute from VRAM (AutoBatchSizer)
  • Step 2: Update constructor to honor batch_size=0 as auto

In constructor.rs, find where hyperparams.batch_size is set from the profile. After the profile applies, if batch_size == 0, let AutoBatchSizer compute it:

// After profile.apply_to(&mut hyperparams):
if hyperparams.batch_size == 0 {
    // Auto-compute: AutoBatchSizer already ran and set the VRAM ceiling.
    // Use the ceiling as the batch size, capped at a sane max (16384).
    hyperparams.batch_size = auto_batch_ceiling.min(16384);
    info!("AutoBatchSizer: batch_size auto-computed to {}", hyperparams.batch_size);
}
  • Step 3: Verify the auto-computed value is reasonable

Run locally (RTX 3050): the sizer should still pick a small batch (64-128) for 4GB VRAM. Check H100 would get: look at the AutoBatchSizer: GPU VRAM ceiling = 2085808 log — that's the theoretical max. The actual batch should be much smaller (constrained by kernel shared memory and practical gradient noise). Cap at 8192 or 16384.

  • Step 4: Commit
git commit -m "perf: dynamic batch sizing — AutoBatchSizer drives batch_size on H100"

Task 4: Experience collection → training phase overlap

Files:

  • Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs:1196

  • Step 1: Replace stream.synchronize() with event after experience collection

At line 1196, change:

// BEFORE:
stream.synchronize()
    .map_err(|e| MLError::DeviceError(format!("experience sync: {e}")))?;

// AFTER:
let experience_done_event = stream.record_event(None)
    .map_err(|e| MLError::DeviceError(format!("experience event record: {e}")))?;

Then before the first training step that samples from the replay buffer, wait on the event:

// Before first replay buffer sample:
if !experience_done_event.is_complete() {
    experience_done_event.synchronize()
        .map_err(|e| MLError::DeviceError(format!("experience event wait: {e}")))?;
}

This allows CPU-side training setup (lock acquisition, pre-sampling initialization) to overlap with the tail end of experience collection.

  • Step 2: Verify + commit
git commit -m "perf: event-based experience→training phase transition"

Task 5: fxcache PVC regeneration

  • Step 1: Regenerate PVC cache with aligned data-dir

The precompute template already uses data-dir/SYMBOL. Just re-run:

./scripts/argo-precompute.sh ES.FUT --pool ci-compile-cpu
  • Step 2: Verify training hits fxcache

Submit a quick 2-epoch baseline and check logs for "fxcache hit":

argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=2

Check: argo logs <name> | grep fxcache

  • Step 3: Commit (if any template changes needed)

Phase 2: Multi-Stream Parallelism

Task 6: Multi-stream branch dispatch in forward pass

Files:

  • Modify: crates/ml/src/cuda_pipeline/batched_forward.rs:265-284

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (add child streams)

  • Step 1: Add 3 child streams to CublasForward

In batched_forward.rs, add to the struct:

/// Child streams for parallel branch dispatch (forked from main training stream)
branch_streams: [Arc<CudaStream>; 3],

Create them in the constructor by forking:

let branch_streams = [
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 0: {e}")))?,
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 1: {e}")))?,
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?,
];
  • Step 2: Dispatch branches on separate streams

In forward_online_raw(), replace the sequential for d in 0..3 loop (lines 270-284):

// Fork: dispatch each branch on its own stream
for d in 0..3 {
    let branch_stream = &self.branch_streams[d];
    // Wait for trunk (h_s2) to be ready on this branch stream
    let trunk_event = self.stream.record_event(None)
        .map_err(|e| MLError::ModelError(format!("trunk event: {e}")))?;
    branch_stream.wait(&trunk_event)
        .map_err(|e| MLError::ModelError(format!("branch wait trunk: {e}")))?;

    let n_d = branch_sizes[d];
    let w_fc_idx = branch_w_base[d];

    // Hidden layer
    self.gemmex_bf16_on_stream(branch_stream, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "h_bd")?;
    self.launch_add_bias_relu_bf16_on_stream(branch_stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;

    // Output layer
    let adv_out_ptr = b_logits_ptr + logit_byte_offset;
    self.gemmex_bf16_to_f32_on_stream(branch_stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?;
    self.launch_add_bias_f32_on_stream(branch_stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;

    logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
}

// Join: main stream waits for all branches to finish
for d in 0..3 {
    let branch_done = self.branch_streams[d].record_event(None)
        .map_err(|e| MLError::ModelError(format!("branch done event: {e}")))?;
    self.stream.wait(&branch_done)
        .map_err(|e| MLError::ModelError(format!("join branch {d}: {e}")))?;
}

Note: gemmex_bf16_on_stream and launch_add_bias_relu_bf16_on_stream are new variants that take an explicit stream parameter. Implement them by extracting the cuBLAS handle → set stream → call GemmEx → restore stream. cuBLAS handles are NOT stream-specific — cublasSetStream() switches the active stream per-call.

  • Step 3: Apply same pattern to f32 forward (line 325) and backward (line 461)

Repeat the fork/join pattern for both the f32 experience collection forward and the backward pass branch loops.

  • Step 4: Verify compile + smoke test

Run smoke tests — output logits must be identical (branches write to non-overlapping memory regions).

  • Step 5: Commit
git commit -m "perf: multi-stream branch dispatch — 3 advantage heads in parallel"

Task 7: Multi-stream target/online forward parallelism

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4307-4323

  • Step 1: Add a second stream for Double DQN pass

Add to GpuDqnTrainer:

/// Second stream for parallel target/online next-state forward passes
double_dqn_stream: Arc<CudaStream>,

Fork in constructor.

  • Step 2: Dispatch Pass 2 and Pass 3 on separate streams

In submit_forward_ops(), after Pass 1 (online forward on current states):

// Pass 2 + Pass 3 are independent: both read next_states, write to different outputs.
// Fork: record event after stochastic depth (Pass 1 done)
let pass1_done = self.stream.record_event(None)?;
self.double_dqn_stream.wait(&pass1_done)?;

// Pass 2: target forward on next_states (main stream)
cublas.forward_target_raw(
    &self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, ...
)?;

// Pass 3: online forward on next_states (double_dqn_stream)
cublas.forward_online_raw(
    &self.double_dqn_stream, self.ptrs.next_states_buf, &on_w_ptrs, ...
)?;

// Join before loss kernel
let target_done = self.stream.record_event(None)?;
let online_done = self.double_dqn_stream.record_event(None)?;
self.stream.wait(&online_done)?;
// main stream now has both pass 2 + pass 3 complete
  • Step 3: Handle CUDA Graph capture with multi-stream

CUDA Graph capture with fork/join requires CU_STREAM_CAPTURE_MODE_THREAD_LOCAL (already used at line 3734). The fork/join events within the graph are captured as graph dependencies. Verify this works with cudarc — may need raw cuStreamWaitEvent calls during capture.

  • Step 4: Verify + smoke test + commit
git commit -m "perf: parallel target/online forward on separate CUDA streams"

Task 8: CUDA Graph for experience collection

Files:

  • Modify: crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1506

This is the most complex task. The 100-iteration host loop dispatches 400+ kernels. Capturing as a graph eliminates all host→GPU latency.

  • Step 1: Make per-timestep parameters GPU-resident

The loop currently passes epsilon, step counter, and state pointers from the host. These must become GPU buffers:

/// GPU-resident parameters for graph-captured timestep loop
epsilon_buf: CudaSlice<f32>,       // [1] — updated via DtoH before graph launch
step_counter_buf: CudaSlice<u32>,  // [1] — incremented by a tiny CUDA kernel per step
  • Step 2: Create a step-increment kernel

Tiny kernel that increments step_counter_buf[0] and computes epsilon = eps_end + (eps_start - eps_end) * exp(-step / decay) on GPU. No host involvement.

  • Step 3: Capture the 100-timestep loop as a single graph
stream.begin_capture(CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
for t in 0..timesteps {
    self.launch_experience_state_gather()?;
    self.launch_cublas_forward_online()?;
    self.launch_compute_expected_q()?;
    self.launch_experience_action_select()?;  // reads epsilon_buf
    self.launch_experience_env_step()?;
    self.launch_step_increment()?;  // increments step_counter, updates epsilon
}
let graph = stream.end_capture(CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;

Before each epoch: memcpy_htod the initial epsilon and step counter, then graph.launch().

  • Step 4: Handle dynamic epsilon

Epsilon decays across episodes within an epoch. The step-increment kernel computes the new epsilon from the step counter using the exponential decay formula, writing to epsilon_buf. The action selection kernel reads epsilon_buf directly — no host involvement.

  • Step 5: Verify collected experiences match sequential path

Run 1 epoch with both paths (graph and sequential), compare states/actions/rewards for the first 100 experiences. Must be numerically identical given the same Philox RNG seed.

  • Step 6: Commit
git commit -m "perf: CUDA Graph capture for 100-timestep experience collection loop"

Task 9: Lock contention removal

Files:

  • Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs:1336-1360

  • Step 1: Pre-sample all batches for the epoch in one lock acquisition

Instead of acquiring/releasing the read lock every 32 steps, sample ALL batches at epoch start:

let all_batches = {
    let agent = self.agent.read().await;
    let buffer = agent.memory();
    (0..num_training_steps)
        .map(|_| buffer.sample(self.current_batch_size))
        .collect::<Vec<_>>()
};
// Lock released. Training loop iterates over pre-sampled batches.

This is safe because the GPU PER buffer is immutable during training (new experiences are only added in Phase 2, which is finished before training starts).

  • Step 2: Remove PREFETCH_K chunking

The for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) loop becomes a simple for batch in all_batches.iter().

  • Step 3: Verify + commit
git commit -m "perf: pre-sample all training batches in single lock acquisition"

Validation

Task 10: H100 benchmark

  • Step 1: Rebuild and upload binary to PVC

Compile in-cluster via the baseline template.

  • Step 2: Run 10-epoch baseline and measure
argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=10

Check:

  • Epoch time: target <10s (from 2 min)

  • nvidia-smi: SM utilization >85%, memory bandwidth >50%, TDP >250W

  • Loss trajectory: finite, converging

  • "fxcache hit" in logs

  • Step 3: Compare loss with pre-optimization baseline

The 7-epoch metrics from the terminated run serve as reference:

  • Epoch 4: Sharpe 0.06, PF 0.94
  • Epoch 6: Sharpe -1.12, PF 1.17

Post-optimization should show similar or better convergence (larger batches may improve gradient quality).

  • Step 4: Run full 50-epoch baseline if metrics look good

  • Step 5: Save metrics to memory