perf(prover): halve GPU continuation proving time - #863
Conversation
- run_profile.sh: nsys export needs --force-overwrite (nsys stats already materializes the sqlite); tolerate runs that produce no timeline JSON - flamegraphs.sh: fixed off-CPU capture window sized from the on-CPU run (SIGINT through sudo is unreliable and produced 0-byte captures); find offcputime-bpfcc in /usr/sbin (Debian) - bench_mode.sh: set the CPU governor via sysfs when cpupower is absent - setup_machine.sh: Debian-aware perf install (linux-perf); extract libnvToolsExt from the cuda-nvtx-12-8 deb into ~/nvtx (CUDA >= 12.9 removed NVTX v2 from the toolkit) with LAMBDA_VM_NVTX_LIB override - docs: benchmark/profiling examples use ethrex 5tx/10tx fixtures only (team convention: never fibonacci); plan status updated
Adds the pieces needed to use the tooling without reading the scripts: column-by-column semantics for phase_table.md and phase_busy.md (including NVML gpu% vs nsys busy% and launch-site attribution), a reference table of every script with its flags, the environment variables the tooling understands plus the pre-existing prover knobs for A/B experiments, and a troubleshooting section (missing NVTX ranges, silent CPU fallback, empty off-CPU captures, jitter, concurrent-thread span nesting).
…ee cache The optimization half of the original campaign commit, without its profiling layer (this branch keeps gpu-profiling-tooling's toolkit as the only instrumentation): - async_dtoh_via/PendingD2H: big D2H copies go through per-worker pinned slabs via raw cuMemcpyDtoHAsync + a reusable completion event, instead of cudarc's memcpy_dtoh whose pageable path blocks the calling thread for all prior stream work (host DtoH blocking 12.4s -> 6.7s on the original ethrex A/B). - GpuLdeBase/GpuLdeExt3 carry a 'ready' event; consumers wait device-side (cuStreamWaitEvent) instead of producers host-synchronizing. - Events are pre-created at backend init plus a reusable pool: a mid-prove cuEventCreate convoys the driver lock (~30ms/call measured under load). - Precomputed-column Merkle trees are cached process-wide keyed by their commitment root, so preprocessed tables (DECODE/BITWISE/range) stop rebuilding identical trees on every prove; only the multiplicity columns are recommitted.
Producer thread executes and builds epoch i+1's traces while epoch i proves; K epoch provers (LAMBDA_VM_EPOCH_CONCURRENCY, default 3) consume prepared epochs concurrently — epoch proofs are mutually independent (label-domain-separated transcripts), results re-ordered by index so proof bytes match the sequential schedule. The DECODE commitment is computed once per continuation prove instead of per epoch. Same as the original campaign commit minus its epoch-timeline instrumentation (this branch keeps the profiling toolkit's spans as the only instrumentation; they are re-homed onto this pipelined flow at the end of the series).
…e slots The constraint interp/composition kernels evaluated every IR node as ext3 and kept one global-memory scratch slot per node, so scratch size and traffic scaled with program length (KECCAK_RND/ECSM/ECDAS at full thread count needed 26-39 GB, failing the alloc and silently falling back to CPU via result.ok()). Lowering (constraint_ir/device.rs) now assigns dim-split slots: - Base-dim nodes compute in the base field (1 mul vs 9 for ext3) and live in u64 slots (8B vs 24B); mixed base*ext ops use mul_base / componentwise shortcuts that are bit-identical to the full ext op on the embedded operand (SUB components keep the literal sub(0, x) form, which is NOT bitwise neg on non-canonical limbs). - Slots are liveness-reused (linear scan, freed at last use, roots pinned), so per-thread scratch is the max-live-set, not the node count: 8-35x smaller across the 26 tables (CPU 14.4KB -> 1.3KB, ECDAS 596KB -> 17KB per thread). Scratch allocs drop the memset. - Row-invariant leaves (constants, RAP challenges, alpha powers, table offset) are propagated into operand encodings (kind<<29|payload) and never touch scratch; they only materialize when a root needs them. The CPU walker eval_device_program mirrors the new walk and stays the pre-GPU parity oracle; the 26-table differential vs the production folder and the on-GPU parity tests (synthetic + all real programs) pass bit-for-bit. ir_stats_dump (ignored) prints per-table node/slot stats to size scratch when tuning. Measured on RTX 5090 (nsys, ethrex): constraint_composition_kernel 814ms -> 267ms (-67%) over the same 29 launches; ethrex 10tx continuations ABBA 15.16s -> 14.77s.
Preprocessed tables (DECODE/BITWISE: precomputed + multiplicity column split) skipped the fused GPU commit entirely — commit_main_trace only tried the GPU when precomputed.is_none() — so they paid the CPU row-major LDE plus two CPU subset Merkle trees (~2.2s thread-time of R1 'Main commit Merkle CPU' on ethrex). - keccak256_leaves_base_row_major_row_pair_range: column-range variant of the row-pair leaf kernel, byte-identical to the CPU commit_rows_bit_reversed_subset layout. - coset_lde_row_major_split_trees: one row-major GPU LDE of all columns plus the two subset trees built on device; both node buffers download to host and rebuild full host trees via from_precomputed_nodes, so the preprocessed opening path, the process-wide precomputed-tree cache and disk-spill work unchanged. The shared expansion stage is factored into expand_row_major_on_stream (same code path as the existing fused commit). - The table now gets a GpuLdeBase handle (column-major LDE + trace snapshot, no device tree), so its rounds 2-4 (composition, DEEP, barycentric) run on GPU too. Preprocessed openings short-circuit to the host trees via is_preprocessed, as before. - REGISTER stays on CPU (LDE below the dispatch threshold). Parity: split_tree_tests pins roots and opening paths against the CPU subset commits on device; cross-binary verification of full ethrex bundles passes both ways. Measured on RTX 5090: ethrex 10tx continuations interleaved 3-way 14.77s -> 14.23s (cumulative -6.1% vs the pre-kernel baseline).
prove_global consumes only execution artifacts — the per-epoch cell boundaries built by the producer, the ELF and the genesis pages — never an epoch proof, yet it ran serially after every epoch prove finished (~0.9s of pure tail on ethrex 10tx). The producer now publishes each epoch's boundary (an Arc share of the one already flowing to the epoch provers — no data copy) on a dedicated channel, in epoch order. A scoped thread drains that channel until the producer hangs up (last epoch prepared) and proves the global memory argument while the tail epochs are still proving. On an epoch failure first_err still wins and the global result is discarded; proof bytes and bundle content are unchanged — only the schedule moves. The epoch timeline confirms the tail is gone: the global prove runs fully inside the window of the last three in-flight epoch proves. Measured on RTX 5090: ethrex 10tx continuations ABBA 14.16s -> 13.66s (-3.5%); cross-binary verification passes both ways. Day cumulative across the three optimizations: -9.4% (15.16s -> 13.66s).
Every epoch's trace build re-parsed the ELF and regenerated the pristine DECODE trace (~1M rows) inside the serial producer chain, plus moved a ~900K-entry pc->row map by value per epoch. DecodeArtifacts (instruction map + pristine DECODE trace + pc->row index) is a pure function of the ELF: prove_continuation builds it once and every epoch's build clones the pristine trace (a memcpy) and fills its own multiplicities; build_traces now borrows the pc->row map. The monolithic entry point delegates and is unchanged. Net work removal with identical trace bytes (cross-binary verification passes). Wall-neutral within noise on a 32-core box; groundwork for pipelining the epoch trace build out of the producer chain, where parallel builders would otherwise each redo the ELF parse.
The continuation producer built every epoch's full trace tables inline, so the serial chain feeding the provers was execute + collect + BUILD per epoch (~95% of it table generation) — 7.2s of a ~18s wall on a 32-core box, with the last epochs' proves gated on it. The epoch trace build is now split at its real sequential boundary: - Traces::collect_epoch (Phases 1-2): op collection over the advancing memory image — stays on the producer, in epoch order. - Traces::build_from_collected (Phases 3-5): table generation — pure epoch-local work, runs on a small builder pool (LAMBDA_VM_TRACE_BUILDERS, default 2) between the producer and the epoch provers, bounded channels capping peak memory. The cross-epoch chain no longer touches traces: the boundary derives from CollectedEpoch::touched_memory_cells (same function, same immutable memory_state as the build) and the next epoch's register init from register::fini_from_final_state — a trace-free mirror of the REGISTER FINI column, pinned by fini_from_final_state_matches_trace. PAGE tables are the build's only image consumers and continuation mode skips them, so builders need no image snapshot. Measured on a 32-core RTX 5090 box (ethrex 10tx continuations): the producer chain drops 7.2s -> 2.9s and the first three proves start ~1s earlier, but the wall ties (~18s) — the box is bound by total CPU work, which this change conserves (proves and the global dilate to absorb the freed schedule). A K/builders sweep confirms K=3/B=2 stays optimal. Expected to pay on wider boxes where idle cores can absorb the parallelism; groundwork for cutting per-epoch CPU work (AIR/capture caching), which is the binding constraint on narrow boxes.
Constructing an AirWithBuses runs every constraint body through a MetaBuilder, and the first constraint_program() runs them again for the IR capture — for ECDAS/ECSM/KECCAK_RND (16-25K IR nodes) that dominates AIR construction (0.78s per VmAirs::new on ethrex). Continuation epochs rebuild the full AIR set per epoch and shard tables build one instance per shard, so the same walks re-ran dozens of times per prove. build_air now keeps a process-wide prototype cache keyed by (table name, proof options): the prototype is built and pre-captured once, and every later request clones it — Clone on AirWithBuses copies the derived meta, LogUp layout and the captured IR inside the OnceLock, never re-running the bodies. PAGE stays correct because its page base is part of its name. with_name/with_preprocessed apply to the caller's clone; the cached prototype stays pristine. Wall-neutral within noise on the 32-core box (the removed work is a few core-seconds against a ~580 core-second prove); cross-binary verification passes both ways. Also cuts AIR construction out of the monolithic path and the test suites.
…flow The toolkit's continuation instrumentation assumed the sequential epoch loop. With the producer/builder/prover pipeline the stages run on different threads, so the spans move to where the work actually happens: - prove_continuation_total root span + timeline reset at entry, drained at the end exactly like the monolithic path (stdout tree + LAMBDA_VM_TIMELINE_JSON for phase_table.py). - epoch_execute / epoch_collect on the producer, epoch_trace_build on the builder pool, epoch_prove on the prove workers — each prove/build/ collect also opens an NVTX range with per-epoch identity (epoch_*[i=N]) for Nsight timelines. - Spans close BEFORE blocking channel sends, so backpressure waits are never booked as work. - prove_global span on the overlapped global-prove thread.
|
/bench-gpu |
GPU Benchmark (ABBA) —
|
Domain and LdeTwiddles are now shared across epochs and concurrent epoch provers via a process-wide cache keyed by (field, trace_length, blowup, coset_offset). The OOD barycentric constants, FRI inverse twiddles, and the d=2 decomposition inverses hang off them as lazy per-domain values instead of being rebuilt (each an LDE-size-order batch inversion or clone) per table per epoch.
Each boundary constraint paid its own LDE-size batch inversion even when sharing the step with its neighbours, and the vectors are identical for every table and epoch on the same domain. The inverted vector now lives in the shared domain, keyed by step, and constraints hold an Arc to it.
Upload each distinct column once (GpuBaseVec, cached keyed by its host Arc — storing the Arc pins the allocation so the key can never alias) and D2D-copy into each dispatch's flat buffer, instead of re-uploading tens of MB per table per epoch over PCIe.
The composition evaluations stay resident after the fused kernel; a pointwise kernel decomposes them into the H0/H1 slabs, the batched slab LDE extends both halves with no H2D, and the parts handle feeds R4 DEEP. One drain of the final evaluations (still read by the commit tree and the query openings) replaces four codeword-sized PCIe trips per table per epoch. Falls back to downloading H and running the host decompose on any device failure.
The fully-resident DEEP arm keeps its output on device, bit-reverses it into FRI order with a permutation kernel, and hands the buffer to the FRI fold state as its working codeword — removing the download / CPU-bit-reverse / re-upload round trip. The commit loop is shared between the host and device entries and restores the transcript on any mid-loop failure so the CPU path reruns cleanly.
The shared domain caches ran the parallel batch inversion inside their OnceLock initializers. A rayon worker that starts such an initialization farms chunks to the pool while sibling workers block on the same cell; with every worker parked the chunks never run and the prove deadlocks (observed as a full-process futex stall). Initializers now use the sequential inversion, and domain construction pre-fills every lazy cell from the setup thread so pool workers never run — or wait on — an initializer mid-prove.
ColoCarletti
changed the title
There was a problem hiding this comment.
⚠️ 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 batch adds an overlapped epoch pipeline for continuations, CUDA LDE/commit work, and a profiling toolkit. The test shape is right — the dim-split slot-file interpreter is differentially tested against the CPU model, the fused preprocessed-commit path has a CPU parity check, and fini_from_final_state pins the register fini to the trace — and the ABBA numbers back the perf claim. Two blocking issues, plus cleanup.
prover/src/continuation.rs:1366 — the new epoch pipeline deadlocks instead of returning Err on any downstream failure. rx and build_rx are Mutex<Receiver> locals declared outside std::thread::scope (1130-1131), so they stay alive for the whole scope and a send can never observe disconnection; tx.send(...).is_err() (1237) and build_tx.send(Ok(job)).is_err() (1366) are therefore dead branches, contradicting their own "the downstream side hung up, stop quietly" comments.
The chain: a worker's prove_epoch Err (1174-1177, reachable from Prover::multi_prove) or a builder error forwarded as tx.send(Err(e)) (1241-1244) sets first_err; every worker then exits at 1148-1150 and every builder at 1196-1198 after at most one more message. Nothing drains build_rx after that, and the producer never consults first_err — its only exits are is_final and the dead send check — so once the ~3 jobs of drain slack are consumed it blocks forever in build_tx.send, producer.join() (1442) never returns, the scope cannot exit, and the overlapped global thread stays wedged in boundary_rx.recv() (1399) because boundary_tx was moved into the stuck producer. Net: a failure with ~3+ epochs still pending turns a clean Err into a permanent hang of the main continuation entry point.
Scope limits: producer-originated errors (e.g. HaltInNonFinalEpoch from Traces::collect_epoch) shut down fine because the producer drops build_tx; a failure on or near the final epoch also returns normally; a single panicking prover with workers >= 2 still completes.
Fix: check first_err at the top of the producer loop and return early (draining in the workers alone unblocks the scope but makes the producer execute the whole remaining program first), and have workers/builders keep draining to close rather than returning on the error path. This path has zero coverage — the only Err-expecting continuation test returns before the pipeline is built, and trace_builder_tests.rs:1009 calls Traces::from_image_and_logs directly — so a negative test is needed: cfg(test) fault injection in prove_epoch or build_from_collected, enough epochs to exceed the slack, asserting prove_continuation returns Err, wrapped in a timeout so a regression fails instead of wedging CI.
crypto/math-cuda/src/device.rs:632 — PendingD2H holds the staging slot's MutexGuard but has no Drop, so every ? between enqueue and wait releases the slot with cuMemcpyDtoHAsync still writing into the pinned slab. Three sites do exactly that: coset_lde_row_major_inner (pending at 592-602, then launch_row_to_col_major(...)? at 606 which allocates, plus take_event()?/record(..)? at 609-610, wait only at 616), coset_lde_row_major_split_trees (759 → 764 → 770), and coset_lde_batch_base_into_with_merkle_tree_inner (1541, held across d2h_bytes_via_pinned_hashes(...)? at 1543, whose nested ensure_capacity can fail on pinned-host OOM).
Slots are per-rayon-worker (529-536), so the next user is usually the same worker on a later call — on a different stream, since next_stream round-robins, so nothing orders them — and slot 0 is shared by every non-rayon thread. That next user either grows the slot (ensure_capacity cuMemFreeHosts the old pointer at 63-69, and the in-flight DMA writes into a freed mapping) or repacks the slab for its own H2D (as_mut_slice + copy_from_slice, lde.rs:1170-1180) and the stale DMA silently corrupts that upload. These Errs are meant to be survivable — gpu_lde.rs falls back to CPU on Err at 467-471, 1042, 1114, 1656, and the crate ships a test-cuda-faults injector — so a recoverable OOM leaves the process running with a live DMA into an unowned slab. ? does not poison the mutex; only the panic path is accidentally safe. Fix: impl Drop for PendingD2H doing a best-effort self.staging.sync_event() (ignoring the error) before the guard releases, so the invariant holds by construction. That also covers the lower-risk holds in deep.rs:249, fri.rs:223, logup.rs:205 and constraint_interp.rs:140, where the intervening ? happens to be a sync/memcpy_dtoh that drains the copy on success.
prover/src/continuation.rs:1120 / :1433 — LAMBDA_VM_EPOCH_CONCURRENCY defaults to 3 at 1120, but the sizing comment at 1433 says "the default is a conservative 2". Both lines are new in this PR, so one is stale on arrival, and it is not only a comment: main proves epochs strictly sequentially, while the new pipeline holds up to ~7 epoch-sized working sets concurrently per the comment at 1103-1104 (1 collected + 2 building + 1 queued + 3 proving), on CPU-only builds too. Nothing sizes this from available memory — unwrap_or(3) is unconditional — and the per-prove VRAM admission control does not coordinate across concurrent proves: crypto/stark/src/prover.rs:2548-2552 hands each prove backend().vram_budget_bytes(), which is 80% of total device memory (crypto/math-cuda/src/device.rs:252-275), so all three plan plan_table_chunks against the whole card. 24 GB cards look safe (the measured ~15 GB peak is already the 3-worker config; the standalone ~9 GB per-epoch peak does not stack on top of it), but ≤16 GB cards, larger EPOCH_SIZE_LOG2, and CPU host RAM are exposed. Pick a number and justify it in one place — default 2 with 3 opt-in, or derive it from the detected budget — and delete whichever comment loses.
scripts/profiling/README.md:149 — four doc sites promise math-cuda entry-point NVTX ranges that do not exist. Nothing under crypto/math-cuda/src/ references nvtx except lib.rs:18 and nvtx.rs itself; the ~40 crate::nvtx::Range::fmt(...) guards present at a48f5d4 (e.g. lde.rs lde_tree_base[n=.. m=.. bf=..], barycentric.rs bary_base[..]) were dropped when 530f423 rewrote those files. They are still claimed at README.md:149 ("every math-cuda public entry point pushes a range with its shape"), crypto/math-cuda/Cargo.toml:40, crypto/stark/Cargo.toml:59, and crypto/math-cuda/src/nvtx.rs:2. The toolkit is not broken — instruments spans mirrored via instruments.rs:89 and the LAMBDA_VM_NSYS_CAPTURE_SPAN gating at 92/112 both work, so the README's "verify ranges appear" check passes — but nsys_phase_busy.py:318's "## math-cuda entry points (innermost NVTX range)" table computes fine_key = base_name(chain[-1]), which now resolves to the innermost instruments span, so that section is mislabeled and mostly duplicates the coarse table above it. Either restore the ranges behind the nvtx feature (zero-cost when disabled) or fix all four doc sites and drop that section from the report.
Nits:
- ~59 statements across
crypto/math-cudaare wrapped in braces that bind nothing and scope no borrow (let out = { stream.clone_dtoh(&out_dev)? };,{ stream.synchronize()?; }, multi-linelet (a, b, c) = { ( ..?, ..?, ..? ) };) where main has the plain form (barycentric.rs:48,:75-77); densest in barycentric.rs (22), merkle.rs (10), lde.rs (~9 inert plus 5 real borrow scopes), logup.rs (8), inverse.rs (5), ntt.rs and fri.rs (3), deep.rs (2), constraint_interp.rs (1) — semantically identical, but they read as if a guard were deleted, andinverse.rs:71,logup.rs:202,deep.rs:246,deep.rs:399still say "the labelled sync" with no label present; unwrap them (keeping lde.rs's five) or fill the holes with thenvtx::Rangeguards this PR already ships. scripts/profiling/__pycache__/{nsys_phase_busy,nvml_sampler,phase_table,timeline_to_perfetto}.cpython-313.pycare checked in (~35 KB, first tracked.pycfiles in the repo) with no__pycache__/*.pycrule in any.gitignore— timestamp-invalidated so no stale-bytecode risk, but they persist in history and show as modified for any 3.13 dev importing these modules;git rm -r --cached scripts/profiling/__pycache__plus a root.gitignoreentry.
…ts senders The prove/build channel receivers live in the outer scope, so a worker that returned on error left the bounded senders parked in send() with no consumer — any mid-run proving error hung prove_continuation forever instead of surfacing. Workers now drain-and-discard until the channels disconnect, the producer stops executing epochs once an error is recorded, and the global-prove thread skips its (whole-prove-sized) run when the bundle can no longer be assembled.
- PendingD2H now synchronizes on drop: an error between enqueue and wait no longer releases the pinned slab to reuse/free while the DMA is in flight. - domain_and_twiddles re-checks the cache under the insert lock so a build race can't pin a duplicate instance's columns in the pointer-keyed device caches. - Hard-assert b_z_inv column length at the D2D copy (a short column left uninitialized VRAM in the kernel's window), mirror the batched-LDE input asserts in the split-trees entry, gate mismatched FRI twiddles to the CPU path, and pin the ext3 tower in the shared FRI drive. - Refresh the event-tracking safety note to the wait_ready_on contract.
A builder-injected fault (keyed by a magic private input, so it is stateless and inert for every real caller and for concurrent tests) fails epoch 3 of a ~9-epoch prove — enough pending work past the bounded channels' slack that a shutdown regression wedges instead of returning. The test runs the prove under a timeout so that regression fails CI rather than hanging it.
- The per-entry-point NVTX shape ranges were dropped when the math-cuda pipelines were rewritten; four doc sites still promised them and the nsys report mislabeled its innermost-range table. Align them with what the nvtx feature actually emits (mirrored instruments spans). - Untrack scripts/profiling/__pycache__ and ignore Python bytecode. - Remove ~86 brace wrappers in math-cuda left inert by the async-DMA refactor (kept the ones that scope real borrows) and reword three comments that referenced a deleted sync label.
|
/bench-gpu |
|
/bench-gpu |
The per-epoch report filtered `epoch[i=N]` ranges, but the pipelined continuation flow emits `epoch_collect[i=N]`, `epoch_trace_build[i=N]` and `epoch_prove[i=N]` — the section silently never printed. Group by stage instead (prove first: a gap between prove instances now means the pipeline stopped feeding the prover, not missing overlap), and update the README names to match.
set_staging_size_hints and its plumbing (per-slot Arc hint, ensure_capacity branch) were dead since the pre-sizing experiment was measured as a regression on the 5090 — remove them so it can't be silently re-enabled. Update comments that drifted from the code (preprocessed tables keep a device handle, boundary zerofiers are D2D from the resident cache, register chaining binds fini_from_final_state, the continuations capture span is epoch_prove), point the profiling examples at the fixture generator (the .bin files are not checked in), gitignore reports/, and untrack the working notes.
…dump An Err between enqueueing a DMA that reads a pinned-staging slab and its sync used to release the slot mutex with the copy in flight (the next locker can realloc the slab mid-DMA): async_dtoh_via now drains the stream when the event record fails, and the batched-LDE upload window holds a drain-on-error guard. try_fri_commit_gpu gains the inv_twiddles length gate its from_dev sibling had (degrade to CPU instead of panicking on a wiring bug), with a debug_assert on the CPU path. ir_stats_dump was never declared in tests/mod.rs (its documented command matched nothing), and capture_env.sh reported every tree as dirty.
3 participants
Result
ethrex 10tx
--continuationson RTX 5090 + 16 cores, interleaved A/B/B/A (12 pairs): main 22.58s → 11.05s = 2.04x (−51.1%, sd 0.88%). Proofs cross-verify in both directions against main's verifier.What changed
Validation
Per-change interleaved ABBA + bidirectional cross-verify + full stark cuda suite (212/212) at every step; clippy clean; 20-prove stress run for the deadlock fix; GPU↔CPU parity bit-exact on the constraint kernels.
Notes
--epoch-size-log2 22+LAMBDA_VM_EPOCH_CONCURRENCY=1gives a further −18% on 32GB/16-core boxes (~27GB RAM per prove) — left at defaults, documented for deployment tuning.