Skip to content

feat: add Stage 1 multi-GPU base proving - #1403

Draft
hero78119 wants to merge 12 commits into
masterfrom
feat/multi_gpu
Draft

feat: add Stage 1 multi-GPU base proving#1403
hero78119 wants to merge 12 commits into
masterfrom
feat/multi_gpu

Conversation

@hero78119

@hero78119 hero78119 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

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.
  • CLI/SDK: explicit device list/count, recursion device, and worker CPU sets; logical GPU 0 remains the default. GPU builds require AOT.

Benchmark / Performance Impact

Operation

Operation master (s) this PR (s) Improve (master -> this PR)
Base proofs ready, before verification/recursion 42.662679 25.914000 39.26% lower wall; 1.6463x throughput
Total application proof 48.488613 33.674897 30.55% lower wall; 1.4399x throughput

Layer

Layer master (s) this PR (s) Improve (master -> this PR)
AOT preflight 1.528036 1.524144 0.25% faster
Base proof vector ready 42.662679 25.914000 39.26% faster
Recursion proving 5.759648 6.146068 6.71% slower

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 MainFrontloadTermAir allocation 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):

# baseline: run 32964242819
target/release/ceno-reth-benchmark-bin --mode prove-stark --block-number 23817600 --output-dir output --cache-dir rpc-cache --chain-id 1

# current dual GPU: run 34038286261
target/release/ceno-reth-benchmark-bin --mode prove-stark --block-number 23817600 --output-dir output --cache-dir rpc-cache --gpu-devices 0,1 --chain-id 1

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):

  • Runner scroll-SEA2-10; CPU model/core count were not captured and affinity was unset.
  • Baseline: one RTX 4090, Ceno 32a12d5131a26f68bea1a97d6da699432299c019.
  • PR: two RTX 4090s, Ceno f7c630927423fce51bded0bdf90851095aa27b96, ceno-gpu 773c74d546238d2faa9ec83faf49eb1eef51cace, CUDA_ARCH=89,120.
  • Rust 1.93.0-nightly (07bdbaedc 2025-11-19); Cargo 1.93.0-nightly (2d4fa1395 2025-11-12).

raw data:

Testing

cargo fmt --all --check
CUDA_ARCH=89,120 cargo test -p ceno_zkvm --features gpu,aot-x86_64 multi_gpu
cargo test -p ceno_emul --release --features aot-x86_64,u16limb_circuit gpu_replay_fast_flight_preserves_state_and_next_compact_shard -- --nocapture
CUDA_ARCH=89,120 cargo check -p ceno_zkvm --features gpu,aot-x86_64
cargo make clippy
  • Selection, ownership, bounded-queue, cancellation, affinity, replay-digest, collection, and CUDA-binding regression tests pass.
  • Local one-GPU two-shard AOT proving, base verification, recursion, and root verification pass.
  • Current-code dual-RTX-4090 E2E passes for all 11 round-robin shards, compact/fast-flight digest pairs, canonical base verification, exactly one measured post-base GPU0 pool handoff, recursion, and root verification, with an empty failure scan.

Risks and Rollout

  • Stage 1 duplicates replay on every GPU and uses fixed round-robin ownership, so uneven shard costs reduce scaling.
  • HAL/context/stream isolation is correctness-sensitive; focused concurrency tests and the dual-GPU E2E cover it.
  • Device capability, HAL construction, and conservative usable memory are validated before preflight.
  • The default remains logical GPU 0. No proof-format or recursion-protocol migration is introduced.

Follow-ups (optional)

  • Measure formal scaling with one frozen current binary on device 0, then 0,1.
  • Stage 2: shared replay and cost-aware assignment. Stage 3: multi-GPU recursion.
  • Retries and P2P transfer remain out of scope.

Copilot Reviewer Directive (keep this section)

When Copilot reviews this PR, apply .github/copilot-instructions.md strictly.

@hero78119
hero78119 marked this pull request as draft September 4, 2026 08:43
@hero78119 hero78119 changed the title perf: [stage 1] multi-gpu base layer pipeline proving feat: add Stage 1 multi-GPU base proving Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant