# GPU Hot Path Performance — Zero-Sync Training Step > **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:** Eliminate all CPU sync points and F32 compute waste from the per-step training hot path, restoring 3s/epoch performance on H100. **Architecture:** Cache all CUDA device pointers at construction/graph-capture time into a `CachedPtrs` struct. Replace per-step `cuStreamSynchronize` + `EventTrackingGuard` + `raw_device_ptr` calls with cached u64 pointer lookups. Convert cuBLAS `cublasSgemm` (F32) to `cublasGemmEx` (BF16 compute with F32 accumulate) for 2-3× tensor core throughput. Batch 8 head spectral norm launches into a single multi-matrix kernel. **Tech Stack:** Rust, CUDA, cudarc 0.19.3, cuBLAS (cublasGemmEx), BF16 tensor cores **Verification requirement:** Every task MUST compile (`SQLX_OFFLINE=true cargo check -p ml`) and pass tests (`SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`). No exceptions. No skipping. --- ## Root Cause Analysis **Measured:** Epoch time regressed from 3s → 30s. Per-step hot path has: | Issue | Per-step cost | Source | Lines | |---|---|---|---| | `cuStreamSynchronize` in `apply_iqn_trunk_gradient` | ~50μs (stalls entire GPU pipeline) | Pre-existing | `gpu_dqn_trainer.rs:708` | | `cuStreamSynchronize` in `apply_ensemble_trunk_gradient` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:911` | | `cuStreamSynchronize` in `run_ensemble_step` | ~50μs | Pre-existing | `fused_training.rs:875` | | `cuStreamSynchronize` in `replay_adam_and_readback` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:2369` | | 99× `raw_device_ptr()` calls | ~1μs each but forces `device_ptr()` event machinery | Per-step | Throughout | | 11× `EventTrackingGuard` create/drop | ~0.1μs each (atomic ops) | Per-step | Throughout | | F32 `cublasSgemm` instead of BF16 `cublasGemmEx` | 2-3× slower SGEMM on H100 | Pre-existing | `batched_forward.rs:513` | | 8× single-block spectral norm launches | ~5μs each | New | `gpu_dqn_trainer.rs:1439-1453` | | 10× DtoD spectral norm sync-back | ~2μs each | New | `gpu_dqn_trainer.rs:1490-1503` | **Total per-step overhead: ~200μs sync stalls + 2-3× slower SGEMM.** Over 10K steps/epoch: 2s sync + 2-3× slower forward/backward = easily 30s epochs. --- ## File Map | File | Action | |---|---| | `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Major: add `CachedPtrs` struct, remove syncs from IQN/ensemble/CQL paths, cache all pointers at capture time | | `crates/ml/src/cuda_pipeline/batched_forward.rs` | Major: replace `cublasSgemm` with `cublasGemmEx` (BF16 input, F32 accumulate) | | `crates/ml/src/cuda_pipeline/batched_backward.rs` | Major: same `cublasSgemm` → `cublasGemmEx` conversion | | `crates/ml/src/trainers/dqn/fused_training.rs` | Moderate: remove `cuStreamSynchronize` from `run_ensemble_step`, dynamic budget fix | | `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Minor: add batched spectral norm kernel | | `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` | Minor: add performance regression test | --- ### Task 1: Cache all device pointers at graph-capture time The root cause of most overhead: `raw_device_ptr()` calls `slice.device_ptr(stream)` which goes through cudarc's event tracking machinery. Even with events disabled, it creates a `SyncOnDrop` guard that's immediately leaked via `ManuallyDrop`. This happens 99 times per step. **Fix:** Compute all u64 pointers ONCE at construction (or re-capture), store them in a `CachedPtrs` struct, and use the cached values in all per-step methods. **Files:** - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` - [ ] **Step 1: Define `CachedPtrs` struct** After the `GpuDqnTrainConfig` struct (~line 219), add: ```rust /// Pre-computed raw CUDA device pointers for all buffers used in the per-step hot path. /// Eliminates 99 `raw_device_ptr()` → `device_ptr()` calls per training step. /// Pointers are invalidated when buffers are reallocated (which never happens after init). struct CachedPtrs { // Flat parameter buffers params: u64, target_params: u64, grad: u64, grad_norm: u64, m: u64, v: u64, t: u64, total_loss: u64, cql_grad_scratch: u64, // Batch input buffers states: u64, next_states: u64, actions: u64, rewards: u64, dones: u64, is_weights: u64, // Activation save buffers save_h_s1: u64, save_h_s2: u64, save_h_v: u64, save_h_b0: u64, save_h_b1: u64, save_h_b2: u64, // Backward scratch bw_d_h_s1: u64, bw_d_h_s2: u64, bw_d_h_v: u64, bw_d_h_b0: u64, bw_d_h_b1: u64, bw_d_h_b2: u64, // IQN trunk scratch iqn_trunk_m: u64, iqn_trunk_grad_norm: u64, // Loss outputs td_errors: u64, // Per-weight pointers (20 online + 20 target, computed from params base) on_weight_ptrs: [u64; 20], tg_weight_ptrs: [u64; 20], } ``` - [ ] **Step 2: Build `CachedPtrs` in constructor** At the end of `new()`, before `Ok(Self { ... })`, compute all pointers. Since event tracking is enabled at this point and all buffers were just allocated on the same stream, `device_ptr()` is safe: ```rust let cached_ptrs = { let _evt_guard = EventTrackingGuard::new(stream.context()); let param_sizes_arr = compute_param_sizes(&config); CachedPtrs { params: raw_device_ptr(¶ms_buf, &stream), target_params: raw_device_ptr(&target_params_buf, &stream), grad: raw_device_ptr(&grad_buf, &stream), grad_norm: raw_device_ptr(&grad_norm_buf, &stream), m: raw_device_ptr(&m_buf, &stream), v: raw_device_ptr(&v_buf, &stream), t: raw_device_ptr_i32(&t_buf, &stream), total_loss: raw_device_ptr(&total_loss_buf, &stream), cql_grad_scratch: raw_device_ptr(&cql_grad_scratch, &stream), states: raw_device_ptr(&states_buf, &stream), next_states: raw_device_ptr(&next_states_buf, &stream), actions: raw_device_ptr_i32(&actions_buf, &stream), rewards: raw_device_ptr(&rewards_buf, &stream), dones: raw_device_ptr(&dones_buf, &stream), is_weights: raw_device_ptr(&is_weights_buf, &stream), save_h_s1: raw_device_ptr(&save_h_s1, &stream), save_h_s2: raw_device_ptr(&save_h_s2, &stream), save_h_v: raw_device_ptr(&save_h_v, &stream), save_h_b0: raw_device_ptr(&save_h_b0, &stream), save_h_b1: raw_device_ptr(&save_h_b1, &stream), save_h_b2: raw_device_ptr(&save_h_b2, &stream), bw_d_h_s1: raw_device_ptr(&bw_d_h_s1, &stream), bw_d_h_s2: raw_device_ptr(&bw_d_h_s2, &stream), bw_d_h_v: raw_device_ptr(&bw_d_h_v, &stream), bw_d_h_b0: raw_device_ptr(&bw_d_h_b0, &stream), bw_d_h_b1: raw_device_ptr(&bw_d_h_b1, &stream), bw_d_h_b2: raw_device_ptr(&bw_d_h_b2, &stream), iqn_trunk_m: raw_device_ptr(&iqn_trunk_m, &stream), iqn_trunk_grad_norm: raw_device_ptr(&iqn_trunk_grad_norm, &stream), td_errors: raw_device_ptr(&td_errors_buf, &stream), on_weight_ptrs: f32_weight_ptrs(¶ms_buf, ¶m_sizes_arr, &stream), tg_weight_ptrs: f32_weight_ptrs(&target_params_buf, ¶m_sizes_arr, &stream), } }; ``` Add `cached_ptrs: CachedPtrs` to the struct and `Ok(Self { ... })`. - [ ] **Step 3: Replace `raw_device_ptr` calls in `apply_iqn_trunk_gradient`** Replace all `raw_device_ptr(&self.xxx, &self.stream)` with `self.cached_ptrs.xxx`. Remove the `cuStreamSynchronize` (line 708) and `EventTrackingGuard` (line 712) — they were only needed because `device_ptr()` could trigger event validation. With cached pointers, no events are involved. **CRITICAL:** The `cuStreamSynchronize` at line 708 has the comment "ensure graph_forward replay completed before we touch buffers." With cached pointers, we don't "touch" buffers via device_ptr — we just use pre-computed addresses. All kernel launches go to the same stream, so ordering is guaranteed by CUDA stream semantics. The sync is unnecessary. - [ ] **Step 4: Replace `raw_device_ptr` calls in `apply_ensemble_trunk_gradient`** Same pattern. Remove the `cuStreamSynchronize` (line 911) and `EventTrackingGuard` (line 914). - [ ] **Step 5: Replace `raw_device_ptr` calls in `apply_cql_gradient` and `apply_cql_clipped_saxpy`** Same pattern for both methods. - [ ] **Step 6: Replace `raw_device_ptr` calls in `clip_grad_buf_inplace` and `read_grad_norm_sync`** For `clip_grad_buf_inplace`: use `self.cached_ptrs.grad` and `self.cached_ptrs.grad_norm`. Remove `EventTrackingGuard`. For `read_grad_norm_sync`: keep the `cuStreamSynchronize` (this is a diagnostic method that explicitly syncs for readback) but use cached pointers. This method only runs every 1000 steps. - [ ] **Step 7: Replace `raw_device_ptr` calls in `apply_spectral_norm`** Use `self.cached_ptrs.params` for `params_base` in the sync-back section. The spectral norm kernel launches on individual weight set buffers (DuelingWeightSet/BranchingWeightSet) — these are passed as references, so their `device_ptr()` calls remain. BUT: spectral norm runs ONCE per step before graph_forward. The sync at the method start is the graph_forward guard — remove it and rely on stream ordering. - [ ] **Step 8: Replace `raw_device_ptr` in kernel launch methods** Update `launch_grad_norm`, `launch_adam_update`, and the SAXPY kernel launches (IQN clipped, ensemble clipped) to use cached pointers. Remove the `EventTrackingGuard` in each. - [ ] **Step 9: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` Expected: Compiles clean, all 129+ tests pass. - [ ] **Step 10: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs git commit -m "perf: cache all CUDA device pointers — eliminate 99 raw_device_ptr + 4 cuStreamSync per step" ``` --- ### Task 2: Remove `cuStreamSynchronize` from `run_ensemble_step` The ensemble step in `fused_training.rs:875` syncs before touching `save_h_s2`. With cached pointers and single-stream execution, this sync is unnecessary — CUDA guarantees ordering on the same stream. **Files:** - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` - [ ] **Step 1: Remove the sync + EventTrackingGuard in `run_ensemble_step`** At ~line 874-876, remove: ```rust unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } let _ = self.stream.context().check_err(); ``` And the local `EvtGuard` struct + its creation (~lines 881-893). The ensemble heads use `device_ptr()` on their own weight buffers — since these are created after event tracking is disabled by the top-level guard in `apply_spectral_norm`, they have no events. Use raw u64 pointers from the ensemble head structs instead. Note: the `device_ptr()` calls in the ensemble loop (on `ensemble_extra_heads[k]` weights) cannot be cached in `CachedPtrs` because ensemble heads are optional and variable-count. For these, use a SINGLE `EventTrackingGuard` scope at the top of `run_ensemble_step` and call `device_ptr()` within it — but do NOT sync the stream. - [ ] **Step 2: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` Expected: All pass. - [ ] **Step 3: Commit** ```bash git add crates/ml/src/trainers/dqn/fused_training.rs git commit -m "perf: remove cuStreamSynchronize from run_ensemble_step" ``` --- ### Task 3: Convert cuBLAS forward from `cublasSgemm` (F32) to `cublasGemmEx` (BF16 tensors, F32 accumulate) The forward pass uses `cublasSgemm` (F32×F32→F32). On H100, `cublasGemmEx` with BF16 inputs and F32 accumulate uses tensor cores at 2-3× the throughput. The BF16 weight buffers (`bf16_params_buf`, `bf16_target_params_buf`) already exist but are never used. **Key constraints:** - Weights: BF16 (pre-converted from F32 via `f32_to_bf16_kernel`) - Activations/states: MUST stay F32 (written by upstream kernels in F32) - Accumulator: F32 (no precision loss in output) - cuBLAS GemmEx can do BF16×F32→F32 using `CUBLAS_COMPUTE_32F` Actually — `cublasGemmEx` requires both A and B to have the same type when using tensor cores. We need BF16×BF16→F32. This means **states and activations must also be BF16** for the forward SGEMM. But: - States come from the experience collector in F32 - Activations (h_s1, h_s2) are written by SGEMM and consumed by the next layer **Practical approach:** Use **mixed-precision GemmEx** with: - A = weights (BF16, from bf16_params_buf) - B = input (F32, states/activations) - C = output (F32, activations/logits) - Compute: `CUBLAS_COMPUTE_32F` (uses tensor cores on H100 for BF16×F32 with TF32 accumulate) On H100, `CUBLAS_COMPUTE_32F` with mixed types uses TF32 tensor cores, which is ~1.5× faster than pure F32. For full BF16 tensor core speed (3×), both inputs must be BF16. That requires converting activations to BF16 too — a larger change. **This task implements the TF32 mixed-precision path (1.5× speedup, minimal changes).** **Files:** - Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` - [ ] **Step 1: Replace `cublasSgemm` with `cublasGemmEx` using TF32 compute** In `sgemm_layer` (~line 495), change from: ```rust cublas_result::sgemm( self.handle.0, cublas_sys::cublasOperation_t::CUBLAS_OP_T, cublas_sys::cublasOperation_t::CUBLAS_OP_N, n as i32, b as i32, k as i32, &alpha, w_ptr as *const f32, k as i32, a_ptr as *const f32, k as i32, &beta, c_ptr as *mut f32, n as i32, ) ``` To: ```rust cublas_sys::cublasGemmEx( self.handle.0, cublas_sys::cublasOperation_t::CUBLAS_OP_T, cublas_sys::cublasOperation_t::CUBLAS_OP_N, n as i32, b as i32, k as i32, &alpha as *const f32 as *const std::ffi::c_void, w_ptr as *const std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_32F, // A type k as i32, a_ptr as *const std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_32F, // B type k as i32, &beta as *const f32 as *const std::ffi::c_void, c_ptr as *mut std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_32F, // C type n as i32, cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32, // TF32 tensor cores cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, ) ``` `CUBLAS_COMPUTE_32F_FAST_TF32` enables TF32 tensor cores on H100/A100. With F32 inputs, cuBLAS internally truncates to TF19 mantissa for the multiply, accumulates in F32. ~1.5× speedup over pure F32 SGEMM. Do the same for `sgemm_layer_raw`. - [ ] **Step 2: Set cuBLAS math mode for tensor ops** In the `CublasForward::new()` constructor, set the math mode: ```rust unsafe { cublas_sys::cublasSetMathMode( handle.0, cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH, ); } ``` This allows `cublasSgemm` to use TF32 even without switching to `cublasGemmEx`. - [ ] **Step 3: Do the same for `batched_backward.rs`** The backward pass also uses `cublasSgemm`. Apply the same TF32 math mode or `cublasGemmEx` conversion. - [ ] **Step 4: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` Expected: All pass. TF32 has ≤0.1% numerical difference from F32 — tests should still pass. - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs git commit -m "perf: enable TF32 tensor cores for cuBLAS SGEMM — 1.5x forward+backward throughput" ``` --- ### Task 4: Batch spectral norm into fewer kernel launches Currently 8 separate kernel launches for head spectral norm, each with `grid=(1,1,1)` — massively underutilized GPU. Batch them into a single kernel that processes all 8 matrices sequentially (same single block, but loop over matrices using a matrix descriptor array). **Files:** - Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` - [ ] **Step 1: Add batched spectral norm kernel** In `dqn_utility_kernels.cu`, add: ```cuda /* ══════════════════════════════════════════════════════════════════════ * BATCHED SPECTRAL NORM KERNEL * * Processes N weight matrices in a single kernel launch. * Grid = (N, 1, 1), Block = (256, 1, 1). * Each block handles one matrix via the same power iteration as * spectral_norm_kernel but indexed by blockIdx.x into descriptor arrays. * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void batched_spectral_norm_kernel( /* Arrays of N pointers, one per matrix */ float** __restrict__ W_arr, /* [N] pointers to weight matrices */ float** __restrict__ u_arr, /* [N] pointers to left singular vectors */ float** __restrict__ v_arr, /* [N] pointers to right singular vectors */ const int* __restrict__ out_dims, /* [N] out_dim per matrix */ const int* __restrict__ in_dims, /* [N] in_dim per matrix */ float sigma_max, int n_matrices ) { int mat_idx = blockIdx.x; if (mat_idx >= n_matrices) return; float* W = W_arr[mat_idx]; float* u = u_arr[mat_idx]; float* v = v_arr[mat_idx]; int out_dim = out_dims[mat_idx]; int in_dim = in_dims[mat_idx]; int n_total = out_dim * in_dim; /* Same power iteration logic as spectral_norm_kernel... */ /* (copy the body of spectral_norm_kernel here, replacing the fixed arguments with the per-matrix values) */ __shared__ float shmem[256]; int tid = threadIdx.x; int bd = blockDim.x; // Step 1: v_new = W^T u for (int col = tid; col < in_dim; col += bd) { float val = 0.0f; for (int row = 0; row < out_dim; row++) val += W[row * in_dim + col] * u[row]; v[col] = val; } __syncthreads(); // Normalize v float local_v2 = 0.0f; for (int col = tid; col < in_dim; col += bd) local_v2 += v[col] * v[col]; shmem[tid] = local_v2; __syncthreads(); for (int s = bd / 2; s > 0; s >>= 1) { if (tid < s) shmem[tid] += shmem[tid + s]; __syncthreads(); } float v_norm = sqrtf(shmem[0] + 1e-12f); for (int col = tid; col < in_dim; col += bd) v[col] /= v_norm; __syncthreads(); // Step 2: u_new = W v_new for (int row = tid; row < out_dim; row += bd) { float val = 0.0f; for (int col = 0; col < in_dim; col++) val += W[row * in_dim + col] * v[col]; u[row] = val; } __syncthreads(); // sigma = ||u_new|| float local_u2 = 0.0f; for (int row = tid; row < out_dim; row += bd) local_u2 += u[row] * u[row]; shmem[tid] = local_u2; __syncthreads(); for (int s = bd / 2; s > 0; s >>= 1) { if (tid < s) shmem[tid] += shmem[tid + s]; __syncthreads(); } float sigma = sqrtf(shmem[0] + 1e-12f); // Normalize u for (int row = tid; row < out_dim; row += bd) u[row] /= (sigma + 1e-12f); __syncthreads(); // Scale W if sigma > sigma_max if (sigma < 1e-6f) return; float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f; if (scale < 1.0f) { for (int i = tid; i < n_total; i += bd) W[i] *= scale; } } ``` - [ ] **Step 2: Load the kernel in Rust** Add `batched_spectral_norm_kernel` loading in `compile_training_kernels`. Add the function field to the struct. - [ ] **Step 3: Allocate descriptor arrays on GPU** Allocate `CudaSlice` arrays for the 10 W/u/v pointer arrays and `CudaSlice` for out_dims/in_dims. Upload once at construction. - [ ] **Step 4: Replace 10 individual spec_norm! calls with single batched launch** In `apply_spectral_norm`, replace the 2 trunk + 8 head launches with: ```rust // Single batched launch: 10 blocks × 256 threads unsafe { self.stream .launch_builder(&self.batched_spectral_norm_kernel) .arg(&self.spec_w_ptrs_ptr) .arg(&self.spec_u_ptrs_ptr) .arg(&self.spec_v_ptrs_ptr) .arg(&self.spec_out_dims_ptr) .arg(&self.spec_in_dims_ptr) .arg(&sigma_max) .arg(&10_i32) .launch(LaunchConfig { grid_dim: (10, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4, })?; } ``` This replaces 10 kernel launches with 1. - [ ] **Step 5: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` Expected: All pass. - [ ] **Step 6: Commit** ```bash git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs git commit -m "perf: batch 10 spectral norm launches into single kernel" ``` --- ### Task 5: Dynamic gradient budget allocation When IQN/ensemble are disabled, their budget is wasted — C51 only gets 70% even as the sole gradient source. **Files:** - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` **NOTE:** This is already partially implemented (dynamic `c51_frac` computation at Step 2b). Verify it's correct and add a test. - [ ] **Step 1: Verify the dynamic budget code** Read the current Step 2b in `fused_training.rs` and confirm: ```rust let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 }; let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 }; let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 }; let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac; ``` This is correct. C51 gets 100% when nothing else is active. - [ ] **Step 2: Add test for dynamic budget** In `gradient_budget.rs`, add a test that constructs a trainer with CQL disabled (cql_alpha=0) and verifies C51 can use the full max_grad_norm: ```rust #[test] fn test_c51_gets_full_budget_when_no_auxiliaries() -> anyhow::Result<()> { let dev = cuda_device(); let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); let mut cfg = test_config(); cfg.cql_alpha = 0.0; // Disable CQL cfg.iqn_lambda = 0.0; // Disable IQN let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?; // Fill grad_buf with large gradient let total_params = trainer.total_params(); stream.memcpy_htod(&vec![1.0_f32; total_params], trainer.grad_buf_mut()) .map_err(|e| anyhow::anyhow!("{e}"))?; // Clip to full max_grad_norm (what C51 should get when solo) trainer.clip_grad_buf_inplace(10.0)?; // max_grad_norm * 1.0 let norm = trainer.read_grad_norm_sync()?; assert!( norm > 9.9, "C51 solo should use full max_grad_norm, got {norm}" ); Ok(()) } ``` - [ ] **Step 3: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget` Expected: All pass including new test. - [ ] **Step 4: Commit** ```bash git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs git commit -m "fix: dynamic gradient budget — C51 gets full max_grad_norm when auxiliaries disabled" ``` --- ### Task 6: Remove `cuStreamSynchronize` from `replay_adam_and_readback` The Adam readback (`replay_adam_and_readback`) syncs the stream and reads back loss + grad_norm (8 bytes DtoH). This is the ONLY legitimate sync in the per-step path — we need the loss for logging. But we can defer it to the NEXT step's start, overlapping GPU work with CPU logging. **Files:** - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` - [ ] **Step 1: Convert to async readback with deferred sync** In `replay_adam_and_readback`, replace: ```rust cuStreamSynchronize(self.stream.cu_stream()); cuMemcpyDtoH_v2(loss_host..., loss_ptr, 4); cuMemcpyDtoH_v2(norm_host..., norm_ptr, 4); ``` With async DtoH into a pinned host buffer, returning the PREVIOUS step's values: ```rust pub fn replay_adam_and_readback(&mut self) -> Result { // Return PREVIOUS step's readback (already synced by this step's graph_forward) let prev_loss = self.scalars_readback_host[0]; let prev_norm = self.scalars_readback_host[1]; // Replay Adam (async) self.replay_adam()?; // Queue async DtoH for THIS step's results (will be ready by next step) unsafe { cudarc::driver::sys::cuMemcpyDtoHAsync_v2( self.scalars_readback_host.as_mut_ptr().cast(), self.cached_ptrs.total_loss, 4, self.stream.cu_stream(), ); cudarc::driver::sys::cuMemcpyDtoHAsync_v2( self.scalars_readback_host.as_mut_ptr().add(1).cast(), self.cached_ptrs.grad_norm, 4, self.stream.cu_stream(), ); } Ok(FusedTrainScalars { total_loss: prev_loss, grad_norm: prev_norm.sqrt(), }) } ``` This eliminates the per-step sync. The readback has 1-step lag (reporting previous step's metrics) which is standard and acceptable for monitoring. - [ ] **Step 2: Allocate pinned host buffer** Replace `scalars_readback_host: Vec` with pinned memory for async DtoH: ```rust // In constructor: let scalars_readback_host = vec![0.0_f32; 2]; // [loss, grad_norm_sq] // cudarc doesn't expose cuMemAllocHost directly — keep Vec but ensure // the async DtoH is followed by a stream sync at epoch boundary. ``` Actually — `cuMemcpyDtoHAsync_v2` requires the destination to be page-locked (pinned). With a regular Vec, this is technically UB. The safe approach: keep the sync but move it to the epoch boundary, NOT per-step. **Revised approach:** Batch the sync. Instead of syncing every step, sync once per N steps (e.g. every 100 steps or at epoch boundary). Use the existing `readback_host` Vec with synchronous DtoH but only when actually needed. - [ ] **Step 3: Move sync to epoch boundary** In `fused_training.rs`, remove the per-step sync from `replay_adam_and_readback`. Instead, sync only when `steps_since_varmap_sync % 100 == 0` or at epoch boundary. For intermediate steps, return the last known values. - [ ] **Step 4: Verify compilation + tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn` Expected: All pass. - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs git commit -m "perf: defer Adam readback sync to every 100 steps — eliminate per-step cuStreamSync" ``` --- ### Task 7: Performance regression smoke test Add a test that measures training step throughput and fails if it drops below a threshold. **Files:** - Modify: `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` - [ ] **Step 1: Add throughput test** ```rust /// Training step throughput must not regress. /// At 3s/epoch with 10K steps: ~3333 steps/sec minimum on RTX 3050. /// On H100: ~10K+ steps/sec. Test uses a conservative threshold. #[test] fn test_training_step_throughput() -> anyhow::Result<()> { let dev = cuda_device(); let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); let cfg = test_config(); let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; let mut dueling = alloc_dueling(&stream, &cfg); let mut branching = alloc_branching(&stream, &cfg); // Warm up trainer.apply_spectral_norm(&mut dueling, &mut branching)?; trainer.clip_grad_buf_inplace(7.0)?; // Measure 1000 iterations of clip + spectral norm (the hot path additions) let n_iters = 1000; let start = std::time::Instant::now(); for _ in 0..n_iters { trainer.apply_spectral_norm(&mut dueling, &mut branching)?; trainer.clip_grad_buf_inplace(7.0)?; } // Sync to ensure all GPU work completes unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); } let elapsed = start.elapsed(); let per_iter_us = elapsed.as_micros() as f64 / n_iters as f64; assert!( per_iter_us < 500.0, // 500μs budget for spectral norm + clip per step "Spectral norm + clip took {per_iter_us:.0}μs/iter — budget is 500μs" ); Ok(()) } ``` - [ ] **Step 2: Verify compilation + run** Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_training_step_throughput --nocapture` Expected: PASS with per-iter time well under 500μs. - [ ] **Step 3: Commit** ```bash git add crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs git commit -m "test: training step throughput regression test" ``` --- ## Execution Dependencies ``` Task 1 (cache pointers) → Task 2 (remove ensemble sync) → Task 6 (defer Adam readback) → Task 3 (TF32 tensor cores) — independent → Task 4 (batch spectral norm) — independent Task 5 (dynamic budget) — independent Task 7 (throughput test) — after Tasks 1-4 ``` Task 1 is the foundation — all other tasks benefit from cached pointers. Tasks 3 and 4 are independent GPU optimizations. Task 6 requires Task 1's cached pointers. Task 7 validates everything. ## Expected Impact | Optimization | Estimated speedup | |---|---| | Eliminate 4× cuStreamSync (200μs/step) | 2s/epoch saved | | TF32 tensor cores (1.5× SGEMM) | 30-40% forward+backward time | | Batch spectral norm (10→1 launch) | ~50μs/step saved | | Eliminate 99 raw_device_ptr calls | ~100μs/step saved | | Defer Adam readback sync | ~50μs/step saved | | **Combined** | **Target: ≤5s/epoch on H100** |