feat: Multi-GPU E2E proving pipeline - #644
Draft
samuelburnham wants to merge 9 commits into
Draft
samuelburnham wants to merge 9 commits into
samuelburnham wants to merge 9 commits into
Conversation
samuelburnham
added this pull request to stack #645
September 22, 2026 18:41
Merged
samuelburnham
force-pushed
the
sb/aiur-gpu-lanes
branch
3 times, most recently
from
October 1, 2026 14:08
76c5dca to
67fe4ed
Compare
Member
|
!benchmark fresh |
Contributor
|
| constant | execute-time (main) | execute-time (PR) | Δ% | prove-time (main) | prove-time (PR) | Δ% | throughput (const/s) (main) | throughput (const/s) (PR) | Δ% | peak-ram (main) | peak-ram (PR) | Δ% | proof-size (main) | proof-size (PR) | Δ% | verify-time (main) | verify-time (PR) | Δ% | fft-cost (main) | fft-cost (PR) | Δ% | trace-shards (main) | trace-shards (PR) | Δ% |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
ByteArray.utf8DecodeChar?_utf8EncodeChar_append |
8.187 s | 8.196 s | +0.1% | 54.946 s | 55.248 s | +0.5% | 50.500 | 50.230 | -0.5% | 73.23 GiB | 73.22 GiB | -0.0% | 5.01 MiB | 5.01 MiB | +0.0% | 33.0 ms | 33.8 ms | +2.4% | 135.66B | 135.66B | +0.0% | 1 | 1 | +0.0% |
Char.ofOrdinal_le_of_le |
5.943 s | 5.999 s | +1.0% | 44.745 s | 44.940 s | +0.4% | 61.750 | 61.480 | -0.4% | 62.33 GiB | 62.19 GiB | -0.2% | 5.01 MiB | 5.01 MiB | +0.0% | 29.5 ms | 30.5 ms | +3.4% |
97.27B | 97.27B | +0.0% | 1 | 1 | +0.0% |
Array.extract_append |
6.131 s | 6.126 s | -0.1% | 40.393 s | 40.686 s | +0.7% | 39.760 | 39.470 | -0.7% | 53.91 GiB | 54.01 GiB | +0.2% | 4.89 MiB | 4.89 MiB | +0.0% | 29.2 ms | 30.1 ms | +3.2% |
97.92B | 97.92B | +0.0% | 1 | 1 | +0.0% |
Std.HashMap |
3.848 s | 3.934 s | +2.2% | 25.309 s | 25.470 s | +0.6% | 80.680 | 80.170 | -0.6% | 36.90 GiB | 36.88 GiB | -0.1% | 4.94 MiB | 4.94 MiB | +0.0% | 29.4 ms | 29.5 ms | +0.5% | 62.29B | 62.29B | +0.0% | 1 | 1 | +0.0% |
_private.Init.Data.Range.Polymorphic.SInt.0.Int64.instRxcHasSize_eq |
3.615 s | 3.563 s | -1.4% | 24.281 s | 24.314 s | +0.1% | 76.890 | 76.790 | -0.1% | 35.41 GiB | 35.37 GiB | -0.1% | 4.91 MiB | 4.91 MiB | +0.0% | 29.9 ms | 30.0 ms | +0.4% | 56.81B | 56.81B | +0.0% | 1 | 1 | +0.0% |
String.append |
734.4 ms | 742.9 ms | +1.2% | 2.443 s | 2.469 s | +1.1% | 133.860 | 132.440 | -1.1% | 5.90 GiB | 5.56 GiB | -5.7% (1.06× smaller) 🟢 | 4.65 MiB | 4.65 MiB | +0.0% | 26.5 ms | 26.8 ms | +1.3% | 3.49B | 3.49B | +0.0% | 1 | 1 | +0.0% |
Nat.add_comm |
524.0 ms | 517.9 ms | -1.2% | 1.061 s | 1.068 s | +0.7% | 43.340 | 43.060 | -0.6% | 4.47 GiB | 4.55 GiB | +1.9% | 4.47 MiB | 4.47 MiB | +0.0% | 24.5 ms | 25.0 ms | +1.9% | 306.71M | 306.71M | +0.0% | 1 | 1 | +0.0% |
FRI verifier on FRI (7 constants)
| constant | execute-time (main) | execute-time (PR) | Δ% | prove-time (main) | prove-time (PR) | Δ% | throughput (const/s) (main) | throughput (const/s) (PR) | Δ% | peak-ram (main) | peak-ram (PR) | Δ% | proof-size (main) | proof-size (PR) | Δ% | verify-time (main) | verify-time (PR) | Δ% | fft-cost (main) | fft-cost (PR) | Δ% |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
ByteArray.utf8DecodeChar?_utf8EncodeChar_append |
3.070 s | 3.047 s | -0.7% | 34.146 s | 34.404 s | +0.8% | 81.270 | 80.660 | -0.8% | 53.63 GiB | 53.44 GiB | -0.3% | 2.67 MiB | 2.67 MiB | +0.0% | 16.6 ms | 17.0 ms | +3.0% | 94.15B | 94.15B | +0.0% |
Char.ofOrdinal_le_of_le |
3.022 s | 3.074 s | +1.7% | 33.506 s | 33.591 s | +0.3% | 82.460 | 82.250 | -0.3% | 52.07 GiB | 52.20 GiB | +0.2% | 2.67 MiB | 2.67 MiB | +0.0% | 15.5 ms | 15.8 ms | +2.2% | 93.40B | 93.40B | +0.0% |
Array.extract_append |
3.119 s | 3.034 s | -2.7% | 33.227 s | 33.337 s | +0.3% | 48.330 | 48.180 | -0.3% | 51.33 GiB | 51.11 GiB | -0.4% | 2.67 MiB | 2.67 MiB | +0.0% | 15.5 ms | 16.2 ms | +4.9% |
92.62B | 92.62B | +0.0% |
Std.HashMap |
2.962 s | 3.041 s | +2.7% | 33.129 s | 33.301 s | +0.5% | 61.640 | 61.320 | -0.5% | 51.42 GiB | 51.40 GiB | -0.0% | 2.67 MiB | 2.67 MiB | +0.0% | 15.5 ms | 15.7 ms | +1.0% | 92.66B | 92.66B | +0.0% |
_private.Init.Data.Range.Polymorphic.SInt.0.Int64.instRxcHasSize_eq |
2.961 s | 3.015 s | +1.8% | 33.280 s | 33.517 s | +0.7% | 56.100 | 55.700 | -0.7% | 51.03 GiB | 51.16 GiB | +0.2% | 2.67 MiB | 2.67 MiB | +0.0% | 15.3 ms | 15.8 ms | +3.3% |
92.06B | 92.06B | +0.0% |
String.append |
2.750 s | 2.858 s | +4.0% |
31.351 s | 31.529 s | +0.6% | 10.430 | 10.370 | -0.6% | 47.23 GiB | 46.90 GiB | -0.7% | 2.67 MiB | 2.67 MiB | +0.0% | 15.3 ms | 15.3 ms | +0.3% | 83.91B | 83.91B | +0.0% |
Nat.add_comm |
2.662 s | 2.615 s | -1.8% | 30.971 s | 31.064 s | +0.3% | 1.490 | 1.480 | -0.7% | 45.94 GiB | 45.90 GiB | -0.1% | 2.67 MiB | 2.67 MiB | +0.0% | 16.2 ms | 16.5 ms | +2.1% | 78.77B | 78.77B | +0.0% |
Aggregate flat join (1 pair)
| pair | execute-time (main) | execute-time (PR) | Δ% | prove-time (main) | prove-time (PR) | Δ% | peak-ram (main) | peak-ram (PR) | Δ% | proof-size (main) | proof-size (PR) | Δ% | verify-time (main) | verify-time (PR) | Δ% | fft-cost (main) | fft-cost (PR) | Δ% |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
Nat.add_comm + String.append |
3.574 s | 3.606 s | +0.9% | 36.570 s | 36.539 s | -0.1% | 55.31 GiB | 55.13 GiB | -0.3% | 2.83 MiB | 2.83 MiB | +0.0% | 17.1 ms | 17.3 ms | +1.2% | 108.96B | 108.96B | +0.0% |
Pipeline total (7 constants)
| constant | total-time (main) | total-time (PR) | Δ% | pipeline-throughput (const/s) (main) | pipeline-throughput (const/s) (PR) | Δ% | pipeline-peak-ram (main) | pipeline-peak-ram (PR) | Δ% |
|---|---|---|---|---|---|---|---|---|---|
ByteArray.utf8DecodeChar?_utf8EncodeChar_append |
1m 29.1s | 1m 29.7s | +0.6% | 31.150 | 30.950 | -0.6% | 73.23 GiB | 73.22 GiB | -0.0% |
Char.ofOrdinal_le_of_le |
1m 18.3s | 1m 18.5s | +0.4% | 35.310 | 35.180 | -0.4% | 62.33 GiB | 62.19 GiB | -0.2% |
Array.extract_append |
1m 13.6s | 1m 14.0s | +0.5% | 21.810 | 21.700 | -0.5% | 53.91 GiB | 54.01 GiB | +0.2% |
Std.HashMap |
58.438 s | 58.771 s | +0.6% | 34.940 | 34.750 | -0.5% | 51.42 GiB | 51.40 GiB | -0.0% |
_private.Init.Data.Range.Polymorphic.SInt.0.Int64.instRxcHasSize_eq |
57.560 s | 57.831 s | +0.5% | 32.440 | 32.280 | -0.5% | 51.03 GiB | 51.16 GiB | +0.2% |
String.append |
33.794 s | 33.998 s | +0.6% | 9.680 | 9.620 | -0.6% | 47.23 GiB | 46.90 GiB | -0.7% |
Nat.add_comm |
32.032 s | 32.133 s | +0.3% | 1.440 | 1.430 | -0.7% | 45.94 GiB | 45.90 GiB | -0.1% |
- Base hardware (base run @ b4fea97 (fresh — bencher bypassed)): CPU:
Intel(R) Xeon(R) 6975P-C· vCPUs:32· RAM:247.7 GiB - PR hardware: CPU:
Intel(R) Xeon(R) 6975P-C· vCPUs:32· RAM:247.7 GiB
…s as a feature Pins multi-stark at `7b3699b`, the head of its `sb/trace-sharding-gpu` branch (PR #81): the sppark-backed NTT/LDE device backend and the resident multi-device prover the lanes build on. `cuda-trace-codegen` is a Cargo feature over `cuda` that compiles the generated CUDA trace writers; `IX_CUDA_TRACE_CODEGEN=1` selects it through Lake, whose archive target now traces the CUDA sources and the trace manifest. `texray` becomes an optional feature of `aiur`, on by default. The Nix Lake source is the union of lean4-nix's `cleanLakeSource` and Crane's `cleanCargoSource`, minus the zkVM guest workspaces, so the files Lake traces and the files Crane compiles cannot drift apart. CI checks the generated trace bundle the way it checks the executors.
…etrics `ix prove --lanes N` builds one resident prover per device from the two Aiur systems and schedules every claim and join over them through one bounded execution queue, joins first. Record reservations follow their storage through both proving rounds; under memory pressure executions wait for capacity, and mutual blocking cancels the youngest waiter. Oversized environment claims are bisected and checkpointed, completed proofs are reused, and `--out-ixes` records the final partition. The GPU trace runtime packs seeds through per-device staging leases, keeps round-two seeds resident, and can build BLAKE3 traces on the device (`AIUR_GPU_TRACE=blake3`). `AIUR_METRICS` writes one JSON line per proving piece and execution with the selected environment knobs, cheap enough to leave on for a whole Mathlib run. Store writes publish through a per-process temporary name so concurrent lanes never expose a partial object.
A trace plan per constrained function, derived from the compiled bytecode: the SSA values each row needs, the seed words an execution records for it, and the column spans it writes. From the plan the compiler emits Rust seed packers and row writers and the matching CUDA row writers, plus a registry the GPU trace runtime binds against after checking the program fingerprint. `ix codegen --trace-bundle` writes the production units for the three programs and `--check` detects stale ones, as the executor check does; `--trace-report` prints the plan inventory. Fixtures and BLAKE3 units for the Rust and CUDA parity tests are generated the same way. The Rust renderer leaves expressions bare where the grammar already delimits them and elides unit return types, which the generated writers need to pass clippy.
How the GPU build is layered and built, the four-command pipeline for a whole environment on several devices, how the lanes scheduler governs memory, what the generated traces are and how to regenerate and test them, every knob a run or a measurement needs, and the benchmark procedure with the Mathlib and FLT results it produced. The design discussions, dated run reports and planning documents it condenses stay on `sb/aiur-gpu-lanes-sp1-compress`, which the document and the handoff point to.
The design document keeps its argument (what is sharded, the shared challenge, memory across shards, soundness, the GPU sizing rule, recursion, proof size, open questions) and gains a status paragraph in place of the plan-and-status table, the parked single-claim design and the one-off measurement runbook, all of which described a branch state that has since merged. The performance audit of PRs 621 to 625, the dated results log, the PR handoff and the bencher proposal were working notes for that merge and are removed; the history keeps them.
The lanes scheduler, the generated trace writers, and the codegen CLI carried instrumentation and variants that nothing in the proving path uses: - The `prover_metrics` tracing subscriber duplicated the texray span tree as flat key/value lines. Its emit sites in the shard prover, the GPU trace bridge, and the lanes scheduler are gone, with the `AIUR_METRICS`, `AIUR_METRICS_RUN_ID`, and `AIUR_PROFILE` knobs. The scheduler's calibration summary (record percentiles and a re-anchored shard count) went with it; the per-unit `[lanes]` lines already carry the record bytes and execution times it aggregated. - `ix codegen --trace-report` rendered a JSON digest of the trace plans that no tool consumed. - The `pack_checked` seed packer re-ran every alias check on the host while the GPU bridge always passed `check_aliases = false`. The plan no longer tracks alias roots, the emitter has three modes (`pack`, `typed`, `row`), and `SeedContext::check_returned` has no caller. The `write` host writer and the schema contract stay: the first is the CPU materialization path, the second catches equal-size layout reorders that the binding's size checks cannot. - `run_manifest` is split into `plan_pass` (plan, host budget, resident provers), `resume_pass` (cached joins and indexed claims), and a `Scheduler` whose `next_work`, `queued`, and `handle` methods carry the dispatch state, leaving the thread scope and the event loop in `run_manifest`. Behaviour is unchanged. Generated programs and the trace-plan tests are regenerated and updated for the three-mode emitter.
multi-stark acfc370a is main after the GPU acceleration PR (#81) was squashed in. The pin moves there; Cargo.lock picks up its sppark revision and the `which`/`rustix`/`home` chain sppark's build script pulls in. The API it dropped is absorbed on the ix side: - `TraceGenerator::host_bytes` left the trait. The CUDA tests' seed accounting keeps it as a test-only inherent method on `CudaTrace`. - `CudaDft::generated_trace_rows`, the tile download the parity tests drove the device writers through, is gone, and `DeviceTraceView` has no constructor outside multi-stark. The writers are now inherent `write_tile` methods behind thin `write_device_rows` adapters, and a test-only `download_tile` runs them on a device buffer from ix's own runtime (`aiur_trace_tile_alloc`/`aiur_trace_tile_download`) and copies the rows back as raw words, which keeps the noncanonical-output check the tests had. - `prepare_trace` and `prepare_memory_trace` return the concrete traces so tests can reach them; `prepare` and `prepare_memory` wrap them. - `generate_coset_lde` is private upstream, so the leak test that drove it through a failing writer goes. - `MULTI_STARK_CUDA_TRACE_FORCE_SPILL` and `MULTI_STARK_CUDA_LOOKUP_TRACE_TILE_ROWS` no longer exist upstream; their row leaves the knob table. Verified with CI's check/clippy/test commands on the CPU path, and the `cuda-trace-codegen` code type-checked and clippied through a stand-in nvcc, since no CUDA toolchain is available here and CI never builds that feature.
CI runs `cargo check` on the `cuda` feature and never builds `cuda-trace-codegen`, so this code had drifted from the rest of the crate: a `pack_row` call in the BLAKE3 timing test still passed the removed checked-packing flag, and the workspace's cast, slice and pass-by-value lints flagged a handful of sites. Found by driving `cargo clippy -p aiur --features cuda-trace-codegen --all-targets` through a stand-in nvcc. `register` now consumes the bytecode it hands to the library instead of cloning it; the row count in the wide-row test is a `usize` like every other count; the release-only guard on the opt-in timing test is part of its ignore reason rather than a constant assertion. The generated program fixtures (`programs/ixvm.rs`, `programs/multi_stark.rs`) still trip `match_same_arms`, `match_single_binding` and `eq_op` on constructs that mirror the bytecode they render; those need the Lean renderer and stay as they are.
The IxVM kernel and aggregator sources come out of the renderer this branch ships, over the Ixon v4 bytecode on main; the generated trace writers, CUDA units and manifest follow the new bytecode fingerprints. `get_tag4` no longer exists in the kernel, so the measured weighted selection names its TagN successor `get_tagn4`; the selection is still the one profiled on the v3 kernel.
samuelburnham
force-pushed
the
sb/aiur-gpu-lanes
branch
from
October 2, 2026 16:20
52eb7d7 to
0379bdd
Compare
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
2 participants
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.
GPU proving: lanes, generated CUDA traces, sppark, SP1 terminal (PR 2 of 2)
After #643. Everything from
sb/aiur-trace-sharding-gputhatneeds a device: multi-GPU lanes, the GPU trace runtime and generated CUDA
trace writers, the sppark-only backend, and the SP1 terminal.
Depends on multi-stark
sb/trace-sharding-gpu(pinned at59df87a, ontop of
sb/batch-proving).Changes
ix prove --lanes Nbuilds one resident prover per device,shares bytecode across the FFI, schedules claims and joins over a shared
execution pool with bisection and checkpoints;
ix aggregate --subtree.round-two seeds, opt-in BLAKE3 seeds, lightweight metrics (
AIUR_METRICS).row writers; CUDA row writers and registries;
ix codegen --trace-bundlewith a manifest CI checks. Emitted Rust is lint-clean by construction.
ix compress-rooton batch roots, Plonk by default,Groth16 available, pinned compatibility guest for the 2026-09-03 root.
READMEs and summaries only.