perf(gpu): overlap replay with shard proving - #1401
Merged
Merged
Conversation
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
GPU shard proving serialized compact CPU replay with the preceding shard proof. Replay ownership also crossed shard boundaries implicitly, and the packed producer priority could not represent the full supported producer domain.
Design Rationale
A capacity-one pipeline prepares shard N+1 on the CPU while shard N is proved. CUDA assignment remains on the proving thread and starts only after the preceding proof completes, preventing assignment kernels or transfers from overlapping sumchecks. Replay owners are retained and recycled explicitly so their storage remains valid across the boundary.
Producer base, count, and order remain separate on the host. The hot device domain contains only base and count; order is read only for terminal emission. A checked 64-bit priority preserves continuation class, instruction order, producer index, and slot through exactly
2^28rows without changing transcript or witness order.Change Highlights
ceno_emul: add compact replay ranges and checked producer-domain metadata.ceno_zkvm: overlap CPU replay with the preceding proof, retain/recycle owners explicitly, and gate each GPU assignment before proving its shard.Benchmark / Performance Impact
Operation
Layer
Benchmark command:
Environment: NVIDIA RTX 4090 (
CUDA_ARCH=89), GPU + jemalloc + AOT benchmark workflow. Baseline used Ceno436a32e2; this PR used6fddf6cb.Raw data:
Testing
cargo fmt --all -- --check cargo check -p ceno_zkvm --all-targets RUST_MIN_STACK=33554432 cargo test --workspace --lib --bins --tests --examples --no-default-features --features goldilocks --no-run RUSTFLAGS=-Dwarnings cargo check --workspace --all-targetsBoth benchmark workflows completed successfully and verified block
23817600.Risks and Rollout
The main risk is cross-shard lifetime or stream-ordering regression. Bounded channels, explicit owner recycling, checked domain bounds, and a publication barrier make these boundaries explicit. Reverting the PR restores serialized replay without data migration or protocol changes.
Copilot Reviewer Directive (keep this section)
When Copilot reviews this PR, apply
.github/copilot-instructions.mdstrictly.