Skip to content

[follow up 1299]: fix gpu build && optimize witness build and avoid GPU OOM #1316

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kunxian-xia merged 3 commits into from
Apr 21, 2026
Merged

[follow up 1299]: fix gpu build && optimize witness build and avoid GPU OOM #1316

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cfdab05
follow up 1299: fix gpu build
hero78119 Apr 20, 2026
a72e4ac
add gkr output class
hero78119 Apr 20, 2026
c8760fd
Merge branch 'master' into bug/fix_gpu_build
hero78119 Apr 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/bitwise_keccakf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -910,6 +910,7 @@ pub fn run_keccakf<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + 'stat
&[],
&[],
&[],
None,
);
exit_span!(span);

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/fptower/fp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,7 @@ mod tests {
&[],
&[],
&challenges,
None,
);

let out_evals = {
Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/fptower/fp2_addsub.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,7 @@ mod tests {
&[],
&[],
&challenges,
None,
);

let out_evals = {
Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/fptower/fp2_mul.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -465,6 +465,7 @@ mod tests {
&[],
&[],
&challenges,
None,
);

let out_evals = {
Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/lookup_keccakf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1128,6 +1128,7 @@ pub fn run_lookup_keccakf<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>
&[],
&[],
&challenges,
None,
);
exit_span!(span);

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/sha256/extend.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -470,6 +470,7 @@ mod tests {
&[],
&[],
&challenges,
None,
);

let out_evals = {
Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/uint256.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -825,6 +825,7 @@ pub fn run_uint256_mul<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + '
&[],
&[],
&challenges,
None,
);
exit_span!(span);

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/weierstrass/weierstrass_add.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,7 @@ pub fn run_weierstrass_add<
&[],
&[],
&challenges,
None,
);
exit_span!(span);

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/weierstrass/weierstrass_decompress.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -659,6 +659,7 @@ pub fn run_weierstrass_decompress<
&[],
&[],
&challenges,
None,
);
exit_span!(span);

Expand Down
1 change: 1 addition & 0 deletions ceno_zkvm/src/precompiles/weierstrass/weierstrass_double.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,6 +681,7 @@ pub fn run_weierstrass_double<
&[],
&[],
&challenges,
None,
);
exit_span!(span);

Expand Down
21 changes: 18 additions & 3 deletions ceno_zkvm/src/scheme/cpu/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,8 @@ use crate::{
hal::{DeviceProvingKey, EccQuarkProver, ProofInput, TowerProverSpec},
septic_curve::{SepticExtension, SepticPoint, SymbolicSepticExtension},
utils::{
assign_group_evals, derive_ecc_bridge_claims, extract_ecc_quark_witness_inputs,
GkrOutputStageMask, assign_group_evals, derive_ecc_bridge_claims,
extract_ecc_quark_witness_inputs, first_layer_output_group_stage_masks,
infer_tower_logup_witness, infer_tower_product_witness, split_rotation_evals,
},
},
Expand Down Expand Up @@ -909,11 +910,13 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<C
panic!("empty gkr circuit")
};
let first_layer = gkr_circuit.layers.first().expect("empty gkr circuit layer");
let group_stage_masks = first_layer_output_group_stage_masks(composed_cs, gkr_circuit);
let selector_ctxs = first_layer
.out_sel_and_eval_exprs
.iter()
.map(|(selector, _)| {
if cs.ec_final_sum.is_empty() {
.zip_eq(group_stage_masks.iter())
.map(|((selector, _), stage_mask)| {
if !stage_mask.contains(GkrOutputStageMask::TOWER) || cs.ec_final_sum.is_empty() {
SelectorContext {
offset: 0,
num_instances,
Expand Down Expand Up @@ -950,6 +953,13 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<C
else {
panic!("rotation proof provided for non-rotation layer")
};
debug_assert!(group_stage_masks[left_group_idx].contains(GkrOutputStageMask::ROTATION));
debug_assert!(
group_stage_masks[right_group_idx].contains(GkrOutputStageMask::ROTATION)
);
debug_assert!(
group_stage_masks[point_group_idx].contains(GkrOutputStageMask::ROTATION)
);

let (left_evals, right_evals, point_evals) =
split_rotation_evals(&rotation.proof.evals);
Expand Down Expand Up @@ -987,6 +997,11 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<C
else {
panic!("ecc proof provided for non-ecc layer")
};
debug_assert!(group_stage_masks[x_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[y_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[slope_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[x3_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[y3_group_idx].contains(GkrOutputStageMask::ECC));

let sample_r = transcript.sample_and_append_vec(b"ecc_gkr_bridge_r", 1)[0];
let claims = derive_ecc_bridge_claims(ecc_proof, sample_r, num_var_with_rotation)
Expand Down
17 changes: 3 additions & 14 deletions ceno_zkvm/src/scheme/gpu/memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::{
scheme::{
constants::{NUM_FANIN, SEPTIC_EXTENSION_DEGREE},
hal::ProofInput,
utils::tower_output_count,
},
structs::ComposedConstrainSystem,
};
Expand Down Expand Up @@ -184,20 +185,9 @@ pub fn estimate_main_witness_bytes<E: ExtensionField>(
composed_cs: &ComposedConstrainSystem<E>,
num_var_with_rotation: usize,
) -> usize {
let cs = &composed_cs.zkvm_v1_css;
let num_reads = cs.r_expressions.len() + cs.r_table_expressions.len();
let num_writes = cs.w_expressions.len() + cs.w_table_expressions.len();
let num_lk_num = cs.lk_table_expressions.len();
let num_lk_den = if !cs.lk_table_expressions.is_empty() {
cs.lk_table_expressions.len()
} else {
cs.lk_expressions.len()
};
let num_records = num_reads + num_writes + num_lk_num + num_lk_den;

let elem_size = std::mem::size_of::<BB31Ext>();
let record_len = 1usize << num_var_with_rotation;
num_records * record_len * elem_size
tower_output_count(composed_cs) * record_len * elem_size
}

pub(crate) fn estimate_main_constraints_bytes<
Expand Down Expand Up @@ -233,8 +223,7 @@ pub(crate) fn estimate_main_constraints_bytes<
// (see ZerocheckLayer verifier: max_degree = self.max_expr_degree + 1)
let main_sumcheck_degree = (layer.max_expr_degree + 1).max(1);

let total_mles =
layer.n_witin + layer.n_structural_witin + layer.n_fixed + layer.n_instance;
let total_mles = layer.n_witin + layer.n_structural_witin + layer.n_fixed;
let main_mle_num_vars_list = vec![num_var_with_rotation; total_mles];
let main_est = estimate_sumcheck_memory(
num_var_with_rotation,
Expand Down
25 changes: 15 additions & 10 deletions ceno_zkvm/src/scheme/gpu/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ use crate::{
DeviceProvingKey, MainSumcheckEvals, ProofInput, RotationProverOutput, TowerProverSpec,
},
utils::{
assign_group_evals, derive_ecc_bridge_claims, extract_ecc_quark_witness_inputs,
GkrOutputStageMask, assign_group_evals, derive_ecc_bridge_claims,
extract_ecc_quark_witness_inputs, first_layer_output_group_stage_masks,
split_rotation_evals,
},
},
Expand Down Expand Up @@ -75,13 +76,7 @@ use util::{

pub struct GpuTowerProver;

use crate::{
scheme::{
constants::{NUM_FANIN, SEPTIC_EXTENSION_DEGREE},
septic_curve::SepticPoint,
},
structs::EccQuarkProof,
};
use crate::scheme::{constants::NUM_FANIN, septic_curve::SepticPoint};
use gkr_iop::{
gpu::{ArcMultilinearExtensionGpu, BB31Base, MultilinearExtensionGpu},
selector::{SelectorContext, SelectorType},
Expand Down Expand Up @@ -321,11 +316,13 @@ pub fn prove_main_constraints_impl<
panic!("empty gkr circuit")
};
let first_layer = gkr_circuit.layers.first().expect("empty gkr circuit layer");
let group_stage_masks = first_layer_output_group_stage_masks(composed_cs, gkr_circuit);
let selector_ctxs = first_layer
.out_sel_and_eval_exprs
.iter()
.map(|(selector, _)| {
if cs.ec_final_sum.is_empty() {
.zip_eq(group_stage_masks.iter())
.map(|((selector, _), stage_mask)| {
if !stage_mask.contains(GkrOutputStageMask::TOWER) || cs.ec_final_sum.is_empty() {
SelectorContext {
offset: 0,
num_instances,
Expand Down Expand Up @@ -362,6 +359,9 @@ pub fn prove_main_constraints_impl<
else {
panic!("rotation proof provided for non-rotation layer")
};
debug_assert!(group_stage_masks[left_group_idx].contains(GkrOutputStageMask::ROTATION));
debug_assert!(group_stage_masks[right_group_idx].contains(GkrOutputStageMask::ROTATION));
debug_assert!(group_stage_masks[point_group_idx].contains(GkrOutputStageMask::ROTATION));

let (left_evals, right_evals, point_evals) = split_rotation_evals(&rotation.proof.evals);

Expand Down Expand Up @@ -398,6 +398,11 @@ pub fn prove_main_constraints_impl<
else {
panic!("ecc proof provided for non-ecc layer")
};
debug_assert!(group_stage_masks[x_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[y_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[slope_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[x3_group_idx].contains(GkrOutputStageMask::ECC));
debug_assert!(group_stage_masks[y3_group_idx].contains(GkrOutputStageMask::ECC));

let sample_r = transcript.sample_and_append_vec(b"ecc_gkr_bridge_r", 1)[0];
let claims = derive_ecc_bridge_claims(ecc_proof, sample_r, num_var_with_rotation)
Expand Down
23 changes: 20 additions & 3 deletions ceno_zkvm/src/scheme/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,8 +429,17 @@ impl<
let num_var_with_rotation = log2_num_instances + cs.rotation_vars().unwrap_or(0);

// build main witness
let records = info_span!("[ceno] build_main_witness")
.in_scope(|| build_main_witness::<E, PCS, PB, PD>(cs, input, challenges));
let records = info_span!("[ceno] build_main_witness").in_scope(|| {
// ECC and rotation have dedicated witness/eval flows. For tower proving we only
// materialize the tower-facing GKR outputs here to avoid keeping unrelated output
// MLEs resident in VRAM during tower prove.
build_main_witness::<E, PCS, PB, PD>(
cs,
input,
challenges,
crate::scheme::utils::WitnessBuildStage::Tower,
)
});

let span = entered_span!("prove_tower_relation", profiling_2 = true);
// prove the product and logup sum relation between layers in tower
Expand Down Expand Up @@ -762,12 +771,20 @@ where
// build main witness
let records =
info_span!("[ceno] build_main_witness").in_scope(|| {
// ECC and rotation have dedicated witness/eval flows. For tower proving we only
// materialize the tower-facing GKR outputs here to avoid keeping unrelated output
// MLEs resident in VRAM during tower prove.
build_main_witness::<
E,
PCS,
GpuBackend<E, PCS>,
gkr_iop::gpu::GpuProver<GpuBackend<E, PCS>>,
>(cs, &input, challenges)
>(
cs,
&input,
challenges,
crate::scheme::utils::WitnessBuildStage::Tower,
)
});

let span = entered_span!("prove_tower_relation", profiling_2 = true);
Expand Down
Loading
Loading