feat(gpu): bytes-aware GPU admission, loud abort inside the device envelope, BALU row-chunking policy - #961
Merged
Merged
Conversation
…pe, BALU row-chunking policy GPU admission was row-only and width-blind: `lde_size = n * blowup >= 2^14` admitted a 2^21 x 449 table whose LDE alone is 14 GiB, and the allocation failure that followed was swallowed by `.ok()?` into a silent host commit. Under the plan's rule that host RAM is a cache and not a compute path, a host fallback in the production pipeline is a failure to report. crypto/stark/src/gpu_lde.rs - One shared predicate, `admit(lde_size, bytes)`: the row floor (`gpu_lde_threshold`) checked first, then a bytes ceiling against the card's admission budget (80% of device memory, or LAMBDA_VM_VRAM_BUDGET_MB). Every former `gpu_lde_threshold()` consumer site now goes through it with the bytes it is about to allocate. The predicate is a pure function of the shape and two process constants — live free memory is deliberately not an input — so R1 admission and every handle-bearing re-derivation agree (LOCKSTEP); FRI re-derives at width 1 and the ceiling, being an upper bound, never degenerates there. - `commit_device_set`: the R1 commit's device set term by term after #956's in-place transpose (one LDE buffer + trace snapshot + full tree + scratch), the model `one_lde_buffer::vram_arm` measures. - Loud abort: an admitted R1 commit that fails on the device, or a table over budget, panics with the stage, shape, device set, budget, live free/total VRAM and the mempool posture; `run_admitted` re-raises the payload. The R2-R4 host recoveries of device-only tables (`materialize_*`) are refused on the same terms. The only way back is LAMBDA_VM_TEST_ONLY_HOST_FALLBACK, implied by the LAMBDA_VM_GPU_FORCE_DOWNGRADE test hook and by the `test-cuda-faults` feature, and announced by a banner. A GPU-less host (no backend) keeps the host path — that is not the production pipeline. - Tests: the admission arithmetic at the brief's shape (2^22 x 612 @ blowup 2 is 57.9 GiB against 25.6) and at LFM_HASH under RPO (21.4 GiB @ b2 fits, 35.6 @ b4 does not); an ignored box test that an over-budget commit aborts with the diagnostic instead of committing on the host. prover/src/lfm/chunking.rs - `BaluChunking`: the row-chunking policy for the tall-narrow chips, sized in its doc — LFM_BALU at 2^27 rows is a 96 GiB whole-prove set; 2^22-row chunks are 3 GiB each and eight prove concurrently inside the budget. Knob LFM_BALU_MAX_CHUNK_ROWS_LOG2, default one table. Blake3Chunking shares the chunk arithmetic. Wiring into the program follows the BLAKE3 template and is not in this change. crypto/math-cuda/src/device.rs - LAMBDA_VM_MEMPOOL_RELEASE_MB wired through a named default and a cached accessor, reported once per process; `Backend::device_mem_info` for the diagnostics and `Backend::trim_mempool_to` for the explicit release. The default stays retain-all pending the sampler measurement.
…retain-all default vram_arm on the RTX 5090 (2^21 x 316 @ blowup 2, five commits): retain-all 1176.9/1057.9/1055.3/1056.2/1055.6 ms vs release-0 1178.5/1057.2/1077.5/ 1081.2/1079.3 ms, in-process peak 15.67 GiB under both. Retention is ~2% faster after the first commit and adds nothing to the peak; the multi-table q=41 wrap rung proves under the same default at a 28,976 MiB VRAM peak with no device decline. The default stays retain-all, now on a measurement.
…E3 builders
The comp-tree cliff test (`gpu_comp_tree_fault_recovers_device_only_parts`)
arms `FAULT_COMP_TREE_STICKY` and expects the first device comp-poly tree
build to fail. The hook lived only in the keccak builders
(`merkle::build_comp_poly_tree_from_{slabs_dev,evals_ext3_keep}`); when the
BLAKE3 twins landed with the by-hash device dispatch (9cb811e) they were not
mirrored, and since `DefaultStarkHash` is BLAKE3 on every build the fault
could never fire — the test has failed at
"injected comp-tree fault never fired" on per-table-gpu since then.
Both BLAKE3 builders now run the same `check_sticky` at their head, under
`#[cfg(feature = "test-faults")]` exactly like the keccak twins, so the
production build is unchanged.
MauroToscano
added a commit
that referenced
this pull request
Sep 7, 2026
…builders #961 gave both BLAKE3 comp-poly tree builders the test-faults-gated sticky hook the keccak twins carry, because the fault suite's comp-tree cliff test arms one counter and must reach it under whichever hash the build pins. The two RPX builders lacked it, so the suite would have gone silent again the day the pinned hash is RPX. Same four lines, first statement of build_comp_poly_tree_from_slabs_dev and build_comp_poly_tree_from_evals_ext3_keep. rpx.cu is untouched.
MauroToscano
added a commit
that referenced
this pull request
Sep 7, 2026
… 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.
MauroToscano
added a commit
that referenced
this pull request
Sep 8, 2026
… 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.
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.
Width-aware GPU admission, loud abort inside the device envelope, BALU row-chunking policy
Per-table redo, lane M (EXECUTION 10b/10c/10e). Cut from
per-table-gpu@7c817051.Why
GPU admission was row-only:
lde_size = n · blowup ≥ 2^14admitted a 2^21 × 449 table whoseLDE alone is 14 GiB, and the allocation failure that followed was swallowed by
.ok()?into asilent host commit — inside and outside the device-only envelope (the hard abort was only
the R4 gather). The pre-existing
VramGateis a concurrency throttle: it admits an oversizedtable alone and never declines. Under the plan's rule that host RAM is a cache and not a
compute path, a host fallback in the production pipeline is a failure to report.
What
crypto/stark/src/gpu_lde.rsadmit(lde_size, bytes): row floor (gpu_lde_threshold) first, then abytes ceiling against the card's admission budget (80% of device memory, or
LAMBDA_VM_VRAM_BUDGET_MB). Every formergpu_lde_threshold()site (16) goes through it withthe transient it allocates. Pure function of the shape and two process constants — live free
memory is deliberately not an input — so R1 admission and every handle-bearing re-derivation
agree (LOCKSTEP); FRI re-derives at width 1 against an upper-bound ceiling.
commit_device_set(n, base_cols, blowup, snapshot): the R1 commit's device set after perf(gpu): transpose the fused commit's LDE in place, one LDE-sized buffer live #956'sin-place transpose (one LDE + trace snapshot + full tree + scratch) — the terms
one_lde_buffer::vram_armmeasures.with stage, shape, device set, budget, live free/total VRAM and the mempool posture;
run_admittedre-raises the payload. The R2–R4 host recoveries of device-only tables(
materialize_*) are refused on the same terms. The only way back isLAMBDA_VM_TEST_ONLY_HOST_FALLBACK=1, implied by theLAMBDA_VM_GPU_FORCE_DOWNGRADEtest hookand by the
test-cuda-faultsfeature, announced by a banner. No CUDA backend (stub cubins)keeps the host path — that is CI, not the production pipeline.
against 25.6) and at LFM_HASH under RPO (2^21 × 449: 21.4 GiB @ b2 fits, 35.6 @ b4 does not);
an ignored box test that an over-budget commit aborts with the diagnostic.
prover/src/lfm/chunking.rs—BaluChunking, the row-chunking policy for tall-narrowchips (knob
LFM_BALU_MAX_CHUNK_ROWS_LOG2, default one table, target 2^22). Sized in its docand pinned by a test: LFM_BALU at 2^27 rows is a 50 GiB R1 set and a ~96 GiB whole-prove set;
a 2^22 chunk is ~3 GiB and eight prove concurrently inside the budget (32 chunks @110q, 64
@219q). Column streaming is rejected in writing: rounds 2–4 read the resident LDE, so streaming
the commit shrinks nothing for a tall-narrow chip.
Blake3Chunkingshares the chunkarithmetic. Wiring into the program (BLAKE3 template) is a follow-up.
crypto/math-cuda/src/device.rs—LAMBDA_VM_MEMPOOL_RELEASE_MBthrough a named default(retain-all) and a cached accessor, reported once per process;
Backend::device_mem_infoforthe diagnostics,
Backend::trim_mempool_tofor the explicit release. The default is now on ameasurement (G6 below), recorded in the doc.
Evidence (box A, RTX 5090, relayed verbatim; @c36eadb9 unless stated)
cargo test -p stark --release --features cuda --lib gpu_lde::admission_tests→ok. 7 passed;cargo test -p lambda-vm-prover --release lfm::chunking::→ok. 14 passed.LAMBDA_VM_VRAM_BUDGET_MB=1024 cargo test -p stark --release --features cuda --lib gpu_lde::admission_box_tests -- --ignored --nocapture→
test result: ok. 1 passed, stderr:[gpu] ABORT: the device path is the production path and it is unavailable — R1 main commit: rows 262144 x 612 base cols @ blowup 2 (LDE 524288); over the VRAM admission budget: needs 4137680864 B (3.85 GiB), budget 107…— the pre-registered line.make test-cuda-integration→ok. 7 passed … 13.82s; mempool lines 1, TEST-ONLY 0,ABORT 0.
make test-cuda-d1→ok. 1 passed.make test-cuda-fallback→ 4 of 5 pass with the predicted banners(
TEST-ONLY host fallback ENABLED (feature test-cuda-faults)once,TEST-ONLY host recovery: …per table).gpu_comp_tree_fault_recovers_device_only_partsfails — pre-existing: thesticky comp-tree hook exists only in the keccak builders (merkle.rs:500,549) while
DefaultStarkHashis BLAKE3 on every build (config.rs:37), so the fault can never fire underBLAKE3. Confirmed pre-existing at the integration tip c34dee4 (per-table-gpu after test(lfm): lever-0 census — the tenant-width bill of a per-table aggregator leg #955, no
hook there either):
make test-cuda-fallback→ the samepanicked at prover/tests/cuda_fallback_tests.rs:170:5: injected comp-tree fault never fired,test result: FAILED. 4 passed; 1 failed … 25.77s. Fixed here by mirroring the keccakbuilders'
test-faults-gated hook into the two BLAKE3 builders (5d2fe24). Re-run at5d2fe24:
make test-cuda-fallback→ exit 0,test result: ok. 5 passed; 0 failed … finished in 26.66s(the comp-tree test now prints[gpu] TEST-ONLY host recovery: host consumer of the composition parts on a device-only table: rows 52428…; the forced-downgradebinary
1 passed … 12.38s).VRAM 31,920 MiB, 0 downgrade lines):
★ WRAP PROVED AND VERIFIED (… blowup 4, 41 queries),prove 270.0s / verify 0.19s, VRAM peak 28,976 MiB, host peak 98.6 GiB (box B's run 97.7 —run-to-run band),
ok. 1 passed … 332.55s, no[gpu] ABORT, 0 downgrade lines.vram_arm×5 (2^21 × 316 @ b2): retain-all 1176.9 / 1057.9 / 1055.3 / 1056.2 / 1055.6 ms,peak 15.67 GiB; release-0 1178.5 / 1057.2 / 1077.5 / 1081.2 / 1079.3 ms, peak 15.67 GiB.
Retention ≈2% faster, peaks equal → retain-all stays the default (recorded in device.rs).
make test-cuda-integration→ exit 0,7 passed; 0 failed … 13.99s, mempool-threshold lines 1, TEST-ONLY 0,[gpu] ABORT0;make test-cuda-d1→ exit 0,1 passed … 2.65s.make lint→ exit 0 at c36eadb and again at 5d2fe24 (no clippy errors, noDiff inlines).Follow-ups (not in this PR)
pt/admission-prover-side: buildestimate_table_vram_bytesfromcommit_device_set; putair.name()in the abort payload; trim the retry/downgrade block that is now test-only.only when over budget.