perf(gpu): transpose the fused commit's LDE in place, one LDE-sized buffer live - #956
Merged
Merged
Conversation
…uffer live The fused row-major R1 commit allocated the row-major LDE for all columns and then, while it was still live, a second full-size column-major buffer for the transpose: peak ~ 2*LDE + trace + tree (24.8 GiB for a 2^21 x 316 table at blowup 2 by the memory model; 44.7 GiB at blowup 4, which no 32 GiB card holds). The transpose now happens inside the one allocation, over the same tiled kernel: a block pass transposes each row block through one spare block of scratch, then the resulting column runs are permuted into column-major order cycle by cycle, issued as batched device copies. Scratch is one block plus one run, capped at 256 MiB, instead of a second LDE. Bytes are unchanged: the passes only move runs of already-computed values, the row-major host D2H is queued ahead of them on the same stream, the trace snapshot transpose is untouched, and no kernel changes. New tests pin the handle's column-major bytes against the row-major host copy the same call returns (base, ext3, split trees), plus an ignored production-shape VRAM/time arm for the sampler. LAMBDA_VM_LDE_TRANSPOSE_UNBATCHED=1 issues the run pass as single copies for measurement.
This was referenced Sep 7, 2026
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.
MauroToscano
added a commit
that referenced
this pull request
Sep 8, 2026
…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.
This was referenced Sep 8, 2026
Merged
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.
Summary
The fused row-major R1 commit allocated the row-major LDE for all columns (
lde.rsexpand_row_major_on_stream,alloc_zeros::<u64>(lde_size * total_cols)) and then, while it was still live, a SECOND full-size column-major buffer for the transpose (launch_row_to_col_major), from bothcoset_lde_row_major_innerandcoset_lde_row_major_split_trees. Peak ≈ 2·LDE + trace snapshot + tree: 24.9 GiB for a 2^21×316 table at blowup 2 by the memory model, 44.7 GiB at blowup 4 and 34.2 GiB for 436 columns — neither of which a 32 GiB card can hold.The LDE is now transposed inside its own allocation, over the same
matrix_transpose_stridedkernel:rows × colsmatrix isblocksrow blocks ofrows_per_blockrows; each block is transposed intocolsruns (one column,rows_per_blockconsecutive rows), ping-ponging through ONE spare block of scratch (block 0 → scratch, block b → slot b−1, scratch → last slot).cuMemcpyBatchAsyncbatches (about2·blocks + 3·cyclesdriver calls — 136 at 2^22×316).Scratch =
(cols + 1) · rows_per_block · 8bytes, capped at 256 MiB (2^16 rows per block for 316 or 436 columns), instead of a second LDE. The trace-snapshot transpose (n×C into its own small buffer) is unchanged.LAMBDA_VM_LDE_TRANSPOSE_UNBATCHED=1issues the run pass as single D2D copies — a measurement/diagnostic knob, same bytes.Only
crypto/math-cuda/src/lde.rschanges (plus the new test file); no kernel, dispatch, caller ormatch hashline moves.Why the bytes cannot change
Nothing computes: both passes only move runs of already-computed values, under exactly the old kernel's map
dst[c·L + r] = src[r·C + c](block b, row r′, column c lands atc·L + b·R + r′). Leaves are hashed before the transpose; the row-major host D2H (retain_host_lde) is enqueued on the same stream ahead of it, so the host copy sees row-major bytes; thereadyevent is recorded after, as before. Same kernels, same values, same roots, same proof bytes — and the measurement below confirms the root byte for byte, old vs new.Tests
lde.rsunit tests (cargo test -p math-cuda --lib inplace_transpose): geometry invariants over 2^1..2^27 rows × {1..65535} cols; a host model of the run pass asserting every batch is independent (distinct destinations, no destination aliasing a source) and that every run lands at its column-major position, for 15 (blocks, cols) shapes incl. 64×316, 64×436, 128×612.tests/one_lde_buffer.rs: for 12 shapes × {keccak, blake3}, the handle's column-major device bytes equal (raw u64) the row-major host copy the same call returns — base, ext3 and split-tree entry points, plus the trace snapshot.vram_arm(#[ignore]) is the production-shape VRAM/time arm for the sampler (envLAMBDA_VM_VRAM_ARM_{LOG_N,COLS,BLOWUP,ITERS,HASH,PREDEV}).Pre-registered predictions (written before any measurement)
Model = LDE + snapshot + tree + scratch + twiddles/weights; GiB = 2^30 bytes. The sampler adds the CUDA context (0.51 GiB on the box) to old AND new.
Time, pre-registered separately: transpose stage +30 ms (b2) / +60 (b4) / +42 (436) batched; +55 / +115 / +80 ms unbatched; commit-level ≤ +5% with host input.
Gate (box A, at d5ba828, verbatim)
cargo test -p math-cuda --release --lib inplace_transpose:test result: ok. 2 passed; 0 failed ... finished in 0.19scargo test -p math-cuda --release --test one_lde_buffer:ok. 3 passed; 0 failed; 1 ignored ... 0.93sblake3_fused_parity:ok. 3 passed·merkle_root_parity:ok. 3 passed·comp_poly_tree:ok. 4 passedmake test-cuda-integration:ok. 7 passed ... 13.50s·make test-cuda-d1:ok. 1 passed ... 2.57smake lint: exit 0VRAM / time arm (
vram_arm, host input, 5 iterations,LAMBDA_VM_MEMPOOL_RELEASE_MB=0, 10 Hz nvidia-smi sampler)LAMBDA_VM_LDE_TRANSPOSE_UNBATCHED=1VRAM: every prediction met within the context term; peak −9.72 GiB at 316/b2 (25.39 → 15.67), and the two shapes the old code cannot commit at all now run.
Time: the miss, stated plainly
Commit-level time at 316/b2 is +123 ms per iteration (+13%, 1065 vs 942 ms), above the pre-registered ±5%. Batched is ~30 ms faster than unbatched, so the run pass is bound by per-copy cost (~20k run copies of 512 KiB ≈ 5 µs each), not by bandwidth; the block pass is bandwidth-bound as predicted. b4 and 436 have no old comparison (the old code cannot allocate them).
Block-level projection (? INFERRED from the campaign's cell counts, not measured on a block): +123 ms per 10.6 GB of committed LDE ≈ +11.6 ms per GB for wide tables. Aggregation layer: 6 proofs per block × ≈19 GiB of LDE (LFM_HASH 9.9 + the ten other chips ≈ 8.9) ≈ +1.3 s. Base layer: 5 epochs × ≈64 GiB (≈8B cells/epoch at blowup 2) ≈ +3.7 s. Block ≈ +5 s ≈ +0.2% of the 42.6-minute record. Narrower tables pay less per byte (longer runs, fewer copies).
Lever if it matters: the copy count is
rows · cols / rows_per_block. Doubling the run to 2^17 rows halves it (≈ −50 ms at 316/b2) for +0.16 GiB (316 cols) / +0.21 GiB (436 cols) of scratch, which would put the RPO shape 0.05 GiB over its ≤ 21 GiB model target. The real fix is a one-launch run-permutation gather kernel in place of the per-copy engine ops, which makes the run pass bandwidth-bound (≈ 15–20 ms) with no extra scratch — a kernel-file change outside this PR.Notes
math-cudahas nocudacargo feature; its tests arecargo test -p math-cuda --release(Makefiletest-math-cuda).match hashline touched. Lane K phase 2: the row-major leaf kernels stay the production layout.mainonpt/one-lde-buffer-main(hash-independent win for keccak and BLAKE3 today).