feat: add Stage 1 multi-GPU base proving - #1403
Draft
hero78119 wants to merge 12 commits into
Draft
Conversation
hero78119
marked this pull request as draft
September 4, 2026 08:43
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
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.
Problem
Ceno's GPU prover relied on process-global CUDA state and could only prove base shards on one device. Device selection, memory validation, replay ownership, and cross-worker failure handling were therefore implicit.
Design Rationale
Each GPU worker now owns an
Arc<CudaHalBB31>and binds its device context and stream to the worker and chip-lane threads. This isolates device-backed state and makes constructor-selected device IDs authoritative.Stage 1 runs one independent AOT replay producer and prover worker per GPU. Deterministic round-robin ownership (
shard_id % device_count) assigns proving work: owned shards emit compact witnesses; unowned shards fast-flight through execution, preserving VM state and replay digests without materializing compact rows, fallback records, or syscall witnesses. Depth-one per-worker FIFOs bound memory and provide backpressure, while failures cancel the operation without retrying on another GPU.Proofs may finish out of order. The coordinator validates ownership and completeness, restores canonical shard order, then performs one full-trace Rust verification. Recursion remains unchanged and starts on one configured GPU only after all base workers finish.
Change Highlights
gkr_iop: worker-owned HALs and scoped context/stream bindings; explicit lane streams; reject unbound fallback.ceno_zkvm:MultiGpuConfig, device discovery/validation, conservative common shard limits, replay/prover workers, round-robin collection, fail-fast diagnostics, metrics, and optional CPU affinity.ceno_emul: switch AOT compact capture at canonical shard boundaries while preserving execution and digest identity.Benchmark / Performance Impact
Operation
Layer
Apparent two-GPU efficiency is 82.32%. Relative to the naive half-baseline of 21.331340 s, the current 25.914000 s proof-ready wall leaves 4.582661 s attributable to serial/non-shard work, six-even/five-odd round-robin imbalance, duplicated replay/resource contention, and build/timing drift. Fast-flight replay largely overlaps proving and is not counted twice.
The current-code rerun performs exactly one measured post-base GPU0 pool handoff after canonical verification and before recursion. The handoff takes 0.424 s and frees 22.0625 GiB. Against the prior dual run, recursion falls from 6.497638 s to 6.146068 s and the first
MainFrontloadTermAirallocation wrapper falls from 572.899 ms to 183.007 ms, confirming the missing-handoff diagnosis. Total application wall is 184 ms higher because the handoff is now paid explicitly and proof-ready varied by 147 ms in this single sample.This is directional, not a strict A/B. Both runs use the same runner, frozen block input, 11-shard configuration, and toolchain, but differ in benchmark/Ceno/ceno-gpu revisions, guest ELF, CUDA packaging, and exact instruction counts and shard boundaries. A same-binary device-0 versus device-0,1 pair remains required for formal scaling.
Benchmark command(s):
Both use
max_cell_per_shard=4500000000, lanes-4, cache level 1, a 3048 MiB booking margin, and jagged reshape height 23.Environment (CPU/GPU, core count, rust toolchain, commit hash):
scroll-SEA2-10; CPU model/core count were not captured and affinity was unset.32a12d5131a26f68bea1a97d6da699432299c019.f7c630927423fce51bded0bdf90851095aa27b96, ceno-gpu773c74d546238d2faa9ec83faf49eb1eef51cace,CUDA_ARCH=89,120.1.93.0-nightly (07bdbaedc 2025-11-19); Cargo1.93.0-nightly (2d4fa1395 2025-11-12).raw data:
Testing
Risks and Rollout
Follow-ups (optional)
0, then0,1.Copilot Reviewer Directive (keep this section)
When Copilot reviews this PR, apply
.github/copilot-instructions.mdstrictly.