From 54aa69c108b74e717ac73e972306c0541cfbeed1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 23:06:28 +0200 Subject: [PATCH] spec(perf): K-loop block-per-batch parallelization design MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Documents the design for fixing the single-SM bottleneck in five backward/K-loop kernels: cfc_step (fwd+bwd), GRN bwd, VSN bwd, and attention_pool bwd. All currently use grid=(1,1,1) with an internal n_batch loop — on L40S (142 SMs) with B=32 this puts <1% of the GPU to work in the K-loop critical path. Architecture: block-per-batch (grid=(B,1,1)) for the kernel body, plus per-batch grad scratch buffers reduced via a single parameterised reduce_axis0 kernel (block tree-reduce, no atomicAdd per feedback_no_atomicadd.md). Same pattern as the existing LayerNorm bwd reducer — CUDA-Graph-safe, debuggable, and consistent with foxhunt's no-cooperative-groups discipline. Target: ≥3× epoch wall speedup (stretch 8-15×). Makes 3-fold CV tractable (10.5h → 2-3h) and unblocks decision-stride / state-dim sweeps that compound the gain. Acceptance gates: (a) all 8 perception_overfit smokes still converge, (b) new B=1 bit-equivalence test asserts the refactored batched bwd kernel at B=1 matches the existing single-sample helper byte-for-byte, (c) cluster A/B vs t6z89 baseline shows AUC trajectory within ±0.005 and epoch wall ≥3× faster. Atomic refactor per kernel — one commit per kernel covering the kernel rewrite, scratch buffer alloc, reducer launch wiring, and smoke. No "_legacy" parallel kernels per feedback_no_legacy_aliases.md + feedback_no_partial_refactor.md. Open implementation-plan decisions: exact memset_zeros ordering inside the captured graph, batch-vs-per-tensor reducer launches, optimal block_dim for reduce_axis0 itself. Co-Authored-By: Claude Opus 4.7 --- ...2026-05-17-kloop-parallelization-design.md | 185 ++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md diff --git a/docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md b/docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md new file mode 100644 index 000000000..afdbf6f31 --- /dev/null +++ b/docs/superpowers/specs/2026-05-17-kloop-parallelization-design.md @@ -0,0 +1,185 @@ +# 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 time 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. Decision-stride sweeps, ablations, and the eventual 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. Target speedup: **≥ 3× per epoch (stretch: 8–15×)**. + +## 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. + +## 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 training step: + +```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 ~10× we get from block-per-batch alone. Costs: CUDA Graph capture risk (`pearl_cudarc_disable_event_tracking_for_graph_capture` already documents fragility here), `-rdc=true` build flag changes, new cudarc API path (`cudaLaunchCooperativeKernel`) not used elsewhere in foxhunt, harder nsys profiling, concurrent-block-resident constraint that forecloses B → 1000+. Scratch+reducer matches the existing LayerNorm-bwd-reducer pattern and is CUDA-Graph-safe. + +### 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 the other grad accumulators), then 64 K-iterations accumulate into it. After the K-loop completes, ONE reducer launch sums B → final param grad. AdamW consumes the final grad as today. Same total work as the current sequential pattern, parallelized across B. + +## Per-Kernel Breakdown + +| Kernel | File | Param-grad scratch (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 (one per param tensor) | +| `multi_horizon_heads_grn_bwd_batched` | `cuda/multi_horizon_heads.cu` | 10 tensors, biggest is `grad_w1[B,5,64,128]` | 8 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** | + +Total additional GPU memory ~12 MB. L40S has 48 GB; scratch grows linearly in B (at B=128 future, still <1% of L40S memory). Reducer launches are microsecond-scale (small grid, simple kernel) so 17 extra launches/step is ~340 µs 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` + +**Step lifecycle changes:** + +1. **Step start:** memset_zeros the new scratch buffers. The OLD final grad buffers (e.g., `grad_w_in_d`) are no longer zeroed at step start — they become reducer outputs (OVERWRITE semantics, written exactly once per step after K-loop). +2. **K-loop bwd:** the 64 cfc_bwd / grn_bwd calls write `+=` into scratch. +3. **After K-loop:** 14 reducer launches collapse cfc + grn scratch → final grad buffers. +4. **Standalone bwd (VSN, attn pool):** write scratch → reducer → final grad buffer (3 reducer launches). +5. **AdamW step:** unchanged — consumes the final grad buffers. + +**Unchanged:** + +- Forward path (except cfc_step_batched fwd's launch config — same kernel signature) +- BCE / ISV lambda +- LayerNorm bwd (already uses this pattern via `layer_norm_reduce_param_grads`) +- Mamba2 (already parallel internally) +- AdamW callsites +- CUDA Graph capture / replay (reducers are plain kernel launches; 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` (covers the kernel rewrites + new reducer cubin). +- No new `-rdc=true` flag, no nvcc API change. + +## Tests + Acceptance + +### Correctness gates + +1. **All 8 existing `perception_overfit` tests pass** — synthetic constant-direction signal still converges 0.30 → 0.00 by step 50 at both stride=1 and stride=4. Catches any chain-rule error in the refactor. +2. **All existing `backward_finite_diff` tests pass** — these test single-sample GPU helpers (`cfc_step_backward_gpu`, `multi_horizon_heads_backward_gpu`) on an orthogonal code path. Unaffected by the batched-kernel refactor, but verifying they still pass confirms the kernel module loading and gradient invariants are intact. +3. **New B=1 bit-equivalence test (`cfc_bwd_b1_matches_single_sample`):** run identical synthetic input through the refactored `cfc_step_backward_batched` kernel at B=1, compare against the existing single-sample helper `cfc_step_backward_gpu` from `crates/ml-alpha/src/cfc/step.rs`. Assert bit-exact equality (`==`, not `relative_eq`). The single-sample helper IS the B=1 oracle. +4. **New scratch-clears test (`cfc_bwd_scratch_clears_between_steps`):** run two `trainer.step(seq_A, labels_A)` calls; assert the scratch buffers contain zero contributions from `seq_A` when `seq_B` is processed. Guards against the "forgot to zero scratch" class of bug. + +### Performance gates + +5. **Local synthetic-overfit step time ≥ 3× faster at B=32:** wall delta on `cargo test stacked_trainer_loss_shrinks_on_constant_signal --release` (release build only — debug build has different timings). Target: pre-refactor 250 ms/step → post-refactor 80 ms/step. +6. **Cluster epoch wall time ≥ 3× faster:** 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`), one on pre-refactor `7a558b88b`, one on perf-branch HEAD. Compare timestamps between successive `epoch complete` log lines. Target: 8 min/epoch → 2.5 min/epoch. +7. **AUC trajectory matches t6z89:** the perf-branch fold-1 run must show mean_auc and h6000 within ±0.005 of t6z89's per-epoch values at the same epoch index. Larger deviation indicates the numerical-equivalence guarantee broke. + +## Risks + Mitigations + +| Risk | Mitigation | +|---|---| +| Numerical drift from cross-batch reduction order | Deterministic reducer (sum in `bi` index order via stride-loop). B=1 bit-equivalence test (acceptance gate #3) catches any inadvertent change. | +| CUDA Graph capture breaks with reducer launches | Reducers are plain `launch_builder.launch()` calls — identical pattern to `layer_norm_reduce_param_grads` which already captures cleanly. The `evaluate_works_after_captured_training_step` test (in `perception_overfit.rs`) guards this. | +| K-loop scratch accumulation bug (forgot to zero, double-zero, wrong index) | `perception_overfit::stacked_trainer_loss_shrinks_on_constant_signal` diverges immediately on any chain-rule break. The new `cfc_bwd_scratch_clears_between_steps` test asserts scratch = 0 at the start of step N+1. | +| Existing single-sample finite-diff tests break | These use single-sample GPU helpers (`cfc_step_backward_gpu` etc.) — orthogonal code path, unchanged by this refactor. If they break, the refactor accidentally touched the wrong file. | +| Memory budget at future B=128 | Scratch grows linearly: B=128 → ~48 MB. Still <0.1% of L40S 48 GB. No mitigation needed at this scope. | +| Refactor lands in pieces, leaving the codebase half-old-half-new | One atomic commit per kernel: kernel rewrite + scratch buffer + reducer launch wiring + smoke validation, all in the same commit. No "_legacy" symbols. `feedback_no_partial_refactor.md` + `feedback_no_legacy_aliases.md`. | + +## Rollout + +Implementation order (one commit per kernel, smoke-validated before moving on): + +1. `reduce_axis0` kernel + build.rs registration + minimal smoke (overwrite-correctness test on a known [B, N] tensor) +2. Refactor `cfc_step_batched` fwd (no scratch needed — just launch-config change) — smoke +3. Refactor `cfc_step_backward_batched` + cfc scratch buffers + 4 reducer launches in trainer + B=1 bit-equivalence test — smoke +4. Refactor `multi_horizon_heads_grn_bwd_batched` + GRN scratch buffers + 10 reducer launches — smoke +5. Refactor `variable_selection_bwd` + VSN scratch + 2 reducer launches — smoke +6. Refactor `attention_pool_bwd` + attn scratch + 1 reducer launch — smoke +7. Local A/B perf benchmark via `perception_overfit` release build wall time +8. Cluster A/B Argo run (acceptance gates #6, #7); merge to `ml-alpha-phase-a` once gates #1–#7 all pass + +Total kernel commits: 7 (steps 1–6 are 6 commits, step 7 is benchmarking/no commit, step 8 is cluster validation). + +## Open Questions + +None at spec time. Implementation plan resolves: + +- Exact memset_zeros ordering inside the captured CUDA Graph region (must come before any kernel that reads the scratch). +- Whether to emit a single `reduce_axis0` launch per scratch tensor or batch them into a multi-tensor sweep (probably one per tensor — the per-launch overhead is negligible and the code stays simple). +- Optimal `block_dim` for `reduce_axis0` itself (likely 256 threads, parameterised via launch config).