foxhunt/docs/superpowers/plans/2026-04-19-kernel-optimization-phase3.md

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:

/* 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:

reward_abs_sharpe_buf: CudaSlice<f32>,   // [max_N] sort keys
reward_indices_buf: CudaSlice<i32>,       // [max_N] original indices
reward_sorted_keys_buf: CudaSlice<f32>,   // [max_N] sorted output
reward_sorted_indices_buf: CudaSlice<i32>,// [max_N] sorted indices
reward_sort_temp_buf: CudaSlice<u8>,      // 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

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

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:

/* 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:

bias_grad_p1_kernel: CudaFunction,
bias_grad_p2_kernel: CudaFunction,
bias_grad_partials_buf: CudaSlice<f32>,  // [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

fn launch_bias_grad(
    &self, stream: &Arc<CudaStream>,
    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

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

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:

/* 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:

bias_grad_p1_func: CudaFunction,
bias_grad_p2_func: CudaFunction,
bias_grad_partials: CudaSlice<f32>,  // [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:

// 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

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

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

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

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.