feat(prover): one device-set model for the VRAM throttle, table names on device aborts, no resident-aux host downgrade - #964
Merged
Conversation
… on device aborts, no resident-aux host downgrade Prover side of the width-aware admission (#961). crypto/stark/src/device_set.rs (new, cfg-free) - The device-set arithmetic moves out of `gpu_lde` (cuda-only) so the scheduler's throttle, which runs on every build, reads the model the dispatch layer admits against: `commit_device_set` (one LDE + snapshot + tree + scratch), the new `table_device_set` for rounds 2-4 (aux LDE and resident aux trace, H and the parts, the R3/R4 inverted denominators, DEEP and the FRI chain, each with its tree), and the pure `admit_bytes`. `gpu_lde` re-exports them. Tests pin LFM_HASH under RPO at 2^21 x 449 (21.4 GiB R1, 23.1 GiB whole table: fits alone, where the retired model said 28.3 GiB) and a 2^22 BALU chunk at 4.9 GiB (five concurrent). crypto/stark/src/prover.rs - `estimate_table_vram_bytes` (two LDE buffers + 256 B per LDE row, stale since #956) is replaced: R1 admits on the commit set, the fused rounds on the table set, both read off the AIR and the domain. - `commit_main_trace` and the aux commit sites pass `air.name()` into the R1 dispatch entry points, so a device abort names its table at the site; `run_admitted` prefixes any string panic payload from a table task with `table <name>: ` before re-raising it. - The resident-aux host downgrade is trimmed: after the drain-and-retry declines, the aux commit returns `ProvingError::DevicePath` with the table, the shape and the live device posture. Host RAM is a cache, not a compute path. `materialize_aux_trace_host` is deleted; `GPU_RESIDENT_AUX_DOWNGRADES` is retired at zero (accessor kept for the integration assertion). prover/src/lfm/chunking.rs - The BALU sizing doc and its test read `stark::device_set::table_device_set` instead of restating the arithmetic: a 2^27 table is 150 GiB, a 2^22 chunk 4.9 GiB, a 2^24 chunk 19 GiB; LFM_LANES at 2^24 is 23 GiB.
…ce-set estimate At TABLE_PARALLELISM=1 the heaviest-first walk order moves nothing on the device (the gate never blocks) and everything about the host allocator's layout. Measured on the q=20 wrap (2^22, blowup 4, 15 tables, RTX 5090): walking by the device-set model peaked at 51.0 GiB max RSS, walking by the retired 2·lde·(8·main + 24·aux) + 256·lde key at 45.2 GiB, prove flat (137.8 vs 137.9 s). The device-set model's 256 MiB scratch floor tied six small tables and moved LFM_HINT two places. The gate keeps the device-set estimates. Both walks now sort by `device_set::host_transient_bytes` — the host bytes the fused task allocates under residency recompute: the recomputed main LDE, the aux copy, the parts, the DEEP codeword and the FRI chain, lde · (8·main + 24·aux + 24·(parts + 3)) — the retired key's per-row weights with 120 in place of 128 in the constant and no floor. The prover prints the walk it took once per prove.
MauroToscano
force-pushed
the
pt/admission-prover-side
branch
from
September 8, 2026 04:22
431f1fc to
6a2fe71
Compare
…del gates, it does not schedule The device-set model this branch introduces is a size model: it says what a stage puts on the card, so the gate can decide whether it fits. It is not a schedule, and sorting the table walk by it cost 5.9 GiB of host peak. The walk and the gate are now two different functions, and only the gate reads the model. The walk goes back to exactly the key it used before this branch: `2·lde·(8·main + 24·aux) + 256·lde`, aux width zero for the R1 main-commit walk and the AIR's aux width for the fused rounds walk, two separate walks as there have always been. The arithmetic is restored verbatim, so the order is byte-for-byte the one that measured well; the function is renamed `table_walk_weight` and its constants renamed with it, because the numbers are kept for the order they produce and are no longer a claim about bytes (the factor of two assumed the second LDE buffer #956's in-place transpose removed). Measured on the q=20 wrap (2^22, blowup 4, RTX 5090, TABLE_PARALLELISM=1), `/usr/bin/time -v` max RSS, prove time and proof bytes identical throughout: this weight, before the device-set model 47,307,284 kB = 45.1 GiB this weight, under the device-set gate 47,365,192 kB = 45.2 GiB the device-set model's own order 53,472,980 kB = 51.0 GiB a fused-phase host-transient order 53,453,344 kB = 51.0 GiB The last row is why this is a restoration and not a re-derivation: an order justified by a truer quantity is still the wrong order. At q=41 the effect is absent, so the mechanism — why reordering only the small tables moves host peak when one table is resident at a time — is open and tracked separately. The prover prints the walk each phase took, one line per phase, so a run that moves host peak can say which order it ran.
…echanism The weight's doc said the mechanism was still open. It is not, and the measurement that closed it is this branch's own. Same walk order, two builds: 45.2 GiB in the diagnostic build, 51.0 GiB here. So the order cannot be the cause. The cause is glibc arena retention — the harness installs no global allocator, and a freed multi-gibibyte buffer goes back to the OS only when the arena top can be trimmed, which depends on what sits above it. The walk order is one input to that layout and the allocations around it are another. Forcing every allocation of a megabyte or more to be mapped and unmapped directly collapses the 5.9 GiB spread to 1.6 MB (46,053,164 kB here against 46,054,788 kB at the pre-#964 tip), so none of it was working set. That makes the allocator the durable fix rather than the walk, and the doc now says so, with the caveat that leaving the retentive state costs ~2.5% prove time. Comments only; no code changed.
MauroToscano
added a commit
that referenced
this pull request
Sep 8, 2026
…uild Two device-memory fixes on the per-table R1 path. Snapshot only when the table builds an aux trace. The fused main commit kept a pre-NTT column-major snapshot of the trace (n * cols * 8 bytes of card) on every device commit. Its only consumer is the LogUp aux build, so a table with no aux trace paid for a buffer nothing read. The two math-cuda entry points gain snapshot-aware siblings (`coset_lde_row_major_with_merkle_tree_keep_snapshot`, `coset_lde_row_major_split_trees_snapshot`); the old signatures are unchanged and keep the snapshot. The stark wrappers take `snapshot_trace`, size the admission's device set with it, and the prover passes `air.has_aux_trace()` at the two commit call sites. With the flag off the handle's `trace_dev` is `None` and `trace_rows` is 0, as its doc says. The size model follows the code: `device_set::table_device_set` and the prover's R1 `main_estimates` take the snapshot on `aux_cols != 0`, which is exactly what `has_aux_trace()` returns (`traits.rs:156`, no override anywhere in the tree), so the gate stops charging aux-less tables for a buffer that is no longer allocated. Bytes admission and a loud abort on the resident LogUp aux build. The build allocated, unchecked, one fingerprint buffer (interactions * rows * 24 B), the batch inverse's prefix, suffix and output (three more of the same size), the term columns and the aux buffer, and on a device error fell to the CPU build through `.ok()?` with no message: a production host compute path outside the sixteen sites #961 covered. This was found by reading the code, not from a failure. `device_set::aux_build_device_set` sizes the set term by term from what `logup_aux_resident` and `batch_inverse_ext3_dev` allocate, with the peak as the larger of the inverse phase and the term/assemble phase; it lives in `device_set` rather than beside its dispatch because the scheduler's throttle reads the same arithmetic on every build, cuda or not. `admit_aux_build` applies the same bytes ceiling as the commits before any device allocation, and both over-budget and a post-admission device error go through the existing abort, with the set printed term by term and the table named (`#964`'s `DispatchShape::table`, threaded from the AIR's own name).
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.
Prover side of width-aware admission: one device-set model, named aborts, no resident-aux host downgrade
Per-table redo, lane M follow-up to #961. Cut from
per-table-gpu@8a34d205, rebased onto 7c311c2 (#963). Three commits: the model + names + clean error, the walk-order fix, then its replacement — the walk restored byte-for-byte to its pre-#964 form.What
crypto/stark/src/device_set.rs(new, cfg-free) — the device-set arithmetic moved out ofgpu_lde(which iscuda-only) so the scheduler's throttle, which runs on every build, readsthe same model the dispatch layer admits against:
commit_device_set(one LDE + snapshot +tree + scratch, the #956 shape), the new
table_device_setfor rounds 2–4 (adds the aux LDEand resident aux trace,
Hand the parts, the R3/R4 inverted denominators, DEEP and the FRIchain, each with its tree), and the pure
admit_bytespredicate.gpu_ldere-exports them.Tests pin the model at the brief's shapes: LFM_HASH under RPO at 2^21 × 449 is 21.4 GiB for
the R1 commit and 23.1 GiB for the whole table (fits alone; the retired 2×-LDE model said
28.3 GiB, over a budget it is under); a 2^22 BALU chunk is 4.9 GiB (five concurrent).
crypto/stark/src/prover.rsestimate_table_vram_bytes(the 2×-LDE + 256 B/row heuristic, stale since perf(gpu): transpose the fused commit's LDE in place, one LDE-sized buffer live #956) stops beingthe ADMISSION model: the R1 phase admits on
commit_device_set(...).total(), the fusedrounds-2–4 phase on
table_device_set(TableShape { n, blowup, main_cols, aux_cols, num_parts, num_eval_points }), both read off the AIR and the domain. Its arithmetic survives verbatim asthe WALK's sort weight — see the third commit, and the host-peak section below for why.
commit_main_traceand the aux commitsites pass
air.name()into the four R1 dispatch entry points (the abort line now startstable LFM_HASH: R1 main commit: …), andrun_admittedprefixes any string panic payloadfrom a table task with
table <name>:before re-raising it, so a device-only recoveryrefusal names its table too.
LDE and continue on the host) is trimmed: after a failed retry the aux commit returns a clean
ProvingError::DevicePathcarrying the table, the shape and the live VRAM / mempoolposture. Host RAM is a cache, not a compute path;
materialize_aux_trace_hostis deletedand
GPU_RESIDENT_AUX_DOWNGRADESis retired at zero (its accessor stays for theintegration assertion). The R2 downgrade comment now names the test-only switch that alone
reaches it.
prover/src/lfm/chunking.rs— the BALU sizing doc and its test readstark::device_set::table_device_setinstead of restating the arithmetic: a 2^27 table is a150 GiB device set (51 GiB at R1), a 2^22 chunk 4.9 GiB (five concurrent, 32 chunks @110q),
a 2^24 chunk 19 GiB (proves alone); LFM_LANES at 2^24 is 23 GiB.
Evidence
Laptop:
cargo test -p stark --release --lib device_set→ 10 passed;cargo test -p lambda-vm-prover --release --lib lfm::chunking→ 14 passed;cargo clippy -p stark --features cuda --all-targets -- -D warnings -A clippy::op_refandthe same without
cuda→ Finished;make fmt→ clean;make lint→ exit 0.Box A (RTX 5090, @431f1fc1, relayed verbatim):
make test-cuda-integration→ exit 0,ok. 7 passed … 13.80s; mempool-threshold line 1,TEST-ONLY 0,
[gpu] ABORT0.make test-cuda-d1→ok. 1 passed … 2.62s.make test-cuda-fallback→ok. 5 passed … 26.83sand the forced-downgrade binaryok. 1 passed … 12.57s; bannersTEST-ONLY host fallback ENABLED (LAMBDA_VM_GPU_FORCE_DOWNGRADE)×1 and
… (feature test-cuda-faults)×1; theTEST-ONLY host recovery: R2/R3/R4 host arm on a device-only trace: rows … main cols … aux cols …set is unchanged (1048576×17/5 ×1, 1048576×29/10×1, 1048576×38/10 ×3, 131072×14/3 ×1).
make lint→ exit 0.LAMBDA_VM_MEMPOOL_RELEASE_MB=0, residency recompute,TABLE_PARALLELISM=1,2^22, 10 Hz VRAM sampler):
★ WRAP PROVED AND VERIFIED (… blowup 4, 20 queries),prove 137.8s / verify 0.18s / proof 51175464 bytes / 2407 published words / 15 sub-proofs,ok. 1 passed … 180.08s, no[gpu] ABORT, 0 downgrade lines; VRAM peak 26,640 MiB; max RSS53,472,980 kB (51.0 GiB).
Against box A's q=20 baseline at 64e1610 (prove 136.6 s, proof 51,175,464 B, VRAM 27,056 MiB,
max RSS 47,306,788 kB = 45.1 GiB): prove +0.9% (predicted ±5%), VRAM −1.5% (predicted ≤ +2%),
proof bytes identical.
Host peak at q=20 reads 51.0 GiB, and it is retention, not working set. The story changed
twice; this is where it landed.
The walk order was attributed as the mechanism, and a
host_transient_byteskey was written tochase it. That key did not reproduce the good order: per LDE row it is
8·main + 24·aux + 24·parts + 72, which equals half the retired key's16·main + 48·aux + 256only where
num_parts == 2, and the residualretired/2 − 8·ldeis not a monotone transformbecause the LDE size varies per table. It measured 53,453,344 kB, unchanged.
The third commit drops that key and restores the pre-#964 walk verbatim instead. The device-set
model feeds ADMISSION only; the walk is a scheduling policy keyed on
table_walk_weight, whosearithmetic is the retired estimate's, character for character, kept for the ORDER it produces and
documented as not a byte claim. Two separate walks again, R1 with aux width zero and the fused
rounds with the AIR's, as before #964. The prover prints the walk each phase took.
Measured at
a6359906, same recipe:q=41 is neutral. q=20 still reads 51.0 with the walk restored byte-for-byte, so the walk order
is not the cause. A four-run ABBA in ONE build, switch OFF/ON/ON/OFF, settles what it is:
Proof bytes identical in all four; repeats agree to 824 kB and 2.9 MB; the ABBA ordering rules out
drift. So order moves 5.9 GiB inside one binary — yet this PR walks the low arm's order and
lands on the high arm's number. Order is therefore a correlate, not a cause: the two builds also
differ in what they allocate around the walk sites, and both inputs feed the same mechanism.
That mechanism is glibc arena retention. The harness installs no
#[global_allocator], so theseruns are glibc, and a freed chunk returns to the OS only if the arena top can be trimmed, which
depends on what sits above it. The page faults are the independent witness: the LOW-memory arm
takes ~1.2M MORE minor faults, which is memory being returned and re-faulted rather than kept. A
change in working set would move them the other way or not at all. The same instrument at q=41
confirms the class: under
MALLOC_MMAP_THRESHOLD_=1048576the tip drops 98.6 → 85.5 GiB at no timecost, and #956's exact parent and the tip then measure 604 kB apart.
The honest trade: the low-retention state is consistently ~2.5% slower (139–140 s against
135–136 s), the price of returning and re-faulting those pages. This branch walks the LOW arm's
order and yet measures the HIGH arm's peak and the HIGH arm's speed — which is the whole reason
the order cannot be the cause. Neither number is a working-set change.
Confirmed: the gap is retention and nothing else. The q=20 pair under
MALLOC_MMAP_THRESHOLD_=1048576, which maps and unmaps every allocation of a megabyte or morerather than leaving it to arena trimming:
a6359906(this PR)7c311c28(tip control)1,624 kB apart, against 5.9 GiB without it, and both land 1.2 GiB below the tip's own no-knob
reading. So this PR costs zero live memory: the 5.9 GiB was retained pages, and the criterion
registered before the run — that the PR should wait if the gap survived the knob — is not met.
The flag is therefore a real host lever, and its price shows up in the same pair: the tip, already
in the low-retention state, goes 140.3 → 140.1 s and pays nothing, while this head goes
136.4 → 140.2 s and pays 2.8%. ⚠ Not yet a universal trade — at q=41 the same flag delivered
13.1 GiB with no reported time cost, so what the move costs is shape-dependent and unpinned. F12b
(q=41 host 86.7 → 98.6 GiB) resolves as this same class and is tracked outside this PR.
Not in this PR