# Phase 3 Kernel Optimization 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:** Eliminate the 3 remaining serial-batch bottleneck kernels found by nsys Phase 2 profiling — `reward_rank_normalize` (18.3%), `bias_grad_reduce_f32_kernel` (11.4%), and `curiosity_bias_grad_reduce` (0.9%) — to reach <40ms/step on H100. **Architecture:** (1) Replace O(N²) rank-normalize with radix-sort-based O(N log N) rank. (2) Convert the main DQN backward bias-grad-reduce from serial batch loop to 2-phase shared-memory reduction (same proven pattern used in IQN/IQL/attention). (3) Convert curiosity bias-grad-reduce to the same 2-phase pattern. **Tech Stack:** CUDA 12.4, CUB radix sort (device-wide), cudarc vendored bindings, existing 2-phase reduce pattern from `iqn_dual_head_kernel.cu`. **Dimensions:** B=8192, SH2=256, out_dim varies (3–256), CUR_OUTPUT=42, CUR_HIDDEN=128. --- ## File Structure | File | Responsibility | |------|---------------| | `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu` | Reward rank-normalize kernel (rewrite to sort-based) | | `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Rust wiring for reward rank-normalize launch | | `crates/ml/src/cuda_pipeline/backward_kernels.cu` | Main DQN backward bias-grad-reduce (rewrite to 2-phase) | | `crates/ml/src/cuda_pipeline/batched_backward.rs` | Rust wiring for bias-grad-reduce launch + partials buffer | | `crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu` | Curiosity bias-grad-reduce (rewrite to 2-phase) | | `crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs` | Rust wiring for curiosity bias-grad-reduce | --- ### Task 1: Rewrite reward_rank_normalize — sort-based O(N log N) rank (18.3%, 10.6s) The current kernel is O(N²): each of N threads scans all N elements to count how many are ≤ its value. With N=~4M (4096 episodes × ~1000 steps), this takes 10.6 seconds for a single call. **Fix:** Sort the absolute Sharpe values using CUB `DeviceRadixSort::SortPairs`, then compute rank from sorted position. O(N log N) via radix sort — should complete in <50ms for 4M elements on H100. **Files:** - Modify: `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu` - Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` - [ ] **Step 1: Write 2 new helper kernels in `reward_shaping_kernel.cu`** Keep the old `reward_rank_normalize` kernel. Add these BEFORE it: ```c /* reward_compute_abs_sharpe — Compute |sharpe[i]| and original index for sort. * Grid: ceil(N/256), Block: 256. */ extern "C" __global__ void reward_compute_abs_sharpe( const float* __restrict__ rewards_in, float* __restrict__ abs_sharpe_out, /* [N] sort keys */ int* __restrict__ indices_out, /* [N] original indices */ int N, float std_ema, float return_mean_ema ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; float raw = rewards_in[i]; float sharpe = (std_ema > 1e-6f) ? (raw - return_mean_ema) / std_ema : raw; abs_sharpe_out[i] = fabsf(sharpe); indices_out[i] = i; } /* reward_scatter_rank — After sort, write rank = position/N to output, preserving sign. * sorted_indices[pos] = original_index. rank = pos/N. * Grid: ceil(N/256), Block: 256. */ extern "C" __global__ void reward_scatter_rank( const float* __restrict__ rewards_in, const int* __restrict__ sorted_indices, float* __restrict__ rewards_out, int N, float std_ema, float return_mean_ema ) { int pos = blockIdx.x * blockDim.x + threadIdx.x; if (pos >= N) return; int orig_idx = sorted_indices[pos]; float raw = rewards_in[orig_idx]; float sharpe = (std_ema > 1e-6f) ? (raw - return_mean_ema) / std_ema : raw; int is_trade = (fabsf(sharpe) > 0.001f) ? 1 : 0; if (!is_trade) { rewards_out[orig_idx] = 0.0f; return; } float rank = (float)(pos + 1) / (float)N; /* 1-based rank / N */ float sign = (sharpe > 0.0f) ? 1.0f : -1.0f; rewards_out[orig_idx] = sign * rank; } ``` - [ ] **Step 2: Add CUB sort buffers and new kernel handles in `gpu_experience_collector.rs`** In the struct, add: ```rust reward_abs_sharpe_buf: CudaSlice, // [max_N] sort keys reward_indices_buf: CudaSlice, // [max_N] original indices reward_sorted_keys_buf: CudaSlice, // [max_N] sorted output reward_sorted_indices_buf: CudaSlice,// [max_N] sorted indices reward_sort_temp_buf: CudaSlice, // CUB sort workspace reward_compute_abs_sharpe_kernel: CudaFunction, reward_scatter_rank_kernel: CudaFunction, ``` max_N = gpu_n_episodes × max_bars (the maximum number of rewards to rank). Allocate CUB temp storage using `cub::DeviceRadixSort::SortPairs` temp size query (via cudarc raw API or pre-computed upper bound of `2 * N * sizeof(f32)` which is safe for radix sort). In the constructor, load `reward_compute_abs_sharpe` and `reward_scatter_rank` from the reward_shaping cubin. Allocate the scratch buffers. - [ ] **Step 3: Rewrite `shape_rewards()` to use sort-based pipeline** Replace the single `reward_rank_normalize` launch with: 1. `reward_compute_abs_sharpe` kernel — compute keys + indices 2. CUB `DeviceRadixSort::SortPairs` — sort keys, permute indices 3. `reward_scatter_rank` kernel — scatter rank to original positions For CUB sort, use `cudarc::driver::sys` to call `cub::DeviceRadixSort::SortPairs` via raw CUDA API. If CUB is not directly available through cudarc, use the alternative: `thrust::sort_by_key` via a small custom kernel that calls `__device__` sort, OR implement a simple GPU merge sort. Simplest practical approach: use `cudarc`'s built-in `CudaSlice` sort if available, otherwise implement bitonic sort as a fallback (O(N log²N), still far better than O(N²)). **Fallback if CUB not available:** Write a 3-kernel bitonic sort: - `bitonic_compare_swap` kernel, launched log²(N) times with varying step/stage parameters - Total kernel launches: ~20 for N=4M (log2(4M) ≈ 22 stages × log2(stage) sub-passes) - Each launch: grid=N/2, block=256, 0.05ms → total ~1ms vs 10.6s - [ ] **Step 4: Run smoke tests** ```bash SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture ``` Expected: 20/20 pass. - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu crates/ml/src/cuda_pipeline/gpu_experience_collector.rs git commit -m "perf: reward_rank_normalize O(N²) → sort-based O(N log N) — 10.6s → <50ms" ``` --- ### Task 2: Rewrite bias_grad_reduce_f32_kernel — 2-phase shared-memory reduce (11.4%, 6.6s) The main DQN backward bias gradient kernel in `backward_kernels.cu` uses the serial anti-pattern: 1 thread per `out_dim` element, loops over `batch_size=8192`. Called 28,321 times during 993 training steps (~28 per step for all FC layers). Total: 6.6 seconds. **Fix:** Same 2-phase pattern proven in IQN/IQL/attention bias grad reduces. **Files:** - Modify: `crates/ml/src/cuda_pipeline/backward_kernels.cu` - Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` - [ ] **Step 1: Write 2-phase kernels in `backward_kernels.cu`** Replace the old `bias_grad_reduce_f32_kernel` with: ```c /* Phase 1: block-level shared-memory reduce. * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float). */ extern "C" __global__ void bias_grad_reduce_f32_p1( const float* __restrict__ dy, float* __restrict__ partials, int out_dim, int batch_size ) { extern __shared__ float sdata[]; int j = blockIdx.y; int tid = threadIdx.x; int gid = blockIdx.x * blockDim.x + tid; float val = (gid < batch_size) ? dy[gid * out_dim + j] : 0.0f; sdata[tid] = val; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) sdata[tid] += sdata[tid + s]; __syncthreads(); } if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0]; } /* Phase 2: final reduce across block partials + accumulate into db. * Grid: ceil(out_dim/256), Block: 256. * NOTE: db[j] += sum (accumulate, not overwrite) to match existing behavior. */ extern "C" __global__ void bias_grad_reduce_f32_p2( const float* __restrict__ partials, float* __restrict__ db, int out_dim, int num_blocks ) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= out_dim) return; float sum = 0.0f; for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j]; db[j] += sum; } ``` Note: The old kernel uses `db[j] += sum` (accumulate), so the p2 kernel must also accumulate. - [ ] **Step 2: Update `batched_backward.rs` struct + constructor** In the `BatchedBackward` struct (or equivalent name), add: ```rust bias_grad_p1_kernel: CudaFunction, bias_grad_p2_kernel: CudaFunction, bias_grad_partials_buf: CudaSlice, // [num_blocks * max_out_dim] bias_grad_num_blocks: u32, ``` `max_out_dim` = maximum `out_dim` across all FC layers = SH2 = 256. `num_blocks` = `(batch_size + 255) / 256` = 32. Partials buffer: 32 × 256 × 4 = 32KB — trivial. In the constructor, load `bias_grad_reduce_f32_p1` and `bias_grad_reduce_f32_p2` from the backward_kernels cubin. - [ ] **Step 3: Rewrite `launch_bias_grad()` to 2-phase pattern** ```rust fn launch_bias_grad( &self, stream: &Arc, dy: u64, db: u64, out_dim: usize, batch: usize, ) -> Result<(), MLError> { let out_dim_i32 = out_dim as i32; let batch_i32 = batch as i32; let num_blocks = self.bias_grad_num_blocks; let partials_ptr = self.bias_grad_partials_buf.raw_ptr(); // Phase 1 unsafe { stream.launch_builder(&self.bias_grad_p1_kernel) .arg(&dy).arg(&partials_ptr) .arg(&out_dim_i32).arg(&batch_i32) .launch(LaunchConfig { grid_dim: (num_blocks, out_dim as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4, }) .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p1: {e}")))?; } // Phase 2 let num_blocks_i32 = num_blocks as i32; let p2_blocks = ((out_dim + 255) / 256) as u32; unsafe { stream.launch_builder(&self.bias_grad_p2_kernel) .arg(&partials_ptr).arg(&db) .arg(&out_dim_i32).arg(&num_blocks_i32) .launch(LaunchConfig { grid_dim: (p2_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p2: {e}")))?; } Ok(()) } ``` - [ ] **Step 4: Run smoke tests** ```bash SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture ``` Expected: 20/20 pass. - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/backward_kernels.cu crates/ml/src/cuda_pipeline/batched_backward.rs git commit -m "perf: bias_grad_reduce_f32 — 2-phase shared-memory reduce (0.24ms×28K → <0.01ms×28K)" ``` --- ### Task 3: Rewrite curiosity_bias_grad_reduce — 2-phase pattern (0.9%, 134ms/call) Same serial anti-pattern: `out_dim` threads loop over N=8192. Called 4 times (2 bias layers × 2 calls). 134ms per call is extreme for a simple reduction. **Fix:** Same 2-phase shared-memory reduce pattern. **Files:** - Modify: `crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu` - Modify: `crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs` - [ ] **Step 1: Write 2-phase kernels in `curiosity_training_kernel.cu`** Replace the old `curiosity_bias_grad_reduce` with: ```c /* Phase 1: block-level shared-memory reduce. * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float). */ extern "C" __global__ void curiosity_bias_grad_reduce_p1( const float* __restrict__ dy, float* __restrict__ partials, int N, int out_dim ) { extern __shared__ float sdata[]; int d = blockIdx.y; int tid = threadIdx.x; int gid = blockIdx.x * blockDim.x + tid; float val = (gid < N) ? dy[gid * out_dim + d] : 0.0f; sdata[tid] = val; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) sdata[tid] += sdata[tid + s]; __syncthreads(); } if (tid == 0) partials[blockIdx.x * out_dim + d] = sdata[0]; } /* Phase 2: final reduce across block partials. * Grid: ceil(out_dim/256), Block: 256. */ extern "C" __global__ void curiosity_bias_grad_reduce_p2( const float* __restrict__ partials, float* __restrict__ grad_b, int out_dim, int num_blocks ) { int d = blockIdx.x * blockDim.x + threadIdx.x; if (d >= out_dim) return; float sum = 0.0f; for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + d]; grad_b[d] = sum; } ``` - [ ] **Step 2: Update `gpu_curiosity_trainer.rs` struct + constructor** In the struct, add: ```rust bias_grad_p1_func: CudaFunction, bias_grad_p2_func: CudaFunction, bias_grad_partials: CudaSlice, // [num_blocks * max(CUR_OUTPUT, CUR_HIDDEN)] bias_grad_num_blocks: u32, ``` In the constructor, load `curiosity_bias_grad_reduce_p1` and `curiosity_bias_grad_reduce_p2`. Allocate partials buffer: `num_blocks * max(CUR_OUTPUT=42, CUR_HIDDEN=128) = 32 * 128 = 4096 floats = 16KB`. - [ ] **Step 3: Replace both `curiosity_bias_grad_reduce` call sites** At line ~810 (b2) and ~850 (b1), replace the single kernel launch with 2-phase pattern: ```rust // Phase 1: block reduce let partials_ptr = self.bias_grad_partials.raw_ptr(); let num_blocks = self.bias_grad_num_blocks; unsafe { stream.launch_builder(&self.bias_grad_p1_func) .arg(&pred_ptr).arg(&partials_ptr) .arg(&n_i32).arg(&cur_output_i32) .launch(LaunchConfig { grid_dim: (num_blocks, CUR_OUTPUT as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4, }) .map_err(|e| MLError::ModelError(format!("curiosity_bias_grad_reduce_p1 b2: {e}")))?; } // Phase 2: final reduce let num_blocks_i32 = num_blocks as i32; unsafe { stream.launch_builder(&self.bias_grad_p2_func) .arg(&partials_ptr).arg(&grad_b2_ptr) .arg(&cur_output_i32).arg(&num_blocks_i32) .launch(LaunchConfig { grid_dim: (((CUR_OUTPUT + 255) / 256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!("curiosity_bias_grad_reduce_p2 b2: {e}")))?; } ``` Repeat the same pattern for the b1 call site with `d_hidden_ptr`, `grad_b1_ptr`, `CUR_HIDDEN`. - [ ] **Step 4: Run smoke tests** ```bash SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture ``` Expected: 20/20 pass. - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs git commit -m "perf: curiosity_bias_grad_reduce — 2-phase shared-memory reduce (134ms → <0.5ms)" ``` --- ### Task 4: Remove dead old kernels + deploy nsys validation **Files:** - Modify: `crates/ml/src/cuda_pipeline/reward_shaping_kernel.cu` (remove old kernel) - Modify: `crates/ml/src/cuda_pipeline/backward_kernels.cu` (remove old kernel) - Modify: `crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu` (remove old kernel) - [ ] **Step 1: Remove old `reward_rank_normalize` kernel body** (keep comment noting replacement) - [ ] **Step 2: Remove old `bias_grad_reduce_f32_kernel`** (replaced by p1+p2) - [ ] **Step 3: Remove old `curiosity_bias_grad_reduce`** (replaced by p1+p2) - [ ] **Step 4: Run smoke tests** ```bash SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture ``` - [ ] **Step 5: Commit, push, deploy with nsys** ```bash git add -A && git commit -m "cleanup: remove old serial kernels replaced by batch-parallel versions" git push origin main ./scripts/argo-train.sh dqn --sanitizer nsys --epochs 1 --baseline ``` Target: `bias_grad_reduce_f32_kernel` gone from top 10. `reward_rank_normalize` from 10.6s → <50ms. Step time <40ms. --- ## Expected Impact | Kernel | Before | After | Savings | |--------|--------|-------|---------| | reward_rank_normalize | 10,652ms (1 call) | <50ms | 10.6s one-time | | bias_grad_reduce_f32_kernel | 6,659ms (28K calls) | <200ms | 6.5s total | | curiosity_bias_grad_reduce | 539ms (4 calls) | <2ms | 537ms total | | **Total savings** | **17,850ms** | **<252ms** | **~17.6s** | Per-step training time (993 steps in 60s capture): current ~60ms/step → target ~40ms/step after bias_grad_reduce fix. The reward_rank_normalize is one-time per epoch, not per-step.