fix: root layer tracegen and proving on the same GPU context#2914
fix: root layer tracegen and proving on the same GPU context#2914stephenh-axiom-xyz wants to merge 1 commit into
Conversation
Note: cells_used metrics omitted because CUDA tracegen does not expose unpadded trace heights. Commit: 958898b |
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. |
| let root_prover = RootProver::new::<RootEngine>( | ||
| internal_recursive_vk, | ||
| internal_recursive_pcs_data.commitment.into(), | ||
| root_system_params(), | ||
| system_config.memory_config.memory_dimensions(), | ||
| system_config.num_public_values, | ||
| None, | ||
| None, | ||
| ); | ||
| let engine = RootEngine::new(root_prover.get_vk().inner.params.clone()); | ||
| let engine = root_prover.create_engine::<RootEngine>(); | ||
| let ctx = root_prover.generate_proving_ctx_no_def::<<RootEngine as StarkEngine>::PB, _>( | ||
| internal_recursive_proof, | ||
| &user_pvs_proof, | ||
| &engine.device().device_ctx, | ||
| ); | ||
| let root_proof = root_prover.root_prove_from_ctx::<RootEngine>(ctx.unwrap())?; | ||
| let root_proof = root_prover.root_prove_from_ctx::<RootEngine>(ctx.unwrap(), &engine)?; |
There was a problem hiding this comment.
shouldn't we just create an engine in the new constructor and use it internally for all the root prover functions. why pass it around?
There was a problem hiding this comment.
not sure about this but if engine is an abstraction common among all types of provers, maybe it should be an argument to the constructor itself and each prover can either store a reference or a copy (if safe)
There was a problem hiding this comment.
I think we try to avoid storing Engines, especially in structs that are potentially long-lived (i.e. the RootProver will exist for as long as the Sdk does). They contain device-specific stuff
There was a problem hiding this comment.
okay, and is the idea that the same prover can be used on multiple devices?
2 participants
Resolves INT-8503.
Overview
Fix root proof generation so root trace generation and proving use the same caller-owned
StarkEngineinstance. This keeps GPU-backed root proving on a single device context instead of constructing a second engine between trace generation and proof creation.Changes include:
RootProver::create_enginehelpers in the continuations prover and SDK wrapper.generate_proving_ctx,prove_from_ctx, andprove.