# Fused moment-screen executor — design / scoping (2026-06-08, Codex-reviewed)

Phase 5a of the AOM/moment accuracy-then-throughput work. **Scoping doc for review
BEFORE any kernel code.** Revised after an adversarial Codex review (raw review:
`fused_moment_executor_design.codex_review_2026-06-08.md`) that corrected several
overclaims in the first draft — see "Codex review incorporated" at the bottom.

## TL;DR — direction confirmed, diagnosis must be measured first

The handoff's deferred "big task" was a *fused many-chain/fold/candidate IKPLS
**scoring** executor*. That targets the wrong layer: the PLS moment scorer is already
on CUDA and is not where large-p wall-clock goes. **Codex concurs with redirecting to
the upstream moment / data-movement layer, but flags that the first draft overclaimed
*which* sublayer dominates and mislabeled the cost as "host CPU GEMM."**

Corrected thesis: on large-p datasets the staged campaign spends its time on the
**host side of the per-candidate moment pipeline** — but the GEMMs themselves already
route to cuBLAS, so the cost is most plausibly **per-candidate data movement
(H2D/D2H), device allocations, and launch/sync overhead, plus genuinely host-side
Ridge eigensolves** — *not* CPU matmul. **This must be confirmed by a phase-timing
profile (5b-0) before any kernel is written.** 5b-0 is a hard gate, not a warm-up.

## 5b-0 MEASURED (2026-06-08, nsys + before/after)

Profiled one BERRY brix screen (n=1434, p=2101, compact, 71.3 s) with `nsys
--sample=cpu` + a correct-but-neutral first change, and the measurement **overturned
the draft hypothesis**:

- **#1 host hotspot: `n4m_householder_qr` = 37.1% of CPU samples** (`libn4m.so`) — the
  **Ridge GCV solve** (`ridge.cpp` / `sweep.cpp` `solve_square_qr`, single-threaded
  O(p³/n³) QR). Note: this runs even though `n_ridge_moment_*` counters are 0 — Ridge
  candidates are scored through a QR path the moment counters don't track.
- **#2: synchronous `cudaMemcpy` = 4.59 s** (98% of CUDA API time, 888 blocking copies,
  one 174 ms) — host-blocking H2D/D2H, ~6% of wall-clock, ~20% of samples land in the
  libcuda copy machinery.
- CPU OpenBLAS `blas_thread_server` ~6%; `cuda_dispatch::gemm` only **2.6%**;
  `cudaMalloc`+`cudaFree` = **14 ms total** (negligible).
- **First change (neutral, kept):** gave `cuda_dispatch::gemm` a thread-local reusable
  device scratch arena (kills per-call cudaMalloc/Free). Result: RMSEP **bit-identical**
  on BERRY/COLZA/LUCAS, speedup **1.00×** — because alloc/GEMM was never the cost. The
  change is correct and foundational for residency, just not the bottleneck.

**Revised priority (data-driven):** the moment-transform GEMM (draft's #1 target) is
minor. The real targets, in order: **(1) move the Ridge GCV solve off the CPU
householder-QR onto the GPU** (cuSOLVER SPD Cholesky/eigh — the ridge normal system is
SPD); **(2) cut the 4.59 s of synchronous cudaMemcpy** (device residency / pinned). The
moment-transform batching matters only after these.

## What is established vs hypothesized

**Established (direct evidence):**
- BERRY brix (n=1434, p=2101) took **1212 s with GPU0 at ~0% util**, one CPU core
  pinned. Across 51 datasets `corr(log fit_time, log n·p²)=0.86`; median fit
  **40.6 s (p≥1500) vs 3.9 s (p<800)** — 10×. → The slow phase is **not active CUDA
  kernels**, and cost grows steeply with size.
- The PLS moment scorer ran fully on device for BERRY
  (`n_screen_pls_moment_cuda_device_cv_fits=120` host 0; refit 20 host 0). → The
  scorer is not the bottleneck.

**Hypothesized (NOT yet proven — 5b-0 must apportion):**
- *Which* host phase dominates: base-moment build, per-candidate operator→moment
  transform, host Ridge eigensolves, or the data-movement/alloc/launch overhead
  around the cuBLAS calls. `corr=0.86` with `n·p²` is **confounded** — large p
  correlates with large n, chain count, refit load, and stage mix, and at BERRY scale
  `p³` and `n·p²` are same-order, so the correlation cannot separate transform from
  build from eigensolve.
- The BERRY `n_ridge_moment_*=0` reading does **not** prove "wide Ridge materialized";
  it could equally mean the Ridge moment route did not run. Needs the
  dual/materialized counters + active-head evidence.

## Corrected code facts (verified by Codex against source)

These replace the first draft's claims:

1. **Base moment build is NOT host-CPU GEMM in CUDA builds.** `compute_moments`
   (`moments.cpp:190-201,264-270`) builds `xtx/xty` via `linalg::gemm`, which in CUDA
   builds dispatches to `cuda_dispatch::gemm` (`common/linalg.hpp:173-178`) — cuBLAS
   with per-call device alloc + H2D + D2H (`cuda_dispatch.cpp:1439-1472`). Buffers are
   **host-owned**; the matmul is on-device but wrapped in per-call copies. So the
   suspect is **copy/alloc/launch overhead**, not CPU matmul.
2. **The dense O(p³) operator→moment transform is only one route.**
   `transform_moments_by_operator` does `opᵀ·xty` (O(p²q)) + `xtx·op` + `opᵀ·temp`
   (2×O(p³)) at `aom_sweep.cpp:674-687`, but the code tries **structured/banded**
   transforms first (`aom_sweep.cpp:1323-1344,1425-1468`) and disables dense for
   `p>min_train` except a tiny cap (`aom_sweep.cpp:2222-2225`). "Every candidate = two
   O(p³) GEMMs" is false in general.
3. **Wide-Ridge does NOT host-materialize via `should_materialize_cpu_wide_ridge` in
   CUDA builds** — it returns false immediately when `uses_cuda_linalg_build()`
   (`aom_sweep.cpp:555-560`). Materialized fallback can still occur via unsupported
   moment routes, but not from this predicate.
4. **Ridge GCV is host-only (confirmed) but amortizes lambdas.** `fit_ridge_from_moments`
   uses host vectors + host `solve_square_qr` (`sweep.cpp:840-868`, `ridge.cpp:104-120`);
   the eigen path is host (`sweep.cpp:956-958`, `svd.c:263-283`). But multiple λ share
   one host eigen prep per fold (`sweep.cpp:2094-2110`, `fit_ridge_from_eigen_path`
   `sweep.cpp:2157-2172`) — not QR/eigh per λ. **cuSOLVER is not linked** (only
   `CUDA::cudart CUDA::cublas`, `n4m_targets.cmake:144-148`).
5. **The PLS scorer is device-resident only *inside* the scorer.** It takes host
   `C/s` pointers and copies them in per tile (`cuda_dispatch.cpp:834-843`), runs
   cuBLAS strided-batched (`854-1072`), copies W/P/Q back (`1104-1113`). Feeding it
   device-resident moments is **not a drop-in** — the scorer would need an internal
   device-pointer entry variant.

## Current batching layout

The scorer `pls1_moment_components_many_batched_tiled` (`cuda_dispatch.cpp:784`)
batches **folds** into device tiles (one cuBLAS strided-batched call per component over
all folds of one candidate), `dcur_yy`/`dQ` device-resident, W/P/Q read back per tile.
It is fp64, called **once per (chain×candidate)** from `sweep.cpp:1636`, fed
host-built moments. **Host-looped today:** chain, candidate, and the entire
moment-production pipeline that precedes the scorer.

## Revised proposed design — kill per-candidate data movement, keep moments on device

Keep `libn4m` single-runtime. The intent is **not** "move CPU GEMM to GPU" (it's
already cuBLAS) but **"stop round-tripping moments through host buffers per
candidate"** and **batch the per-candidate launches** so the device stays busy.
Invariants preserved: fold-clean stateful heads, P-dependent detrend projector Q,
`MomentStats` raw+centered fields, fp64-for-scores, train-CV-only selection.

Candidate components (each gated on 5b-0 evidence):
1. **Device-resident base moments per fold-config** — compute `XtX/XtY/yy` once and
   keep on device, instead of host-owned buffers re-copied per candidate.
2. **Batched operator→moment transform on device, route-aware** — for the routes that
   actually run (dense / banded / structured), apply the operator transform to
   device-resident base moments in **batched** cuBLAS calls over candidates. Do **not**
   force dense p×p batches for large p where the code already uses banded/structured.
3. **Device-pointer scorer entry** — add an internal variant of the PLS scorer that
   accepts device-resident `C/s/yy`, eliminating the per-tile H2D of `hC/hs`.
4. **(Conditional) device Ridge GCV** — only if 5b-0 shows host Ridge is a material
   fraction; would add a cuSOLVER dependency (currently unlinked).
5. **Batch chain×candidate into the tile dimension** — so the executor is one call per
   stage-chunk, not one per candidate, eliminating the host launch/copy loop.

## Internal device-state contract (new — per Codex)

Any device-resident state must be defined before code:
- **Memory budget.** One p×p fp64 base moment at p=4200 (LUCAS) = **141 MB**; raw +
  centered + transformed + temp + operator matrices multiply that. Must fit the
  existing PLS tile budget (`choose_pls_batch_tile_jobs`, `cuda_dispatch.cpp:433-469`)
  — chain×candidate×fold batching **collides** with the current per-job p×p budget and
  must share one accounting.
- **No persistent cross-call cache.** Device state stays **strictly internal to one
  public call** (so "no ABI change" holds). Persisting base moments in `Context` or
  across calls changes lifetime/thread-safety/failure semantics even with symbols at
  ABI 1.22.0 (`n4m_version.h:20-22`) — out of scope unless explicitly designed.
- **fp64 + equivalence.** Keep fp64; assert batched-device transform == host transform
  within tol; a degenerate-component batch must preserve the existing per-job
  failure/recovery semantics (don't let one bad candidate poison the batch).
- **Fold-local state.** MSC/EMSC re-fit per fold; the detrend projector/inverse is
  P-dependent (`aom_sweep.cpp:904-928`) — must not be hoisted across folds/P.
- **Device selection.** CUDA state currently hardcodes device 0
  (`cuda_dispatch.cpp:50-56`); the executor must honor `CUDA_VISIBLE_DEVICES` / the
  chosen device, not assume 0.
- **cuSOLVER gating.** Only linked if 5b-3 is greenlit by profiling.

## Increments (5b-0 is a HARD GATE)

- **5b-0 — Phase-timing profile (mandatory, gates everything).** Native phase timers
  (or NVTX ranges) around `compute_moments*`, `transform_chain`/operator-transform,
  Ridge eigen/direct fits, PLS scoring, and the Python chunk/checkpoint overhead, run
  on **BERRY / COLZA / LUCAS** with checkpoints off, scale_x grid fixed, plus Nsight
  Systems CUDA kernel-vs-copy time and the dual/materialized counters. **Deliverable: a
  reproducible artifact apportioning the large-p host time.** No kernel work until this
  names the dominant sublayer.
- **5b-1..5b-5** (device-resident base moments → batched route-aware transform →
  device-pointer scorer entry → conditional cuSOLVER Ridge → chain×candidate batching):
  **scope and order to be set by 5b-0's measured split**, biggest measured win first.
  Each behind the full green gate (`n4m_tests` + `n4m_internal_tests` CUDA &
  dev-release, targeted CUDA pytest, catalog `split/--check-references/--strict-abi` +
  `reconcile_abi --check`, `git diff --check`), with host-equivalence asserted and ABI
  held at 1.22.0.

## Open questions — resolved (Codex + author agree)

1. **Scope:** target the upstream moment / data-movement layer, not the PLS scorer —
   **but narrow only after 5b-0**; do not assume dense-transform+Ridge dominates. ✓
2. **cuSOLVER:** **not yet** — keep Ridge host-side until 5b-0 shows it's material. ✓
3. **Precision:** **fp64 first**; fp32-screen/fp64-verify is a later experiment needing
   a recall/top-k-miss audit (the accuracy track already showed selection is variance-
   sensitive). ✓
4. **5b-0 first:** **mandatory hard gate** with native timers + route/materialized
   counters + CUDA copy/kernel time. ✓

## Codex review incorporated (2026-06-08)

Changed vs first draft: (a) the cost is reframed from "host CPU GEMM" to "per-candidate
data-movement/alloc/launch overhead + host eigensolves" — the GEMMs already route to
cuBLAS; (b) the `should_materialize_cpu_wide_ridge` and "every candidate = two O(p³)
GEMMs" and "Ridge QR/eigh per λ" claims are corrected/refuted with line evidence;
(c) the empirical section now separates *established* from *hypothesized* and flags the
`n·p²` confound and the counter over-reading; (d) 5b-0 is promoted to a hard gate that
must apportion the host time before any kernel; (e) a device-state contract (memory at
p=4200, tiling collision, fold-local state, device selection, no cross-call cache) is
added; (f) feeding the scorer device moments is acknowledged to need a new device-
pointer entry, not a drop-in. Raw review preserved alongside.

## Why this is the right investment

The accuracy track established the *accuracy* ceiling is model expressiveness, not
selection (`RECALL_FIX_RESULTS.md`). The remaining engine work is *throughput*, and the
honest lever is the host-bound large-p moment pipeline (10× slower than small-p) — but
**what exactly to fuse is a measurement question, answered by 5b-0, not assumed.**