foxhunt/docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md

# K-loop Parallelization — Block-per-Batch Kernel Refactor

**Status:** Design
**Date:** 2026-05-17
**Affects:** `crates/ml-alpha/cuda/{cfc_step.cu, multi_horizon_heads.cu, variable_selection.cu, attention_pool.cu}` + `crates/ml-alpha/src/trainer/perception.rs`

## Problem

Five backward / K-loop kernels in `ml-alpha` use the "ONE block per launch, internal `n_batch` loop" pattern:

| Kernel | Calls per training step |
|---|---|
| `cfc_step_batched` (fwd) | 64 (K-loop) |
| `cfc_step_backward_batched` (bwd) | 64 (K-loop) |
| `multi_horizon_heads_grn_bwd_batched` | 64 (K-loop) |
| `variable_selection_bwd` | 1 |
| `attention_pool_bwd` | 1 |

Each runs as a single block (`grid_dim=(1,1,1)`, `block_dim=(C,1,1)` where C is the channel/feature dim) and iterates over the batch sequentially inside the kernel body. On an L40S with 142 SMs and B=32 this puts **<1% of the GPU to work** in the K-loop's critical path.

Empirical evidence from t6z89 (Phase 1+2+3 fold-1, stride=4): epoch wall 8 min for 750 steps → 640 ms/step. Estimated real compute per step is < 10 ms based on FLOP count. The gap (~630 ms/step idle) is the single-SM bottleneck.

Architectural impact: 3-fold CV at 30 epochs × 3 folds × 8 min = 12 hours wall. With the fix, target is 3 min/epoch → 4.5 hours. Future decision-stride sweeps, state-dim sweeps, and the Phase 3 CV all compound this.

## Goal

Parallelize the five kernels across the batch dimension so they use B simultaneous thread blocks instead of one.

**Realistic target:** ≥ 3× per-epoch speedup (8 min → ≤ 2.5 min).
**Stretch:** ≥ 5× (8 min → ≤ 1.5 min).

Higher multiples (e.g., 10×+) would require either bigger batch (B=128+), H100 hardware, or mixed-precision — all out of scope here. The K-loop block-per-batch refactor alone is bounded by Amdahl on the non-K-loop work (Mamba2 scan, LayerNorm, cuBLAS calls) which is already parallel.

## Non-goals

- Mamba2's selective scan kernel (already parallel via `backward_with_slices_into`).
- LayerNorm fwd/bwd (already parallel across rows).
- Forward GRN, snap_feature_assemble, BCE (already parallel).
- Fusing the K-loop into a single launch — would require parallel-scan reformulation of CfC's recurrence; out of scope.
- Mixed-precision (BF16) — orthogonal, considered in a future spec.
- **Hardware change.** L40S remains the target pool (`ci-training-l40s`). H100 is intentionally excluded from this round — the goal is to fix the architectural bottleneck on the hardware we're already running on, so the speedup is attributable to the refactor (not a hardware bump). H100 may revisit in a later spec once the L40S floor is established.

## Architecture Choice

For every refactored kernel, the same three changes:

**1. Launch config**

```rust
- LaunchConfig { grid_dim: (1, 1, 1), block_dim: (C, 1, 1), shared_mem_bytes: ... }
+ LaunchConfig { grid_dim: (b_sz as u32, 1, 1), block_dim: (C, 1, 1), shared_mem_bytes: ... }
```

**2. Kernel body**

```cuda
- int i = blockIdx.x * blockDim.x + threadIdx.x;
- if (i >= C) return;
- for (int bi = 0; bi < n_batch; ++bi) {
-     /* per-batch work */
- }
+ int bi = blockIdx.x;
+ int i = threadIdx.x;
+ if (bi >= n_batch || i >= C) return;
+ /* process exactly (bi, i) — no batch loop */
```

**3. Param-grad writes (bwd kernels only)**

```cuda
- grad_w[i * n_in + k] += d_pre * x_b[k];                        // safe when sequential
+ grad_w_scratch[(bi * C + i) * n_in + k] += d_pre * x_b[k];     // sole writer per (bi, i)
```

Followed by one reducer kernel call per (bwd-kernel × param-tensor) pair:

```cuda
// reduce_axis0 — single parameterised reducer, ONE cubin, ONE symbol.
extern "C" __global__ void reduce_axis0(
    const float* per_batch,   // [B, N]
    int n_batch,
    int n_tail,
    float* out                // [N]   — OVERWRITE
);
```

Block tree-reduce across `bi` (no atomicAdd, per `feedback_no_atomicadd.md`), grid over the output tail dim.

### Why scratch+reducer, not cooperative_groups grid sync

Cooperative grid sync (`cooperative_groups::this_grid().sync()`) would skip the separate reducer launch. Net additional saving: ~64 launches/step × ~20 µs = ~1.3 ms/step ≈ 1% on top of the block-per-batch win. Costs: CUDA Graph capture risk (`pearl_cudarc_disable_event_tracking_for_graph_capture` already documents fragility here), `-rdc=true` build flag change, new cudarc API (`cudaLaunchCooperativeKernel`) not used elsewhere in foxhunt, harder nsys profiling, concurrent-block-resident constraint. Scratch+reducer matches the existing LayerNorm-bwd-reducer pattern and is CUDA-Graph-safe.

### Large shared-mem note

`cfc_step_backward_batched` uses `extern __shared__ float smem[]` sized `(B+1) * n_hid * 4` bytes — currently 16.5 KiB at B=32, n_hid=128. The refactored version has **one block per batch** with shared mem sized `(1+1) * n_hid * 4` = 1 KiB per block (one row of `sd_pre` for the block's single `bi`, plus `sdecay[n_hid]`). Lower per-block shared mem, no `cudaFuncSetAttribute` call needed. If a future expansion increases shared mem above 48 KiB per block we'd need `cudaFuncAttributeMaxDynamicSharedMemorySize` — flagged for awareness, not action.

### Per-K-iteration accumulation invariant

The K-loop calls `cfc_bwd` and `grn_bwd` 64 times per training step. Each call's scratch writes are `+=`. The scratch is `memset_zeros` ONCE at step start (alongside other grad accumulators), then 64 K-iterations accumulate into it. After the K-loop completes, one reducer launch per param tensor sums B → final param grad. Same total work as the current sequential pattern, parallelized across B.

### AdamW-after-reducer invariant

The final grad buffer (e.g., `grad_w_in_d`) was previously zeroed at step start, accumulated by the bwd kernel, then read by AdamW. After the refactor:

- Final grad buffer is **OVERWRITTEN by the reducer**, not zeroed at step start.
- Final grad buffer holds meaningful data **only after the reducer for it has run in the current step**.
- **AdamW must run after the reducer.** Any future code path that reads a final grad without running the reducer first (diagnostic dump, alternative optimizer path, checkpoint serialiser) will read stale data from the previous step.

This invariant is asserted via the ordering inside `dispatch_train_step`: reducer launches sit between the bwd kernels and the AdamW step calls, all captured in the same CUDA Graph region.

## Per-Kernel Breakdown

| Kernel | File | Param-grad scratch at B=32 | Scratch bytes | Reducer launches |
|---|---|---|---|---|
| `cfc_step_batched` (fwd) | `cuda/cfc_step.cu` | none | 0 | 0 |
| `cfc_step_backward_batched` | `cuda/cfc_step.cu` | `grad_w_in[B,128,128]` + `grad_w_rec[B,128,128]` + `grad_b[B,128]` + `grad_tau[B,128]` | 4.1 MB | 4 |
| `multi_horizon_heads_grn_bwd_batched` | `cuda/multi_horizon_heads.cu` | 10 tensors, biggest `grad_w1[B,5,64,128]` | 8.0 MB | 10 |
| `variable_selection_bwd` | `cuda/variable_selection.cu` | `grad_W_vsn[B,40,40]` + `grad_b_vsn[B,40]` | 210 KB | 2 |
| `attention_pool_bwd` | `cuda/attention_pool.cu` | `grad_Q[B,128]` | 16 KB | 1 |
| **Totals** | — | — | **~12 MB** | **17 per step** |

**Scratch sizing formula:** for kernel K with parameters {Pᵢ} of trailing-dim size {Nᵢ}, scratch is

```
scratch_bytes(K, B) = B × Σᵢ (Nᵢ × 4)
```

Linear in B. At B=128 → ~48 MB total (still < 0.1 % of L40S 48 GB). At B=512 → ~190 MB. Memory ceiling for B at this architecture is bound by other allocators long before scratch matters.

Reducer launches (17/step) are microsecond-scale (small grid, simple kernel) → ~340 µs total overhead, rounding error against the K-loop savings.

## Trainer Wiring

**New `PerceptionTrainer` fields (~17):**

- `cfc_grad_w_in_scratch_d`, `cfc_grad_w_rec_scratch_d`, `cfc_grad_b_scratch_d`, `cfc_grad_tau_scratch_d`
- `grn_grad_w1_scratch_d`, `grn_grad_b1_scratch_d`, `grn_grad_w2_scratch_d`, `grn_grad_b2_scratch_d`, `grn_grad_w_gate_scratch_d`, `grn_grad_b_gate_scratch_d`, `grn_grad_w_main_scratch_d`, `grn_grad_b_main_scratch_d`, `grn_grad_w_skip_scratch_d`, `grn_grad_b_skip_scratch_d`
- `vsn_grad_w_scratch_d`, `vsn_grad_b_scratch_d`
- `attn_grad_q_scratch_d`
- `reduce_axis0_fn: CudaFunction`
- `_reduce_module: Arc<CudaModule>`

**Step lifecycle:**

1. **Step start:** `memset_zeros` the new scratch buffers. The OLD final grad buffers are no longer zeroed at step start — they are reducer outputs (OVERWRITE semantics).
2. **K-loop bwd:** 64 cfc_bwd + 64 grn_bwd calls write `+=` into scratch.
3. **After K-loop:** 14 reducer launches (4 cfc + 10 grn) collapse scratch → final grad buffers.
4. **VSN bwd + reducer (2) + attn pool bwd + reducer (1):** standalone path. For these, scratch zero + write + reducer all happen in one logical step (the kernel runs once, so the K-accumulation invariant is moot — but we still memset scratch at step start for uniformity with the rest of the pattern).
5. **AdamW step:** unchanged. Consumes final grad buffers. Runs strictly after all reducers per the AdamW-after-reducer invariant.

**Unchanged:**
- Forward path (except cfc_step_batched fwd's launch config — same kernel signature, no scratch).
- BCE / ISV lambda.
- LayerNorm bwd (already uses scratch+reducer via `layer_norm_reduce_param_grads`).
- Mamba2 (already parallel internally).
- CUDA Graph capture / replay (reducers are plain `launch_builder.launch()` — same pattern as LN reducer which captures cleanly).

## Build System

- `crates/ml-alpha/build.rs` KERNELS list gets `"reduce_axis0"`.
- Cache-bust comment bumped to `v11`.
- No new `-rdc=true` flag, no nvcc API change.

## Tests + Acceptance

### Correctness gates

**1. All 8 existing `perception_overfit` tests pass** (B=1 path).
Synthetic constant-direction signal still converges 0.30 → 0.00 by step 50 at both stride=1 and stride=4. Catches any chain-rule break.

**2. NEW: `stacked_trainer_loss_shrinks_at_batch_32` test.**
Same synthetic signal, but `n_batch: 32`. Asserts loss converges from baseline to < 0.1 within 200 steps. This is the **first test that actually exercises the cross-batch reduction code path** — the existing perception_overfit suite uses B=1 throughout, so without this we'd never validate the reducer at real B.

**3. All existing `backward_finite_diff` tests pass** (single-sample GPU helpers).
Orthogonal code path, but verifying confirms the kernel module loading + gradient invariants are intact.

**4. NEW: `cfc_bwd_b1_matches_single_sample` test (`cfc_step` only).**
Run identical synthetic input through the refactored `cfc_step_backward_batched` at B=1, compare against the single-sample helper `cfc_step_backward_gpu`. Assert `relative_eq` ≤ 1e-6 (NOT bit-exact — FP addition is non-associative and the kernels may sum in different orders even with identical math).

**5. NEW: `cfc_bwd_scratch_clears_between_steps` test.**
Run two `trainer.step(seq_A, labels_A)` calls in succession. After step A, read each cfc scratch buffer; after step B, read each scratch buffer; assert step B's scratch reflects ONLY seq_B's contribution (i.e., scratch was zeroed at start of step B, not left holding step A's accumulation).

**Note on B=1 oracle asymmetry.** Single-sample GPU helpers exist only for `cfc_step` (via `crates/ml-alpha/src/cfc/step.rs::cfc_step_backward_gpu`). There is NO single-sample helper for the GRN bwd, VSN bwd, or attention pool bwd kernels, and `feedback_no_cpu_test_fallbacks.md` forbids writing a CPU reference oracle. Validation for the other three refactored kernels rests on:
- The new B=32 synthetic smoke (gate #2) — sees them work end-to-end.
- The fact that the refactor preserves chain rule structure unchanged — only the parallelization changes.
- The cluster A/B AUC trajectory match (gate #7) — catches numerical regressions on real data.

Writing single-sample helpers for GRN / VSN / attn could be done, but the engineering cost (~3 kernels × 100 LoC each + binding code) exceeds the marginal validation value given the convergence and A/B safety nets.

### Performance gates

**6. Local synthetic-overfit step time ≥ 3× faster at B=32 (release build).**
Pre-refactor: ~250 ms/step. Post-refactor target: ~80 ms/step. Measured by wall delta on `cargo test stacked_trainer_loss_shrinks_at_batch_32 --release`.

**7. Cluster epoch wall time ≥ 3× faster (stretch: ≥ 5×).**
A/B Argo submission. Same config as t6z89 (`--decision-stride 4 --seq-len 64 --mamba2-state-dim 32 --epochs 5 --cv-fold 1 --cv-n-folds 3 --cv-train-window 4`). Compare timestamps between successive `epoch complete` log lines.
Pre: 8 min/epoch (t6z89). Post target: ≤ 2.5 min/epoch.

**8. AUC trajectory matches t6z89 within ±0.005.**
The perf-refactor run must show mean_auc and h6000 within ±0.005 of t6z89's per-epoch values at the same epoch index. Larger deviation indicates a numerical regression.

### Post-refactor bottleneck analysis (informational, not a gate)

After block-per-batch lands on L40S the next limiter is expected to be one of:
- **Mamba2's selective scan** — already parallel across (batch, channels) but K-sequential by construction. ~5-10% of step time today; possibly more after K-loop is fast.
- **Kernel launch latency floor** — ~350 kernel launches per step at ~5 µs each = ~2 ms minimum step time (with CUDA Graph) or much higher without.
- **cuBLAS calls inside Mamba2's W_in / W_a / W_b / W_c projections** — pre-allocated workspaces help but small matmuls don't saturate the GPU.

These set the floor for any future optimization round on L40S. Once that floor is known, a follow-up spec can decide whether to spend on H100 hardware, BF16 trunk, or larger batch — but this spec stays on L40S so the speedup is attributable to the refactor alone.

## Risks + Mitigations

| Risk | Mitigation |
|---|---|
| FP order drift from cross-batch reduction | Deterministic reducer (sum in `bi` index order via stride-loop). Gate #4 (`relative_eq` 1e-6) catches drift large enough to matter. |
| CUDA Graph capture breaks with reducer launches | Reducers are plain `launch_builder.launch()` — identical pattern to `layer_norm_reduce_param_grads`. `evaluate_works_after_captured_training_step` guards this. |
| K-loop scratch accumulation bug | Gate #5 (`cfc_bwd_scratch_clears_between_steps`) explicitly asserts scratch = 0 at step N+1. Gate #2 (B=32 convergence) diverges immediately on chain-rule break. |
| Existing single-sample finite-diff tests break | These use single-sample GPU helpers — orthogonal code path, unaffected by this refactor. If they break, the refactor accidentally touched the wrong file. |
| Memory budget at future B | `scratch_bytes ≈ B × Σ(param tensor sizes)`. At B=512 → ~190 MB. No mitigation needed at any plausible scale. |
| Refactor lands in pieces leaving half-old-half-new state | Each kernel refactor is one atomic commit covering kernel + scratch + reducer wiring + smoke. No `_legacy` symbols. `feedback_no_partial_refactor.md` + `feedback_no_legacy_aliases.md`. |
| AdamW reads stale grad | The AdamW-after-reducer invariant is structurally enforced by the order of calls in `dispatch_train_step`. Documented in the trainer-wiring section; any future diagnostic / checkpoint code that needs final grads must run reducers first. |

## Rollout

Each kernel refactor lands as a single atomic commit on `ml-alpha-phase-a`. No worktree, no merge step — commits go straight to the active branch (per session decision).

**Commit 1: `reduce_axis0` kernel + first consumer (cfc_step refactor).**
Bundles three things to avoid the orphan-kernel anti-pattern (`feedback_wire_everything_up.md`):
- New `cuda/reduce_axis0.cu` + build.rs registration
- Refactor `cfc_step_batched` fwd (launch config only — no scratch)
- Refactor `cfc_step_backward_batched` + add 4 cfc scratch fields + zero scratch at step start + 4 reducer launches after K-loop bwd
- Add `cfc_bwd_b1_matches_single_sample` + `cfc_bwd_scratch_clears_between_steps` tests
- All existing smokes pass + new B=32 smoke (`stacked_trainer_loss_shrinks_at_batch_32`)

**Commit 2: GRN bwd refactor.**
Refactor `multi_horizon_heads_grn_bwd_batched` + 10 scratch fields + 10 reducer launches. All smokes pass.

**Commit 3: VSN bwd refactor.**
Refactor `variable_selection_bwd` + 2 scratch fields + 2 reducer launches. All smokes pass.

**Commit 4: attention pool bwd refactor.**
Refactor `attention_pool_bwd` + 1 scratch field + 1 reducer launch. All smokes pass.

**Step 5 (no commit): local A/B perf benchmark.**
`cargo test --release` wall-time delta old vs new. Verify gate #6 (≥ 3× local).

**Step 6 (no commit): cluster A/B benchmark.**
Argo submit on commit 4's HEAD with t6z89-identical config. Verify gates #7 (≥ 3× wall) and #8 (AUC within ±0.005).

## Rollback

If a gate fails after a commit lands:

- **Single broken commit:** `git revert <sha>` on `ml-alpha-phase-a`. Atomic-commit discipline (one kernel per commit) means each revert touches one kernel cleanly.
- **Cascade breakage across multiple commits:** revert the range. E.g., if commits 2 + 3 both broke and only commit 1 is salvageable, `git revert <sha2>..<sha4>`. The architectural Phase 1+2+3 work (everything before commit 1) and t6z89's reference commit `7a558b88b` remain in history regardless.
- **Worst case (perf is fine but AUC trajectory diverges by >0.005):** revert everything; we've kept the architectural state and gained data about which kernel changed semantics. Then write a follow-up spec that adds the missing oracle (e.g., GRN single-sample helper) before retrying.

`ml-alpha-phase-a`'s history before commit 1 (currently `54aa69c10` + the spec commit) remains the safe rollback point for the entire refactor.

## Open Questions

- Optimal `block_dim` for `reduce_axis0` itself (likely 256 threads; can be parameterised via launch config — decided in implementation).
- Whether to fuse multiple reducer launches into a multi-tensor sweep kernel (probably one per tensor — per-launch overhead is negligible vs code-complexity savings).