GPU trace-gen (part 1): device-resident fills for 7 tables

[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.}