feat(stark): full GPU trace residency — skip the round-1 D2H#799
Draft
MauroToscano wants to merge 7 commits into
Draft
feat(stark): full GPU trace residency — skip the round-1 D2H#799MauroToscano wants to merge 7 commits into
MauroToscano wants to merge 7 commits into
Conversation
gather_rows_base / gather_rows_ext3: read full rows (all columns) from a device-resident LDE buffer at given row indices, returning them row-major — the building block for gathering R4 query-opening values from the resident LDE instead of the host trace (toward eliminating the round-1 trace D2H). Wrappers gather_rows_base_on_device / gather_rows_ext3_on_device + a standalone parity test vs direct buffer indexing.
…oss-checked) Read each query's row-pair straight off the resident device LDE in open_deep_composition_poly instead of indexing the host trace, keeping the round-1 host D2H this stage. For each query, gather LDE rows reverse_index(2c) and reverse_index(2c+1) via the Stage-1 kernels, lift u64 -> FieldElement through TypeId-gated base_u64_to_field / ext3_u64_to_field (inverse of the *_slice_to_u64 seam), and feed them into a new open_polys_from_values. Every device row is cross-checked (release-active assert_eq!) against gather_main_row / gather_aux_row before it is trusted in the proof. Value gathers are gated on the corresponding device-tree proofs so the two device arms stay aligned and no rows are gathered for a non-resident tree. Preprocessed and host-fallback arms are untouched. De-risks Stage 3 (dropping the host D2H). Non-cuda build unaffected (all new code cuda-gated).
Add GPU_OPENING_GATHER_CALLS, incremented once per main/aux trace whose R4 query rows are read off the device LDE in open_deep_composition_poly, and a gpu_opening_gather_fires_and_verifies e2e test that proves fib_iterative_1M, asserts the counter fired (no silent revert to the host gather), and verifies the proof. Mirrors the existing GPU dispatch counters (e.g. GPU_COMPOSITION_CALLS) so a future regression that quietly drops the residency path fails loudly.
Scaffolding for the device-only round-1 residency path (D2H skip lands next). Adds a per-table `host_trace_empty` flag on LDETraceTable (default false, so behavior is identical) and hard-abort `assert!(!host_trace_empty)` guards on every host-LDE-trace fallback that would otherwise index the (future) empty buffers: R2 composition CPU path (boundary + transition), R3 barycentric main/aux host fallbacks, R4 DEEP host loop, and the R4 opening host-fallback arms. The Stage-2 device-gather cross-checks are gated on `!host_trace_empty` so they are skipped when there is no host copy to check against. All guards are inert until step 2 wires the `device_only` gate.
…ce D2H Wire the full-residency win: when a table's `device_only` gate holds, keep its round-1 LDE device-resident and skip the host D2H (the big transfer/alloc the prover paid every round-1). Threads a `retain_host_lde` bool through the math-cuda keep wrappers (`coset_lde_row_major_inner` returns an empty host Vec when false) and the three stark keep entries. `device_only_gate` (gpu_lde.rs) is the strict AND of every precondition that guarantees the R2 composition, R3 barycentric, R4 DEEP and R4 opening GPU paths all fire and read the device LDE: Goldilocks+ext3, lde_size>=gpu_lde_threshold, n>=gpu_bary_threshold, power-of-two, !preprocessed, contiguous offsets, uniform zerofier (all end_exemptions==0), and neither composition nor device-only disabled; forced off under debug-checks so reconstruct_round1 keeps a host trace. `Prover::device_only_for` derives it from the AIR so the main (Phase A) and aux commit closures compute the identical value and skip both host buffers together. build_round1 infers `host_trace_empty` from the ACTUAL buffer state (empty host buffer where data was expected) — a safety property: if the gate held but the GPU keep fell back to CPU, the buffer is populated and the proof runs on the host trace as normal. num_rows is recovered from the device handle's lde_size when the main buffer is empty. Adds GPU_DEVICE_ONLY_CALLS + a LAMBDA_VM_DISABLE_DEVICE_ONLY A/B toggle + a gpu_device_only_residency_fires_and_verifies e2e test. Step-1 host_trace_empty guards are now live.
…ncy A/B) Add CONTINUATIONS=1 + EPOCH_SIZE_LOG2 (default 20) + TX_COUNT (default 5) env knobs so the two-ref ABBA harness can prove with --continuations on both sides. Monolithic proving OOMs on large workloads (e.g. ethrex-20tx) and can't surface a per-epoch GPU-residency change at a realistic trace size; continuation mode proves epoch-by-epoch with flat peak memory. Backward compatible (default off = monolithic as before). Also fixes the missing-fixture generator to honor TX_COUNT instead of hardcoding 5. Usage for the device-only residency win vs main: CONTINUATIONS=1 EPOCH_SIZE_LOG2=20 TX_COUNT=20 \ BENCH_FEATURES=jemalloc-stats,prover/cuda \ scripts/bench_abba.sh HEAD origin/main 5
…nce + gather fallback - device_only_for: require has_aux_trace() + a non-empty constraint set (R2 composition needs a device aux handle and a uniform zerofier with >=1 group). - build_round1: recover num_rows from the device handle whenever host_trace_empty, covering the aux-only (num_main_cols==0) case a main_empty-only guard missed. - open_deep_composition_poly: a device row-gather Err falls back to the host gather unless the trace is device-only (was .expect() → panic on a transient error). - merkle_root_parity: pass the retain_host_lde arg (test target no longer compiled).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR 2 of 3, stacked on #798. Review this diff against
feat/gpu-constraint-eval.What
Keep the LDE trace on the GPU across rounds 2–4 and skip the round-1 device→host copy (the big transfer the prover paid every proof). Round 4 openings gather query rows straight off the device LDE; rounds 2/3/4 read the resident trace.
Safety (the important part)
A per-table
device_onlygate skips the D2H only when every downstream round is guaranteed to take its GPU path. Ahost_trace_emptyflag then makes every host-read fallback hard-abort — so a mis-gate or an unexpected GPU fallback fails loudly, and can never produce a wrong proof (a mis-gathered value is also Merkle-bound, so it fails verification). An adversarial soundness audit found no wrong-proof path.Win
−12.5% incremental (device-only A/B toggle) / −15% cumulative vs
mainon the original box; −12.75% on a 384-core box (faster CPU baseline → smaller GPU share). All confirmed, all-negative pairs.Validation
cuda_path_integration7/7 incl. a residency-fires-and-verifies counter test; gather-kernel parity; e2e prove→verify. Multi-agent review (soundness + kernel specialists + high/max) — findings addressed, no wrong-proof bug.