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

Custom 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: Rewrite 7 custom CUDA kernels that consume 95% of GPU time — replace inner matmul loops with cuBLAS GEMMs, replace batch-serial patterns with batch-parallel grids. Target: 1950ms/step → <50ms/step.

Architecture: Split each "one-thread-per-weight-loops-over-batch" kernel into: (1) cuBLAS GEMMs for matrix projections (tensor cores), (2) lightweight element-wise/scan kernels for non-GEMM ops. The key insight from NVIDIA's official Mamba implementation: grid=(batch, dim) not grid=(params).

Tech Stack: CUDA 13, cublasLtMatmul (TF32), cudarc (vendored), existing IqnGemmDesc pattern for cached GEMM descriptors.

Dimensions: SH2=256, STATE_D=16, K=8, batch=8192, IQN_Q=64, adv_h=128.

---

Task 1: Mamba2 Forward — cuBLAS Projections + Lightweight Scan (1.1%)

Files:

• Modify: crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

The forward scan currently does per-sample inner matmuls: for j in 0..SH2: a_val += W_A[j,s] * h_t[j]. This is a GEMM [STATE_D, SH2] × [SH2, B] that should use cuBLAS.

• Step 1: Allocate scratch buffers for pre-projected values

In gpu_dqn_trainer.rs constructor, allocate:

// Mamba2 pre-projection scratch: A_vals[B*K, STATE_D], B_vals[B*K, STATE_D]
let mamba2_a_proj = alloc_f32(&stream, batch_size * MAMBA2_HISTORY_K * MAMBA2_STATE_DIM, "mamba2_a_proj")?;
let mamba2_b_proj = alloc_f32(&stream, batch_size * MAMBA2_HISTORY_K * MAMBA2_STATE_DIM, "mamba2_b_proj")?;
let mamba2_c_proj = alloc_f32(&stream, batch_size * MAMBA2_STATE_DIM, "mamba2_c_proj")?;

• Step 2: Create cuBLAS GEMM descriptors for Mamba2 projections

Using the existing create_cached_fwd_gemm_desc pattern, create 3 GEMM descriptors:

// A_proj = W_A^T @ h_history_flat: m=STATE_D(16), n=B*K(65536), k=SH2(256)
// B_proj = W_B^T @ h_history_flat: same dims
// C_proj = W_C^T @ h_s2:          m=STATE_D(16), n=B(8192), k=SH2(256)

Load these from a new CUmodule (graph-captured, must be isolated).

• Step 3: Write lightweight scan kernel mamba2_scan_projected

Replace mamba2_temporal_scan with a kernel that reads pre-projected values:

extern "C" __global__ void mamba2_scan_projected(
    const float* __restrict__ a_proj,    // [B*K, STATE_D] pre-projected
    const float* __restrict__ b_proj,    // [B*K, STATE_D] pre-projected
    const float* __restrict__ c_proj,    // [B, STATE_D] pre-projected
    const float* __restrict__ h_s2,      // [B, SH2] current trunk
    float* __restrict__ h_enriched,      // [B, SH2] output
    int B, int K, int sh2, int state_d,
    const float* __restrict__ isv_signals,
    const float* __restrict__ temporal_weight)
{
    int b = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (b >= B || j >= sh2) return;

    // Scan over K=8 timesteps — reads from pre-projected buffers
    float context = 0.0f;
    for (int s = 0; s < state_d; s++) {
        float x = 0.0f;
        for (int t = 0; t < K; t++) {
            int idx = (b * K + t) * state_d + s;
            float gate = 1.0f / (1.0f + expf(-a_proj[idx]));
            if (isv_signals) gate *= isv_signals[11];
            x = gate * x + b_proj[idx];
        }
        context += c_proj[b * state_d + s] * x;
    }

    // Apply temporal routing and residual
    float h_val = h_s2[b * sh2 + j];
    float weight = temporal_weight ? temporal_weight[j] : 1.0f;
    float blended = h_val + weight * context * (1.0f / (float)state_d);
    h_enriched[b * sh2 + j] = blended;
}

Grid: (B, ceil(SH2/256)), Block: (256). Each thread handles one (sample, feature). K=8 inner loop — trivial. STATE_D=16 inner loop — 16 iterations, no memory bottleneck.

• Step 4: Update mamba2_forward in gpu_dqn_trainer.rs

Replace the old kernel launch with:

1. cuBLAS GEMM: W_A^T @ h_history_flat → mamba2_a_proj
2. cuBLAS GEMM: W_B^T @ h_history_flat → mamba2_b_proj
3. cuBLAS GEMM: W_C^T @ h_s2 → mamba2_c_proj
4. Launch mamba2_scan_projected kernel

Note: h_history is [B, K, SH2] — reshape to [B*K, SH2] for the GEMM by treating it as a flat contiguous buffer (it already is in memory).

• Step 5: Build and test

SQLX_OFFLINE=true cargo check -p ml
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 6: Commit

git commit -m "perf: mamba2 forward — cuBLAS projections + lightweight scan (22ms → <1ms)"

---

Task 2: Mamba2 Backward — cuBLAS Projections + Lightweight Reverse Scan (75%)

Files:

• Modify: crates/ml/src/cuda_pipeline/experience_kernels.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

This is THE bottleneck — 1567ms/call, 75% of total GPU time. Same approach as forward but with reverse scan + weight gradient accumulation via cuBLAS.

• Step 1: Allocate backward scratch buffers

// Backward needs: d_gate[B*K, STATE_D], d_x[B*K, STATE_D], d_context[B, STATE_D]
let mamba2_d_gate = alloc_f32(&stream, batch_size * MAMBA2_HISTORY_K * MAMBA2_STATE_DIM, "mamba2_d_gate")?;
let mamba2_d_x = alloc_f32(&stream, batch_size * MAMBA2_HISTORY_K * MAMBA2_STATE_DIM, "mamba2_d_x")?;
let mamba2_d_context = alloc_f32(&stream, batch_size * MAMBA2_STATE_DIM, "mamba2_d_context")?;

• Step 2: Write mamba2_scan_backward_projected kernel

This kernel does the reverse scan using pre-projected values. No inner matmuls. Grid: (B, STATE_D), each thread handles one (sample, state_dim) pair.

extern "C" __global__ void mamba2_scan_backward_projected(
    const float* __restrict__ a_proj,      // [B*K, STATE_D] forward A projections
    const float* __restrict__ b_proj,      // [B*K, STATE_D] forward B projections
    const float* __restrict__ c_proj,      // [B, STATE_D] forward C projections
    const float* __restrict__ d_h_enriched,// [B, SH2] upstream gradient
    const float* __restrict__ w_c,         // [SH2, STATE_D]
    float* __restrict__ d_gate,            // [B*K, STATE_D] output: sigmoid_deriv * d_x * x_prev
    float* __restrict__ d_x_out,           // [B*K, STATE_D] output: d_x per timestep
    float* __restrict__ d_context,         // [B, STATE_D] output: d_context = d_out @ W_C
    int B, int K, int sh2, int state_d,
    const float* __restrict__ isv_signals)
{
    int b = blockIdx.x;
    int s = blockIdx.y * blockDim.x + threadIdx.x;
    if (b >= B || s >= state_d) return;

    // Forward replay: recompute x_states[0..K] for this (b, s)
    float x_states[9]; // K+1 max
    x_states[0] = 0.0f;
    float a_raw[8];
    for (int t = 0; t < K; t++) {
        int idx = (b * K + t) * state_d + s;
        a_raw[t] = a_proj[idx];
        float gate = 1.0f / (1.0f + expf(-a_raw[t]));
        if (isv_signals) gate *= isv_signals[11];
        x_states[t + 1] = gate * x_states[t] + b_proj[idx];
    }

    // d_context[b, s] = sum_j(d_h_enriched[b, j] * w_c[j, s])
    float d_ctx = 0.0f;
    for (int j = 0; j < sh2; j++) {
        d_ctx += d_h_enriched[b * sh2 + j] * w_c[j * state_d + s];
    }
    d_context[b * state_d + s] = d_ctx;

    // d_x[K] = d_context * c_proj was the last step contribution
    float d_x_s = d_ctx;

    // Reverse scan
    for (int t = K - 1; t >= 0; t--) {
        int idx = (b * K + t) * state_d + s;
        float a_val = a_raw[t];
        float gate = 1.0f / (1.0f + expf(-a_val));
        float sigmoid_deriv = gate * (1.0f - gate);
        float stability = isv_signals ? isv_signals[11] : 1.0f;

        // d_gate for W_A gradient: d_x * x_prev * sigmoid'
        d_gate[idx] = d_x_s * x_states[t] * sigmoid_deriv * stability;
        // d_x for W_B gradient
        d_x_out[idx] = d_x_s;
        // Propagate d_x backward
        d_x_s = d_x_s * gate * stability;
    }
}

• Step 3: Create cuBLAS GEMM descriptors for weight gradient accumulation

After the scan kernel produces d_gate[B*K, STATE_D] and d_x[B*K, STATE_D]:

d_W_A[SH2, STATE_D] = h_history^T[SH2, B*K] @ d_gate[B*K, STATE_D]
                       GEMM: transa=N, transb=N, m=STATE_D(16), n=SH2(256), k=B*K(65536)
d_W_B[SH2, STATE_D] = h_history^T[SH2, B*K] @ d_x[B*K, STATE_D]
                       same dims
d_W_C[SH2, STATE_D] = h_s2^T[SH2, B] @ d_context[B, STATE_D]
                       GEMM: m=STATE_D(16), n=SH2(256), k=B(8192)

• Step 4: Update mamba2_backward in gpu_dqn_trainer.rs

Replace old kernel launch with:

1. Reuse forward A/B projections (already computed in forward, save them)
2. Launch mamba2_scan_backward_projected kernel
3. cuBLAS GEMM: h_history^T @ d_gate → d_W_A
4. cuBLAS GEMM: h_history^T @ d_x → d_W_B
5. cuBLAS GEMM: h_s2^T @ d_context → d_W_C

• Step 5: Build and test

SQLX_OFFLINE=true cargo check -p ml
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 6: Run local smoke test

SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib --release --features ml/cuda -- smoke_tests --ignored --nocapture

• Step 7: Commit

git commit -m "perf: mamba2 backward — cuBLAS projections + reverse scan (1567ms → <5ms)"

---

Task 3: Q-Denoise Backward — cuBLAS GEMMs (7.3%)

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

Replace the custom q_denoise_backward kernel (2-layer MLP backward, one thread per weight, loops over batch) with 2× cuBLAS backward GEMMs + bias_grad_reduce. The MLP is 12→24→12 with batch=8192.

• Step 1: Create cuBLAS GEMM descriptors for denoise backward

// Layer 2 backward dX: m=24, n=8192, k=12 (dX = W2^T @ dY2)
// Layer 2 backward dW: m=12, n=24, k=8192 (dW = dY2^T @ h1)
// Layer 1 backward dX: m=24, n=8192, k=24 (dX = W1^T @ dY1) -- not needed (no upstream)
// Layer 1 backward dW: m=24, n=24, k=8192 (dW = dY1^T @ Q_input)

• Step 2: Rewrite launch_q_denoise_backward to use cuBLAS

Replace single kernel with: loss gradient → cuBLAS dW2 → ReLU mask → cuBLAS dW1 → bias_grad_reduce × 2.

• Step 3: Build, test, commit

git commit -m "perf: q_denoise_backward — cuBLAS GEMMs (152ms → <1ms)"

---

Task 4: IQN Bias Grad Reduce — Warp-Parallel Reduction (3.9%)

Files:

• Modify: crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu

Current: one thread per output neuron (out_dim=3), each loops over B×Q=524,288. 3 threads doing 524K iterations.

• Step 1: Rewrite iqn_bias_grad_reduce with batch-parallel grid

// Phase 1: each block reduces 256 elements
// Grid: (ceil(B*Q/256), out_dim), Block: 256
// Each thread reads one element, warp shuffle reduces to partial sum
// Block-level reduction via shared memory → one partial per block

// Phase 2: reduce block partials → final db[out_dim]
// Grid: (out_dim), Block: 256
// Each thread reads one block partial, warp reduce to final sum

• Step 2: Build, test, commit

git commit -m "perf: iqn_bias_grad_reduce — warp-parallel (16ms×130 → 0.1ms×130)"

---

Task 5: KAN Grad Reduce — Batch-Parallel Reduction (3.2%)

Files:

• Modify: crates/ml/src/cuda_pipeline/experience_kernels.cu

Same anti-pattern as IQN — one thread per param, loops over batch. Fix: batch-parallel 2-phase reduction.

• Step 1: Rewrite kan_grad_reduce with batch-parallel grid

Same pattern as Task 4 — phase 1 block reduce + phase 2 final reduce.

• Step 2: Build, test, commit

git commit -m "perf: kan_grad_reduce — batch-parallel (8ms×208 → 0.1ms×208)"

---

Task 6: Curiosity Fwd+Bwd — cuBLAS GEMMs (2.1%)

Files:

• Modify: crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

The training kernel is a single launch doing forward + backward for the entire batch. Replace with cuBLAS GEMMs for the curiosity MLP.

• Step 1: Replace curiosity_fwd_bwd_per_block with cuBLAS pipeline

Forward: 2 cuBLAS GEMMs (state→hidden→pred). Backward: 2 cuBLAS GEMMs (dW). Loss: element-wise kernel.

• Step 2: Build, test, commit

git commit -m "perf: curiosity training — cuBLAS GEMMs (1085ms → <5ms)"

---

Task 7: Risk Budget Forward — cuBLAS GEMMs (1.0%)

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• Modify: crates/ml/src/cuda_pipeline/experience_kernels.cu

Per-sample inner matmul over SH2=256. Replace with cuBLAS.

• Step 1: Replace risk_budget_forward with cuBLAS pipeline

Concat [h_s2; ISV; pred_error] → cuBLAS GEMM1 → ReLU → cuBLAS GEMM2 → sigmoid. Same pattern as main forward.

• Step 2: Build, test, commit

git commit -m "perf: risk_budget_forward — cuBLAS GEMMs (20ms → <0.1ms)"

---

Task 8: Deploy + nsys Validation

• Step 1: Push and deploy with nsys

git push origin main
./scripts/argo-train.sh dqn --baseline --epochs 1 --sanitizer nsys

• Step 2: Verify step timing Expected: STEP_TIMING step=50 step_ms=<50 (was 1950ms)

• Step 3: Compare nsys kernel breakdown Verify all 7 kernels are now <5ms each. cuBLAS GEMMs should dominate.

• Step 4: Run smoke tests on H100 All 20 smoke tests pass. Training metrics within 5% of baseline.

• Step 5: Commit validation results

git commit -m "docs: kernel optimization validation — step time Xms (was 1950ms)"