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

- 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

- 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)

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

// 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).