From d751d2ad4f5fe9c09aa6c733be7100da4b306636 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 12:15:04 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20kernel=20optimization=20=E2=80=94=208?= =?UTF-8?q?=20tasks,=207=20kernels=20=E2=86=92=20cuBLAS=20+=20batch-parall?= =?UTF-8?q?el?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-19-kernel-optimization.md | 355 ++++++++++++++++++ 1 file changed, 355 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-19-kernel-optimization.md diff --git a/docs/superpowers/plans/2026-04-19-kernel-optimization.md b/docs/superpowers/plans/2026-04-19-kernel-optimization.md new file mode 100644 index 000000000..a02e67b05 --- /dev/null +++ b/docs/superpowers/plans/2026-04-19-kernel-optimization.md @@ -0,0 +1,355 @@ +# 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: +```rust +// 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: +```c +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** +```bash +SQLX_OFFLINE=true cargo check -p ml +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +- [ ] **Step 6: Commit** +```bash +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** + +```rust +// 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. + +```c +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** +```bash +SQLX_OFFLINE=true cargo check -p ml +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +- [ ] **Step 6: Run local smoke test** +```bash +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** +```bash +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** + +```rust +// 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** +```bash +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** + +```c +// 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** +```bash +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** +```bash +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** +```bash +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** +```bash +git commit -m "perf: risk_budget_forward — cuBLAS GEMMs (20ms → <0.1ms)" +``` + +--- + +### Task 8: Deploy + nsys Validation + +- [ ] **Step 1: Push and deploy with nsys** +```bash +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** +```bash +git commit -m "docs: kernel optimization validation — step time Xms (was 1950ms)" +```