[ColoCarletti]

What this PR does

Moves the first slice of trace generation onto the GPU. Seven table fills now build their main trace directly in device memory and feed the LDE/commit from there — no host→device upload for those tables.

Tables on GPU: CPU, MEMW_R, MEMW_A, LOAD, STORE, LT, SHIFT

Core pieces:

• A device-resident LDE seam (*_keep_dev): a trace already on the GPU feeds the LDE device-to-device instead of being uploaded.
• The commit path (commit_main_trace) takes the device buffer directly when present, and frees it after commit to keep VRAM flat on large blocks.
• The LogUp aux build reads the same resident buffer.
• Transfer accounting that reports H2D bytes vs device-resident builds.

Measured (ethrex 10-tx, RTX 5090): main-trace H2D drops from ~4.5 GiB to
~166 MiB (~96% less). Proofs verify.

What's left for future PRs

• Device register/memory-model walk (trace-gen stage 3).
• Remaining table fills (MEMW general, MUL, DVRM, EQ, BYTEWISE, BRANCH, CPU32,
  BITWISE, DECODE, PAGE, precompiles).
• Precompile crypto on GPU.
• Remaining trace-gen stages (decode, op-collection, derive, bitwise-collect).

Total time is still FFT/proving-math dominated, so this part is correctness + plumbing groundwork; the bulk of the trace-gen speedup lands as the remaining stages move over in follow-ups.

[ColoCarletti]

/ai-review

[claude]

Review: GPU trace-gen (part 1)

Reviewed the full diff, focusing on kernel↔Rust correctness, the trace-builder integration, VRAM lifetime, and the LDE/commit seam. Overall this is careful, well-documented work and I found no correctness bugs. The seven fill kernels each match their Rust reference table-generator; I verified column indices and bit-slicing against cols and the set_dword_* layouts for CPU, MEMW_R, MEMW_A, LOAD, STORE, SHIFT (including the compute_aux/compute_shifted recomputation), and LT (DWordHHW/DWordHL decomposition, signed compare, out = lt ^ invert).

Notes (no blockers)

1. CPU-table kernel lacks a parity test — see inline comment on kernels/trace_cpu.cu. Six of the seven tables have a gpu_*_fill_matches_cpu test; the CPU table (the most intricate kernel) is the exception. Worth adding one for symmetry and regression safety since the GPU path is on by default.

2. CI validation of the GPU kernels — every parity test early-returns when backend().is_err(), and gpu_fill_tests is gated on feature = cuda. So unless a CI job builds the cuda feature on a GPU runner, none of the new kernels are exercised; with LAMBDA_VM_CPU_TRACE defaulting to the GPU path, a kernel regression would only surface as a proof-verification failure on a GPU box. Worth confirming a GPU CI job runs these.

Things I checked that are correct

• The u32::MAX / BLOCK_SIZE → u32::MAX guard relaxation in inverse.rs/logup.rs: grid.x max (2^31-1) covers total.div_ceil(256) for any total <= u32::MAX, and the flat index tops out at exactly u32::MAX with no overflow. The old bound was 256x too strict.
• MemwBuckets/RegRow refactor and classify_memw shared between the walk and count_table_lengths — routing order (Register→Aligned→General) preserved; build_reg_fallback cannot re-enter MEMW_R.
• RegRow::fill_soa reconstructing the old_ts high limb from timestamp>>32 is sound (routing enforces ts_hi == old_ts_hi), and the kernel only consumes the low limb anyway.
• Device-input LDE seam always runs the GPU LDE for a resident input (no threshold decline that would silently commit the zeroed host placeholder), matching the parallel change in logup_gpu.rs.
• Device buffers are freed via clear_main_input_dev after Round-1 Phase A, before the aux/DEEP/FRI VRAM peak.
• [u32; 8] shrink of MemwOperation value/old — all element domains (bytes, register halves) fit u32; timestamps stay u64.
• zeroed_fe_vec swap is guarded by the zeroed_fe_vec_matches_fe_zero test.

[ColoCarletti]

/bench

[ColoCarletti]

/bench-gpu

[github-actions]

GPU Benchmark (ABBA) — fd95330def vs main (14 pairs)

_{RTX 5090 · Vast.ai datacenter @ $0.6064814814814815/hr · prover/cuda · drift-free A/B/B/A}

=== ABBA paired result  (improvement: - = PR faster) ===
  pairs: 14   mean A (PR): 23.642s   mean B (base): 25.269s

  [parametric] paired-t   mean -6.42%   sd 1.87%   se 0.50%
               95% CI: [-7.50%, -5.34%]   (t df=13 = 2.16)
  [robust]     median -6.38%   Wilcoxon W+=0 W-=105  p(exact)=0.0001  (z=-3.26)

  --- server stability (this run; compare across servers) ---
  run-to-run jitter:    A CV 1.51%   B CV 1.50%        (lower = steadier)
  within-session drift: +0.29% over the run, 1st->2nd half +0.20%
    (jitter -> Tier-1 cached gate floor; drift -> whether the cached baseline can be trusted)

  VERDICT: REAL IMPROVEMENT - PR faster by ~6.42% (t-CI and Wilcoxon agree)

  raw pairs: /tmp/abba_run/pairs.csv

_{- = PR faster. Trust the verdict when paired-t and Wilcoxon agree.}

[MauroToscano]

⚠️ AI-generated review. Produced by an automated agent and posted from this account — not written by hand. Findings below are leads to verify, not confirmed conclusions.

This PR moves main-trace generation for seven tables onto the GPU, keeping the trace device-resident across the LDE seam and the aux build, and fixes a u32::MAX / BLOCK_SIZE overflow. The plumbing (LDE seam, resident-aux reuse, clear_main_input_dev lifetime) is coherent, but the host main_table is now a zeroed placeholder for those seven tables, which converts several previously-safe "decline and fall back" paths into silent wrong-proof paths, and the H2D headline number does not count what the new path uploads.

crypto/stark/src/prover.rs:818 — when try_expand_leaf_and_tree_row_major_keep_dev returns None, control falls through to the host GPU path and then the CPU path, both of which read trace.main_data_row_major() — the zeroed_fe_vec placeholder. The doc comment reasons only about the threshold decline; the error path commits zeros. The function ends in coset_lde_row_major_with_merkle_tree_keep_dev(...).ok()?, so any CUDA error inside the LDE (most plausibly alloc_zeros::<u64>(lde_size * total_cols) under VRAM pressure) yields a committed, transcript-absorbed all-zero main trace with no error raised. The symptom is a bus imbalance at verify time with nothing pointing at the GPU trace path. Fix: when trace.main_input_dev().is_some(), treat None as fatal (return Err(...)) rather than falling through, or guard the two host branches with a hard runtime check that main_input_dev is None.

prover/src/tables/gpu_trace.rs:31 — gpu_trace_disabled() gates only on LAMBDA_VM_CPU_TRACE, and the seven use_gpu predicates in trace_builder.rs only exclude StorageMode::Disk. Nothing excludes debug-checks, yet the residency mechanism these tables depend on is compiled out under it: prover.rs:2292 (#[cfg(all(feature = "cuda", not(feature = "debug-checks")))], which threads main_trace_dev) and the try_build_aux_resident_gpu block at lookup.rs:1116/1136. Under --features cuda,debug-checks the aux is therefore built from trace.columns_main() — zeros — and run_debug_checks reconstructs Round 1 and prints the bus-balance report against the same zeros. That is exactly the workflow documented at README:271 for diagnosing bus imbalances, and no CI job builds the combination (make test-prover-debug is CPU-only). One-line fix: gpu_trace_disabled() should also return true when cfg!(feature = "debug-checks"), documented next to the disk-spill rationale. The same applies to LAMBDA_VM_NO_GPU_LOGUP (logup_gpu.rs:445): a formerly safe A/B switch that now silently corrupts device-born tables — fold it into the same kill-switch, or make the non-resident aux paths error out when the main is a placeholder.

crypto/stark/src/prover.rs:2144 — main_chunks is planned from estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) against vram_budget, which detect_vram_budget_bytes (device.rs:254) fixes once at init as 80% of total device memory. The device-born main inputs are allocated during trace-gen and held until clear_main_input_dev() at prover.rs:2245, so all seven tables' num_rows * ncols * 8 are resident for the whole of Phase A and none of it enters estimates. The planner admits chunks as if that memory were free, and chunking cannot bound it because later tables already hold their buffers while chunk 0 runs. Subtract total resident main-input bytes from vram_budget before plan_table_chunks, and consider freeing per chunk (right after each chunk's roots are absorbed) instead of after the whole loop. This matters more given the exhaustion path is the silent zero-commit above.

crypto/stark/src/instruments.rs:294 and prover/src/lib.rs:976 — MAIN_H2D_BYTES is fed only from the host branch at prover.rs:851; the device branch records zero. Every device-born build still uploads: gpu_build_cpu_trace, build_interleaved_on and fill_memw_register_on all clone_htod the packed ops, uncounted. Per row: CPU 88 B vs 304 B, MEMW_A 96 vs 232, LOAD 56 vs 144, STORE 32 vs 128, SHIFT 24 vs 232, LT 32 vs 136, and MEMW_R — the largest table — 45 B (reg_addr 4 + ts 8 + value 8 + is_read 1 + old_value 8 + old_ts 8 + row_index 8) against a 10×8 = 80 B/row matrix, only ~44% less. So ~4.5 GiB → ~166 MiB (~96% less) is the delta in bytes uploaded by commit_main_trace, not in PCIe traffic; the real reduction across these seven tables is roughly half that. Add a counter for the packed uploads (each gpu_build_* knows its byte count), print both, and restate the claim as total H2D. The doc comment calling the counted bytes "the transfer the on-GPU trace generator removes" should say "replaces with a smaller upload".

prover/src/tables/gpu_trace.rs:61 — gpu_build_tables runs for chunk in chunks { tables.push(build_chunk(chunk)?) }, and every gpu_build_* ends in a blocking stream.synchronize(); the chunk_and_generate it replaces uses op_chunks.into_par_iter().map(generate). At production max_rows (MEMW_R = 2^20) a ~15M-row MEMW_R is ~15 chunks: previously 15 fills across rayon workers, now 15 serial pack→H2D→kernel→sync round-trips on one worker, each chunk's host-side pack marshalling 1M RegRows into six fresh Vecs (~45 MB) before it can start. Because these closures run inside the rayon::scope at trace_builder.rs:3536, blocking in synchronize() also starves the ~18 concurrent CPU table builds. Nothing forces the serialization — buffers are retained to end of Phase A either way, and next_stream() round-robins a 32-stream pool for exactly this. Use chunks.into_par_iter().map(build_chunk).collect::<Option<Vec<_>>>() (add Sync to the bound), or better, hoist synchronize() out of each gpu_build_* and sync once after all chunks are enqueued.

prover/tests/cuda_path_integration.rs:1 — that file exists to assert "every dispatch counter introduced by the cuda backend actually fired… catches silent CPU-fallback regressions (GPU path skipped while the proof still verifies)", and every prior dispatch got a sibling test. This is the largest new dispatch of the series and has none. Each of the seven builders returns Option and falls through to chunk_and_generate on failure, and the seam declines silently on .ok()?; if the kernels fail to load or LAMBDA_VM_CPU_TRACE leaks into the environment, the whole feature is dead while the GPU suite stays green because the proofs still verify. gpu_lde_calls() cannot cover it — gpu_lde.rs:526 (host) and :597 (dev) bump the same GPU_LDE_CALLS. Add a GPU_MAIN_DEV_LDE_CALLS next to the others in gpu_lde.rs, bumped only inside the dev seam, plus a gpu_device_trace_gen_fires_and_verifies test asserting it is > 0 on fib_iterative_1M.

Nits:

• crypto/math-cuda/src/trace_cpu.rs:481 — row_index: Vec<i64> = (0..n).collect() is built on the host and uploaded every chunk purely to be the identity (per the kernel's own comment): 8 B/row of pure-waste PCIe on the largest table, ~18% of the 80 B/row matrix it replaces. Use a null/identity sentinel (null pointer or n_acc == 0 flag) and use i directly.
• crypto/stark/src/gpu_lde.rs:584 — let _lde_size = n.saturating_mul(blowup_factor); is dead; drop the binding.
• prover/src/tables/trace_builder.rs:3196 and six more — the identical seven-line use_gpu disk-spill block is copy-pasted per closure; pull it into a helper (or fold storage_mode into gpu_trace_disabled()), which also makes the debug-checks fix a single edit.
• crypto/stark/src/lookup.rs:1108 — main_segment_cols = trace.columns_main() is materialized before the resident check, so on a device-born table it is a full transposed copy of a zeroed matrix (and faults in the whole calloc'd placeholder) while only being consumed for is_empty()/column count; hoist it below the resident branch.