feat: GPU trace gen and sppark NTT backend - #81
Merged
Merged
Conversation
samuelburnham
added this pull request to stack #82
September 22, 2026 18:36
samuelburnham
force-pushed
the
sb/trace-sharding-gpu
branch
2 times, most recently
from
September 29, 2026 12:31
ed8eefa to
7b3699b
Compare
…nesses Bind a GoldilocksBlake3Config to a CUDA device (GoldilocksBlake3Config::with_device, CudaMmcs::with_device, CudaDft::new taking the device) instead of reading MULTI_STARK_CUDA_DEVICE once per process, so one process can drive one prover per GPU. The host-side pools that were process-wide (the pinned host pool, the upload staging slots, the resident-LDE control-block slab) are now one per device, indexed by the calling thread's current device, which every kernel entry point sets first. Add the witness module (TraceSource, PreparedWitness, TraceGenerator) and the device trace view so a circuit can generate its trace on the GPU, with the lookup-graph LDE writing device rows directly. build.rs: compile the kernels' host code with the POSIX and default feature sets instead of _GNU_SOURCE, so glibc 2.38+ keeps the plain strto*/clock_gettime names; the ix binary is linked by Lean's bundled clang against an older sysroot that lacks the __isoc23_* variants. Claude-Session: https://claude.ai/code/session_01WD1WGwm74eqamotwg34pZt
…ssion The bounded tree cache (TreeCheckpoint, Retention::MerkleTrees, checkpoint_main, prove_stage_1_restored, trim_tree_cache, tree_cache_bytes and the cache_headroom reservation) is removed. Its headroom rule reserved four times the main LDE plus the main trace plus the minimum-free allowance, about 95 GiB of a 96 GiB device for one 2^20-row BLAKE3 matrix, so it evicted every retained tree before reuse; with the accounting fixed (each trace/LDE construction admitted on its own, existing cache bytes not charged twice, optional trees released before required LDEs spill) a 24 GiB allowance reused all 91 trees on the eight-claim Init fixture and still did not move end-to-end time (290.06 s against a 284.54 s regenerate median), because hashing is a small slice of the proof. Regenerate is the one cross-round policy. What stays: the generated main commitment admits each trace and LDE construction separately from Merkle hashing, and LDE spills are counted per device (record_lde_spill) so a run can report them. Measurements and the archived experiment: ix bench/tree-cache-2026-09-15 and bench/gpu-trace-regenerate-2026-09-15.
The BLAKE3 leaf kernel hashes one message of at most 32 KiB per row, and
a mixed tree's leaf row at a height is every matrix of that height side
by side, so a height group wider than 4,096 Goldilocks columns made
multi_stark_cuda_mixed_merkle_create reject the tree ("invalid argument")
and the prover panic in stage-one commit. ix's Aiur test programs hit it
(a 9,282-column BLAKE3 hashing circuit at height 4); the production IxVM
circuits are narrower, so no proof run had.
Such a group is now hashed on the host and enters the tree as a digest
group, the way a group spilled for memory already does:
- commit_cuda_resident, the sink of every all-resident and spillable
commit, copies the group's LDEs back to the host and delegates to
commit_cuda_hybrid with hash_host_only_height_groups' digests;
- the PCS placement under memory pressure never makes such a group
durable or transient, so its LDEs are computed on the host directly;
- the generated-trace commit spills the group whole, like a memory spill.
host_hashed_heights (MAX_DEVICE_LEAF_ROW_BYTES) is the one rule all three
share, with a unit test; the CUDA suite (72 tests) passes.
Export the field helpers in a shared CUDA header and expose its include directory through Cargo dependency metadata. Keep backend kernels on the same arithmetic implementation.
Lookup admission budgeted the graph kernel for every resident LDE while execution required the raw trace as well, so a spilled or trace-released matrix was planned as regenerable and then evaluated on the slow path. `CudaMmcsData::has_resident_with_trace` is now the one predicate behind admission, the spill sort key and the re-budget after a spill, and the chosen path is passed into `evaluate` instead of being recomputed there. Under `MULTI_STARK_CUDA_TRACE_FORCE_SPILL` every lookup job now logs `graph=true`.
…up job A generator may hold device memory that serves its tiles faster, such as seeds kept resident between a shard's commitment and the lookup pass that regenerates rows from them. `TraceGenerator::release_device` frees it; the lookup path calls it after each LDE's job, and admission calls it before any LDE is spilled or evicted, in the commit loop and in both headroom paths, since a cache is rebuilt from host seeds on demand while a spilled LDE has to be uploaded again.
Every resident LDE construction records its kind, height, width and blowup; the commit loop records each source and the Merkle build; the streamed FRI opening records preparation, interpolation, observation, reduction, the per-round commitments and the query phase. A CUPTI profile joined to these spans attributes kernel time per transform shape and idle time per phase.
The streamed FRI opening rebuilt, bit-reversed and converted the coset of the largest LDE height for every shard and uploaded it again for the denominators: 0.32 s of host time per 2^26 shard. The bit-reversed coset of a smaller height is a prefix of a larger one's, so one vector per process serves every height, and the workspace takes it from the device constant cache instead of owning a copy. A staged upload copied each 64 MiB chunk into the pinned slot, transferred it and waited before the next. The slot is now a ring of 16 MiB chunks with per-chunk events, so the host copy of one chunk overlaps the transfers before it, and the host copy of a chunk is split across four pthreads since one thread fills pinned memory well below the PCIe rate. On the profiled Init claim the proof went from 44.9 s to 40.3 s: FRI opening 5.7 s to 2.9 s, host-trace commit idle 5.8 s to 4.2 s.
Add opt-in process/device counters for staged uploads, coset caching, memory samples and NTT shapes. Emit cumulative snapshots at piece boundaries and annotate coarse DFT/LDE operations and admission decisions. Metrics introduce no CUDA events or synchronization. Validated with the focused aiur release check using cuda-trace-codegen and the local backend override.
`cuda-sppark` adds the pinned sppark crate and compiles a device-pointer adapter with the upstream GPU runtime into the kernel archive, in their own nvcc step because that runtime needs libstdc++'s `<mutex>`, which the kernels' host flags would hide. Each transform runs on an upstream stream fenced to the caller's per-thread stream by events, so the completion contract matches the first-party kernels. The contract tests pin what the prover will rely on: the natural forward transform is the CPU DFT, the reversed orders are its bit-reversal, the inverse is normalized upstream so an inverse-forward pair needs no scaling, the coset variant uses the field generator, canonical extremes come back canonical, and the compiled domain limit covers 2^26 rows. Upstream takes canonical words only; representatives at or above the modulus transform to different values.
The fork's `SPPARK_NO_CXX_RUNTIME` mode records a failing CUDA call in a per-thread status instead of throwing, and drops the thread pool from the GPU handle, so the sppark objects reference nothing from libstdc++ and the archive links into the Lean executable, which carries libc++. The adapter reads the recorded status after each transform. The dependency now names the organization's fork; the pinned revision moves to the fork's branch once it is pushed.
… releasing its stream Upstream keeps only the devices it supports, so its logical index can differ from the CUDA ordinal; the adapter now finds the entry by ordinal and rejects an unknown one before touching the buffer, which a test covers. The caller's stream waits on the transform while the private stream is still alive, and if that wait cannot be installed the private stream is drained before returning, so no queued kernel outlives the call. With the fork stopping at the first failing CUDA call, the recorded-status read is gone. The pin moves to the fork's fail-fast commit once pushed.
With `cuda-sppark`, MULTI_STARK_CUDA_NTT=sppark (or a runtime selection for one-process comparisons) routes `coset_lde_create` through a column panel: a gather that canonicalizes into column-major scratch, an inverse transform per column, one pass that restores natural coefficient order and applies the coset powers into a second panel with a zeroed tail, a forward transform per column, and a scatter into the bit-reversed row order the commitment expects. Panels are sized from MULTI_STARK_SPPARK_PANEL_BYTES. The stored words match the first-party kernels bit for bit across the tested heights, widths, blowups and shifts, raw representatives included. Medians of five warm resident LDEs on one RTX PRO 6000: 2^24 x 6 at blowup 4 from 170 to 71 ms, the 2^22 width-2 codeword from 9.7 to 4.5 ms, the 2^20 x 533 BLAKE3 shape from 386 to 315 ms; 2^24 x 17 even; short wide matrices slower, 2^16 x 925 from 25 to 57 ms, on per-column launches.
…s per backend The commit loop adds the two panels the sppark path allocates to a source's admission. Below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (20) a resident LDE stays on the first-party kernels, where the per-column baseline is launch-bound and slower; the selection has a third state that takes sppark at every height for one-process comparisons. The metrics snapshot counts taken and declined dispatches and keeps transform shapes per backend.
…shold The lookup LDEs (graph, direct and partitioned), the quotient's two forward transforms, the general DFT and the host coset LDE dispatch through two helpers, one for the inverse-shift-forward sequence in place and one for a forward transform in place, so the backend rule lives in one place: sppark's column panel above MULTI_STARK_SPPARK_MIN_LOG_HEIGHT, the first-party stages otherwise. The adapter gains the forward panel entry and an ungated transform counter. A batch proof and the compatibility example produce identical bytes on the first-party kernels, on sppark above the threshold and on sppark at every height; the general DFT and host coset LDE entries are compared bit for bit as well.
nvcc's host compiler builds the adapter unit without the C standard pin the crate's C units get, so glibc redirected strtoul to a C23 symbol the Lean toolchain's libc lacks and the ix executable no longer linked.
…n spans The dispatch rule takes the complete shape: tall enough, the extended height within upstream's compiled domain, one column's scratch within the panel budget; a shape outside any of these stays on the first-party kernels instead of failing after allocation. Lookup and quotient admission add the sppark panels their transforms allocate, as the commit loop already did. The paths that go through sppark no longer build or upload the first-party twiddle tables, in the prewarm and at the lookup and quotient sites. The backend selector is an atomic with a single publication of the environment setting, the comparison tests serialize on a lock so both constructions of a comparison see the backend they selected, and the LDE, lookup and quotient spans label the backend they ran on.
The fork's batched entry runs a group of panel columns through one launch sequence per stage. A group is sized to the device's L2 cache (MULTI_STARK_SPPARK_BATCH_BYTES overrides it, 0 launches every column on its own): short columns share launches, and a group of tall columns keeps the reuse between stages that one column at a time had. The coefficient panel is compact, so an LDE's scratch is (N + M) words per column rather than 2M. The build script reruns when the fork's NTT or runtime sources change. On one RTX PRO 6000 the resident LDE medians against the unbatched baseline: 2^16 x 925 x2 52.4 to 20.8 ms, 2^18 x 128 x2 15.8 to 11.7 ms, 2^20 x 533 x4 343 to 336 ms, the tall narrow shapes unchanged; every shape now at or ahead of the first-party kernels.
…ench With the columns batched, the resident LDE through sppark is ahead of the first-party kernels from 2^18 input rows up (0.97 to 1.19x at 2^18, 1.07 to 1.22x at 2^19) and behind below (0.64 to 0.93x from 2^8 to 2^17), so the default threshold comes down from 2^20. The bench covers those short shapes now.
…ured The gather, expansion and scatter passes around the transforms were per-element kernels with a 64-bit division each and strided accesses; they took 227 of the 349 ms of the BLAKE3 shape's LDE against 122 ms of transforms. The gather and scatter are shared-memory tiled transposes now, or row-per-thread kernels for matrices narrower than eight columns, and the expansion runs one column per grid row. The admission counts the reversed coset powers. The fused expansion of the plan (bit-reversed coefficients spread with the coset powers straight into the forward transform in RN order, the row permutation folded into the scatter) is implemented behind MULTI_STARK_SPPARK_FUSED and measured: equal to the restoring path on the wide shape and 15 to 18 percent behind on the tall ones, since its reversing scatter costs what the restoring pass saves, so the restoring path stays the default. MULTI_STARK_SPPARK_STAGE_TIMING prints the per-stage times of every LDE; the resident bench takes a shape list. Resident LDE medians on one RTX PRO 6000 against the first-party kernels: 2^20 x 533 x4 385 to 284 ms, 2^24 x 6 x4 169 to 63 ms, 2^22 x 2 x4 9.5 to 4.3 ms, 2^18 x 128 x2 13.9 to 11.4 ms; the shapes below 2^18 rows stay on the first-party kernels.
…end only The Rust side decides a transform's backend once, by the shape, and builds the first-party twiddle tables only for the first-party path; the kernels take the sppark path exactly when no table was handed to them, so a backend switch between the two sides cannot split a transform, and the generic DFT and LDE entries neither build nor upload tables they will not use. Their spans label the backend that ran. The batched host entry checks its word count and stride representability before the FFI, and the adapter rejects a stride beyond upstream's 32-bit index. The bit-reversed feed's reversed coset powers are allocated only when that expansion is selected and count against the panel budget with the panels.
The panel scratch was a cudaMalloc per LDE freed asynchronously: a device synchronization per transform, and a pairing on which CUPTI 2026.2.1's memory tracking faults (an async free of memory that did not come from a pool dereferences the pool it expects, with the MEMORY2 activity kind on, whatever the stream). Pool allocations on the caller's stream remove both: 2^24 x 6 x4 61 to 55 ms, 2^22 x 2 x4 4.1 to 3.5 ms, and the prover profile collector records the sppark path again.
Share immutable panel plans between admission and execution, and run the fork on borrowed caller streams. Remove the first-party GPU NTT, backend selection, twiddle plumbing, and unused expansion experiment. Keep the generic CPU DFT and compare resident results and proof bytes against CPU references.
Reject mismatched heights, widths, expansions, and forward/coset plans before the shared NTT helpers launch. Keep raw plans private and non-copyable, borrow them from their owner, and test that rejected calls preserve host output and do not return resident storage. Refresh the backend descriptions and restore-stage timer label.
Nothing read the per-device counters, the prover_metrics tracing events or the sppark transform counter, so they go, together with the AIUR_METRICS, MULTI_STARK_SPPARK_BATCH_BYTES, MULTI_STARK_SPPARK_STAGE_TIMING, MULTI_STARK_CUDA_TRACE_FORCE_SPILL and MULTI_STARK_CUDA_LOOKUP_TRACE_TILE_ROWS settings. Launch groups always follow the device's L2 size. The two dated experiment docs described the first-party NTT this branch replaced, and the standalone BLAKE3 row bench no longer linked once kernels.cu called into sppark. The short-row dispatch boundary is noted in cuda/README.md instead. Also removes the unused generated-tile download, the device-pointer NTT entry, a dead bit-reversal helper and the TraceGenerator::host_bytes method, parses MULTI_STARK_CUDA_MIN_FREE_BYTES in one place, and makes the sppark module crate-private with its host contract entries gated to tests.
samuelburnham
force-pushed
the
sb/trace-sharding-gpu
branch
from
October 2, 2026 13:31
7b3699b to
c9532e3
Compare
arthurpaulino
approved these changes
Oct 2, 2026
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.
CUDA: per-device provers, GPU trace witnesses, metrics, sppark transforms
After #80.
The device side of trace-shard batch proving, consumed by ix's multi-GPU lanes and generated trace writers.
witnesses that keep device state through the lookup job; lookup graphs
admitted only for resident LDEs that kept their trace.
double-buffered; the FRI coset cached per process; the Merkle tree cache
removed in favour of per-construction admission; wide height groups
hashed on the host; fused wide NTTs; single-chunk BLAKE3 rows hashed one
thread per message; Goldilocks arithmetic shared with generated kernels.
sources and the FRI opening phases; device snapshots under event-only
tracing filters.
the stream's memory pool, dispatch by whole shape, batched panel
launches, then sppark for all CUDA transforms with plans validated
against buffer dimensions.