Skip to content

perf(prover): device-only preprocessed tables and GPU commits for mid-size tables - #888

Open
ColoCarletti wants to merge 54 commits into
mainfrom
gpu-opt-round4-residency
Open

perf(prover): device-only preprocessed tables and GPU commits for mid-size tables#888
ColoCarletti wants to merge 54 commits into
mainfrom
gpu-opt-round4-residency

Conversation

@ColoCarletti

@ColoCarletti ColoCarletti commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator
  • Preprocessed tables (BITWISE/DECODE) were excluded from the device-only gate, so every epoch drained their full LDE (and their aux) to host. The split-tree commit now honors retain_host_lde, R4 openings serve both subsets (multiplicity range and precomputed range) from the device row gather, and the gate no longer excludes preprocessed tables.
  • Mid-size tables sat below the GPU-commit threshold, committed on CPU, and — having no device handle — re-uploaded their LDE on every R2-R4 dispatch (~40 GB per prove). The default threshold drops from 2^19 to 2^14, the sweep optimum: 2^14 beats 2^15..2^19, and also beats "everything on GPU", where sub-2^14 tables lose to launch overhead.

Also in this round, wall-neutral on this box (the scheduler already hides transfer time) but cutting per-prove PCIe from 187 GB to ~70 GB and freeing GPU busy for smaller-lane setups:

  • R3 OOD barycentric kernels evaluate every OOD point in one pass over the LDE on a cols x chunks grid (the old kernels re-read the LDE per point from a single block per column).
  • Domain coset points are cached on device instead of re-uploaded per table per epoch (~19 GB of identical data).
  • The epoch builder thread pre-uploads the big main traces during its idle slack, so the R1 commit D2D-copies instead of paying the H2D in its chain. BITWISE is excluded (prove_epoch edits its multiplicities post-build) and update_multiplicities drops any stale pre-upload defensively.

ColoCarletti and others added 30 commits July 23, 2026 15:52
- 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.
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.
…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.
ColoCarletti and others added 6 commits July 29, 2026 14:53
Grid-stride the fused row-major NTT past gridDim.y (lde >= 2^24 silently
fell back to CPU), assert the device-only contract in the R2 composition
commit and preprocessed opening fallbacks, validate htod_via bounds, retain
FRI device evals only under device-only, and move the inverse fault-injection
hook so every batch-inverse entry is covered.
…eduler

Fiat-Shamir only requires the main roots absorbed in index order before the
shared challenges; past that fork every table's chain is independent. Phase A
now runs all main commits under a byte-budget admission gate (no chunk
barriers), and aux build, aux commit and rounds 2-4 run fused as one task per
table, heaviest first — while a big table works through a host-bound stretch,
the other tables' GPU stages fill the device. GPU builds default
TABLE_PARALLELISM to 2/3 of the cores (swept flat at 10 on a 16-core RTX 5090).

ethrex 10tx continuations on RTX 5090: 10.64s -> 8.54s (-19.7%, 8 ABBA pairs).
… hygiene

- Extend the device-only gate to preprocessed tables (BITWISE/DECODE): the
  split-trees path takes retain_host_lde, R4 openings serve both subsets from
  the device row gather (multiplicity range + precomputed range), and the
  is_preprocessed exclusion is gone.
- Default GPU LDE threshold 2^19 -> 2^14: CPU-committed mid tables had no
  device handle, so every R2-R4 dispatch re-uploaded their LDE per round.
- Multi-eval-point chunked barycentric kernels for R3 OOD (one pass over the
  LDE for all eval points, cols x chunks grid) with per-point fallback.
- Device cache for domain coset points keyed by (len, p0, p1).
- Pre-upload big main traces from the epoch builder thread; the R1 commit
  D2D-copies instead of paying the H2D in its chain. BITWISE is excluded
  (prove_epoch edits its multiplicities post-build) and update_multiplicities
  drops any stale pre-upload defensively.
- scripts/profiling/h2d_histo.py: memcpy attribution histogram by NVTX phase
  and transfer size from an nsys sqlite export.
@ColoCarletti ColoCarletti changed the title Gpu opt round4 residency perf(prover): device-only preprocessed tables and GPU commits for mid-size tables Aug 3, 2026
@ColoCarletti
ColoCarletti marked this pull request as ready for review August 3, 2026 21:38
@ColoCarletti

Copy link
Copy Markdown
Collaborator Author

/bench-gpu

@ColoCarletti

Copy link
Copy Markdown
Collaborator Author

/ai-review

@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown

GPU Benchmark (ABBA) — 9c65293e28 vs main (14 pairs)

RTX 5090 · AMD Ryzen 9 7945HX with Radeon Graphics (32 threads) · Vast.ai datacenter @ $0.7956481481481481/hr · prover/cuda · ethrex real block, continuations · drift-free A/B/B/A

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

  [parametric] paired-t   mean -15.53%   sd 1.20%   se 0.32%
               95% CI: [-16.22%, -14.83%]   (t df=13 = 2.16)
  [robust]     median -15.61%   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.88%   B CV 0.93%        (lower = steadier)
  within-session drift: +3.03% over the run, 1st->2nd half +1.18%
    (jitter -> Tier-1 cached gate floor; drift -> whether the cached baseline can be trusted)

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

  raw pairs: /tmp/abba_run/pairs.csv

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

@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown

Codex Code Review

  • Mediumprover.rs:1039: Enabling device-only mode for preprocessed tables can create a mixed host/device state. If the new split-tree GPU main commit fails and falls back to CPU, the aux commit still uses the static device_only gate and skips its D2H. build_round1 then marks the whole trace device-only, but no main GPU handle exists, causing later rounds to panic instead of using the valid host main LDE. Base aux retention on whether the main GPU commit actually succeeded, or disable device-only when it fell back.

@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown

AI Review

PR #888 · 13 changed files

Findings

Status Sev Location Finding Found by
confirmed medium crypto/math-cuda/kernels/barycentric.cu:206 CUDA BARY_MAX_K and Rust BARY_MAX_EVAL_POINTS are manually coupled kimi
openrouter/moonshotai/kimi-k2.7-code
nemotron
openrouter/nvidia/nemotron-3-ultra-550b-a55b
confirmed low crypto/stark/src/gpu_lde.rs:2398 coset_points_device_cache relies on an unchecked geometric-coset invariant glm
openrouter/z-ai/glm-5.2
confirmed low prover/src/tables/trace_builder.rs:3773 Pre-uploaded main-trace VRAM is outside admission control glm
openrouter/z-ai/glm-5.2

Status column reflects the verdict from the verifier: deepseek-verifier (openrouter/deepseek/deepseek-v4-pro).

AI-001: CUDA BARY_MAX_K and Rust BARY_MAX_EVAL_POINTS are manually coupled
  • Status: confirmed
  • Severity: medium
  • Location: crypto/math-cuda/kernels/barycentric.cu:206
  • Found by: kimi:openrouter/moonshotai/kimi-k2.7-code, nemotron:openrouter/nvidia/nemotron-3-ultra-550b-a55b
  • Verified by: deepseek-verifier:openrouter/deepseek/deepseek-v4-pro
  • Rejected by: -

Claim

The CUDA kernel allocates ext3::Fe3 acc[BARY_MAX_K] on each thread's stack and loops over k_points without runtime bounds checking. Rust asserts k_points <= BARY_MAX_EVAL_POINTS (8). These two compile-time limits are defined independently in different languages/files. If a future change raises BARY_MAX_EVAL_POINTS without updating BARY_MAX_K, the kernel will be launched with k_points larger than the local array, causing stack/linear-memory corruption and incorrect OOD values.

Evidence

barycentric.cu:206 defines #define BARY_MAX_K 8 and uses it for the per-thread acc[] array in barycentric_base_strided_multi (line 226) and barycentric_ext3_strided_multi (line 274). barycentric.rs:356 defines pub const BARY_MAX_EVAL_POINTS: usize = 8 and asserts k_points <= BARY_MAX_EVAL_POINTS before launching. There is no generated/shared constant or static check keeping them in sync.

Suggested fix

Generate the CUDA constant from the Rust value at build time (e.g. pass it via the build script as a -D flag), or add a prominent comment in both files warning that they must be identical. A runtime kernel-argument cap would also work: pass max_k as a launch argument and assert k_points <= BARY_MAX_K at the top of each kernel.

AI-007: coset_points_device_cache relies on an unchecked geometric-coset invariant
  • Status: confirmed
  • Severity: low
  • Location: crypto/stark/src/gpu_lde.rs:2398
  • Found by: glm:openrouter/z-ai/glm-5.2
  • Verified by: deepseek-verifier:openrouter/deepseek/deepseek-v4-pro
  • Rejected by: -

Claim

The process-wide device coset cache keys only on (len, points[0], points[1]) and never validates the full buffer contents; correctness depends entirely on every caller passing a geometric coset fully determined by its first two terms.

Evidence

coset_points_device_handle builds key = (coset_u64.len(), coset_u64[0], coset_u64[1]) and returns a cached Arc<CudaSlice> on a hit with no content re-check. All current callers (dc.points for R3, coset_base for R4) do pass geometric cosets so the invariant holds today, but a future caller passing a non-geometric point set that collides on len+first two limbs would silently receive a wrong buffer.

Suggested fix

Either assert/verify the geometric property (point[i] == point[0]*generator^i for a derived generator) on insert, or key the cache on a content hash / the full slice. Alternatively, document the invariant as a hard contract on coset_points_device_handle and have callers pass an explicit Coset type rather than a raw &[u64].

AI-010: Pre-uploaded main-trace VRAM is outside admission control
  • Status: confirmed
  • Severity: low
  • Location: prover/src/tables/trace_builder.rs:3773
  • Found by: glm:openrouter/z-ai/glm-5.2
  • Verified by: deepseek-verifier:openrouter/deepseek/deepseek-v4-pro
  • Rejected by: -

Claim

preupload_main_traces allocates up to LAMBDA_VM_TRACE_PREUPLOAD_MB (default 4 GiB) of device buffers from the builder thread on pool streams, but these riding-ahead buffers are not accounted for by the per-table vram_budget_bytes admission control that gates the prover's LDE allocations.

Evidence

preupload_main_traces caps total bytes via budget but never consults Backend::vram_budget_bytes(); the buffers live in trace.main_rowmajor_dev until clear_main_rowmajor_dev runs in the aux stage, so while epoch i is being proven, epoch i+1's preuploaded buffers (up to 4 GiB) can be resident concurrently with epoch i's LDE. The only protection is that clone_htod fails (returns 0, skips) if VRAM is exhausted, which silently demotes the table to the H2D path; an LDE alloc failure similarly falls back to CPU — a large perf regression rather than a correctness break.

Suggested fix

Either have the preupload consult the same VRAM budget/admission accounting used by the prover (so riding-ahead bytes reduce the admitted set), or document explicitly that the budget env var must be set with headroom for the prove peak. At minimum, log/warn when an upload is skipped due to failure so the silent H2D fallback is observable.

Reviewer Lanes

Lane Model Prompt Status Findings
glm openrouter/z-ai/glm-5.2 general success 3
kimi openrouter/moonshotai/kimi-k2.7-code general success 2
minimax minimax/MiniMax-M3 general error: opencode failed (provider/auth/runtime error) and no findings were submitted 0
moonmath zro/minimax-m3 general error: opencode failed (provider/auth/runtime error) and no findings were submitted 0
nemotron openrouter/nvidia/nemotron-3-ultra-550b-a55b general success 5

Verification Lanes

Lane Model Status Confirmed Rejected Uncertain
deepseek-verifier openrouter/deepseek/deepseek-v4-pro success 3 6 0

Native Codex and Claude reviews run separately and post their own comments. They are not included in this structured provenance report.

Discarded candidates (6) — rejected by the verifier
  • Unbounded growth of coset_points_device_cache (crypto/stark/src/gpu_lde.rs:2155, found by nemotron:openrouter/nvidia/nemotron-3-ultra-550b-a55b) — The cache comment (lines 2147-2152) explicitly documents the policy: 'a handful of sizes, ~2-16 MiB each, never evicted — same policy as the host-side domain caches.' The keys are (len, points[0], points[1]) which for geometric cosets with a fixed offset produce one entry per domain size. In a STARK prover, the number of distinct domain sizes is bounded (typically 1-2 per epoch, 1 per process). Total VRAM is well under 100 MiB. This is a deliberate design tradeoff, not a leak.
  • Stream synchronization in coset_points_device_handle blocks producer stream (crypto/stark/src/gpu_lde.rs:2178, found by nemotron:openrouter/nvidia/nemotron-3-ultra-550b-a55b) — The stream.synchronize() at line 2180 is necessary for correctness: the buffer must be resident before publishing to the cache because consumers (on other streams) will read it without synchronization. The sync occurs only once per distinct coset per process lifetime (on cache miss, line 2178). The comment explicitly states the design: 'Settle the copy before publishing: consumers run on other streams.' This is not a performance bug — it's a required fence, and the cost is amortized over all subsequent cache hits.
  • Pre-uploaded main trace can become stale if tables mutate it after preupload (crypto/stark/src/trace.rs:91, found by kimi:openrouter/moonshotai/kimi-k2.7-code) — The PR explicitly excludes BITWISE from pre-upload (trace_builder.rs line 3870-3872) with the comment 'BITWISE is excluded: prove_epoch mutates its multiplicities in place (L2G range-check lookups) after the build.' The finding speculates that other tables may also mutate their main trace between build and R1 commit but provides no concrete example. The list of pre-uploaded tables (CPU, LT, shift, memw, etc.) are stable after build. Pre-upload happens immediately after Traces::build_from_collected, and there are no interleaving mutations before the prover reads them. The only known mutable table is properly handled.
  • bary_num_chunks heuristic may produce suboptimal chunk count for edge cases (crypto/math-cuda/src/barycentric.rs:361, found by nemotron:openrouter/nvidia/nemotron-3-ultra-550b-a55b) — The bary_num_chunks function (barycentric.rs lines 361-365) is explicitly documented as a heuristic: 'Row-chunk count for the multi kernels: enough cols * chunks blocks to occupy the device, without shrinking a chunk's row range below the point where launch + combine overhead dominates.' The constants (2048, 8192) are tuned design parameters, not bugs. The finding itself concedes 'This is reasonable but the constants are heuristic.' A heuristic that may be suboptimal for untested edge cases is not a defect — it's the nature of heuristics.
  • Stale dispatch-threshold comment in split_tree test (crypto/stark/src/gpu_lde.rs:2705, found by glm:openrouter/z-ai/glm-5.2) — The test comment at line 2754 reads 'Above the dispatch threshold (2^19 LDE) so the GPU path must engage.' The parenthetical '(2^19 LDE)' describes the LDE size of the test (n=2^18, blowup=2 → lde_size=2^19), not the threshold value. The test LDE of 2^19 is indeed above the current threshold of 2^14, so the statement is still factually correct. The comment never claimed the threshold itself was 2^19.
  • Pre-uploaded trace in main_rowmajor_dev not cleared on all error paths (crypto/stark/src/prover.rs:962, found by nemotron:openrouter/nvidia/nemotron-3-ultra-550b-a55b) — commit_main_trace at line 1083-1115: the GPU path reads trace.main_rowmajor_dev() and passes it as predev. If try_expand_leaf_and_tree_row_major_keep returns None (GPU path declined), the function falls through to the CPU path and returns Ok. No error propagates — the function returns successfully. The pre-uploaded buffer remains allocated until explicitly cleared at line 3417 (clear_main_rowmajor_dev) during the aux stage. There is no error path within commit_main_trace that would leak the buffer permanently; the only 'failure' is a GPU-decline which is normal control flow returning Ok.

Raw lane outputs, candidates, final issues, and model metrics are uploaded as workflow artifacts.

A static device-only gate on the aux commit could mark the trace device-only
with no main GPU handle to serve it, turning a recoverable CPU fallback of
the main commit into a hard abort downstream.
BARY_MAX_K (kernel accumulator array) and BARY_MAX_EVAL_POINTS (dispatch
assert) were defined independently; build.rs now defines both from one
constant, so they cannot drift into kernel stack corruption.
…RAM budget

The device coset cache keys on (len, p0, p1), which only determines the
contents for a geometric sequence — verify it at sampled indices on insert.
The builder's pre-uploaded traces ride ahead of the admission gate, so cap
them to a slice of the device budget instead of competing with the prove
peak on small cards.
…hold

Lowering the commit threshold to 2^14 silently widened device-only to every
mid-size table. The gate cannot mirror kernel-side dispatch eligibility, so
a single R2 decline on one of those tables hard-aborts the prove (seen at
100tx once main's keccak rework landed) and deadlocks the epoch pipeline.
GPU commits and resident handles keep paying from 2^14; dropping the host
copy stays at the proven 2^19 envelope (LAMBDA_VM_GPU_DEVICE_ONLY_THRESHOLD
overrides).
…le in the R2 abort

The device-resident R2 path only exists for the d=2 quotient decomposition.
DECODE proves with a single part, so admitting it to device-only skipped the
whole device path and hard-aborted into the empty host trace, deadlocking the
epoch pipeline at 100tx. Mirror the parts count in the gate, and include the
table identity in the abort message — finding this one took a live-process
backtrace because the message did not say which table died.
Wall-neutral on the 5090 (the scheduler already hides the H2D) and its
riding-ahead buffers sit outside the VRAM admission gate: at epoch 2^22 the
real-block prove peaks at ~23 GiB and the extra 4 GiB pushed it into
CUDA_ERROR_OUT_OF_MEMORY. Opt-in via LAMBDA_VM_TRACE_PREUPLOAD_MB.
… on an R2 miss

The device-only gate is a static predicate over a dynamic dispatch: it cannot
mirror every reason the device R2 path might decline (parts count, kernel
eligibility, transient errors, shapes a new workload brings), and each miss
was a hard abort that deadlocked the epoch pipeline — DECODE on the synthetic
workload, then a second table on the real-block bench. Instead of excluding
tables one by one, treat the resident handles as the source of truth: on a
miss, download the main/aux LDEs back to host, clear the device-only flag,
and continue on the host path. Slower for that table, never wrong; the abort
remains only when the handles themselves cannot serve the data.
gpu_device_only_downgrades() counts recoveries so a persistently-missing
condition still gets mirrored into the gate.
LAMBDA_VM_GPU_XCHECK runs the verifier's composition consistency check
inside the prover after round 3, per table at negligible cost; on a
failure a post-mortem recomputes each device stage on host, reports the
corruption shape, reruns the device chain to tell a transient race from
a corrupted resident input, and aborts. LAMBDA_VM_GPU_FORCE_DOWNGRADE
exercises the device-only R2 recovery end to end. The R2 downgrade path
now names the table it recovered. A proof_diff ignored test structurally
diffs two continuation bundles.
…declines

A transient CUDA OOM on the resident aux LDE was a hard prove failure:
the resident build leaves no host aux trace to fall back to. A device
drain releases the concurrent VRAM peaks, so one retry usually keeps the
table fully resident; if it still declines, download the resident aux
trace (and the main LDE when the table is device-only) and continue
host-backed. The drain before dropping the resident buffer also keeps
kernels enqueued by the failed attempt from reading pool memory reused
by a concurrent table.
…ption race

Concurrent device R2 windows under VRAM pressure can transiently produce
a fully wrong H for one or two tables while every input stays correct
(rerunning the same chain on the same resident inputs matches the host),
yielding a proof that fails the composition check. Serializing only the
constraint-eval + decompose window across tables eliminates it; commits
and host arms stay parallel, and the windows overlap rarely enough that
the lock is near-free. LAMBDA_VM_GPU_SERIALIZE_R2=0 lifts the lock to
bisect further or once the underlying race is found.
@ColoCarletti

ColoCarletti commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bench-gpu

1 similar comment
@ColoCarletti

Copy link
Copy Markdown
Collaborator Author

/bench-gpu

A mixed state (one commit fell back to CPU while the other stayed
device-only) left the recovery refusing to proceed: it treated a missing
device handle as fatal even when that side already had a valid host
copy. Only the missing side is downloaded now, and the R3 host-arm
guards check the buffer they are about to read instead of the
table-wide flag.
LAMBDA_VM_GPU_FORCE_DOWNGRADE declines every device R2 path so each
device-only table goes through materialize_lde_trace_host and finishes
on the host evaluator; the test proves a small ethrex fixture with a
lowered device-only threshold, asserts the downgrade counter moved and
that the proof verifies. Wired into the test-cuda-fallback group.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants