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:

```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<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**

```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<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**

```rust
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**

```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<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:

```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.